Mojo struct
MMASmemDescriptor
@register_passable(trivial)
struct MMASmemDescriptor
Descriptor for shared memory operands tcgen05 mma instructions.
This struct represents a descriptor that encodes information about shared memory layout and access patterns for warp group matrix multiply operations. The descriptor contains the following bit fields:
bits layout:
Bit-field | size | Description 0-13 | 14 | Base address in shared memory 14-15 | 2 | Unused, 0 16-29 | 14 | LBO: leading dim byte offset 30-31 | 2 | Unused, 0 32-45 | 14 | SBO: stride dim byte offset 46-48 | 3 | Unused, 0 49-51 | 3 | Matrix Base offset, 0 for canonical layouts 52 | 1 | LBO mode, only matters for 48B K tile 53-60 | 8 | fixed, 0 61-63 | 3 | Swizzle mode
- Start address, LBO, SBO ingnores 4 LSBs.
See https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#tcgen05-shared-memory-desc-layout
Fields
- desc (
SIMD[uint64, 1]): The 64-bit descriptor encodes shared memory operand information.
Implemented traits
AnyType,
Copyable,
Movable,
UnknownDestructibility
Methods
__init__
@implicit
__init__(val: SIMD[uint64, 1]) -> Self
Initialize descriptor with raw 64-bit value.
This constructor allows creating a descriptor directly from a 64-bit integer that already contains the properly formatted bit fields for the descriptor.
The implicit attribute enables automatic conversion from UInt64 to MMASmemDescriptor.
Args:
- val (
SIMD[uint64, 1]): A 64-bit integer containing the complete descriptor bit layout.
__add__
__add__(self, offset: Int) -> Self
Add offset to descriptor's base address.
Args:
- offset (
Int): Byte offset to add to base address.
Returns:
New descriptor with updated base address.
__iadd__
__iadd__(mut self, offset: Int)
Add offset to descriptor's base address in-place.
Args:
- offset (
Int): Byte offset to add to base address.
create
static create[stride_byte_offset: Int, leading_byte_offset: Int, swizzle_mode: TensorMapSwizzle = TensorMapSwizzle(__init__[__mlir_type.!pop.int_literal](0))](smem_ptr: UnsafePointer[type, address_space=AddressSpace(3), alignment=alignment, mut=mut, origin=origin]) -> Self
Create a descriptor for shared memory operand.
Parameters:
- stride_byte_offset (
Int): Stride dimension offset in bytes. - leading_byte_offset (
Int): Leading dimension stride in bytes. - swizzle_mode (
TensorMapSwizzle): Memory access pattern mode.
Args:
- smem_ptr (
UnsafePointer[type, address_space=AddressSpace(3), alignment=alignment, mut=mut, origin=origin]): Pointer to shared memory operand.
Returns:
Initialized descriptor for the shared memory operand.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!