Skip to content

feat(cuda-oxide): shared-memory i32/addrspace(3) codegen and CTA-local TMA without SharedPtr #1656

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
nihalpasham merged 3 commits into from
Oct 27, 2026
Merged

feat(cuda-oxide): shared-memory i32/addrspace(3) codegen and CTA-local TMA without SharedPtr #1656

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
5cf2507
feat(cuda-oxide): use i32 shared-memory offsets in SM120 cute export …
nihalpasham Oct 26, 2026
4c57760
feat(cuda-oxide): propagate shared address space across pointer offse…
nihalpasham Oct 26, 2026
32fb58b
refactor(tma): remove SharedPtr<T> and take *mut u8 in CTA-local TMA G2S
nihalpasham Oct 26, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 15 additions & 15 deletions cuda-oxide/crates/cuda-device/src/generated/tma.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -237,11 +237,11 @@ pub unsafe fn cp_async_bulk_tensor_1d_g2s(
///
/// # Safety
/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.
/// `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.
/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations.
#[inline(never)]
pub unsafe fn cp_async_bulk_tensor_1d_g2s_cta(
dst: SharedPtr<u8>,
dst: *mut u8,
tensor_map: *const TmaDescriptor,
coord0: i32,
barrier: *mut Barrier,
Expand Down Expand Up @@ -284,11 +284,11 @@ pub unsafe fn cp_async_bulk_tensor_2d_g2s(
///
/// # Safety
/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.
/// `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.
/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations.
#[inline(never)]
pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta(
dst: SharedPtr<u8>,
dst: *mut u8,
tensor_map: *const TmaDescriptor,
coord0: i32,
coord1: i32,
Expand Down Expand Up @@ -368,11 +368,11 @@ pub unsafe fn cp_async_bulk_tensor_3d_g2s(
///
/// # Safety
/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.
/// `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.
/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations.
#[inline(never)]
pub unsafe fn cp_async_bulk_tensor_3d_g2s_cta(
dst: SharedPtr<u8>,
dst: *mut u8,
tensor_map: *const TmaDescriptor,
coord0: i32,
coord1: i32,
Expand Down Expand Up @@ -421,11 +421,11 @@ pub unsafe fn cp_async_bulk_tensor_4d_g2s(
///
/// # Safety
/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.
/// `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.
/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations.
#[inline(never)]
pub unsafe fn cp_async_bulk_tensor_4d_g2s_cta(
dst: SharedPtr<u8>,
dst: *mut u8,
tensor_map: *const TmaDescriptor,
coord0: i32,
coord1: i32,
Expand Down Expand Up @@ -480,12 +480,12 @@ pub unsafe fn cp_async_bulk_tensor_5d_g2s(
///
/// # Safety
/// `dst`, `tensor_map`, and `barrier` must remain valid until the copy completes.
/// `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.
/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other
/// `mbarrier` operations.
#[inline(never)]
#[allow(clippy::too_many_arguments)]
pub unsafe fn cp_async_bulk_tensor_5d_g2s_cta(
dst: SharedPtr<u8>,
dst: *mut u8,
tensor_map: *const TmaDescriptor,
coord0: i32,
coord1: i32,
Expand Down
2 changes: 1 addition & 1 deletion cuda-oxide/crates/cuda-device/src/lib.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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, SharedPtr};
pub use shared::{DynamicSharedArray, SharedArray};
pub use tcgen05::{
TensorMemoryHandle, TmemAddress, TmemDeallocated, TmemF32x4, TmemF32x32, TmemGuard, TmemReady,
TmemUninit,
Expand Down
178 changes: 0 additions & 178 deletions cuda-oxide/crates/cuda-device/src/shared.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -233,34 +233,6 @@ 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.
///
/// # Safety
///
/// `shared` must point to a `static mut SharedArray` in the current CUDA
/// kernel, normally obtained with `&raw mut`. The returned pointer is valid
/// only within that kernel. This requirement is what makes the type's
/// claim true: a `SharedPtr` built from any other pointer hands
/// [`SharedPtr::addr_u32`] an address outside the shared window, where the
/// conversion yields an unspecified value. Callers must keep concurrent
/// accesses disjoint and use the CUDA synchronization required before
/// reading data written by another thread.
///
/// ```rust,ignore
/// static mut TILE: SharedArray<u8, 4096, 128> = SharedArray::UNINIT;
/// let tile: SharedPtr<u8> = unsafe { SharedArray::shared_ptr(&raw mut TILE) };
/// ```
#[inline(always)]
pub unsafe 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> {
Expand Down Expand Up @@ -548,24 +520,6 @@ 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.
Expand Down Expand Up @@ -633,136 +587,4 @@ 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 the `unsafe`
/// [`SharedArray::shared_ptr`] or the safe [`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 = unsafe { 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");
1 change: 0 additions & 1 deletion cuda-oxide/crates/cuda-device/src/tma.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -113,7 +113,6 @@
//! - **sm_120 (Blackwell)**: Latest TMA capabilities

use crate::barrier::Barrier;
use crate::shared::SharedPtr;

// =============================================================================
// TMA Descriptor Type
Expand Down
11 changes: 2 additions & 9 deletions cuda-oxide/crates/cuda-intrinsics-gen/src/render/compat.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -944,7 +944,7 @@ pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String {
);
if operation.is_cta_g2s() {
output.push_str(
"/// `barrier` must point into the issuing CTA's shared memory, as for the other\n/// `mbarrier` operations. `dst` is a [`SharedPtr`], which points there by its type.\n",
"/// `dst` and `barrier` must point into the issuing CTA's shared memory, as for the other\n/// `mbarrier` operations.\n",
);
}
} else if is_s2g {
Expand Down Expand Up @@ -973,15 +973,8 @@ pub(super) fn render_compat_tma(catalog: &CatalogFile, hash: &str) -> String {
output.push_str("#[inline(never)]\n");
if is_g2s {
let dimensions = dimensions.unwrap();
// The CTA-local forms take a typed CTA-shared destination; the
// raw ABI keeps `*mut u8`, as `*mut Barrier` keeps `*mut u64`.
let destination = if operation.is_cta_g2s() {
"dst: SharedPtr<u8>"
} else {
"dst: *mut u8"
};
let mut arguments = vec![
destination.to_owned(),
"dst: *mut u8".to_owned(),
"tensor_map: *const TmaDescriptor".to_owned(),
];
let mut values = vec!["dst".to_owned(), "tensor_map".to_owned()];
Expand Down
4 changes: 2 additions & 2 deletions cuda-oxide/crates/cuda-intrinsics-gen/src/render/tests/tensor.rs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -420,11 +420,11 @@ fn tma_rendering_preserves_api_and_injects_backend_defaults() {
"pub unsafe fn cp_async_bulk_tensor_reduce_xor_im2col_5d(src: *const u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, coord2: i32, coord3: i32, coord4: i32)"
));
assert!(compatibility.contains(
"pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta(dst: SharedPtr<u8>, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, barrier: *mut Barrier)"
"pub unsafe fn cp_async_bulk_tensor_2d_g2s_cta(dst: *mut u8, tensor_map: *const TmaDescriptor, coord0: i32, coord1: i32, barrier: *mut Barrier)"
));
for dimensions in 1..=5 {
assert!(compatibility.contains(&format!(
"pub unsafe fn cp_async_bulk_tensor_{dimensions}d_g2s_cta(dst: SharedPtr<u8>,"
"pub unsafe fn cp_async_bulk_tensor_{dimensions}d_g2s_cta(dst: *mut u8,"
)));
assert!(compatibility.contains(&format!(
"pub unsafe fn cp_async_bulk_prefetch_tensor_{dimensions}d_l2("
Expand Down
Loading
Loading