PGO: new -profile-sample-use= frontend option

This is meant to match clang's `-fprofile-sample-use=`

rdar://135443278

Author: kavon

include/swift/AST/IRGenOptions.h

@@ -500,6 +500,10 @@ 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 = "";
+
   /// List of backend command-line options for -embed-bitcode.
   std::vector<uint8_t> CmdArgs;
 

include/swift/Option/Options.td

@@ -1481,6 +1481,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 an external tool 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,9 @@ 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.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,6 +200,23 @@ 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,
+      /*DebugInfoForProfiling=*/ false
+    );
+    return;
+  }
+}
+
 void swift::performLLVMOptimizations(const IRGenOptions &Opts,
                                      llvm::Module *Module,
                                      llvm::TargetMachine *TargetMachine,
@@ -222,6 +239,7 @@ void swift::performLLVMOptimizations(const IRGenOptions &Opts,
     PTO.LoopVectorization = true;
     PTO.SLPVectorization = true;
     PTO.MergeFunctions = true;
+    populatePGOOptions(PGOOpt, Opts);
     // Splitting trades code size to enhance memory locality, avoid in -Osize.
     DoHotColdSplit = Opts.EnableHotColdSplit && !Opts.optimizeForSize();
     level = llvm::OptimizationLevel::Os;

lib/IRGen/IRGenSIL.cpp

@@ -1939,6 +1939,13 @@ IRGenSILFunction::IRGenSILFunction(IRGenModule &IGM, SILFunction *f)
     // LLVM doesn't have an attribute for -O
   }
 
+  const bool optimizing = IGM.IRGen.Opts.shouldOptimize() &&
+                          !IGM.IRGen.Opts.DisableLLVMOptzns;
+  if (optimizing && !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/Profiler/samplepgo.swift

@@ -0,0 +1,49 @@
+// 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 optimizations are requested, or no data is provided
+
+// RUN: %target-swift-frontend %t/program.swift -module-name test -emit-ir \
+// RUN:            -profile-sample-use=%t/profile.txt -o %t/no-optimizations.ll
+
+// RUN: %FileCheck --check-prefix CHECK-NODATA %s < %t/no-optimizations.ll
+
+
+// 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