Mojo struct
NVIDIASharedMemoryManager
struct NVIDIASharedMemoryManager[name: StringSlice[StaticConstantOrigin] = "extern_ptr_syml", alignment: Int = 128, memory_alignment: Int = 8]
Manager for allocating and organizing shared memory on NVIDIA GPUs.
This struct provides a unified interface for allocating tiles, tile arrays, and regular arrays in shared memory with proper alignment.
Parameters
- name (
StringSlice): Symbol name for external shared memory. - alignment (
Int): Default alignment for memory allocations (default: 128). - memory_alignment (
Int): Base memory alignment (default: 8).
Fields
- base_ptr (
UnsafePointer[Int8, address_space=AddressSpace(3)]): - offset (
Int):
Implemented traits
AnyType,
UnknownDestructibility
Aliases
__del__is_trivial
alias __del__is_trivial = True
Array
alias Array[type: AnyType, num_tiles: Int] = SMemArray[type, num_tiles]
Parameters
Tile
alias Tile[dtype: DType, layout: Layout] = SMemTile[dtype, layout, alignment]
Parameters
TileArray
alias TileArray[dtype: DType, layout: Layout, num_tiles: Int] = SMemTileArray[dtype, layout, alignment, num_tiles]
Parameters
Methods
__init__
__init__(out self)
Initialize the shared memory manager.
tile
tile[dtype: DType, layout: Layout](mut self) -> LayoutTensor[dtype, layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=alignment]
Allocate a single tile in shared memory.
Returns:
LayoutTensor: A shared memory tile with the specified configuration.
tile_array
tile_array[dtype: DType, layout: Layout, num_tiles: Int](mut self) -> SMemTileArrayType[dtype, layout, num_tiles, alignment]
Allocate an array of tiles in shared memory.
Returns:
SMemTileArrayType: A shared memory tile array with the specified configuration.
array
array[type: AnyType, size: Int](mut self) -> UnsafePointer[type, address_space=AddressSpace(3)]
Allocate a regular array in shared memory.
Returns:
UnsafePointer: A pointer to the allocated shared memory array.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!