[CUDA] Allow selecting nvcc's host compiler, and invalidate on change

Plugins/CudaBuild/plugin.swift

@@ -16,6 +16,7 @@ struct CudaBuild: BuildToolPlugin {
         var cppLanguageStandard: String? = nil
         var verbose: Bool = false
         var codeGeneration: [CodeGeneration] = []
+        var hostCompiler: String? = nil
 
         init() {}
 
@@ -28,9 +29,29 @@ struct CudaBuild: BuildToolPlugin {
             verbose = try c.decodeIfPresent(Bool.self, forKey: .verbose) ?? false
             codeGeneration =
                 try c.decodeIfPresent([CodeGeneration].self, forKey: .codeGeneration) ?? []
+            hostCompiler = try c.decodeIfPresent(String.self, forKey: .hostCompiler)
         }
     }
 
+    /// Host compiler nvcc should use, in priority order: the
+    /// `MLX_CUDA_HOST_COMPILER` environment variable, the `hostCompiler` key in
+    /// CudaBuild.json, then the toolchain's own clang++.
+    ///
+    /// It must still be a clang: nvcc preprocesses with it and SwiftPM compiles
+    /// the generated `.cpp`, so a GCC-preprocessed translation unit will not build.
+    func resolveHostCompiler(context: PluginContext, settings: Settings) -> String {
+        if let env = ProcessInfo.processInfo.environment["MLX_CUDA_HOST_COMPILER"], !env.isEmpty {
+            return env
+        }
+        if let configured = settings.hostCompiler, !configured.isEmpty {
+            return configured
+        }
+        guard let clangUrl = try? context.tool(named: "clang++") else {
+            fatalError("clang++ not found")
+        }
+        return clangUrl.url.path
+    }
+
     func createBuildCommands(context: PluginContext, target: Target) async throws -> [Command] {
 
         print("CUDA Build Plugin")
@@ -40,12 +61,6 @@ struct CudaBuild: BuildToolPlugin {
             return []
         }
 
-        guard let clangUrl = try? context.tool(named: "clang++") else {
-            fatalError("clang++ not found")
-        }
-
-        print("Use clang++ at: \(clangUrl.url.path)")
-
         let sourceDir = target.directoryURL
 
         print("Source directory: \(sourceDir.path)")
@@ -61,6 +76,10 @@ struct CudaBuild: BuildToolPlugin {
             settings = Settings()
         }
 
+        let hostCompiler = resolveHostCompiler(context: context, settings: settings)
+
+        print("Use clang++ at: \(hostCompiler)")
+
         // Scan source directory for .cu files
 
         let sourceDirPath = sourceDir.path.hasSuffix("/") ? sourceDir.path : sourceDir.path + "/"
@@ -141,7 +160,7 @@ struct CudaBuild: BuildToolPlugin {
                         "Compiling \(inputFile.lastPathComponent) to \(outputCpp.lastPathComponent)",
                     executable: encuda.url,
                     arguments: ["compile"] + verboseFlag + stdArgs + incrementalFlag + [
-                        "--clangpp", clangUrl.url.path,
+                        "--clangpp", hostCompiler,
                         "-I", sourceDir.path,
                     ] + headerSearchPathArgs + [
                         inputFile.path,
@@ -168,7 +187,7 @@ struct CudaBuild: BuildToolPlugin {
                 displayName: "Linking CUDA objects",
                 executable: encuda.url,
                 arguments: ["link"] + verboseFlag + stdArgs + incrementalFlag + [
-                    "--clangpp", clangUrl.url.path,
+                    "--clangpp", hostCompiler,
                 ] + outputCpps.map { $0.path } + ["-o", linkOutput.path],
                 inputFiles: outputCpps,
                 outputFiles: [linkOutput]

Source/Encuda/encuda-compile.swift

@@ -72,12 +72,37 @@ extension Encuda {
                 guard process.terminationStatus == 0 else {
                     throw EncudaError.nvccFailed(process.terminationStatus)
                 }
+                writeStamp()
             #endif
         }
 
+        /// Inputs to the generated `.cpp` that are not themselves input *files*.
+        /// A modification-time comparison cannot see any of these change, so the
+        /// signature is stamped beside the output and compared on the next run.
+        private var configurationSignature: String {
+            [
+                "nvcc=\(nvccPath ?? "")",
+                "ccbin=\(clangppPath ?? "")",
+                "std=\(std ?? "")",
+                "arch=\(ProcessInfo.processInfo.environment["CUDA_ARCH"] ?? "")",
+                "include=\(includeDirs.joined(separator: ":"))",
+            ].joined(separator: "\n")
+        }
+
+        private var stampURL: URL {
+            URL(fileURLWithPath: output + ".encuda-stamp")
+        }
+
+        private func writeStamp() {
+            try? configurationSignature.write(to: stampURL, atomically: true, encoding: .utf8)
+        }
+
         private func isUpToDate() -> Bool {
             let fm = FileManager.default
             let outputURL = URL(fileURLWithPath: output)
+            guard let stamp = try? String(contentsOf: stampURL, encoding: .utf8),
+                stamp == configurationSignature
+            else { return false }
             guard fm.fileExists(atPath: output),
                 let outputMod =
                     (try? outputURL.resourceValues(forKeys: [.contentModificationDateKey]))?