Update Nvcc compiler to support distributed compilation

* Adds top-level `Cicc` and `Ptxas` compiler types
* Updates `Nvcc` compiler implementation to decompose nvcc calls into its constituent subcompiler invocations via the `nvcc --dryrun` flag
* Bumps the `sccache-dist` request timeout from 5 to 10 minutes, because nvcc compilations can take a while
* Updates the CUDA tests, separates into tests for nvcc and clang since now their behavior is different
* Fixes lint

Cargo.toml

@@ -54,6 +54,7 @@ hyper-util = { version = "0.1.3", optional = true, features = [
   "server",
 ] }
 is-terminal = "0.4.12"
+itertools = "0.12"
 jobserver = "0.1"
 jwt = { package = "jsonwebtoken", version = "9", optional = true }
 libc = "0.2.153"
@@ -118,6 +119,7 @@ object = "0.32"
 rouille = { version = "3.6", optional = true, default-features = false, features = [
   "ssl",
 ] }
+shlex = "1.3.0"
 syslog = { version = "6", optional = true }
 version-compare = { version = "0.1.1", optional = true }
 

src/bin/sccache-dist/build.rs

@@ -231,7 +231,7 @@ impl OverlayBuilder {
                     for (tc, _) in entries {
                         warn!("Removing old un-compressed toolchain: {:?}", tc);
                         assert!(toolchain_dir_map.remove(tc).is_some());
-                        fs::remove_dir_all(&self.dir.join("toolchains").join(&tc.archive_id))
+                        fs::remove_dir_all(self.dir.join("toolchains").join(&tc.archive_id))
                             .context("Failed to remove old toolchain directory")?;
                     }
                 }

src/compiler/c.rs

@@ -153,8 +153,12 @@ pub enum CCompilerKind {
     Diab,
     /// Microsoft Visual C++
     Msvc,
-    /// NVIDIA cuda compiler
+    /// NVIDIA CUDA compiler
     Nvcc,
+    /// NVIDIA CUDA optimizer and PTX generator
+    Cicc,
+    /// NVIDIA CUDA PTX assembler
+    Ptxas,
     /// NVIDIA hpc c, c++ compiler
     Nvhpc,
     /// Tasking VX
@@ -1160,6 +1164,7 @@ impl<T: CommandCreatorSync, I: CCompilerImpl> Compilation<T> for CCompilation<I>
             ref env_vars,
             ..
         } = *self;
+
         compiler.generate_compile_commands(
             path_transformer,
             executable,
@@ -1379,6 +1384,10 @@ impl pkg::ToolchainPackager for CToolchainPackager {
                 add_named_file(&mut package_builder, "liblto_plugin.so")?;
             }
 
+            CCompilerKind::Cicc => {}
+
+            CCompilerKind::Ptxas => {}
+
             CCompilerKind::Nvcc => {
                 // Various programs called by the nvcc front end.
                 // presumes the underlying host compiler is consistent

src/compiler/cicc.rs

@@ -0,0 +1,314 @@
+// Copyright 2016 Mozilla Foundation
+// SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#![allow(unused_imports, dead_code, unused_variables)]
+
+use crate::compiler::args::*;
+use crate::compiler::c::{ArtifactDescriptor, CCompilerImpl, CCompilerKind, ParsedArguments};
+use crate::compiler::{
+    CCompileCommand, Cacheable, ColorMode, CompileCommand, CompilerArguments, Language,
+    SingleCompileCommand,
+};
+use crate::{counted_array, dist};
+
+use crate::mock_command::{CommandCreator, CommandCreatorSync, RunCommand};
+
+use async_trait::async_trait;
+
+use std::collections::HashMap;
+use std::ffi::OsString;
+use std::fs;
+use std::path::{Path, PathBuf};
+use std::process;
+
+use crate::errors::*;
+
+/// A unit struct on which to implement `CCompilerImpl`.
+#[derive(Clone, Debug)]
+pub struct Cicc {
+    pub version: Option<String>,
+}
+
+#[async_trait]
+impl CCompilerImpl for Cicc {
+    fn kind(&self) -> CCompilerKind {
+        CCompilerKind::Cicc
+    }
+    fn plusplus(&self) -> bool {
+        true
+    }
+    fn version(&self) -> Option<String> {
+        self.version.clone()
+    }
+    fn parse_arguments(
+        &self,
+        arguments: &[OsString],
+        cwd: &Path,
+    ) -> CompilerArguments<ParsedArguments> {
+        parse_arguments(arguments, cwd, Language::Ptx, &ARGS[..])
+    }
+    #[allow(clippy::too_many_arguments)]
+    async fn preprocess<T>(
+        &self,
+        _creator: &T,
+        _executable: &Path,
+        parsed_args: &ParsedArguments,
+        cwd: &Path,
+        _env_vars: &[(OsString, OsString)],
+        _may_dist: bool,
+        _rewrite_includes_only: bool,
+        _preprocessor_cache_mode: bool,
+    ) -> Result<process::Output>
+    where
+        T: CommandCreatorSync,
+    {
+        preprocess(cwd, parsed_args).await
+    }
+    fn generate_compile_commands<T>(
+        &self,
+        path_transformer: &mut dist::PathTransformer,
+        executable: &Path,
+        parsed_args: &ParsedArguments,
+        cwd: &Path,
+        env_vars: &[(OsString, OsString)],
+        _rewrite_includes_only: bool,
+    ) -> Result<(
+        Box<dyn CompileCommand<T>>,
+        Option<dist::CompileCommand>,
+        Cacheable,
+    )>
+    where
+        T: CommandCreatorSync,
+    {
+        generate_compile_commands(path_transformer, executable, parsed_args, cwd, env_vars).map(
+            |(command, dist_command, cacheable)| {
+                (CCompileCommand::new(command), dist_command, cacheable)
+            },
+        )
+    }
+}
+
+pub fn parse_arguments<S>(
+    arguments: &[OsString],
+    cwd: &Path,
+    language: Language,
+    arg_info: S,
+) -> CompilerArguments<ParsedArguments>
+where
+    S: SearchableArgInfo<ArgData>,
+{
+    let mut args = arguments.to_vec();
+    let input_loc = arguments.len() - 3;
+    let input = args.splice(input_loc..input_loc + 1, []).next().unwrap();
+
+    let mut take_next = false;
+    let mut extra_inputs = vec![];
+    let mut outputs = HashMap::new();
+
+    let mut common_args = vec![];
+    let mut unhashed_args = vec![];
+
+    for arg in ArgsIter::new(args.iter().cloned(), arg_info) {
+        match arg {
+            Ok(arg) => {
+                let args = match arg.get_data() {
+                    Some(PassThrough(_)) => {
+                        take_next = false;
+                        &mut common_args
+                    }
+                    Some(Output(o)) => {
+                        take_next = false;
+                        let path = cwd.join(o);
+                        outputs.insert(
+                            "obj",
+                            ArtifactDescriptor {
+                                path,
+                                optional: false,
+                            },
+                        );
+                        continue;
+                    }
+                    Some(UnhashedInput(o)) => {
+                        take_next = false;
+                        let path = cwd.join(o);
+                        if !path.exists() {
+                            continue;
+                        }
+                        extra_inputs.push(path);
+                        &mut unhashed_args
+                    }
+                    Some(UnhashedOutput(o)) => {
+                        take_next = false;
+                        let path = cwd.join(o);
+                        if let Some(flag) = arg.flag_str() {
+                            outputs.insert(
+                                flag,
+                                ArtifactDescriptor {
+                                    path,
+                                    optional: false,
+                                },
+                            );
+                        }
+                        &mut unhashed_args
+                    }
+                    Some(UnhashedFlag) | Some(Unhashed(_)) => {
+                        take_next = false;
+                        &mut unhashed_args
+                    }
+                    None => match arg {
+                        Argument::Raw(ref p) => {
+                            if take_next {
+                                take_next = false;
+                                &mut common_args
+                            } else {
+                                continue;
+                            }
+                        }
+                        Argument::UnknownFlag(ref p) => {
+                            let s = p.to_string_lossy();
+                            take_next = s.starts_with('-');
+                            &mut common_args
+                        }
+                        _ => unreachable!(),
+                    },
+                };
+                args.extend(arg.iter_os_strings());
+            }
+            _ => continue,
+        };
+    }
+
+    CompilerArguments::Ok(ParsedArguments {
+        input: input.into(),
+        outputs,
+        double_dash_input: false,
+        language,
+        compilation_flag: OsString::new(),
+        depfile: None,
+        dependency_args: vec![],
+        preprocessor_args: vec![],
+        common_args,
+        arch_args: vec![],
+        unhashed_args,
+        extra_dist_files: extra_inputs,
+        extra_hash_files: vec![],
+        msvc_show_includes: false,
+        profile_generate: false,
+        color_mode: ColorMode::Off,
+        suppress_rewrite_includes_only: false,
+        too_hard_for_preprocessor_cache_mode: None,
+    })
+}
+
+pub async fn preprocess(cwd: &Path, parsed_args: &ParsedArguments) -> Result<process::Output> {
+    // cicc and ptxas expect input to be an absolute path
+    let input = if parsed_args.input.is_absolute() {
+        parsed_args.input.clone()
+    } else {
+        cwd.join(&parsed_args.input)
+    };
+    std::fs::read(input)
+        .map_err(anyhow::Error::new)
+        .map(|s| process::Output {
+            status: process::ExitStatus::default(),
+            stdout: s,
+            stderr: vec![],
+        })
+}
+
+pub fn generate_compile_commands(
+    path_transformer: &mut dist::PathTransformer,
+    executable: &Path,
+    parsed_args: &ParsedArguments,
+    cwd: &Path,
+    env_vars: &[(OsString, OsString)],
+) -> Result<(
+    SingleCompileCommand,
+    Option<dist::CompileCommand>,
+    Cacheable,
+)> {
+    // Unused arguments
+    #[cfg(not(feature = "dist-client"))]
+    {
+        let _ = path_transformer;
+    }
+
+    trace!("compile");
+
+    let lang_str = &parsed_args.language.as_str();
+    let out_file = match parsed_args.outputs.get("obj") {
+        Some(obj) => &obj.path,
+        None => return Err(anyhow!("Missing {:?} file output", lang_str)),
+    };
+
+    let mut arguments: Vec<OsString> = vec![];
+    arguments.extend_from_slice(&parsed_args.common_args);
+    arguments.extend_from_slice(&parsed_args.unhashed_args);
+    arguments.extend(vec![
+        (&parsed_args.input).into(),
+        "-o".into(),
+        out_file.into(),
+    ]);
+
+    let command = SingleCompileCommand {
+        executable: executable.to_owned(),
+        arguments,
+        env_vars: env_vars.to_owned(),
+        cwd: cwd.to_owned(),
+    };
+
+    #[cfg(not(feature = "dist-client"))]
+    let dist_command = None;
+    #[cfg(feature = "dist-client")]
+    let dist_command = (|| {
+        let mut arguments: Vec<String> = vec![];
+        arguments.extend(dist::osstrings_to_strings(&parsed_args.common_args)?);
+        arguments.extend(dist::osstrings_to_strings(&parsed_args.unhashed_args)?);
+        arguments.extend(vec![
+            path_transformer.as_dist(&parsed_args.input)?,
+            "-o".into(),
+            path_transformer.as_dist(out_file)?,
+        ]);
+        Some(dist::CompileCommand {
+            executable: path_transformer.as_dist(executable.canonicalize().unwrap().as_path())?,
+            arguments,
+            env_vars: dist::osstring_tuples_to_strings(env_vars)?,
+            cwd: path_transformer.as_dist_abs(cwd)?,
+        })
+    })();
+
+    Ok((command, dist_command, Cacheable::Yes))
+}
+
+ArgData! { pub
+    Output(PathBuf),
+    UnhashedInput(PathBuf),
+    UnhashedOutput(PathBuf),
+    UnhashedFlag,
+    PassThrough(OsString),
+    Unhashed(OsString),
+}
+
+use self::ArgData::*;
+
+counted_array!(pub static ARGS: [ArgInfo<ArgData>; _] = [
+    take_arg!("--gen_c_file_name", PathBuf, Separated, UnhashedOutput),
+    take_arg!("--gen_device_file_name", PathBuf, Separated, UnhashedOutput),
+    flag!("--gen_module_id_file", UnhashedFlag),
+    take_arg!("--include_file_name", OsString, Separated, PassThrough),
+    take_arg!("--module_id_file_name", PathBuf, Separated, UnhashedInput),
+    take_arg!("--stub_file_name", PathBuf, Separated, UnhashedOutput),
+    take_arg!("-o", PathBuf, Separated, Output),
+]);