feat(cuda-device): add SharedPtr, a CTA-shared pointer typed addrspace(3)

crates/cuda-device/src/lib.rs

@@ -82,7 +82,7 @@ pub use cusimd::{CuSimd, Float2, Float4, TmemRegs4, TmemRegs32};
 pub use disjoint::{__LaunchContractDisjointSlice, __LaunchContractDisjointSliceAbi};
 pub use disjoint::{DisjointSlice, SpaceLayout};
 pub use fence::*;
-pub use shared::{DynamicSharedArray, SharedArray};
+pub use shared::{DynamicSharedArray, SharedArray, SharedPtr};
 pub use tcgen05::{
     TensorMemoryHandle, TmemAddress, TmemDeallocated, TmemF32x4, TmemF32x32, TmemGuard, TmemReady,
     TmemUninit,

crates/cuda-device/src/shared.rs

@@ -220,6 +220,23 @@ impl<T, const N: usize, const ALIGN: usize> SharedArray<T, N, ALIGN> {
         let _ = shared;
         unreachable!("SharedArray::as_raw_mut_ptr called outside CUDA kernel context")
     }
+
+    /// Returns a [`SharedPtr`] to the first element of a shared array.
+    ///
+    /// Unlike [`Self::as_raw_mut_ptr`], which returns a generic pointer, the
+    /// result keeps the array's CTA-shared address space in its type, so it
+    /// stays a shared-memory pointer when it is passed to helper functions.
+    ///
+    /// ```rust,ignore
+    /// static mut TILE: SharedArray<u8, 4096, 128> = SharedArray::UNINIT;
+    /// let tile: SharedPtr<u8> = SharedArray::shared_ptr(&raw mut TILE);
+    /// ```
+    #[inline(always)]
+    pub fn shared_ptr(shared: *mut Self) -> SharedPtr<T> {
+        SharedPtr {
+            ptr: shared as *mut SharedSlot<T>,
+        }
+    }
 }
 
 impl<T, const N: usize, const ALIGN: usize> Index<usize> for SharedArray<T, N, ALIGN> {
@@ -493,6 +510,24 @@ impl<T, const ALIGN: usize> DynamicSharedArray<T, ALIGN> {
         let _ = byte_offset;
         unreachable!("DynamicSharedArray::offset called outside CUDA kernel context")
     }
+
+    /// Like [`Self::get`], but returns a [`SharedPtr`], which keeps the
+    /// CTA-shared address space in its type.
+    #[inline(always)]
+    pub fn shared_ptr() -> SharedPtr<T> {
+        SharedPtr {
+            ptr: Self::get() as *mut SharedSlot<T>,
+        }
+    }
+
+    /// Like [`Self::offset`], but returns a [`SharedPtr`], which keeps the
+    /// CTA-shared address space in its type.
+    #[inline(always)]
+    pub fn shared_ptr_at(byte_offset: usize) -> SharedPtr<T> {
+        SharedPtr {
+            ptr: Self::offset(byte_offset) as *mut SharedSlot<T>,
+        }
+    }
 }
 
 /// Convert a generic-address pointer into its raw `.shared` window offset.
@@ -546,4 +581,136 @@ pub unsafe fn cvta_generic_to_shared_u32(ptr: *const u8) -> u32 {
     unreachable!("cvta_generic_to_shared_u32 called outside CUDA kernel context")
 }
 
+// ============================================================================
+// SharedPtr - a pointer into the issuing CTA's shared memory, typed as such
+// ============================================================================
+
+/// The pointee of [`SharedPtr`]'s field. A pointer to `SharedSlot<T>` is a
+/// pointer into CTA-shared memory (`addrspace(3)`) by its type, as a pointer
+/// to [`SharedArray`] is; `repr(transparent)` gives it `T`'s size, so pointer
+/// arithmetic strides by `T`.
+#[doc(hidden)]
+#[repr(transparent)]
+pub struct SharedSlot<T>(T);
+
+/// A raw pointer into the shared memory of the issuing CTA (`.shared::cta`).
+///
+/// A `*mut T` in device code is a generic pointer, which may address global,
+/// local or shared memory, including another CTA's shared memory in the
+/// cluster. `SharedPtr<T>` instead carries its address space in its type: the
+/// compiler knows it is a CTA-shared pointer wherever it goes, including
+/// through calls to helper functions, and emits the `.shared::cta` forms of
+/// instructions that take it without converting a generic address.
+///
+/// It is created from shared memory the CTA owns, with
+/// [`SharedArray::shared_ptr`] or [`DynamicSharedArray::shared_ptr`], or from
+/// a generic pointer with the `unsafe` [`SharedPtr::from_raw`]. Like a raw
+/// pointer, it may dangle; dereferencing it goes through [`SharedPtr::as_ptr`]
+/// and `unsafe` code.
+///
+/// Pointers from [`crate::cluster::map_shared_rank`], which may name another
+/// CTA's shared memory, are not `SharedPtr`s.
+///
+/// # Example
+///
+/// ```rust,ignore
+/// static mut TILE: SharedArray<f32, 256> = SharedArray::UNINIT;
+///
+/// fn store(tile: SharedPtr<f32>, i: usize, value: f32) {
+///     // `tile` is still a shared pointer here, although `store` is a
+///     // separate function.
+///     unsafe { tile.add(i).as_ptr().write(value) };
+/// }
+///
+/// let tile = SharedArray::shared_ptr(&raw mut TILE);
+/// store(tile, thread::threadIdx_x() as usize, 1.0);
+/// ```
+#[repr(transparent)]
+pub struct SharedPtr<T> {
+    ptr: *mut SharedSlot<T>,
+}
+
+impl<T> Clone for SharedPtr<T> {
+    #[inline(always)]
+    fn clone(&self) -> Self {
+        *self
+    }
+}
+
+impl<T> Copy for SharedPtr<T> {}
+
+impl<T> SharedPtr<T> {
+    /// Converts a generic pointer to a `SharedPtr`, as CUDA C++'s
+    /// `__cvta_generic_to_shared` or CuTe's `make_smem_ptr` do.
+    ///
+    /// # Safety
+    ///
+    /// `ptr` must point into the shared memory of the issuing CTA. A pointer
+    /// to global or local memory, or to another CTA's shared memory, gives an
+    /// unspecified address.
+    #[inline(always)]
+    pub unsafe fn from_raw(ptr: *mut T) -> Self {
+        Self {
+            ptr: ptr as *mut SharedSlot<T>,
+        }
+    }
+
+    /// The generic pointer to the same memory, for ordinary loads and stores
+    /// and for APIs that take raw pointers.
+    #[inline(always)]
+    pub fn as_ptr(self) -> *mut T {
+        self.ptr as *mut T
+    }
+
+    /// The pointer's 32-bit `.shared::cta` address, as taken by the
+    /// `_shared_u32` forms of CTA-shared instructions.
+    #[inline(always)]
+    pub fn addr_u32(self) -> u32 {
+        // SAFETY: `self` points into the issuing CTA's shared memory.
+        unsafe { cvta_generic_to_shared_u32(self.as_ptr() as *const u8) }
+    }
+
+    /// Reinterprets the pointee type, keeping the address space.
+    #[inline(always)]
+    pub fn cast<U>(self) -> SharedPtr<U> {
+        SharedPtr {
+            ptr: self.ptr as *mut SharedSlot<U>,
+        }
+    }
+
+    /// The pointer `count` elements of `T` further on, as `<*mut T>::add`.
+    ///
+    /// # Safety
+    ///
+    /// As for `<*mut T>::add`.
+    #[inline(always)]
+    pub unsafe fn add(self, count: usize) -> Self {
+        Self {
+            ptr: unsafe { self.ptr.add(count) },
+        }
+    }
+
+    /// The pointer `count` elements of `T` away, as `<*mut T>::offset`.
+    ///
+    /// # Safety
+    ///
+    /// As for `<*mut T>::offset`.
+    #[inline(always)]
+    pub unsafe fn offset(self, count: isize) -> Self {
+        Self {
+            ptr: unsafe { self.ptr.offset(count) },
+        }
+    }
+
+    /// The pointer `count` bytes further on, as `<*mut T>::byte_add`.
+    ///
+    /// # Safety
+    ///
+    /// As for `<*mut T>::byte_add`.
+    #[inline(always)]
+    pub unsafe fn byte_add(self, count: usize) -> Self {
+        unsafe { self.cast::<u8>().add(count).cast::<T>() }
+    }
+}
+
 include!("generated/shared_sreg.rs");

crates/mir-importer/src/translator/facts.rs

@@ -389,6 +389,15 @@ mod pointer_origin {
         }
     }
 
+    /// ABI RULE: `cuda_device::SharedPtr<T>` stores a `*mut SharedSlot<T>`, so
+    /// the shared pointer it lowers to is always `RawMut`.
+    pub(crate) fn abi_shared_ptr_field() -> PointerOrigin {
+        PointerOrigin {
+            kind: MirPointerKind::RawMut,
+            mutable: true,
+        }
+    }
+
     /// ABI RULE: `cuda_device::DisjointSlice<'_, T>` stores its data pointer
     /// as `*mut T`, so its data/element pointers are always `RawMut`.
     pub(crate) fn abi_disjoint_slice_data_ptr() -> PointerOrigin {

crates/mir-importer/src/translator/rvalue/place_addr.rs

@@ -64,6 +64,7 @@ pub(crate) fn translate_place_address(
     prev_op: Option<Ptr<Operation>>,
     loc: Location,
 ) -> TranslationResult<Option<(Value, Option<Ptr<Operation>>)>> {
+    let place = &*super::place_read::strip_shared_ptr_field_projections(body, place);
     let Some(slot) = value_map.get_slot(place.local) else {
         return Ok(None);
     };

crates/mir-importer/src/translator/rvalue/place_read.rs

@@ -80,6 +80,7 @@ pub fn translate_place(
     prev_op: Option<Ptr<Operation>>,
     loc: Location,
 ) -> TranslationResult<(Value, Option<Ptr<Operation>>)> {
+    let place = &*strip_shared_ptr_field_projections(body, place);
     match classify_place_read_strategy(ctx, place, value_map)? {
         PlaceReadStrategy::Address => {
             if let Some((value, last_op)) = translate_place_load_from_address(
@@ -109,6 +110,51 @@ pub fn translate_place(
     }
 }
 
+/// Drop the projections of a place onto the field of a `cuda_device::SharedPtr`.
+///
+/// `SharedPtr<T> { ptr: *mut SharedSlot<T> }` is a scalar-lowered newtype: its
+/// value is its one field, a CTA-shared pointer (see `types::translate_type`).
+/// Projecting that field is therefore the identity. Without this, the place
+/// walkers would read the pointer value of a `SharedPtr` nested in another
+/// place, such as `(region.base).ptr`, as the address of an aggregate.
+pub(crate) fn strip_shared_ptr_field_projections<'a>(
+    body: &mir::Body,
+    place: &'a mir::Place,
+) -> std::borrow::Cow<'a, mir::Place> {
+    use rustc_public::ty::{RigidTy, TyKind};
+
+    if !place
+        .projection
+        .iter()
+        .any(|elem| matches!(elem, mir::ProjectionElem::Field(0, _)))
+    {
+        return std::borrow::Cow::Borrowed(place);
+    }
+    let mut place_ty = body.locals()[place.local].ty;
+    let mut kept = Vec::with_capacity(place.projection.len());
+    for elem in &place.projection {
+        let is_shared_ptr = matches!(
+            place_ty.kind(),
+            TyKind::RigidTy(RigidTy::Adt(adt_def, _)) if facts::is_cuda_device_adt(&adt_def, "SharedPtr")
+        );
+        let Ok(next_ty) = elem.ty(place_ty) else {
+            return std::borrow::Cow::Borrowed(place);
+        };
+        if !(is_shared_ptr && matches!(elem, mir::ProjectionElem::Field(0, _))) {
+            kept.push(elem.clone());
+        }
+        place_ty = next_ty;
+    }
+    if kept.len() == place.projection.len() {
+        std::borrow::Cow::Borrowed(place)
+    } else {
+        std::borrow::Cow::Owned(mir::Place {
+            local: place.local,
+            projection: kept,
+        })
+    }
+}
+
 // ============================================================================
 // Place Read Strategy and Address Path
 // ============================================================================

crates/mir-importer/src/translator/types.rs

@@ -364,6 +364,14 @@ fn translate_pointer_like(
             let elem = shared_array_element_type(ctx, &substs, "SharedArray")?;
             Ok(facts::mint_shared_ptr_type(ctx, elem, origin).into())
         }
+        rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, substs))
+            if is_cuda_device_adt(&adt_def, "SharedSlot") =>
+        {
+            // `*mut SharedSlot<T>` is the field of `SharedPtr<T>`: a pointer
+            // to a `T` in CTA-shared memory, by its type.
+            let elem = shared_array_element_type(ctx, &substs, "SharedSlot")?;
+            Ok(facts::mint_shared_ptr_type(ctx, elem, origin).into())
+        }
         rustc_public::ty::TyKind::RigidTy(rustc_public::ty::RigidTy::Adt(adt_def, _substs))
             if is_cuda_device_adt(&adt_def, "Barrier") =>
         {
@@ -718,6 +726,14 @@ pub fn translate_type(
                 // ThreadIndex is a newtype around usize - translate to usize
                 // The type safety is enforced at the Rust level, not the IR level
                 Ok(get_usize_type(ctx).into())
+            } else if is_cuda_device_adt(&adt_def, "SharedPtr") {
+                // `SharedPtr<T> { ptr: *mut SharedSlot<T> }` is a scalar-lowered
+                // newtype: the value is its field, a CTA-shared (`addrspace(3)`)
+                // pointer to `T`. Params, returns, locals and fields of type
+                // `SharedPtr<T>` are therefore shared pointers by type, across
+                // function boundaries.
+                let elem = shared_array_element_type(ctx, &substs, "SharedPtr")?;
+                Ok(facts::mint_shared_ptr_type(ctx, elem, facts::abi_shared_ptr_field()).into())
             } else if is_cuda_device_adt(&adt_def, "RangeToken") {
                 Ok(dialect_iket::types::IketRangeTokenType::get(ctx).into())
             } else if is_cuda_device_adt(&adt_def, "SharedArray") {

crates/rustc-codegen-cuda/examples/shared_memory_helper/README.md

@@ -13,6 +13,26 @@ unsafe fn helper(smem: *mut u8, offset: usize) -> u64 {
 }
 ```
 
+Every build also runs the same probe through `SharedPtr`, a CTA-shared
+pointer typed as such:
+
+```rust
+let smem = DynamicSharedArray::<u32, 1024>::shared_ptr();
+shared_ptr_probe(Region { base: smem, offset });
+
+#[inline(never)]
+unsafe fn shared_ptr_probe<T>(region: Region<T>) -> u64 {
+    let full = region.base.cast::<u8>().byte_add(1024).cast::<Barrier>().as_ptr();
+    // ...
+}
+```
+
+`SharedPtr<T>` lowers to an `addrspace(3)` pointer by its type, so the
+generic helper's parameter, nested in a struct, is one too, and the call
+needs no address-space adaptation. `verify-code-shape.sh` checks that the
+helper takes `ptr addrspace(3)` and that no PTX converts a generic address to a
+shared one.
+
 The launch supplies 2048 dynamic shared bytes. Four independent CTAs each use
 one thread. Nine launches cover three data offsets and three input values,
 including wrapping arithmetic. Every CTA checks both a volatile shared-memory
@@ -79,6 +99,9 @@ RTX 5090, NVIDIA driver 580.173.02, and CUDA tools 13.3.
 | Local helper, all three inline policies | MIR verification rejects the call |
 | Dependency helper, all three inline policies | Same rejection |
 
+With no features, the `SharedPtr` probe runs after the direct control and
+reports `PASS shared-ptr: 9 launches, 36 CTA data/barrier checks`.
+
 Both Compute Sanitizer runs of the direct control reported:
 
 ```text