Feat: experimental dynamic shared mem support, cuda_std 0.2.1

Author: RDambrosio016

crates/cuda_std/CHANGELOG.md

@@ -4,6 +4,12 @@ Notable changes to this project will be documented in this file.
 
 ## Unreleased
 
+## 0.2.1 - 12/8/21
+
+- Fixed `shared_array!` not using fully qualified MaybeUninit.
+- Fixed `shared_array!` working on the CPU.
+- Added experimental dynamic shared memory support through `shared::dynamic_shared_memory`.
+
 ## 0.2.0 - 12/5/21
 
 - Added `#[externally_visible]` in conjunction with cg_nvvm dead code elimination changes to mark that

crates/cuda_std/Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "cuda_std"
-version = "0.2.0"
+version = "0.2.1"
 edition = "2018"
 license = "MIT OR Apache-2.0"
 description = "Standard library for CUDA with rustc_codegen_nvvm"

crates/cuda_std/src/shared.rs

@@ -1,4 +1,6 @@
-//! Shared memory handling. Currently only macros.
+//! Static and Dynamic shared memory handling.
+
+use crate::gpu_only;
 
 /// Statically allocates a buffer large enough for `len` elements of `array_type`, yielding
 /// a `*mut array_type` that points to uninitialized shared memory. `len` must be a constant expression.
@@ -42,9 +44,47 @@
 #[macro_export]
 macro_rules! shared_array {
     ($array_type:ty; $len:expr) => {{
-        // the initializer is discarded when declaring shared globals, so it is unimportant.
-        #[$crate::address_space(shared)]
-        static mut SHARED: MaybeUninit<[$array_type; $len]> = MaybeUninit::uninit();
-        SHARED.as_mut_ptr() as *mut $array_type
+        #[$crate::gpu_only]
+        #[inline(always)]
+        fn shared_array() -> *mut $array_type {
+            use ::core::{cell::UnsafeCell, mem::MaybeUninit};
+            struct SyncWrapper(UnsafeCell<MaybeUninit<[$array_type; $len]>>);
+            // SAFETY: it is up to the user to verify sound shared memory usage, we cannot
+            // fundamentally check it for soundness.
+            unsafe impl Send for SyncWrapper {}
+            // SAFETY: see above
+            unsafe impl Sync for SyncWrapper {}
+
+            // the initializer is discarded when declaring shared globals, so it is unimportant.
+            #[$crate::address_space(shared)]
+            static SHARED: SyncWrapper = SyncWrapper(UnsafeCell::new(MaybeUninit::uninit()));
+
+            SHARED.0.get() as *mut $array_type
+        }
+        shared_array()
     }};
 }
+
+/// Gets a pointer to the dynamic shared memory that was allocated by the caller of the kernel. The
+/// data is left uninitialized.
+///
+/// **Calling this function multiple times will yield the same pointer**.  
+#[gpu_only]
+pub fn dynamic_shared_mem<T>() -> *mut T {
+    // it is unclear whether an alignment of 16 is actually required for correctness, however,
+    // it seems like nvcc always generates the global with .align 16 no matter the type, so we just copy
+    // nvcc's behavior for now.
+    extern "C" {
+        // need to use nvvm_internal and not address_space because address_space only parses
+        // static definitions, not extern static definitions.
+        #[nvvm_internal(addrspace(3))]
+        #[allow(improper_ctypes)]
+        // mangle it a bit to make sure nobody makes the same thing
+        #[link_name = "_Zcuda_std_dyn_shared"]
+        static DYN_SHARED: ::core::cell::UnsafeCell<u128>;
+    }
+
+    // SAFETY: extern statics is how dynamic shared mem is done in CUDA. This will turn into
+    // an extern variable decl in ptx, which is the same thing nvcc does if you dump the ptx from a cuda file.
+    unsafe { DYN_SHARED.get() as *mut T }
+}