For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).
Mojo function
load_cluster_smem
def load_cluster_smem[dtype: DType, width: Int](local_ptr: Pointer[Scalar[dtype], address_space=AddressSpace.SHARED], peer_rank: UInt32) -> SIMD[dtype, width]
Load width elements from peer peer_rank's shared memory at local_ptr.
local_ptr is a pointer into this CTA's shared memory; the returned vector is
the value of the same shared object as it exists in CTA peer_rank. Must be
called after a cluster_sync() so the peer's writes are visible. Restricted to
32-bit element dtypes (covers f32/u32, all the split-K combine needs); moved
with the widest vectorized ld.shared::cluster.{v4,v2,b32} that fits width
(16 B groups first), so a width-element read costs ceil(width/4) memory ops.
Parameters:
- βdtype (
DType): Element dtype of the shared buffer; must be a 32-bit dtype (inferred). - βwidth (
Int): Number of elements to load (inferred).
Args:
- βlocal_ptr (
Pointer[Scalar[dtype], address_space=AddressSpace.SHARED]): Pointer into this CTA's shared memory identifying the shared object to read. - βpeer_rank (
UInt32): Cluster rank of the CTA whose copy of the shared object is read.
Returns:
SIMD[dtype, width]: The width elements of the shared object as they exist in CTA
peer_rank's shared memory.