Add ir and cs instrumentation options

Author: chiragramani

Sources/SwiftDriver/Driver/Driver.swift

@@ -3334,24 +3334,87 @@ extension Driver {
     }
   }
 
+  static private func validateProfilingGenerateArgs(
+    _ parsedOptions: inout ParsedOptions,
+    diagnosticEngine: DiagnosticsEngine
+  ) {
+    let genFlags: [Option] = [
+      .profileGenerate,
+      .irProfileGenerate,
+      .csProfileGenerate,
+      .csProfileGenerateEq,
+    ]
+
+    var providedGen = genFlags.filter { parsedOptions.hasArgument($0) }
+    if parsedOptions.hasArgument(.csProfileGenerate),
+      parsedOptions.hasArgument(.csProfileGenerateEq)
+    {
+      // If both forms were specified, report a clear conflict.
+      diagnosticEngine.emit(
+        .error(Error.conflictingOptions(.csProfileGenerate, .csProfileGenerateEq)),
+        location: nil
+      )
+      providedGen.removeAll { $0 == .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,
+    ]
+    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

@@ -243,6 +243,9 @@ 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(.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,32 @@ 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(.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(.csProfileGenerate) ||
+    parsedOptions.contains(.csProfileGenerateEq)
+  }
 }
 
 @_spi(Testing) public enum ToolchainError: Swift.Error {

Sources/SwiftOptions/Options.swift

@@ -804,6 +804,9 @@ extension Option {
   public static let printZeroStats: Option = Option("-print-zero-stats", .flag, attributes: [.helpHidden, .frontend], helpText: "Prints all stats even if they are zero")
   public static let profileCoverageMapping: Option = Option("-profile-coverage-mapping", .flag, attributes: [.frontend, .noInteractive], helpText: "Generate coverage data for use with profiled execution counts")
   public static let profileGenerate: Option = Option("-profile-generate", .flag, attributes: [.frontend, .noInteractive], helpText: "Generate instrumented code to collect execution counts")
+  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 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 csProfileGenerateEq: Option = Option("-cs-profile-generate=", .joined, attributes: [.frontend, .noInteractive, .argumentIsPath], metaVar: "<dir>", helpText: "Generate instrumented code to collect context sensitive execution counts into <directory>/default.profraw (overridden by LLVM_PROFILE_FILE env var)")
   public static let profileSampleUse: Option = Option("-profile-sample-use=", .joined, attributes: [.frontend, .noInteractive, .argumentIsPath], metaVar: "<profile data>", helpText: "Supply sampling-based profiling data from llvm-profdata to enable profile-guided optimization")
   public static let profileStatsEntities: Option = Option("-profile-stats-entities", .flag, attributes: [.helpHidden, .frontend], helpText: "Profile changes to stats in -stats-output-dir, subdivided by source entity")
   public static let profileStatsEvents: Option = Option("-profile-stats-events", .flag, attributes: [.helpHidden, .frontend], helpText: "Profile changes to stats in -stats-output-dir")
@@ -1815,6 +1818,9 @@ extension Option {
       Option.printZeroStats,
       Option.profileCoverageMapping,
       Option.profileGenerate,
+      Option.irProfileGenerate,
+      Option.csProfileGenerate,
+      Option.csProfileGenerateEq,
       Option.profileSampleUse,
       Option.profileStatsEntities,
       Option.profileStatsEvents,