Adds IR level PGO Instrumentation options

Sources/SwiftDriver/Driver/Driver.swift

@@ -3307,24 +3307,92 @@ extension Driver {
     }
   }
 
+  static private func validateProfilingGenerateArgs(
+    _ parsedOptions: inout ParsedOptions,
+    diagnosticEngine: DiagnosticsEngine
+  ) {
+    let genFlags: [Option] = [
+      .profileGenerate,
+      .irProfileGenerate,
+      .irProfileGenerateEQ,
+      .csProfileGenerate,
+      .csProfileGenerateEQ,
+    ]
+    func resolveDualFormConflict(_ plain: Option, _ equalsForm: Option) {
+       if parsedOptions.hasArgument(plain),
+          parsedOptions.hasArgument(equalsForm)
+       {
+         diagnosticEngine.emit(
+           .error(Error.conflictingOptions(plain, equalsForm)),
+           location: nil
+         )
+         providedGen.removeAll { $0 == equalsForm }
+       }
+     }
+
+    var providedGen = genFlags.filter { parsedOptions.hasArgument($0) }
+    resolveDualFormConflict(.irProfileGenerate, .irProfileGenerateEQ)
+    resolveDualFormConflict(.csProfileGenerate, .csProfileGenerateEQ)
+
+    guard providedGen.count >= 2 else { return }
+    for i in 1..<providedGen.count {
+      let error = Error.conflictingOptions(providedGen[i - 1], providedGen[i])
+      diagnosticEngine.emit(.error(error), location: nil)
+    }
+  }
+
+  static private func validateProfilingUseArgs(
+    _ parsedOptions: inout ParsedOptions,
+    diagnosticEngine: DiagnosticsEngine
+  ) {
+    let conflictingGenFlags: [Option] = [
+      .profileGenerate,
+      .irProfileGenerate,
+      .irProfileGenerateEQ,
+    ]
+    let useProfArgs: [Option] = [
+      .profileUse,
+      .profileSampleUse,
+    ]
+    let providedUse = useProfArgs.filter { parsedOptions.hasArgument($0) }
+    guard !providedUse.isEmpty else { return }
+
+    // At most one *use* option allowed
+    if providedUse.count > 1 {
+      for i in 0..<(providedUse.count - 1) {
+        for j in (i + 1)..<providedUse.count {
+          diagnosticEngine.emit(
+            .error(Error.conflictingOptions(providedUse[i], providedUse[j])),
+            location: nil
+          )
+        }
+      }
+    }
+
+    // If no generate flags, we're good.
+    let providedGen = conflictingGenFlags.filter { parsedOptions.hasArgument($0) }
+    guard !providedGen.isEmpty else { return }
+
+    // We already diagnosed if the user passed more than one "use" option
+    // (e.g. both `-profile-use` and `-profile-sample-use`). To avoid
+    // spamming diagnostics, we now treat the first provided "use" flag
+    // as the canonical representative.
+    let canonicalUse = providedUse[0]
+
+    // Generate vs Use are mutually exclusive
+    for g in providedGen {
+      diagnosticEngine.emit(.error(Error.conflictingOptions(g, canonicalUse)), location: nil)
+    }
+  }
+
+
   static func validateProfilingArgs(_ parsedOptions: inout ParsedOptions,
                                     fileSystem: FileSystem,
                                     workingDirectory: AbsolutePath?,
                                     diagnosticEngine: DiagnosticsEngine) {
-    let conflictingProfArgs: [Option] = [.profileGenerate,
-                                         .profileUse,
-                                         .profileSampleUse]
-
     // Find out which of the mutually exclusive profiling arguments were provided.
-    let provided = conflictingProfArgs.filter { parsedOptions.hasArgument($0) }
-
-    // If there's at least two of them, there's a conflict.
-    if provided.count >= 2 {
-      for i in 1..<provided.count {
-        let error = Error.conflictingOptions(provided[i-1], provided[i])
-        diagnosticEngine.emit(.error(error), location: nil)
-      }
-    }
+    validateProfilingGenerateArgs(&parsedOptions, diagnosticEngine: diagnosticEngine)
+    validateProfilingUseArgs(&parsedOptions, diagnosticEngine: diagnosticEngine)
 
     // Ensure files exist for the given paths.
     func checkForMissingProfilingData(_ profileDataArgs: [String]) {

Sources/SwiftDriver/Jobs/DarwinToolchain+LinkerSupport.swift

@@ -258,9 +258,7 @@ extension DarwinToolchain {
       fileSystem: fileSystem
     )
 
-    if parsedOptions.hasArgument(.profileGenerate) {
-      commandLine.appendFlag("-fprofile-generate")
-    }
+    commandLine.appendFlags(mapInstrumentationTypeToClangArgs(from: &parsedOptions))
 
     // These custom arguments should be right before the object file at the
     // end.

Sources/SwiftDriver/Jobs/FrontendJobHelpers.swift

@@ -237,6 +237,10 @@ extension Driver {
     try commandLine.appendLast(.RpassMissedEQ, from: &parsedOptions)
     try commandLine.appendLast(.suppressWarnings, from: &parsedOptions)
     try commandLine.appendLast(.profileGenerate, from: &parsedOptions)
+    try commandLine.appendLast(.irProfileGenerate, from: &parsedOptions)
+    try commandLine.appendLast(.irProfileGenerateEQ, from: &parsedOptions)
+    try commandLine.appendLast(.csProfileGenerate, from: &parsedOptions)
+    try commandLine.appendLast(.csProfileGenerateEQ, from: &parsedOptions)
     try commandLine.appendLast(.profileUse, from: &parsedOptions)
     try commandLine.appendLast(.profileCoverageMapping, from: &parsedOptions)
     try commandLine.appendLast(.debugInfoForProfiling, from: &parsedOptions)

Sources/SwiftDriver/Jobs/GenericUnixToolchain+LinkerSupport.swift

@@ -303,7 +303,7 @@ extension GenericUnixToolchain {
         }
       }
 
-      if parsedOptions.hasArgument(.profileGenerate) {
+      if needsInstrumentedProfile(from: &parsedOptions) {
         let environment = (targetTriple.environment == .android) ? "-android" : ""
         let libProfile = VirtualPath.lookup(targetInfo.runtimeResourcePath.path)
           .appending(components: "clang", "lib", targetTriple.osNameUnversioned,

Sources/SwiftDriver/Jobs/WebAssemblyToolchain+LinkerSupport.swift

@@ -176,7 +176,7 @@ extension WebAssemblyToolchain {
         commandLine.appendFlag("-fsanitize=\(sanitizerNames)")
       }
 
-      if parsedOptions.hasArgument(.profileGenerate) {
+      if needsInstrumentedProfile(from: &parsedOptions) {
         let libProfile = VirtualPath.lookup(targetInfo.runtimeResourcePath.path)
           .appending(components: "clang", "lib", targetTriple.osName,
                                  "libclang_rt.profile-\(targetTriple.archName).a")

Sources/SwiftDriver/Jobs/WindowsToolchain+LinkerSupport.swift

@@ -50,7 +50,7 @@ extension WindowsToolchain {
       // 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 parsedOptions.hasArgument(.profileGenerate) { return true }
+      if needsInstrumentedProfile(from: &parsedOptions) { return true }
 
       return false
     }()
@@ -228,11 +228,12 @@ extension WindowsToolchain {
       commandLine.appendFlag("-fsanitize=\(sanitize)")
     }
 
-    if parsedOptions.contains(.profileGenerate) {
+    if needsInstrumentedProfile(from: &parsedOptions) {
       assert(bForceLLD,
              "LLD is currently required for profiling (rdar://131295678)")
 
-      commandLine.appendFlag("-fprofile-generate")
+      commandLine.appendFlags(mapInstrumentationTypeToClangArgs(from: &parsedOptions))
+
       // FIXME(rdar://131295678): Currently profiling requires the ability to
       // emit duplicate weak symbols. Assume we're using lld and pass
       // `-lld-allow-duplicate-weak` to enable this behavior.

Sources/SwiftDriver/Toolchains/Toolchain.swift

@@ -376,6 +376,38 @@ extension Toolchain {
     }
     return clangArg
   }
+
+  internal func mapInstrumentationTypeToClangArgs(from options: inout ParsedOptions) -> [String] {
+    var args: [String] = []
+
+    if options.contains(.profileGenerate) || options.contains(.irProfileGenerate) {
+      args.append("-fprofile-generate")
+    }
+
+    if options.contains(.irProfileGenerateEQ),
+       let path = options.getLastArgument(.irProfileGenerateEQ)?.asSingle {
+      args.append("-fprofile-generate=\(path)")
+    }
+
+    if options.contains(.csProfileGenerate) {
+      args.append("-fcs-profile-generate")
+    }
+
+    if options.contains(.csProfileGenerateEQ),
+       let path = options.getLastArgument(.csProfileGenerateEQ)?.asSingle {
+      args.append("-fcs-profile-generate=\(path)")
+    }
+
+    return args
+  }
+
+  internal func needsInstrumentedProfile(from parsedOptions: inout ParsedOptions) -> Bool {
+    parsedOptions.contains(.profileGenerate) ||
+    parsedOptions.contains(.irProfileGenerate) ||
+    parsedOptions.contains(.irProfileGenerateEQ) ||
+    parsedOptions.contains(.csProfileGenerate) ||
+    parsedOptions.contains(.csProfileGenerateEQ)
+  }
 }
 
 @_spi(Testing) public enum ToolchainError: Swift.Error {

Sources/SwiftOptions/Options.swift

@@ -105,6 +105,8 @@ extension Option {
   public static let coveragePrefixMap: Option = Option("-coverage-prefix-map", .separate, attributes: [.frontend], metaVar: "<prefix=replacement>", helpText: "Remap source paths in coverage info")
   public static let CrossModuleOptimization: Option = Option("-cross-module-optimization", .flag, attributes: [.helpHidden, .frontend], helpText: "Perform cross-module optimization")
   public static let crosscheckUnqualifiedLookup: Option = Option("-crosscheck-unqualified-lookup", .flag, attributes: [.frontend, .noDriver], helpText: "Compare legacy DeclContext- to ASTScope-based unqualified name lookup (for debugging)")
+  public static let csProfileGenerateEQ: Option = Option("-cs-profile-generate=", .joined, attributes: [.frontend, .noInteractive], metaVar: "<directory>", helpText: "Generate instrumented code to collect context sensitive execution counts into <directory>/default.profraw (overridden by LLVM_PROFILE_FILE env var)")
+  public static let csProfileGenerate: Option = Option("-cs-profile-generate", .flag, attributes: [.frontend, .noInteractive], helpText: "Generate instrumented code to collect context sensitive execution counts into default.profraw (overridden by LLVM_PROFILE_FILE env var)")
   public static let cxxInteropGettersSettersAsProperties: Option = Option("-cxx-interop-getters-setters-as-properties", .flag, attributes: [.helpHidden, .frontend, .noDriver], helpText: "Import getters and setters as computed properties in Swift")
   public static let cxxInteropUseOpaquePointerForMoveonly: Option = Option("-cxx-interop-use-opaque-pointer-for-moveonly", .flag, attributes: [.helpHidden, .frontend, .noDriver], helpText: "Testing flag that will be eliminated soon. Do not use.")
   public static let cxxInteroperabilityMode: Option = Option("-cxx-interoperability-mode=", .joined, attributes: [.frontend, .synthesizeInterface], helpText: "Enables C++ interoperability; pass 'default' to enable or 'off' to disable")
@@ -643,6 +645,9 @@ extension Option {
   public static let swiftinterfaceCompilerVersion: Option = Option("-interface-compiler-version", .separate, attributes: [.helpHidden, .frontend], metaVar: "<intcvers>", helpText: "The version of the Swift compiler used to generate a .swiftinterface file")
   public static let internalizeAtLink: Option = Option("-internalize-at-link", .flag, attributes: [.helpHidden, .frontend, .noDriver], helpText: "Allow internalizing public symbols and vtables at link time (assume all client code of public types is part of the same link unit, or that external symbols are explicitly requested via -exported_symbols_list)")
   public static let interpret: Option = Option("-interpret", .flag, attributes: [.helpHidden, .frontend, .noDriver], helpText: "Immediate mode", group: .modes)
+  public static let irProfileGenerateEQ: Option = Option("-ir-profile-generate=", .joined, attributes: [.frontend, .noInteractive], metaVar: "<directory>", helpText: "Generate instrumented code to collect execution counts into <directory>/default.profraw (overridden by LLVM_PROFILE_FILE env var)")
+  public static let irProfileGenerate: Option = Option("-ir-profile-generate", .flag, attributes: [.frontend, .noInteractive], helpText: "Generate instrumented code to collect execution counts into default.profraw (overridden by LLVM_PROFILE_FILE env var)")
+  public static let irProfileUse: Option = Option("-ir-profile-use=", .commaJoined, attributes: [.frontend, .noInteractive, .argumentIsPath], metaVar: "<profdata>", helpText: "Supply an IR-level PGO profdata file to enable profile-guided optimization")
   public static let Isystem: Option = Option("-Isystem", .separate, attributes: [.frontend, .synthesizeInterface, .argumentIsPath], helpText: "Add directory to the system import search path")
   public static let I: Option = Option("-I", .joinedOrSeparate, attributes: [.frontend, .synthesizeInterface, .argumentIsPath], helpText: "Add directory to the import search path")
   public static let i: Option = Option("-i", .flag, group: .modes)
@@ -1078,6 +1083,8 @@ extension Option {
       Option.coveragePrefixMap,
       Option.CrossModuleOptimization,
       Option.crosscheckUnqualifiedLookup,
+      Option.csProfileGenerateEQ,
+      Option.csProfileGenerate,
       Option.cxxInteropGettersSettersAsProperties,
       Option.cxxInteropUseOpaquePointerForMoveonly,
       Option.cxxInteroperabilityMode,
@@ -1616,6 +1623,9 @@ extension Option {
       Option.swiftinterfaceCompilerVersion,
       Option.internalizeAtLink,
       Option.interpret,
+      Option.irProfileGenerateEQ,
+      Option.irProfileGenerate,
+      Option.irProfileUse,
       Option.Isystem,
       Option.I,
       Option.i,