Add supporting for remaining UseIRProfile cases

Author: chiragramani

include/swift/AST/DiagnosticsIRGen.def

@@ -41,6 +41,11 @@ ERROR(too_few_output_filenames,none,
 ERROR(no_input_files_for_mt,none,
       "no swift input files for multi-threaded compilation", ())
 
+ERROR(ir_profile_read_failed, none,
+      "error reading profile '%0': %1", (StringRef, StringRef))
+ERROR(ir_profile_invalid, none,
+      "invalid ir profile '%0'", (StringRef))
+
 ERROR(alignment_dynamic_type_layout_unsupported,none,
       "'@_alignment' is not supported on types with dynamic layout", ())
 ERROR(alignment_less_than_natural,none,

lib/IRGen/IRGen.cpp

@@ -72,6 +72,7 @@
 #include "llvm/Passes/PassBuilder.h"
 #include "llvm/Passes/PassPlugin.h"
 #include "llvm/Passes/StandardInstrumentations.h"
+#include "llvm/ProfileData/InstrProfReader.h"
 #include "llvm/Remarks/Remark.h"
 #include "llvm/Remarks/RemarkStreamer.h"
 #include "llvm/Support/CommandLine.h"
@@ -214,8 +215,23 @@ static void align(llvm::Module *Module) {
     }
 }
 
+static std::unique_ptr<llvm::IndexedInstrProfReader>
+getProfileReader(const Twine &ProfileName, llvm::vfs::FileSystem &FS,
+                 DiagnosticEngine &Diags) {
+  auto ReaderOrErr = llvm::IndexedInstrProfReader::create(ProfileName, FS);
+  if (auto E = ReaderOrErr.takeError()) {
+    llvm::handleAllErrors(std::move(E), [&](const llvm::ErrorInfoBase &EI) {
+      Diags.diagnose(SourceLoc(), diag::ir_profile_read_failed,
+                     ProfileName.str(), EI.message());
+    });
+    return nullptr;
+  }
+  return std::move(*ReaderOrErr);
+}
+
 static void populatePGOOptions(std::optional<PGOOptions> &Out,
-                               const IRGenOptions &Opts) {
+                               const IRGenOptions &Opts,
+                               DiagnosticEngine &Diags) {
   if (!Opts.UseSampleProfile.empty()) {
     Out = PGOOptions(
       /*ProfileFile=*/ Opts.UseSampleProfile,
@@ -234,31 +250,60 @@ static void populatePGOOptions(std::optional<PGOOptions> &Out,
   if (Opts.EnableCSIRProfileGen) {
     const bool hasUse = !Opts.UseIRProfile.empty();
     Out = PGOOptions(
-      /*ProfileFile=*/ Opts.UseIRProfile,
-      /*CSProfileGenFile=*/ Opts.InstrProfileOutput,
-      /*ProfileRemappingFile=*/ "",
-      /*MemoryProfile=*/ "",
-      /*FS=*/ llvm::vfs::getRealFileSystem(),
-      /*Action=*/ hasUse ? PGOOptions::IRUse : PGOOptions::NoAction,
-      /*CSPGOAction=*/ PGOOptions::CSIRInstr,
-      /*ColdType=*/ PGOOptions::ColdFuncOpt::Default,
-      /*DebugInfoForProfiling=*/ Opts.DebugInfoForProfiling
-    );
+        /*ProfileFile=*/Opts.UseIRProfile,
+        /*CSProfileGenFile=*/Opts.InstrProfileOutput,
+        /*ProfileRemappingFile=*/"",
+        /*MemoryProfile=*/"",
+        /*FS=*/llvm::vfs::getRealFileSystem(),
+        /*Action=*/hasUse ? PGOOptions::IRUse : PGOOptions::NoAction,
+        /*CSPGOAction=*/PGOOptions::CSIRInstr,
+        /*ColdType=*/PGOOptions::ColdFuncOpt::Default,
+        /*DebugInfoForProfiling=*/Opts.DebugInfoForProfiling);
     return;
   }
 
   if (Opts.EnableIRProfileGen) {
     Out = PGOOptions(
-      /*ProfileFile=*/ Opts.InstrProfileOutput,
-      /*CSProfileGenFile=*/ "",
-      /*ProfileRemappingFile=*/ "",
-      /*MemoryProfile=*/ "",
-      /*FS=*/ llvm::vfs::getRealFileSystem(),
-      /*Action=*/ PGOOptions::IRInstr,
-      /*CSPGOAction=*/ PGOOptions::NoCSAction,
-      /*ColdType=*/ PGOOptions::ColdFuncOpt::Default,
-      /*DebugInfoForProfiling=*/ Opts.DebugInfoForProfiling
-    );
+        /*ProfileFile=*/Opts.InstrProfileOutput,
+        /*CSProfileGenFile=*/"",
+        /*ProfileRemappingFile=*/"",
+        /*MemoryProfile=*/"",
+        /*FS=*/llvm::vfs::getRealFileSystem(),
+        /*Action=*/PGOOptions::IRInstr,
+        /*CSPGOAction=*/PGOOptions::NoCSAction,
+        /*ColdType=*/PGOOptions::ColdFuncOpt::Default,
+        /*DebugInfoForProfiling=*/Opts.DebugInfoForProfiling);
+    return;
+  }
+
+  if (!Opts.UseIRProfile.empty()) {
+    auto FS = llvm::vfs::getRealFileSystem();
+    std::unique_ptr<llvm::IndexedInstrProfReader> Reader =
+        getProfileReader(Opts.UseIRProfile.c_str(), *FS, Diags);
+    if (!Reader)
+      return;
+
+    // Currently memprof profiles are only added at the IR level. Mark the
+    // profile type as IR in that case as well and the subsequent matching
+    // needs to detect which is available (might be one or both).
+    const bool IsIR = Reader->isIRLevelProfile() || Reader->hasMemoryProfile();
+    const bool IsCS = Reader->hasCSIRLevelProfile();
+    if (!IsIR) {
+      Diags.diagnose(SourceLoc(), diag::ir_profile_invalid,
+                     Opts.UseIRProfile.c_str());
+      return;
+    }
+
+    Out = PGOOptions(
+        /*ProfileFile=*/Opts.UseIRProfile,
+        /*CSProfileGenFile=*/"",
+        /*ProfileRemappingFile=*/"",
+        /*MemoryProfile=*/"",
+        /*FS=*/FS,
+        /*Action=*/PGOOptions::IRUse,
+        /*CSPGOAction=*/IsCS ? PGOOptions::CSIRUse : PGOOptions::NoCSAction,
+        /*ColdType=*/PGOOptions::ColdFuncOpt::Default,
+        /*DebugInfoForProfiling=*/Opts.DebugInfoForProfiling);
     return;
   }
 
@@ -297,7 +342,7 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
                                      llvm::TargetMachine *TargetMachine,
                                      llvm::raw_pwrite_stream *out) {
   std::optional<PGOOptions> PGOOpt;
-  populatePGOOptions(PGOOpt, Opts);
+  populatePGOOptions(PGOOpt, Opts, Diags);
 
   PipelineTuningOptions PTO;