Add support for SamplePGO

include/swift/AST/IRGenOptions.h

@@ -500,6 +500,13 @@ class IRGenOptions {
   /// Path to the profdata file to be used for PGO, or the empty string.
   std::string UseProfile = "";
 
+  /// Path to the data file to be used for sampling-based PGO,
+  /// or the empty string.
+  std::string UseSampleProfile = "";
+
+  /// Controls whether DWARF discriminators are added to the IR.
+  unsigned DebugInfoForProfiling : 1;
+
   /// List of backend command-line options for -embed-bitcode.
   std::vector<uint8_t> CmdArgs;
 
@@ -584,7 +591,8 @@ class IRGenOptions {
         ColocateTypeDescriptors(true), UseRelativeProtocolWitnessTables(false),
         UseFragileResilientProtocolWitnesses(false), EnableHotColdSplit(false),
         EmitAsyncFramePushPopMetadata(false), EmitYieldOnce2AsYieldOnce(true),
-        AsyncFramePointerAll(false), UseProfilingMarkerThunks(false), CmdArgs(),
+        AsyncFramePointerAll(false), UseProfilingMarkerThunks(false),
+        DebugInfoForProfiling(false), CmdArgs(),
         SanitizeCoverage(llvm::SanitizerCoverageOptions()),
         TypeInfoFilter(TypeInfoDumpFilter::All),
         PlatformCCallingConvention(llvm::CallingConv::C), UseCASBackend(false),

include/swift/Option/Options.td

@@ -1477,6 +1477,10 @@ def profile_generate : Flag<["-"], "profile-generate">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   HelpText<"Generate instrumented code to collect execution counts">;
 
+def debug_info_for_profiling : Flag<["-"], "debug-info-for-profiling">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  HelpText<"Emit extra debug info (DWARF discriminators) to make sampling-based profiling more accurate">;
+
 def profile_use : CommaJoined<["-"], "profile-use=">,
   Flags<[FrontendOption, NoInteractiveOption, ArgumentIsPath]>,
   MetaVarName<"<profdata>">,
@@ -1486,6 +1490,11 @@ def profile_coverage_mapping : Flag<["-"], "profile-coverage-mapping">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   HelpText<"Generate coverage data for use with profiled execution counts">;
 
+def profile_sample_use : CommaJoined<["-"], "profile-sample-use=">,
+  Flags<[FrontendOption, NoInteractiveOption, ArgumentIsPath]>,
+  MetaVarName<"<profile data>">,
+  HelpText<"Supply sampling-based profiling data from llvm-profdata to enable profile-guided optimization">;
+
 def embed_bitcode : Flag<["-"], "embed-bitcode">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   HelpText<"Embed LLVM IR bitcode as data">;

lib/Frontend/CompilerInvocation.cpp

@@ -3177,6 +3177,11 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
   const Arg *ProfileUse = Args.getLastArg(OPT_profile_use);
   Opts.UseProfile = ProfileUse ? ProfileUse->getValue() : "";
 
+  const Arg *ProfileSampleUse = Args.getLastArg(OPT_profile_sample_use);
+  Opts.UseSampleProfile = ProfileSampleUse ? ProfileSampleUse->getValue() : "";
+
+  Opts.DebugInfoForProfiling |= Args.hasArg(OPT_debug_info_for_profiling);
+
   Opts.PrintInlineTree |= Args.hasArg(OPT_print_llvm_inline_tree);
   // Always producing all outputs when caching is enabled.
   Opts.AlwaysCompile |= Args.hasArg(OPT_always_compile_output_files) ||

lib/IRGen/IRGen.cpp

@@ -200,11 +200,45 @@ static void align(llvm::Module *Module) {
     }
 }
 
+static void populatePGOOptions(std::optional<PGOOptions> &Out,
+                               const IRGenOptions &Opts) {
+  if (!Opts.UseSampleProfile.empty()) {
+    Out = PGOOptions(
+      /*ProfileFile=*/ Opts.UseSampleProfile,
+      /*CSProfileGenFile=*/ "",
+      /*ProfileRemappingFile=*/ "",
+      /*MemoryProfile=*/ "",
+      /*FS=*/ llvm::vfs::getRealFileSystem(), // TODO: is this fine?
+      /*Action=*/ PGOOptions::SampleUse,
+      /*CSPGOAction=*/ PGOOptions::NoCSAction,
+      /*ColdType=*/ PGOOptions::ColdFuncOpt::Default,
+      /*DebugInfoForProfiling=*/ Opts.DebugInfoForProfiling
+    );
+    return;
+  }
+
+  if (Opts.DebugInfoForProfiling) {
+    Out = PGOOptions(
+        /*ProfileFile=*/ "",
+        /*CSProfileGenFile=*/ "",
+        /*ProfileRemappingFile=*/ "",
+        /*MemoryProfile=*/ "",
+        /*FS=*/ nullptr,
+        /*Action=*/ PGOOptions::NoAction,
+        /*CSPGOAction=*/ PGOOptions::NoCSAction,
+        /*ColdType=*/ PGOOptions::ColdFuncOpt::Default,
+        /*DebugInfoForProfiling=*/ true
+    );
+    return;
+  }
+}
+
 void swift::performLLVMOptimizations(const IRGenOptions &Opts,
                                      llvm::Module *Module,
                                      llvm::TargetMachine *TargetMachine,
                                      llvm::raw_pwrite_stream *out) {
   std::optional<PGOOptions> PGOOpt;
+  populatePGOOptions(PGOOpt, Opts);
 
   PipelineTuningOptions PTO;
 

lib/IRGen/IRGenDebugInfo.cpp

@@ -839,6 +839,7 @@ class IRGenDebugInfoImpl : public IRGenDebugInfo {
                               TheCU->getProducer(), true, StringRef(), 0,
                               RemappedASTFile, llvm::DICompileUnit::FullDebug,
                               Signature);
+        // NOTE: not setting DebugInfoForProfiling here
         DIB.finalize();
       }
     }
@@ -2504,7 +2505,7 @@ IRGenDebugInfoImpl::IRGenDebugInfoImpl(const IRGenOptions &Opts,
           ? llvm::DICompileUnit::FullDebug
           : llvm::DICompileUnit::LineTablesOnly,
       /* DWOId */ 0, /* SplitDebugInlining */ true,
-      /* DebugInfoForProfiling */ false,
+      /* DebugInfoForProfiling */ Opts.DebugInfoForProfiling,
       llvm::DICompileUnit::DebugNameTableKind::Default,
       /* RangesBaseAddress */ false, DebugPrefixMap.remapPath(Sysroot), SDK);
 

lib/IRGen/IRGenSIL.cpp

@@ -1939,6 +1939,11 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
     // LLVM doesn't have an attribute for -O
   }
 
+  if (!IGM.IRGen.Opts.UseSampleProfile.empty()) {
+    // This attribute helps in LTO situations: https://reviews.llvm.org/D79959
+    CurFn->addFnAttr("use-sample-profile");
+  }
+
   // Emit the thunk that calls the previous implementation if this is a dynamic
   // replacement.
   if (f->getDynamicallyReplacedFunction()) {

test/DebugInfo/discriminators_for_profiling.swift

@@ -0,0 +1,16 @@
+// RUN: %target-swift-frontend %s -emit-ir -g -debug-info-for-profiling -o - | %FileCheck %s
+
+// CHECK: distinct !DICompileUnit(language: DW_LANG_Swift
+// CHECK-SAME: debugInfoForProfiling: true
+
+// CHECK: !DILexicalBlockFile({{.*}} discriminator: 2)
+
+public func lobster(_ n: Int) -> Bool {
+  guard n > 0 else { fatalError("too cold!") }
+
+  if n > 100 {
+    print("warm but ok")
+  }
+
+  return n < 120
+}

test/Profiler/samplepgo.swift

@@ -0,0 +1,43 @@
+// RUN: %empty-directory(%t)
+// RUN: split-file %s %t
+
+// ----------------------------------------
+// Test the -profile-sample-use= flag using bogus data, to ensure it's actually
+// reaching LLVM in the expected way.
+
+// RUN: %target-swift-frontend %t/program.swift -module-name test -emit-ir \
+// RUN:                  -O -profile-sample-use=%t/profile.txt -o %t/has-data.ll
+
+// RUN: %FileCheck %s < %t/has-data.ll
+
+// CHECK: define{{.*}} @"$s4test8anythingyyF"() #[[ATTRID:[0-9]+]]
+// CHECK: attributes #[[ATTRID]] = {{.*}} "use-sample-profile"
+
+// CHECK-LABEL: !llvm.module.flags
+// CHECK: !{!"ProfileFormat", !"SampleProfile"}
+// CHECK: !{!"TotalCount", i64 2001}
+
+// ----------------------------------------
+// Now, test cases where there there should not be any profile metadata,
+// such as when no data is provided
+
+// RUN: %target-swift-frontend %t/program.swift -module-name test -emit-ir \
+// RUN:                        -O -o %t/no-data.ll
+
+// RUN: %FileCheck --check-prefix CHECK-NODATA %s < %t/no-data.ll
+
+
+// CHECK-NODATA: define{{.*}} @"$s4test8anythingyyF"() #[[ATTRID:[0-9]+]]
+// CHECK-NODATA-NOT: attributes #[[ATTRID]] = {{.*}} "use-sample-profile"
+
+// CHECK-NODATA-LABEL: !llvm.module.flags
+// CHECK-NODATA-NOT: Profile
+
+
+//--- program.swift
+public func anything() {}
+
+
+//--- profile.txt
+bar:100:100
+ 1: 2001