Add IRPGO and CSIRPGO options to Swift

include/swift/AST/IRGenOptions.h

@@ -552,6 +552,16 @@ class IRGenOptions {
   /// or the empty string.
   std::string UseSampleProfile = "";
 
+  /// Path to the context-sensitive profile file to be used for CS-IR PGO,
+  /// or the default string.
+  std::string CSProfileGenFile = "default_%m.profraw";
+
+  /// Whether to enable context-sensitive IR PGO generation.
+  bool EnableCSIRProfileGen = false;
+
+  /// Whether to enable IR level instrumentation.
+  bool EnableIRProfileGen = false;
+
   /// Controls whether DWARF discriminators are added to the IR.
   unsigned DebugInfoForProfiling : 1;
 

include/swift/Option/Options.td

@@ -1608,6 +1608,19 @@ def profile_sample_use : Joined<["-"], "profile-sample-use=">,
   MetaVarName<"<profile data>">,
   HelpText<"Supply sampling-based profiling data from llvm-profdata to enable profile-guided optimization">;
 
+def cs_profile_generate : Flag<["-"], "cs-profile-generate">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  HelpText<"Generate instrumented code to collect context sensitive execution counts into default.profraw (overridden by LLVM_PROFILE_FILE env var)">;
+
+def cs_profile_generate_EQ : Joined<["-"], "cs-profile-generate=">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  MetaVarName<"<directory>">,
+  HelpText<"Generate instrumented code to collect context sensitive execution counts into <directory>/default.profraw (overridden by LLVM_PROFILE_FILE env var)">;
+
+def ir_profile_generate: Flag<["-"], "ir-profile-generate">,
+  Flags<[FrontendOption, NoInteractiveOption]>,
+  HelpText<"Generate instrumented code to collect execution counts into default.profraw (overridden by LLVM_PROFILE_FILE env var)">;
+
 def embed_bitcode : Flag<["-"], "embed-bitcode">,
   Flags<[FrontendOption, NoInteractiveOption]>,
   HelpText<"Embed LLVM IR bitcode as data">;

lib/Driver/DarwinToolChains.cpp

@@ -42,6 +42,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string
 toolchains::Darwin::findProgramRelativeToSwiftImpl(StringRef name) const {
@@ -470,7 +471,7 @@ void
 toolchains::Darwin::addProfileGenerationArgs(ArgStringList &Arguments,
                                              const JobContext &context) const {
   const llvm::Triple &Triple = getTriple();
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     SmallString<128> LibProfile;
     getClangLibraryPath(context.Args, LibProfile);
 

lib/Driver/Driver.cpp

@@ -175,14 +175,58 @@ static void validateWarningControlArgs(DiagnosticEngine &diags,
   }
 }
 
+/// Validates only *generate* profiling flags and their mutual conflicts.
+static void validateProfilingGenerateArgs(DiagnosticEngine &diags,
+                                          const ArgList &args) {
+  const Arg *ProfileGenerate = args.getLastArg(options::OPT_profile_generate);
+  const Arg *IRProfileGenerate =
+      args.getLastArg(options::OPT_ir_profile_generate);
+  const Arg *CSProfileGenerate =
+      args.getLastArg(options::OPT_cs_profile_generate);
+  const Arg *CSProfileGenerateEQ =
+      args.getLastArg(options::OPT_cs_profile_generate_EQ);
+
+  // If both CS Profile forms were specified, report a clear conflict.
+  if (CSProfileGenerate && CSProfileGenerateEQ) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-cs-profile-generate", "-cs-profile-generate=");
+    CSProfileGenerateEQ = nullptr;
+  }
+
+  llvm::SmallVector<std::pair<const Arg *, const char *>, 3> gens;
+  if (ProfileGenerate)
+    gens.push_back({ProfileGenerate, "-profile-generate"});
+  if (IRProfileGenerate)
+    gens.push_back({IRProfileGenerate, "-ir-profile-generate"});
+  if (CSProfileGenerate)
+    gens.push_back({CSProfileGenerate, "-cs-profile-generate"});
+  else if (CSProfileGenerateEQ)
+    gens.push_back({CSProfileGenerateEQ, "-cs-profile-generate="});
+
+  // Emit pairwise conflicts if more than one generate-mode was selected
+  for (size_t i = 0; i + 1 < gens.size(); ++i) {
+    for (size_t j = i + 1; j < gens.size(); ++j) {
+      diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                     gens[i].second, gens[j].second);
+    }
+  }
+}
+
 static void validateProfilingArgs(DiagnosticEngine &diags,
                                   const ArgList &args) {
+  validateProfilingGenerateArgs(diags, args);
   const Arg *ProfileGenerate = args.getLastArg(options::OPT_profile_generate);
   const Arg *ProfileUse = args.getLastArg(options::OPT_profile_use);
+  const Arg *IRProfileGenerate =
+      args.getLastArg(options::OPT_ir_profile_generate);
   if (ProfileGenerate && ProfileUse) {
     diags.diagnose(SourceLoc(), diag::error_conflicting_options,
                    "-profile-generate", "-profile-use");
   }
+  if (IRProfileGenerate && ProfileUse) {
+    diags.diagnose(SourceLoc(), diag::error_conflicting_options,
+                   "-ir-profile-generate", "-profile-use");
+  }
 
   // Check if the profdata is missing
   if (ProfileUse && !llvm::sys::fs::exists(ProfileUse->getValue())) {

lib/Driver/ToolChains.cpp

@@ -160,6 +160,15 @@ bool containsValue(
 
 }
 
+namespace swift::driver::toolchains {
+bool needsInstrProfileRuntime(const llvm::opt::ArgList &Args) {
+  return Args.hasArg(options::OPT_profile_generate) ||
+         Args.hasArg(options::OPT_cs_profile_generate) ||
+         Args.hasArg(options::OPT_cs_profile_generate_EQ) ||
+         Args.hasArg(options::OPT_ir_profile_generate);
+}
+} // namespace swift::driver::toolchains
+
 void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
                                       const CommandOutput &output,
                                       const ArgList &inputArgs,
@@ -323,6 +332,8 @@ void ToolChain::addCommonFrontendArgs(const OutputInfo &OI,
   inputArgs.AddLastArg(arguments, options::OPT_PackageCMO);
   inputArgs.AddLastArg(arguments, options::OPT_profile_generate);
   inputArgs.AddLastArg(arguments, options::OPT_profile_use);
+  inputArgs.AddLastArg(arguments, options::OPT_ir_profile_generate);
+  inputArgs.AddLastArg(arguments, options::OPT_cs_profile_generate);
   inputArgs.AddLastArg(arguments, options::OPT_profile_coverage_mapping);
   inputArgs.AddAllArgs(arguments, options::OPT_warning_treating_Group);
   inputArgs.AddLastArg(arguments, options::OPT_sanitize_EQ);

lib/Driver/ToolChains.h

@@ -25,6 +25,11 @@ class DiagnosticEngine;
 namespace driver {
 namespace toolchains {
 
+/// True if any *generation* mode of instrumentation-based profile is enabled.
+///
+/// This is used to determine if the profiler runtime should be linked.
+bool needsInstrProfileRuntime(const llvm::opt::ArgList &Args);
+
 class LLVM_LIBRARY_VISIBILITY Darwin : public ToolChain {
 protected:
 

lib/Driver/UnixToolChains.cpp

@@ -40,6 +40,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string
 toolchains::GenericUnix::sanitizerRuntimeLibName(StringRef Sanitizer,
@@ -330,7 +331,7 @@ toolchains::GenericUnix::constructInvocation(const DynamicLinkJobAction &job,
     }
   }
 
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     SmallString<128> LibProfile(SharedResourceDirPath);
     llvm::sys::path::remove_filename(LibProfile); // remove platform name
     llvm::sys::path::append(LibProfile, "clang", "lib");

lib/Driver/WebAssemblyToolChains.cpp

@@ -38,6 +38,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string
 toolchains::WebAssembly::sanitizerRuntimeLibName(StringRef Sanitizer,
@@ -168,7 +169,7 @@ toolchains::WebAssembly::constructInvocation(const DynamicLinkJobAction &job,
         "-fsanitize=" + getSanitizerList(context.OI.SelectedSanitizers)));
   }
 
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     SmallString<128> LibProfile(SharedResourceDirPath);
     llvm::sys::path::remove_filename(LibProfile); // remove platform name
     llvm::sys::path::append(LibProfile, "clang", "lib");

lib/Driver/WindowsToolChains.cpp

@@ -36,6 +36,7 @@
 using namespace swift;
 using namespace swift::driver;
 using namespace llvm::opt;
+using namespace swift::driver::toolchains;
 
 std::string toolchains::Windows::sanitizerRuntimeLibName(StringRef Sanitizer,
                                                          bool shared) const {
@@ -89,7 +90,7 @@ toolchains::Windows::constructInvocation(const DynamicLinkJobAction &job,
     // for now, which supports the behavior via a flag.
     // TODO: Once we've changed coverage to no longer rely on emitting
     // duplicate weak symbols (rdar://131295678), we can remove this.
-    if (context.Args.getLastArg(options::OPT_profile_generate)) {
+    if (swift::driver::toolchains::needsInstrProfileRuntime(context.Args)) {
       return true;
     }
     return false;
@@ -186,7 +187,7 @@ toolchains::Windows::constructInvocation(const DynamicLinkJobAction &job,
                         sanitizerRuntimeLibName("ubsan"));
   }
 
-  if (context.Args.hasArg(options::OPT_profile_generate)) {
+  if (needsInstrProfileRuntime(context.Args)) {
     Arguments.push_back(context.Args.MakeArgString("-Xlinker"));
     Arguments.push_back(context.Args.MakeArgString(
         Twine({"-include:", llvm::getInstrProfRuntimeHookVarName()})));

lib/Frontend/CompilerInvocation.cpp

@@ -3497,6 +3497,12 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
   const Arg *ProfileSampleUse = Args.getLastArg(OPT_profile_sample_use);
   Opts.UseSampleProfile = ProfileSampleUse ? ProfileSampleUse->getValue() : "";
 
+  Opts.EnableIRProfileGen = Args.hasArg(OPT_ir_profile_generate);
+  Opts.EnableCSIRProfileGen = Args.hasArg(OPT_cs_profile_generate) ||
+                              Args.hasArg(OPT_cs_profile_generate_EQ);
+  if (auto V = Args.getLastArgValue(OPT_cs_profile_generate_EQ); !V.empty())
+    Opts.CSProfileGenFile = V.str();
+
   Opts.DebugInfoForProfiling |= Args.hasArg(OPT_debug_info_for_profiling);
 
 

lib/IRGen/IRGen.cpp

@@ -231,6 +231,37 @@ static void populatePGOOptions(std::optional<PGOOptions> &Out,
     return;
   }
 
+  if (Opts.EnableCSIRProfileGen) {
+    const bool hasUse = !Opts.UseProfile.empty();
+    Out = PGOOptions(
+      /*ProfileFile=*/ hasUse ? Opts.UseProfile : "",
+      /*CSProfileGenFile=*/ Opts.CSProfileGenFile,
+      /*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=*/ "",
+      /*CSProfileGenFile=*/ "",
+      /*ProfileRemappingFile=*/ "",
+      /*MemoryProfile=*/ "",
+      /*FS=*/ llvm::vfs::getRealFileSystem(),
+      /*Action=*/ PGOOptions::IRInstr,
+      /*CSPGOAction=*/ PGOOptions::NoCSAction,
+      /*ColdType=*/ PGOOptions::ColdFuncOpt::Default,
+      /*DebugInfoForProfiling=*/ Opts.DebugInfoForProfiling
+    );
+    return;
+  }
+
   if (Opts.DebugInfoForProfiling) {
     Out = PGOOptions(
         /*ProfileFile=*/ "",