Enable code for dynamic parallelism

crates/cuda_std/src/lib.rs

@@ -36,17 +36,16 @@
 
 extern crate alloc;
 
+pub mod atomic;
+pub mod cfg;
 pub mod float;
 #[allow(warnings)]
 pub mod intrinsics;
 pub mod io;
 pub mod mem;
 pub mod misc;
-// WIP
-// pub mod rt;
-pub mod atomic;
-pub mod cfg;
 pub mod ptr;
+pub mod rt;
 pub mod shared;
 pub mod thread;
 pub mod warp;

crates/cuda_std/src/rt/error.rs

@@ -132,7 +132,7 @@ pub enum CudaError {
 /// Result type for most CUDA functions.
 pub type CudaResult<T> = Result<T, CudaError>;
 
-pub(crate) trait ToResult {
+pub trait ToResult {
     fn to_result(self) -> CudaResult<()>;
 }
 impl ToResult for cudaError_t {

crates/cuda_std/src/rt/mod.rs

@@ -31,10 +31,15 @@ bitflags::bitflags! {
 
 #[derive(Debug)]
 pub struct Stream {
-    raw: cuda::cudaStream_t,
+    pub raw: cuda::cudaStream_t,
 }
 
 impl Stream {
+    // /// Creates a new stream with flags.
+    // pub fn new(flags: StreamFlags) -> Self {
+    //     Self {}
+    // }
+
     /// Creates a new stream with flags.
     pub fn new(flags: StreamFlags) -> CudaResult<Self> {
         let mut stream = MaybeUninit::uninit();
@@ -47,10 +52,11 @@ impl Stream {
         }
     }
 
-    #[doc(hidden)]
-    pub fn launch(&self, param_buf: *mut c_void) -> CudaResult<()> {
-        unsafe { cuda::cudaLaunchDeviceV2(param_buf, self.raw).to_result() }
-    }
+    // #[doc(hidden)]
+    // pub fn launch(&self, param_buf: *mut c_void) -> CudaResult<()> {
+    //     unsafe { cuda::cudaLaunchDeviceV2(param_buf, core::ptr::null_mut()).to_result() }
+    //     // unsafe { cuda::cudaLaunchDeviceV2(param_buf, self.raw).to_result() }
+    // }
 }
 
 impl Drop for Stream {
@@ -63,39 +69,108 @@ impl Drop for Stream {
 
 #[macro_export]
 macro_rules! launch {
-    ($func:ident<<<$grid_dim:expr, $block_dim:expr, $smem_size:expr, $stream:ident>>>($($param:expr),* $(,)?)) => {{
-        let grid_dim = ::$crate::rt::GridDim::from($grid_dim);
-        let block_dim = ::$crate::rt::BlockDim::from($block_dim);
-        let mut buf = ::$crate::rt::sys::cudaGetParameterBufferV2(
-            &$func as *const _ as *const ::core::ffi::c_void,
-            ::$crate::rt::sys::dim3 {
+    // ($func:ident<<<$grid_dim:expr, $block_dim:expr, $smem_size:expr, $stream:ident>>>($($param:expr),* $(,)?)) => {{
+    ($func:ident<<<$grid_dim:expr, $block_dim:expr, ($smem_size:expr)>>>($($param:expr),* $(,)?)) => {{
+        use $crate::rt::ToResult;
+        use $crate::float::GpuFloat;
+        let grid_dim = $crate::rt::GridSize::from($grid_dim);
+        let block_dim = $crate::rt::BlockSize::from($block_dim);
+
+        // Get a device buffer for kernel launch.
+        let fptr = $func as *const ();
+        let mut buf = $crate::rt::sys::cudaGetParameterBufferV2(
+            fptr as *const ::core::ffi::c_void,
+            $crate::rt::sys::dim3 {
                 x: grid_dim.x,
                 y: grid_dim.y,
                 z: grid_dim.z
             },
-            ::$crate::rt::sys::dim3 {
+            $crate::rt::sys::dim3 {
                 x: block_dim.x,
                 y: block_dim.y,
                 z: block_dim.z
             },
-            $smem_size
-        ) as *mut u8;
-        unsafe {
-            let mut offset = 0;
-            $(
-                let param = $param;
-                let size = ::core::mem::size_of_val(&param)
-                let mut buf_idx = (offset as f32 / size as f32).ceil() as usize + 1;
-                offset = buf_idx * size;
-                let ptr = &param as *const _ as *const u8;
-                let dst = buf.add(offset);
-                ::core::ptr::copy_nonoverlapping(&param as *const _ as *const u8, dst, size);
-            )*
+            $smem_size,
+        );
+
+        // Ensure buffer is not a nil ptr.
+        if buf.is_null() {
+            return;
         }
+
+        // Load data into buffer.
+        let mut offset = 0;
+        $(
+            let param = $param;
+            let size = ::core::mem::size_of_val(&param);
+            let param_ptr = &param as *const _ as *const ::core::ffi::c_void;
+            let dst = buf.add(offset).copy_from(param_ptr, size);
+            offset += size;
+        )*
         if false {
             $func($($param),*);
         }
-        $stream.launch(buf as *mut ::core::ffi::c_void).to_result()
+        // unsafe {
+        //     let mut offset = 0;
+        //     $(
+        //         let param = $param;
+        //         let size = ::core::mem::size_of_val(&param);
+        //         let mut buf_idx = (offset as f32 / size as f32).ceil() as usize + 1;
+        //         offset = buf_idx * size;
+        //         let ptr = &param as *const _ as *const u8;
+        //         let dst = buf.add(offset);
+        //         ::core::ptr::copy_nonoverlapping(&param as *const _ as *const u8, dst, size);
+        //     )*
+        // }
+        // if false {
+        //     $func($($param),*);
+        // }
+
+        // Launch the kernel.
+        $crate::rt::sys::cudaLaunchDeviceV2(buf as *mut ::core::ffi::c_void, ::core::ptr::null_mut() as *mut _)
+
+        // let mut buf = $crate::rt::sys::cudaGetParameterBuffer(alignment, size) as *mut u8;
+
+        // // Populate the buffer with given arguments.
+        // let mut offset = 0;
+        // $(
+        //     let param = $param;
+        //     let size = ::core::mem::size_of_val(&param);
+        //     let buf_bytes_ptr = (buf as *mut u8).add(offset);
+        //     ::core::ptr::copy_nonoverlapping($param as *const _, buf_bytes_ptr.into(), size);
+        //     offset += size;
+        // )*
+
+        // let mut offset = 0;
+        // $(
+        //     let param = $param;
+        //     let size = ::core::mem::size_of_val(&param);
+        //     let mut buf_idx = (offset as f32 / size as f32).ceil() as usize + 1;
+        //     offset = buf_idx * size;
+        //     let ptr = &param as *const _ as *const u8;
+        //     let dst = buf.add(offset);
+        //     ::core::ptr::copy_nonoverlapping(&param as *const _ as *const u8, dst, size);
+        // )*
+
+        // // Launch the kernel.
+        // let fptr = $func as *const ();
+        // $crate::rt::sys::cudaLaunchDevice(
+        //     fptr as *const ::core::ffi::c_void,
+        //     buf as *mut ::core::ffi::c_void,
+        //     $crate::rt::sys::dim3 {
+        //         x: grid_dim.x,
+        //         y: grid_dim.y,
+        //         z: grid_dim.z
+        //     },
+        //     $crate::rt::sys::dim3 {
+        //         x: block_dim.x,
+        //         y: block_dim.y,
+        //         z: block_dim.z
+        //     },
+        //     $smem_size,
+        //     ::core::ptr::null_mut() as *mut _,
+        //     // $stream.raw,
+        // )
     }};
 }
 

crates/cuda_std/src/rt/sys.rs

@@ -14,6 +14,16 @@ pub use crate::rt::driver_types_sys::*;
 // to share this stuff with cust.
 
 extern "C" {
+    pub fn cudaGetParameterBuffer(alignment: usize, size: usize) -> *mut c_void;
+    pub fn cudaLaunchDevice(
+        func: *const c_void,
+        parameterBuffer: *const c_void,
+        gridDimension: dim3,
+        blockDimension: dim3,
+        sharedMemSize: c_uint,
+        stream: cudaStream_t,
+    ) -> cudaError_t;
+
     pub fn cudaDeviceGetAttribute(
         value: *mut c_int,
         attr: cudaDeviceAttr,

crates/cust/src/link.rs

@@ -114,6 +114,28 @@ impl Linker {
         }
     }
 
+    /// Link device runtime lib.
+    pub fn add_libcudadevrt(&mut self) -> CudaResult<()> {
+        let mut bytes = std::fs::read("/usr/local/cuda-11/lib64/libcudadevrt.a")
+            .expect("could not read libcudadevrt.a");
+
+        unsafe {
+            cuda::cuLinkAddData_v2(
+                self.raw,
+                cuda::CUjitInputType::CU_JIT_INPUT_LIBRARY,
+                // cuda_sys wants *mut but from the API docs we know we retain ownership so
+                // this cast is sound.
+                bytes.as_mut_ptr() as *mut _,
+                bytes.len(),
+                UNNAMED.as_ptr().cast(),
+                0,
+                std::ptr::null_mut(),
+                std::ptr::null_mut(),
+            )
+            .to_result()
+        }
+    }
+
     /// Runs the linker to generate the final cubin bytes. Also returns a duration
     /// for how long it took to run the linker.
     pub fn complete(self) -> CudaResult<Vec<u8>> {

crates/cust/src/module.rs

@@ -338,7 +338,7 @@ impl Module {
     /// ```
     #[deprecated(
         since = "0.3.0",
-        note = "load_from_string was an inconsistent name with inconsistent params, use from_ptx/from_ptx_cstr, passing 
+        note = "load_from_string was an inconsistent name with inconsistent params, use from_ptx/from_ptx_cstr, passing
     an empty slice of options (usually)
     "
     )]

crates/rustc_codegen_nvvm/build.rs

@@ -21,7 +21,7 @@ fn main() {
     // this is set by cuda_builder, but in case somebody is using the codegen
     // manually, default to 520 (which is what nvvm defaults to).
     if option_env!("CUDA_ARCH").is_none() {
-        println!("cargo:rustc-env=CUDA_ARCH=520")
+        println!("cargo:rustc-env=CUDA_ARCH=750")
     }
 }