Add a library for dumping kernels arguments before and after launch (v…

…osen#18)

.github/workflows/rust.yml

@@ -48,7 +48,7 @@ jobs:
         sudo apt update
         sudo apt install ocl-icd-opencl-dev
     - name: Build
-      run: cargo build --verbose
+      run: cargo build --workspace --verbose
     # TODO(take-cheeze): Support testing
     # - name: Run tests
     #   run: cargo test --verbose

Cargo.toml

@@ -6,6 +6,7 @@ members = [
     "level_zero",
     "spirv_tools-sys",
     "zluda",
+    "zluda_dump",
     "zluda_lib",
     "zluda_inject",
     "zluda_redirect",

detours-sys/build.rs

@@ -1,10 +1,8 @@
-use std::error::Error;
-
 #[cfg(not(target_os = "windows"))]
 fn main() {}
 
 #[cfg(target_os = "windows")]
-fn main() -> Result<(), Box<dyn Error>> {
+fn main() -> Result<(), Box<dyn std::error::Error>> {
     windows::main_impl()
 }
 
@@ -37,18 +35,15 @@ mod windows {
         .try_compile("detours")?;
         Ok(())
     }
-
-    #[cfg(target_env = "msvc")]
-    fn add_target_options(build: &mut cc::Build) -> &mut cc::Build {
-        build
-    }
-
-    #[cfg(not(target_env = "msvc"))]
     fn add_target_options(build: &mut cc::Build) -> &mut cc::Build {
-        build
-            .compiler("clang")
-            .cpp(true)
-            .flag("-fms-extensions")
-            .flag("-Wno-everything")
+        if std::env::var("CARGO_CFG_TARGET_ENV").unwrap() != "msvc" {
+            build
+                .compiler("clang")
+                .cpp(true)
+                .flag("-fms-extensions")
+                .flag("-Wno-everything")
+        } else {
+            build
+        }
     }
 }

ptx/src/ast.rs

@@ -294,6 +294,17 @@ pub enum KernelArgumentType {
     Shared,
 }
 
+impl From<KernelArgumentType> for Type {
+    fn from(this: KernelArgumentType) -> Self {
+        match this {
+            KernelArgumentType::Normal(typ) => typ.into(),
+            KernelArgumentType::Shared => {
+                Type::Pointer(PointerType::Scalar(ScalarType::B8), LdStateSpace::Shared)
+            }
+        }
+    }
+}
+
 impl FnArgumentType {
     pub fn to_type(&self, is_kernel: bool) -> Type {
         if is_kernel {

ptx/src/translate.rs

@@ -6039,7 +6039,7 @@ impl ast::Type {
         }
     }
 
-    fn size_of(&self) -> usize {
+    pub fn size_of(&self) -> usize {
         match self {
             ast::Type::Scalar(typ) => typ.size_of() as usize,
             ast::Type::Vector(typ, len) => (typ.size_of() as usize) * (*len as usize),
@@ -6253,18 +6253,6 @@ impl<'a> ast::Instruction<ast::ParsedArgParams<'a>> {
     }
 }
 
-impl From<ast::KernelArgumentType> for ast::Type {
-    fn from(this: ast::KernelArgumentType) -> Self {
-        match this {
-            ast::KernelArgumentType::Normal(typ) => typ.into(),
-            ast::KernelArgumentType::Shared => ast::Type::Pointer(
-                ast::PointerType::Scalar(ast::ScalarType::B8),
-                ast::LdStateSpace::Shared,
-            ),
-        }
-    }
-}
-
 impl<T: ArgParamsEx> ast::Arg1<T> {
     fn map<U: ArgParamsEx, V: ArgumentMapVisitor<T, U>>(
         self,

zluda/Cargo.toml

@@ -13,7 +13,7 @@ level_zero = { path = "../level_zero" }
 level_zero-sys = { path = "../level_zero-sys" }
 lazy_static = "1.4"
 num_enum = "0.4"
-lz4 = "1.23"
+lz4-sys = "1.9"
 
 [dev-dependencies]
 cuda-driver-sys = "0.3.0"

zluda/src/impl/export_table.rs

@@ -5,12 +5,12 @@ use crate::{
 };
 
 use super::{context, context::ContextData, device, module, Decuda, Encuda, GlobalState};
-use std::mem;
 use std::os::raw::{c_uint, c_ulong, c_ushort};
 use std::{
     ffi::{c_void, CStr},
-    ptr, slice,
+    ptr,
 };
+use std::{mem, os::raw::c_int};
 
 pub fn get(table: *mut *const std::os::raw::c_void, id: *const CUuuid) -> CUresult {
     if table == ptr::null_mut() || id == ptr::null_mut() {
@@ -177,6 +177,7 @@ const FATBIN_FILE_HEADER_VERSION_CURRENT: c_ushort = 0x101;
 
 // assembly file header is a bit different, but we don't care
 #[repr(C)]
+#[derive(Debug)]
 struct FatbinFileHeader {
     kind: c_ushort,
     version: c_ushort,
@@ -221,12 +222,10 @@ unsafe extern "C" fn get_module_from_cubin(
     let mut ptx_files = get_ptx_files(file, end);
     ptx_files.sort_unstable_by_key(|f| c_uint::max_value() - (**f).sm_version);
     for file in ptx_files {
-        let slice = slice::from_raw_parts(
-            (file as *const u8).add((*file).header_size as usize),
-            (*file).payload_size as usize,
-        );
-        let kernel_text =
-            lz4::block::decompress(slice, Some((*file).uncompressed_payload as i32)).unwrap();
+        let kernel_text = match decompress_kernel_module(file) {
+            None => continue,
+            Some(vec) => vec,
+        };
         let kernel_text_string = match CStr::from_bytes_with_nul(&kernel_text) {
             Ok(c_str) => match c_str.to_str() {
                 Ok(s) => s,
@@ -264,6 +263,33 @@ unsafe fn get_ptx_files(file: *const u8, end: *const u8) -> Vec<*const FatbinFil
     result
 }
 
+const MAX_PTX_MODULE_DECOMPRESSION_BOUND: usize = 16 * 1024 * 1024;
+
+unsafe fn decompress_kernel_module(file: *const FatbinFileHeader) -> Option<Vec<u8>> {
+    let decompressed_size = usize::max(1024, (*file).uncompressed_payload as usize);
+    let mut decompressed_vec = vec![0u8; decompressed_size];
+    loop {
+        match lz4_sys::LZ4_decompress_safe(
+            (file as *const u8).add((*file).header_size as usize) as *const _,
+            decompressed_vec.as_mut_ptr() as *mut _,
+            (*file).payload_size as c_int,
+            decompressed_vec.len() as c_int,
+        ) {
+            error if error < 0 => {
+                let new_size = decompressed_vec.len() * 2;
+                if new_size > MAX_PTX_MODULE_DECOMPRESSION_BOUND {
+                    return None;
+                }
+                decompressed_vec.resize(decompressed_vec.len() * 2, 0);
+            }
+            real_decompressed_size => {
+                decompressed_vec.truncate(real_decompressed_size as usize);
+                return Some(decompressed_vec);
+            }
+        }
+    }
+}
+
 unsafe extern "C" fn cudart_interface_fn6(_: u64) {}
 
 const TOOLS_TLS_GUID: CUuuid = CUuuid {

zluda/src/lib.rs

@@ -4,7 +4,6 @@ extern crate level_zero_sys as l0_sys;
 extern crate lazy_static;
 #[cfg(test)]
 extern crate cuda_driver_sys;
-extern crate lz4;
 #[cfg(test)]
 #[macro_use]
 extern crate paste;

zluda_dump/Cargo.toml

@@ -0,0 +1,22 @@
+[package]
+name = "zluda_dump"
+version = "0.0.0"
+authors = ["Andrzej Janik <[email protected]>"]
+edition = "2018"
+
+[lib]
+name = "zluda_dump"
+crate-type = ["cdylib"]
+
+[dependencies]
+ptx = { path = "../ptx" }
+lz4-sys = "1.9"
+regex = "1.4"
+
+[target.'cfg(windows)'.dependencies]
+winapi = { version = "0.3", features = ["libloaderapi", "debugapi"] }
+wchar = "0.6"
+detours-sys = { path = "../detours-sys" }
+
+[target.'cfg(not(windows))'.dependencies]
+libc = "0.2"