Skip to content

Commit f736d4c

Browse files
trxcllntsylvestre
trxcllnt
authored and
sylvestre
committed
Rework nvcc name rewriting to ensure more cache hits
These changes ensure cache hits for compilations which are subsets of previously cached compilations * Normalize cudafe++, ptx, and cubin names regardless of whether the compilation flag is `-c`, `-ptx`, `-cubin`, or whether there are one or many `-gencode` flags * Include the compiler `hash_key` in the output dir for internal nvcc files to guarantee stability and uniqueness * Fix cache error due to hash collision from not hashing all the PTX and cubin flags
1 parent a43cade commit f736d4c

File tree

16 files changed

+723
-409
lines changed

16 files changed

+723
-409
lines changed

‎Cargo.lock‎

Lines changed: 43 additions & 2 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

‎Cargo.toml‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -69,6 +69,7 @@ opendal = { version = "0.52.0", optional = true, default-features = false }
6969
openssl = { version = "0.10.72", optional = true }
7070
rand = "0.8.4"
7171
regex = "1.10.3"
72+
regex_static = "0.1.1"
7273
reqsign = { version = "0.16.0", optional = true }
7374
reqwest = { version = "0.12", features = [
7475
"json",

‎src/compiler/c.rs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -200,6 +200,7 @@ pub trait CCompilerImpl: Clone + fmt::Debug + Send + Sync + 'static {
200200
T: CommandCreatorSync;
201201
/// Generate a command that can be used to invoke the C compiler to perform
202202
/// the compilation.
203+
#[allow(clippy::too_many_arguments)]
203204
fn generate_compile_commands<T>(
204205
&self,
205206
path_transformer: &mut dist::PathTransformer,
@@ -208,6 +209,7 @@ pub trait CCompilerImpl: Clone + fmt::Debug + Send + Sync + 'static {
208209
cwd: &Path,
209210
env_vars: &[(OsString, OsString)],
210211
rewrite_includes_only: bool,
212+
hash_key: &str,
211213
) -> Result<(
212214
Box<dyn CompileCommand<T>>,
213215
Option<dist::CompileCommand>,
@@ -1157,6 +1159,7 @@ impl<T: CommandCreatorSync, I: CCompilerImpl> Compilation<T> for CCompilation<I>
11571159
&self,
11581160
path_transformer: &mut dist::PathTransformer,
11591161
rewrite_includes_only: bool,
1162+
hash_key: &str,
11601163
) -> Result<(
11611164
Box<dyn CompileCommand<T>>,
11621165
Option<dist::CompileCommand>,
@@ -1178,6 +1181,7 @@ impl<T: CommandCreatorSync, I: CCompilerImpl> Compilation<T> for CCompilation<I>
11781181
cwd,
11791182
env_vars,
11801183
rewrite_includes_only,
1184+
hash_key,
11811185
)
11821186
}
11831187

‎src/compiler/cicc.rs‎

Lines changed: 31 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -85,6 +85,7 @@ impl CCompilerImpl for Cicc {
8585
cwd: &Path,
8686
env_vars: &[(OsString, OsString)],
8787
_rewrite_includes_only: bool,
88+
_hash_key: &str,
8889
) -> Result<(
8990
Box<dyn CompileCommand<T>>,
9091
Option<dist::CompileCommand>,
@@ -118,7 +119,6 @@ where
118119
let mut take_next = false;
119120
let mut outputs = HashMap::new();
120121
let mut extra_dist_files = vec![];
121-
let mut gen_module_id_file = false;
122122
let mut module_id_file_name = Option::<PathBuf>::None;
123123

124124
let mut common_args = vec![];
@@ -128,6 +128,20 @@ where
128128
match arg {
129129
Ok(arg) => {
130130
let args = match arg.get_data() {
131+
Some(ExtraOutput(o)) => {
132+
take_next = false;
133+
let path = cwd.join(o);
134+
if let Some(flag) = arg.flag_str() {
135+
outputs.insert(
136+
flag,
137+
ArtifactDescriptor {
138+
path,
139+
optional: false,
140+
},
141+
);
142+
}
143+
&mut common_args
144+
}
131145
Some(PassThrough(_)) => {
132146
take_next = false;
133147
&mut common_args
@@ -146,7 +160,6 @@ where
146160
}
147161
Some(GenModuleIdFileFlag) => {
148162
take_next = false;
149-
gen_module_id_file = true;
150163
&mut common_args
151164
}
152165
Some(ModuleIdFileName(o)) => {
@@ -158,24 +171,6 @@ where
158171
take_next = false;
159172
&mut unhashed_args
160173
}
161-
Some(UnhashedOutput(o)) => {
162-
take_next = false;
163-
let path = cwd.join(o);
164-
if let Some(flag) = arg.flag_str() {
165-
outputs.insert(
166-
flag,
167-
ArtifactDescriptor {
168-
path,
169-
optional: false,
170-
},
171-
);
172-
}
173-
&mut unhashed_args
174-
}
175-
Some(UnhashedFlag) => {
176-
take_next = false;
177-
&mut unhashed_args
178-
}
179174
None => match arg {
180175
Argument::Raw(ref p) => {
181176
if take_next {
@@ -200,17 +195,16 @@ where
200195
}
201196

202197
if let Some(module_id_path) = module_id_file_name {
203-
if gen_module_id_file {
204-
outputs.insert(
205-
"--module_id_file_name",
206-
ArtifactDescriptor {
207-
path: module_id_path,
208-
optional: true,
209-
},
210-
);
211-
} else {
212-
extra_dist_files.push(module_id_path);
198+
if module_id_path.exists() {
199+
extra_dist_files.push(module_id_path.clone());
213200
}
201+
outputs.insert(
202+
"--module_id_file_name",
203+
ArtifactDescriptor {
204+
path: module_id_path,
205+
optional: false,
206+
},
207+
);
214208
}
215209

216210
CompilerArguments::Ok(ParsedArguments {
@@ -236,13 +230,7 @@ where
236230
}
237231

238232
pub async fn preprocess(cwd: &Path, parsed_args: &ParsedArguments) -> Result<process::Output> {
239-
// cicc and ptxas expect input to be an absolute path
240-
let input = if parsed_args.input.is_absolute() {
241-
parsed_args.input.clone()
242-
} else {
243-
cwd.join(&parsed_args.input)
244-
};
245-
std::fs::read(input)
233+
std::fs::read(cwd.join(&parsed_args.input))
246234
.map_err(anyhow::Error::new)
247235
.map(|s| process::Output {
248236
status: process::ExitStatus::default(),
@@ -329,23 +317,22 @@ pub fn generate_compile_commands(
329317
}
330318

331319
ArgData! { pub
332-
Output(PathBuf),
333-
PassThrough(OsString),
334-
UnhashedFlag,
320+
ExtraOutput(PathBuf),
335321
GenModuleIdFileFlag,
336322
ModuleIdFileName(PathBuf),
323+
Output(PathBuf),
324+
PassThrough(OsString),
337325
UnhashedPassThrough(OsString),
338-
UnhashedOutput(PathBuf),
339326
}
340327

341328
use self::ArgData::*;
342329

343330
counted_array!(pub static ARGS: [ArgInfo<ArgData>; _] = [
344-
take_arg!("--gen_c_file_name", PathBuf, Separated, UnhashedOutput),
345-
take_arg!("--gen_device_file_name", PathBuf, Separated, UnhashedOutput),
331+
take_arg!("--gen_c_file_name", PathBuf, Separated, ExtraOutput),
332+
take_arg!("--gen_device_file_name", PathBuf, Separated, ExtraOutput),
346333
flag!("--gen_module_id_file", GenModuleIdFileFlag),
347334
take_arg!("--include_file_name", OsString, Separated, PassThrough),
348335
take_arg!("--module_id_file_name", PathBuf, Separated, ModuleIdFileName),
349-
take_arg!("--stub_file_name", PathBuf, Separated, UnhashedOutput),
336+
take_arg!("--stub_file_name", PathBuf, Separated, ExtraOutput),
350337
take_arg!("-o", PathBuf, Separated, Output),
351338
]);

‎src/compiler/clang.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -154,6 +154,7 @@ impl CCompilerImpl for Clang {
154154
cwd: &Path,
155155
env_vars: &[(OsString, OsString)],
156156
rewrite_includes_only: bool,
157+
_hash_key: &str,
157158
) -> Result<(
158159
Box<dyn CompileCommand<T>>,
159160
Option<dist::CompileCommand>,

‎src/compiler/compiler.rs‎

Lines changed: 8 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -577,6 +577,7 @@ where
577577
compilation,
578578
weak_toolchain_key,
579579
out_pretty.clone(),
580+
&key,
580581
)
581582
.await?;
582583
let duration_compilation = start.elapsed();
@@ -671,6 +672,7 @@ where
671672
}
672673

673674
#[cfg(not(feature = "dist-client"))]
675+
#[allow(clippy::too_many_arguments)]
674676
async fn dist_or_local_compile<T>(
675677
service: &server::SccacheService<T>,
676678
_dist_client: Option<Arc<dyn dist::Client>>,
@@ -679,13 +681,14 @@ async fn dist_or_local_compile<T>(
679681
compilation: Box<dyn Compilation<T>>,
680682
_weak_toolchain_key: String,
681683
out_pretty: String,
684+
hash_key: &str,
682685
) -> Result<(Cacheable, DistType, process::Output)>
683686
where
684687
T: CommandCreatorSync,
685688
{
686689
let mut path_transformer = dist::PathTransformer::new();
687690
let (compile_cmd, _dist_compile_cmd, cacheable) = compilation
688-
.generate_compile_commands(&mut path_transformer, true)
691+
.generate_compile_commands(&mut path_transformer, true, hash_key)
689692
.context("Failed to generate compile commands")?;
690693

691694
debug!("[{}]: Compiling locally", out_pretty);
@@ -696,6 +699,7 @@ where
696699
}
697700

698701
#[cfg(feature = "dist-client")]
702+
#[allow(clippy::too_many_arguments)]
699703
async fn dist_or_local_compile<T>(
700704
service: &server::SccacheService<T>,
701705
dist_client: Option<Arc<dyn dist::Client>>,
@@ -704,6 +708,7 @@ async fn dist_or_local_compile<T>(
704708
compilation: Box<dyn Compilation<T>>,
705709
weak_toolchain_key: String,
706710
out_pretty: String,
711+
hash_key: &str,
707712
) -> Result<(Cacheable, DistType, process::Output)>
708713
where
709714
T: CommandCreatorSync,
@@ -716,7 +721,7 @@ where
716721
};
717722
let mut path_transformer = dist::PathTransformer::new();
718723
let (compile_cmd, dist_compile_cmd, cacheable) = compilation
719-
.generate_compile_commands(&mut path_transformer, rewrite_includes_only)
724+
.generate_compile_commands(&mut path_transformer, rewrite_includes_only, hash_key)
720725
.context("Failed to generate compile commands")?;
721726

722727
let dist_client = match dist_compile_cmd.clone().and(dist_client) {
@@ -921,6 +926,7 @@ where
921926
&self,
922927
path_transformer: &mut dist::PathTransformer,
923928
rewrite_includes_only: bool,
929+
hash_key: &str,
924930
) -> Result<(
925931
Box<dyn CompileCommand<T>>,
926932
Option<dist::CompileCommand>,

‎src/compiler/cudafe.rs‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,7 @@ impl CCompilerImpl for CudaFE {
8686
cwd: &Path,
8787
env_vars: &[(OsString, OsString)],
8888
_rewrite_includes_only: bool,
89+
_hash_key: &str,
8990
) -> Result<(
9091
Box<dyn CompileCommand<T>>,
9192
Option<dist::CompileCommand>,
@@ -182,7 +183,7 @@ pub fn generate_compile_commands(
182183
use cicc::ArgData::*;
183184

184185
counted_array!(pub static ARGS: [ArgInfo<cicc::ArgData>; _] = [
185-
take_arg!("--gen_c_file_name", PathBuf, Separated, UnhashedOutput),
186+
take_arg!("--gen_c_file_name", PathBuf, Separated, ExtraOutput),
186187
flag!("--gen_module_id_file", GenModuleIdFileFlag),
187188
take_arg!("--module_id_file_name", PathBuf, Separated, Output),
188189
take_arg!("--stub_file_name", OsString, Separated, UnhashedPassThrough),

‎src/compiler/diab.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -87,6 +87,7 @@ impl CCompilerImpl for Diab {
8787
cwd: &Path,
8888
env_vars: &[(OsString, OsString)],
8989
_rewrite_includes_only: bool,
90+
_hash_key: &str,
9091
) -> Result<(
9192
Box<dyn CompileCommand<T>>,
9293
Option<dist::CompileCommand>,

‎src/compiler/gcc.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -104,6 +104,7 @@ impl CCompilerImpl for Gcc {
104104
cwd: &Path,
105105
env_vars: &[(OsString, OsString)],
106106
rewrite_includes_only: bool,
107+
_hash_key: &str,
107108
) -> Result<(
108109
Box<dyn CompileCommand<T>>,
109110
Option<dist::CompileCommand>,

‎src/compiler/msvc.rs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -101,6 +101,7 @@ impl CCompilerImpl for Msvc {
101101
cwd: &Path,
102102
env_vars: &[(OsString, OsString)],
103103
_rewrite_includes_only: bool,
104+
_hash_key: &str,
104105
) -> Result<(
105106
Box<dyn CompileCommand<T>>,
106107
Option<dist::CompileCommand>,

0 commit comments

Comments
 (0)