Skip to content

Commit

Permalink
add cicc and ptxas compiler impls
Browse files Browse the repository at this point in the history
  • Loading branch information
trxcllnt committed Jul 24, 2024
1 parent 9958e7c commit f1f27b1
Show file tree
Hide file tree
Showing 6 changed files with 428 additions and 2 deletions.
6 changes: 5 additions & 1 deletion src/compiler/c.rs
Original file line number Diff line number Diff line change
Expand Up @@ -151,8 +151,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
Expand Down
266 changes: 266 additions & 0 deletions src/compiler/cicc.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,266 @@
// Copyright 2024 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::{Cacheable, ColorMode, CompileCommand, CompilerArguments, Language};
use crate::compiler::c::{ArtifactDescriptor, CCompilerImpl, CCompilerKind, ParsedArguments};
use crate::{counted_array, dist};

use crate::mock_command::{CommandCreator, CommandCreatorSync, RunCommand};

use async_trait::async_trait;

use std::fs;
use std::collections::HashMap;
use std::ffi::OsString;
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, 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(parsed_args).await
}
fn generate_compile_commands(
&self,
path_transformer: &mut dist::PathTransformer,
executable: &Path,
parsed_args: &ParsedArguments,
cwd: &Path,
env_vars: &[(OsString, OsString)],
_rewrite_includes_only: bool,
) -> Result<(CompileCommand, Option<dist::CompileCommand>, Cacheable)> {
generate_compile_commands(
path_transformer,
executable,
parsed_args,
cwd,
env_vars
)
}
}

pub fn parse_arguments<S>(
arguments: &[OsString],
language: Language,
arg_info: S,
) -> CompilerArguments<ParsedArguments>
where
S: SearchableArgInfo<ArgData>,
{
let mut take_next = false;
let mut input = OsString::new();
let mut output = PathBuf::new();

let mut common_args = vec![];
let mut unhashed_args = vec![];

for arg in ArgsIter::new(arguments.iter().cloned(), arg_info) {
match arg {
Ok(arg) => {
let args = match arg.get_data() {
Some(PassThrough(_)) => {
take_next = false;
&mut common_args
},
Some(Output(p)) => {
take_next = false;
output.clone_from(p);
continue
},
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 {
input.clone_from(p);
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,
};
}

let mut outputs = HashMap::new();
outputs.insert(
language.as_str(),
ArtifactDescriptor {
path: output,
optional: false
}
);

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_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(
parsed_args: &ParsedArguments,
) -> Result<process::Output>
{
std::fs::read(&parsed_args.input)
.map_err(|e| { anyhow::Error::new(e) })
.map(|s| {
process::Output {
status: process::ExitStatus::default(),
stdout: s.to_vec(),
stderr: vec![],
}
})
}

pub fn generate_compile_commands(
path_transformer: &mut dist::PathTransformer,
executable: &Path,
parsed_args: &ParsedArguments,
cwd: &Path,
env_vars: &[(OsString, OsString)],
) -> Result<(CompileCommand, 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(lang_str) {
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 = CompileCommand {
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)?,
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
PassThrough(OsString),
Output(PathBuf),
Unhashed(OsString),
}

use self::ArgData::*;

counted_array!(pub static ARGS: [ArgInfo<ArgData>; _] = [
// These are always randomly generated/different between nvcc invocations
take_arg!("--gen_c_file_name", OsString, Separated, Unhashed),
take_arg!("--gen_device_file_name", OsString, Separated, Unhashed),
take_arg!("--include_file_name", OsString, Separated, Unhashed),
take_arg!("--module_id_file_name", OsString, Separated, Unhashed),
take_arg!("--stub_file_name", OsString, Separated, Unhashed),
take_arg!("-o", PathBuf, Separated, Output),
]);
58 changes: 57 additions & 1 deletion src/compiler/compiler.rs
Original file line number Diff line number Diff line change
Expand Up @@ -24,6 +24,8 @@ use crate::compiler::msvc::Msvc;
use crate::compiler::nvcc::Nvcc;
use crate::compiler::nvcc::NvccHostCompiler;
use crate::compiler::nvhpc::Nvhpc;
use crate::compiler::cicc::Cicc;
use crate::compiler::ptxas::Ptxas;
use crate::compiler::rust::{Rust, RustupProxy};
use crate::compiler::tasking_vx::TaskingVX;
#[cfg(feature = "dist-client")]
Expand Down Expand Up @@ -112,6 +114,8 @@ pub enum Language {
ObjectiveC,
ObjectiveCxx,
Cuda,
Ptx,
Cubin,
Rust,
Hip,
}
Expand Down Expand Up @@ -152,6 +156,8 @@ impl Language {
Language::ObjectiveC => "objc",
Language::ObjectiveCxx => "objc++",
Language::Cuda => "cuda",
Language::Ptx => "ptx",
Language::Cubin => "cubin",
Language::Rust => "rust",
Language::Hip => "hip",
}
Expand All @@ -169,6 +175,8 @@ impl CompilerKind {
| Language::ObjectiveC
| Language::ObjectiveCxx => "C/C++",
Language::Cuda => "CUDA",
Language::Ptx => "PTX",
Language::Cubin => "CUBIN",
Language::Rust => "Rust",
Language::Hip => "HIP",
}
Expand All @@ -181,8 +189,10 @@ impl CompilerKind {
CompilerKind::C(CCompilerKind::Diab) => textual_lang + " [diab]",
CompilerKind::C(CCompilerKind::Gcc) => textual_lang + " [gcc]",
CompilerKind::C(CCompilerKind::Msvc) => textual_lang + " [msvc]",
CompilerKind::C(CCompilerKind::Nvhpc) => textual_lang + " [nvhpc]",
CompilerKind::C(CCompilerKind::Nvcc) => textual_lang + " [nvcc]",
CompilerKind::C(CCompilerKind::Cicc) => textual_lang + " [cicc]",
CompilerKind::C(CCompilerKind::Ptxas) => textual_lang + " [ptxas]",
CompilerKind::C(CCompilerKind::Nvhpc) => textual_lang + " [nvhpc]",
CompilerKind::C(CCompilerKind::TaskingVX) => textual_lang + " [taskingvx]",
CompilerKind::Rust => textual_lang,
}
Expand Down Expand Up @@ -984,6 +994,28 @@ fn is_rustc_like<P: AsRef<Path>>(p: P) -> bool {
)
}

/// Returns true if the given path looks like cicc
fn is_nvidia_cicc<P: AsRef<Path>>(p: P) -> bool {
matches!(
p.as_ref()
.file_stem()
.map(|s| s.to_string_lossy().to_lowercase())
.as_deref(),
Some("cicc")
)
}

/// Returns true if the given path looks like ptxas
fn is_nvidia_ptxas<P: AsRef<Path>>(p: P) -> bool {
matches!(
p.as_ref()
.file_stem()
.map(|s| s.to_string_lossy().to_lowercase())
.as_deref(),
Some("ptxas")
)
}

/// Returns true if the given path looks like a c compiler program
///
/// This does not check c compilers, it only report programs that are definitely not rustc
Expand Down Expand Up @@ -1045,6 +1077,30 @@ where

let rustc_executable = if let Some(ref rustc_executable) = maybe_rustc_executable {
rustc_executable
} else if is_nvidia_cicc(executable) {
debug!("Found cicc");
return CCompiler::new(
Cicc {
// TODO: Use nvcc --version
version: Some(String::new()),
},
executable.to_owned(),
&pool,
)
.await
.map(|c| (Box::new(c) as Box<dyn Compiler<T>>, None));
} else if is_nvidia_ptxas(executable) {
debug!("Found ptxas");
return CCompiler::new(
Ptxas {
// TODO: Use nvcc --version
version: Some(String::new()),
},
executable.to_owned(),
&pool,
)
.await
.map(|c|( Box::new(c) as Box<dyn Compiler<T>>, None));
} else if is_known_c_compiler(executable) {
let cc = detect_c_compiler(creator, executable, args, env.to_vec(), pool).await;
return cc.map(|c| (c, None));
Expand Down
2 changes: 2 additions & 0 deletions src/compiler/gcc.rs
Original file line number Diff line number Diff line change
Expand Up @@ -645,6 +645,8 @@ fn language_to_gcc_arg(lang: Language) -> Option<&'static str> {
Language::ObjectiveC => Some("objective-c"),
Language::ObjectiveCxx => Some("objective-c++"),
Language::Cuda => Some("cu"),
Language::Ptx => None,
Language::Cubin => None,
Language::Rust => None, // Let the compiler decide
Language::Hip => Some("hip"),
Language::GenericHeader => None, // Let the compiler decide
Expand Down
Loading

0 comments on commit f1f27b1

Please sign in to comment.