Mojo struct
GroupedTensormapSmem
@register_passable(trivial)
struct GroupedTensormapSmem
Shared memory pointers for tensormap descriptors.
Points to 5 TMA descriptors (128 bytes each) in SMEM for dynamic updates:
- A, B, SFA, SFB for input loading
- C for output storing
These pointers should come from the main SMEM struct (GroupedBlockScaledSmem) to ensure all warps access the same SMEM locations.
Fields
- desc_a (
UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED]): - desc_b (
UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED]): - desc_sfa (
UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED]): - desc_sfb (
UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED]): - desc_c (
UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED]):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
__copyinit__is_trivial
comptime __copyinit__is_trivial = True
__del__is_trivial
comptime __del__is_trivial = True
__moveinit__is_trivial
comptime __moveinit__is_trivial = True
Methods
from_smem
static from_smem(ptr_a: UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED], ptr_b: UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED], ptr_sfa: UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED], ptr_sfb: UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED], ptr_c: UnsafePointer[TMADescriptor, MutAnyOrigin, address_space=AddressSpace.SHARED]) -> Self
Create tensormap pointers from explicit SMEM pointers.
Args:
- ptr_a (
UnsafePointer): Pointer to A tensormap in SMEM. - ptr_b (
UnsafePointer): Pointer to B tensormap in SMEM. - ptr_sfa (
UnsafePointer): Pointer to SFA tensormap in SMEM. - ptr_sfb (
UnsafePointer): Pointer to SFB tensormap in SMEM. - ptr_c (
UnsafePointer): Pointer to C tensormap in SMEM.
Returns:
Self: GroupedTensormapSmem with the provided pointers.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!