Mojo struct
TmemAllocation
struct TmemAllocation[cta_group: Int, max_cols: Int = 512]
Handle to allocated Tensor Memory.
Lifecycle: allocate() → use → release_lock() → wait → deallocate()
Parameters
Fields
- addr (
UInt32):
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDestructible,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
SmemAddrStorage
comptime SmemAddrStorage = SMemArray[UInt32, 1]
Methods
__init__
__init__(addr: UInt32) -> Self
allocate
static allocate(smem_addr: SMemArray[UInt32, 1]) -> Self
Allocate TMEM (MMA warp). Address stored in smem for epilogue.
from_shared
static from_shared(smem_addr: SMemArray[UInt32, 1]) -> Self
Get handle from existing allocation (epilogue warp).
release_lock
release_lock(self)
Release allocation lock before waiting for epilogue.
deallocate
deallocate(self)
Free TMEM after epilogue completion.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!