Implement runtime compiler. (nvrtc)

Author: lshqqytiger

Cargo.toml

@@ -14,6 +14,7 @@ members = [
     "hipblaslt-sys",
     "hipfft-sys",
     "hiprt-sys",
+    "hiprtc-sys",
     "miopen-sys",
     "offline_compiler",
     "optix_base",
@@ -39,6 +40,7 @@ members = [
     "zluda_ml",
     "zluda_redirect",
     "zluda_rt",
+    "zluda_rtc",
     "zluda_sparse",
 ]
 

Makefile.toml

@@ -22,6 +22,7 @@ args = [
     "-p", "zluda_lib",
     "-p", "zluda_ml",
     "-p", "zluda_sparse",
+    "-p", "zluda_rtc",
     "-p", "zluda_redirect",
 ]
 
@@ -38,6 +39,7 @@ args = [
     "-p", "zluda_fft",
     "-p", "zluda_lib",
     "-p", "zluda_ml",
+    "-p", "zluda_rtc",
     "-p", "zluda_sparse",
 ]
 
@@ -55,6 +57,7 @@ args = [
     "-p", "zluda_fft",
     "-p", "zluda_lib",
     "-p", "zluda_ml",
+    "-p", "zluda_rtc",
     "-p", "zluda_sparse",
 ]
 

hiprtc-sys/Cargo.toml

@@ -0,0 +1,8 @@
+[package]
+name = "hiprtc-sys"
+version = "0.0.0"
+authors = ["Seunghoon Lee <op@lsh.sh>"]
+edition = "2018"
+links = "hiprtc"
+
+[lib]

hiprtc-sys/README

@@ -0,0 +1 @@
+bindgen $Env:HIP_PATH/include/hip/hiprtc.h -o src/hiprtc.rs --no-layout-tests --default-enum-style=newtype --no-derive-debug --allowlist-function "hiprtc.*" --must-use-type hiprtcResult_t -- -I$Env:HIP_PATH/include -D__HIP_PLATFORM_AMD__

hiprtc-sys/build.rs

@@ -0,0 +1,14 @@
+use std::env::VarError;
+use std::{env, path::PathBuf};
+
+fn main() -> Result<(), VarError> {
+    println!("cargo:rustc-link-lib=dylib=hiprtc");
+    if cfg!(windows) {
+        let mut path = PathBuf::from(env::var("HIP_PATH")?);
+        path.push("lib");
+        println!("cargo:rustc-link-search=native={}", path.display());
+    } else {
+        println!("cargo:rustc-link-search=native=/opt/rocm/lib/");
+    }
+    Ok(())
+}

hiprtc-sys/src/hiprtc.rs

@@ -0,0 +1,3 @@
+#[allow(warnings)]
+mod hiprtc;
+pub use hiprtc::*;

hiprtc-sys/src/lib.rs

@@ -23,11 +23,8 @@ use winapi::um::{
 use winapi::um::winbase::{INFINITE, WAIT_FAILED};
 
 static REDIRECT_DLL: &'static str = "zluda_redirect.dll";
-static CUBLAS_DLL: &'static str = "cublas.dll";
-static CUDNN_DLL: &'static str = "cudnn.dll";
-static CUFFT_DLL: &'static str = "cufft.dll";
-static CUSPARSE_DLL: &'static str = "cusparse.dll";
 static NCCL_DLL: &'static str = "nccl.dll";
+static NVRTC_DLL: &'static str = "nvrtc.dll";
 static NVCUDA_DLL: &'static str = "nvcuda.dll";
 static NVML_DLL: &'static str = "nvml.dll";
 static NVAPI_DLL: &'static str = "nvapi64.dll";
@@ -38,26 +35,14 @@ include!("../../zluda_redirect/src/payload_guid.rs");
 #[derive(FromArgs)]
 /// Launch application with custom CUDA libraries
 struct ProgramArguments {
-    /// DLL to be injected instead of system cublas.dll. If not provided {0}, will use cublas.dll from its own directory
-    #[argh(option)]
-    cublas: Option<PathBuf>,
-
-    /// DLL to be injected instead of system cudnn.dll. If not provided {0}, will use cudnn.dll from its own directory
-    #[argh(option)]
-    cudnn: Option<PathBuf>,
-
-    /// DLL to be injected instead of system cufft.dll. If not provided {0}, will use cufft.dll from its own directory
-    #[argh(option)]
-    cufft: Option<PathBuf>,
-
-    /// DLL to be injected instead of system cusparse.dll. If not provided {0}, will use cusparse.dll from its own directory
-    #[argh(option)]
-    cusparse: Option<PathBuf>,
-
     /// DLL to be injected instead of system nccl.dll. If not provided {0}, will use nccl.dll from its own directory
     #[argh(option)]
     nccl: Option<PathBuf>,
 
+    /// DLL to be injected instead of system nvrtc.dll. If not provided {0}, will use nvrtc.dll from its own directory
+    #[argh(option)]
+    nvrtc: Option<PathBuf>,
+
     /// DLL to be injected instead of system nvcuda.dll. If not provided {0}, will use nvcuda.dll from its own directory
     #[argh(option)]
     nvcuda: Option<PathBuf>,
@@ -90,11 +75,8 @@ pub fn main_impl() -> Result<(), Box<dyn Error>> {
     let mut startup_info = unsafe { mem::zeroed::<detours_sys::_STARTUPINFOW>() };
     let mut proc_info = unsafe { mem::zeroed::<detours_sys::_PROCESS_INFORMATION>() };
     let mut dlls_to_inject = vec![
-        environment.cublas_path_zero_terminated.as_ptr() as _,
-        //environment.cudnn_path_zero_terminated.as_ptr() as _,
-        environment.cufft_path_zero_terminated.as_ptr() as _,
-        environment.cusparse_path_zero_terminated.as_ptr() as _,
         environment.nccl_path_zero_terminated.as_ptr() as _,
+        environment.nvrtc_path_zero_terminated.as_ptr() as _,
         environment.nvcuda_path_zero_terminated.as_ptr() as _,
         environment.nvml_path_zero_terminated.as_ptr() as *const i8,
         environment.redirect_path_zero_terminated.as_ptr() as _,
@@ -176,11 +158,8 @@ pub fn main_impl() -> Result<(), Box<dyn Error>> {
 }
 
 struct NormalizedArguments {
-    cublas_path: PathBuf,
-    cudnn_path: PathBuf,
-    cufft_path: PathBuf,
-    cusparse_path: PathBuf,
     nccl_path: PathBuf,
+    nvrtc_path: PathBuf,
     nvcuda_path: PathBuf,
     nvml_path: PathBuf,
     nvapi_path: Option<PathBuf>,
@@ -192,16 +171,10 @@ struct NormalizedArguments {
 impl NormalizedArguments {
     fn new(prog_args: ProgramArguments) -> Result<Self, Box<dyn Error>> {
         let current_exe = env::current_exe()?;
-        let cublas_path =
-            Self::get_absolute_path_or_default(&current_exe, prog_args.cublas, CUBLAS_DLL)?;
-        let cudnn_path =
-            Self::get_absolute_path_or_default(&current_exe, prog_args.cudnn, CUDNN_DLL)?;
-        let cufft_path =
-            Self::get_absolute_path_or_default(&current_exe, prog_args.cufft, CUFFT_DLL)?;
-        let cusparse_path =
-            Self::get_absolute_path_or_default(&current_exe, prog_args.cusparse, CUSPARSE_DLL)?;
         let nccl_path =
             Self::get_absolute_path_or_default(&current_exe, prog_args.nccl, NCCL_DLL)?;
+        let nvrtc_path =
+            Self::get_absolute_path_or_default(&current_exe, prog_args.nvrtc, NVRTC_DLL)?;
         let nvcuda_path =
             Self::get_absolute_path_or_default(&current_exe, prog_args.nvcuda, NVCUDA_DLL)?;
         let nvml_path = Self::get_absolute_path_or_default(&current_exe, prog_args.nvml, NVML_DLL)?;
@@ -212,11 +185,8 @@ impl NormalizedArguments {
         let mut redirect_path = current_exe.parent().unwrap().to_path_buf();
         redirect_path.push(REDIRECT_DLL);
         Ok(Self {
-            cublas_path,
-            cudnn_path,
-            cufft_path,
-            cusparse_path,
             nccl_path,
+            nvrtc_path,
             nvcuda_path,
             nvml_path,
             nvapi_path,
@@ -274,11 +244,8 @@ impl NormalizedArguments {
 }
 
 struct Environment {
-    cublas_path_zero_terminated: String,
-    cudnn_path_zero_terminated: String,
-    cufft_path_zero_terminated: String,
-    cusparse_path_zero_terminated: String,
     nccl_path_zero_terminated: String,
+    nvrtc_path_zero_terminated: String,
     nvcuda_path_zero_terminated: String,
     nvml_path_zero_terminated: String,
     nvapi_path_zero_terminated: Option<String>,
@@ -294,31 +261,16 @@ struct Environment {
 impl Environment {
     fn setup(args: NormalizedArguments) -> io::Result<Self> {
         let _temp_dir = TempDir::new()?;
-        let cublas_path_zero_terminated = Self::zero_terminate(Self::copy_to_correct_name(
-            args.cublas_path,
-            &_temp_dir,
-            CUBLAS_DLL,
-        )?);
-        let cudnn_path_zero_terminated = Self::zero_terminate(Self::copy_to_correct_name(
-            args.cudnn_path,
-            &_temp_dir,
-            CUDNN_DLL,
-        )?);
-        let cufft_path_zero_terminated = Self::zero_terminate(Self::copy_to_correct_name(
-            args.cufft_path,
-            &_temp_dir,
-            CUFFT_DLL,
-        )?);
-        let cusparse_path_zero_terminated = Self::zero_terminate(Self::copy_to_correct_name(
-            args.cusparse_path,
-            &_temp_dir,
-            CUSPARSE_DLL,
-        )?);
         let nccl_path_zero_terminated = Self::zero_terminate(Self::copy_to_correct_name(
             args.nccl_path,
             &_temp_dir,
             NCCL_DLL,
         )?);
+        let nvrtc_path_zero_terminated = Self::zero_terminate(Self::copy_to_correct_name(
+            args.nvrtc_path,
+            &_temp_dir,
+            NVRTC_DLL,
+        )?);
         let nvcuda_path_zero_terminated = Self::zero_terminate(Self::copy_to_correct_name(
             args.nvcuda_path,
             &_temp_dir,
@@ -349,11 +301,8 @@ impl Environment {
             .transpose()?;
         let redirect_path_zero_terminated = Self::zero_terminate(args.redirect_path);
         Ok(Self {
-            cublas_path_zero_terminated,
-            cudnn_path_zero_terminated,
-            cufft_path_zero_terminated,
-            cusparse_path_zero_terminated,
             nccl_path_zero_terminated,
+            nvrtc_path_zero_terminated,
             nvcuda_path_zero_terminated,
             nvml_path_zero_terminated,
             nvapi_path_zero_terminated,

zluda_inject/src/bin.rs

@@ -52,10 +52,6 @@ use winapi::{
 include!("payload_guid.rs");
 
 const WIN_MAX_PATH: usize = 260;
-const CUBLAS_UTF8: &'static str = "CUBLAS.DLL";
-const CUBLAS_UTF16: &[u16] = wch!("CUBLAS.DLL");
-const CUDNN_UTF8: &'static str = "CUDNN.DLL";
-const CUDNN_UTF16: &[u16] = wch!("CUDNN.DLL");
 const NVCUDA1_UTF8: &'static str = "NVCUDA.DLL";
 const NVCUDA1_UTF16: &[u16] = wch!("NVCUDA.DLL");
 const NVCUDA2_UTF8: &'static str = "NVCUDA.DLL";
@@ -68,10 +64,6 @@ const NVOPTIX_UTF8: &'static str = "OPTIX.6.6.0.DLL";
 const NVOPTIX_UTF16: &[u16] = wch!("OPTIX.6.6.0.DLL");
 static mut ZLUDA_PATH_UTF8: Option<&'static [u8]> = None;
 static mut ZLUDA_PATH_UTF16: Vec<u16> = Vec::new();
-static mut ZLUDA_BLAS_PATH_UTF8: Option<&'static [u8]> = None;
-static mut ZLUDA_BLAS_PATH_UTF16: Vec<u16> = Vec::new();
-static mut ZLUDA_DNN_PATH_UTF8: Option<&'static [u8]> = None;
-static mut ZLUDA_DNN_PATH_UTF16: Vec<u16> = Vec::new();
 static mut ZLUDA_ML_PATH_UTF8: Option<&'static [u8]> = None;
 static mut ZLUDA_ML_PATH_UTF16: Vec<u16> = Vec::new();
 static mut ZLUDA_API_PATH_UTF8: Option<&'static [u8]> = None;
@@ -207,11 +199,7 @@ unsafe fn get_library_name_utf8(raw_library_name: *const u8) -> *const u8 {
             }
         }
     }
-    if is_cublas_dll_utf8(library_name) {
-        return ZLUDA_BLAS_PATH_UTF8.unwrap().as_ptr();
-    } /*else if is_cudnn_dll_utf8(library_name) {
-        return ZLUDA_DNN_PATH_UTF8.unwrap().as_ptr();
-    }*/ else if is_nvcuda_dll_utf8(library_name) {
+    if is_nvcuda_dll_utf8(library_name) {
         return ZLUDA_PATH_UTF8.unwrap().as_ptr();
     } else if is_nvml_dll_utf8(library_name) {
         return ZLUDA_ML_PATH_UTF8.unwrap().as_ptr();
@@ -249,11 +237,7 @@ unsafe fn get_library_name_utf16(raw_library_name: *const u16) -> *const u16 {
             }
         }
     }
-    if is_cublas_dll_utf16(library_name) {
-        return ZLUDA_BLAS_PATH_UTF16.as_ptr();
-    } /*else if is_cudnn_dll_utf16(library_name) {
-        return ZLUDA_DNN_PATH_UTF16.as_ptr();
-    }*/ else if is_nvcuda_dll_utf16(library_name) {
+    if is_nvcuda_dll_utf16(library_name) {
         return ZLUDA_PATH_UTF16.as_ptr();
     } else if is_nvml_dll_utf16(library_name) {
         return ZLUDA_ML_PATH_UTF16.as_ptr();
@@ -329,22 +313,6 @@ unsafe fn is_driverstore_utf16(lib: &[u16]) -> bool {
     starts_with_ignore_case(lib, &DRIVERSTORE_UTF16, utf16_to_ascii_uppercase)
 }
 
-fn is_cublas_dll_utf8(lib: &[u8]) -> bool {
-    is_dll_utf8(lib, CUBLAS_UTF8.as_bytes())
-}
-
-fn is_cublas_dll_utf16(lib: &[u16]) -> bool {
-    is_dll_utf16(lib, CUBLAS_UTF16)
-}
-
-fn is_cudnn_dll_utf8(lib: &[u8]) -> bool {
-    is_dll_utf8(lib, CUDNN_UTF8.as_bytes())
-}
-
-fn is_cudnn_dll_utf16(lib: &[u16]) -> bool {
-    is_dll_utf16(lib, CUDNN_UTF16)
-}
-
 fn is_nvcuda_dll_utf8(lib: &[u8]) -> bool {
     is_dll_utf8(lib, NVCUDA1_UTF8.as_bytes()) || is_dll_utf8(lib, NVCUDA2_UTF8.as_bytes())
 }

zluda_redirect/src/lib.rs

@@ -0,0 +1,17 @@
+[package]
+name = "zluda_rtc"
+version = "0.0.0"
+authors = ["Seunghoon Lee <op@lsh.sh>"]
+edition = "2018"
+
+[lib]
+name = "nvrtc"
+crate-type = ["cdylib"]
+
+[dependencies]
+hip_common = { path = "../hip_common" }
+hiprtc-sys = { path = "../hiprtc-sys" }
+
+[package.metadata.zluda]
+linux_names = ["libnvrtc.so.10", "libnvrtc.so.11"]
+dump_names = ["libnvrtc.so"]

zluda_rtc/Cargo.toml

@@ -0,0 +1,3 @@
+bindgen include/nvrtc.h -o src/nvrtc.rs --allowlist-function="^nvrtc.*" --default-enum-style=newtype --no-layout-tests --no-derive-debug -- -Iinclude
+sed -i -e 's/extern "C" {//g' -e 's/-> nvrtcResult;/-> nvrtcResult { crate::unsupported()/g' -e 's/pub fn /#[no_mangle] pub extern "system" fn /g' src/nvrtc.rs
+rustfmt src/nvrtc.rs

zluda_rtc/README

@@ -0,0 +1,53 @@
+mod nvrtc;
+pub use nvrtc::*;
+
+use hiprtc_sys::*;
+
+#[cfg(debug_assertions)]
+fn unsupported() -> nvrtcResult {
+    unimplemented!()
+}
+
+#[cfg(not(debug_assertions))]
+fn unsupported() -> nvrtcResult {
+    nvrtcResult::NVRTC_ERROR_INTERNAL_ERROR
+}
+
+fn to_nvrtc(status: hiprtc_sys::hiprtcResult) -> nvrtcResult {
+    match status {
+        hiprtc_sys::hiprtcResult::HIPRTC_SUCCESS => nvrtcResult::NVRTC_SUCCESS,
+        err => panic!("[ZLUDA] HIPRTC failed: {}", err.0),
+    }
+}
+
+unsafe fn create_program(
+    prog: *mut nvrtcProgram,
+    src: *const std::ffi::c_char,
+    name: *const std::ffi::c_char,
+    num_headers: i32,
+    headers: *const *const std::ffi::c_char,
+    include_names: *const *const std::ffi::c_char,
+) -> nvrtcResult {
+    to_nvrtc(hiprtcCreateProgram(
+        prog.cast(),
+        src,
+        name,
+        num_headers,
+        headers.cast_mut(),
+        include_names.cast_mut(),
+    ))
+}
+
+unsafe fn destroy_program(
+    prog: *mut nvrtcProgram,
+) -> nvrtcResult {
+    to_nvrtc(hiprtcDestroyProgram(prog.cast()))
+}
+
+unsafe fn compile_program(
+    prog: nvrtcProgram,
+    num_options: i32,
+    options: *const *const std::ffi::c_char,
+) -> nvrtcResult {
+    to_nvrtc(hiprtcCompileProgram(prog.cast(), num_options, options.cast_mut()))
+}