Allow overriding device compute version in dumper

This commit is contained in:
Andrzej Janik 2021-02-27 20:55:19 +01:00
parent 4d3e37befc
commit bfae2e0d21
2 changed files with 34 additions and 4 deletions

View file

@ -2352,12 +2352,13 @@ extern_redirect! {
dev: CUdevice,
) -> CUresult;
}
extern_redirect! {
extern_redirect_with! {
pub fn cuDeviceGetAttribute(
pi: *mut ::std::os::raw::c_int,
attrib: CUdevice_attribute,
dev: CUdevice,
) -> CUresult;
super::cuDeviceGetAttribute;
}
extern_redirect! {
pub fn cuDeviceGetNvSciSyncAttributes(

View file

@ -13,7 +13,10 @@ use std::{
};
use std::{fs::File, ptr};
use cuda::{CUdeviceptr, CUfunction, CUjit_option, CUmodule, CUresult, CUstream, CUuuid};
use cuda::{
CUdevice, CUdevice_attribute, CUdeviceptr, CUfunction, CUjit_option, CUmodule, CUresult,
CUstream, CUuuid,
};
use ptx::ast;
use regex::Regex;
@ -68,6 +71,7 @@ pub static mut KERNELS: Option<HashMap<CUfunction, KernelDump>> = None;
pub static mut BUFFERS: Vec<(usize, usize)> = Vec::new();
pub static mut LAUNCH_COUNTER: usize = 0;
pub static mut KERNEL_PATTERN: Option<Regex> = None;
pub static mut OVERRIDE_COMPUTE_CAPABILITY_MAJOR: Option<i32> = None;
pub struct ModuleDump {
content: Rc<String>,
@ -90,9 +94,18 @@ pub unsafe fn init_libcuda_handle() {
match env::var("ZLUDA_DUMP_KERNEL") {
Ok(kernel_filter) => match Regex::new(&kernel_filter) {
Ok(r) => KERNEL_PATTERN = Some(r),
Err(err) => {
eprintln!("[ZLUDA_DUMP] Error parsing ZLUDA_DUMP_KERNEL: {}", err);
}
},
Err(_) => (),
}
match env::var("ZLUDA_OVERRIDE_COMPUTE_CAPABILITY_MAJOR") {
Ok(cc_override) => match str::parse::<i32>(&cc_override) {
Ok(ver) => OVERRIDE_COMPUTE_CAPABILITY_MAJOR = Some(ver),
Err(err) => {
eprintln!(
"[ZLUDA_DUMP] Env variable ZLUDA_DUMP_KERNEL is not a regex: {}",
"[ZLUDA_DUMP] Error parsing ZLUDA_OVERRIDE_COMPUTE_CAPABILITY_MAJOR: {}",
err
);
}
@ -393,7 +406,7 @@ unsafe fn create_dump_dir(
fs::create_dir_all(&dump_dir)?;
Ok(Some((dump_dir, kernel_dump)))
}
None => Err("Unknown kernel: {:?}")?,
None => Err(format!("Unknown kernel: {:?}", f))?,
}
}
@ -668,3 +681,19 @@ unsafe fn decompress_kernel_module(file: *const FatbinFileHeader) -> Option<Vec<
}
}
}
#[allow(non_snake_case)]
pub unsafe fn cuDeviceGetAttribute(
pi: *mut ::std::os::raw::c_int,
attrib: CUdevice_attribute,
dev: CUdevice,
cont: impl FnOnce(*mut ::std::os::raw::c_int, CUdevice_attribute, CUdevice) -> CUresult,
) -> CUresult {
if attrib == CUdevice_attribute::CU_DEVICE_ATTRIBUTE_COMPUTE_CAPABILITY_MAJOR {
if let Some(ver) = OVERRIDE_COMPUTE_CAPABILITY_MAJOR {
*pi = ver;
return CUresult::CUDA_SUCCESS;
}
}
cont(pi, attrib, dev)
}