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 struct
UMMADescriptorTS
struct UMMADescriptorTS[operand_type: DType, num_m_mmas: Int, num_n_mmas: Int, *, MMA_M: Int, MMA_N: Int, MMA_K: Int, consumer_group_size: Int]
Pairs a TMEM A-operand descriptor with a shared-memory B-operand descriptor for a tensor-shared (TS) UMMA.
Parameters
- operand_type (
DType): The element type of the A and B operands. - num_m_mmas (
Int): Number of UMMA tiles along the M dimension. - num_n_mmas (
Int): Number of UMMA tiles along the N dimension. - MMA_M (
Int): The M dimension of a single UMMA instruction tile. - MMA_N (
Int): The N dimension of a single UMMA instruction tile. - MMA_K (
Int): The K dimension of a single UMMA instruction tile. - consumer_group_size (
Int): Number of threads in the consumer warp group.
Fields
- a (
UMMADescriptorTS[operand_type, num_m_mmas, num_n_mmas, MMA_M=MMA_M, MMA_N=MMA_N, MMA_K=MMA_K, consumer_group_size=consumer_group_size].a_t): - b (
UMMADescriptorTS[operand_type, num_m_mmas, num_n_mmas, MMA_M=MMA_M, MMA_N=MMA_N, MMA_K=MMA_K, consumer_group_size=consumer_group_size].b_t):
Implemented traits
AnyType,
Copyable,
Deinitable,
DescriptorPairTS,
ImplicitlyCopyable,
Movable,
RegisterPassable,
TrivialRegisterPassable
comptime members
a_t
comptime a_t = TMemOperand[operand_type, num_m_mmas, num_n_mmas, MMA_M, MMA_N, MMA_K, consumer_group_size]
b_t
comptime b_t = MMASmemDescriptor
operand_t
comptime operand_t = operand_type
Methods
__init__
def __init__(a: TMemOperand[operand_type, num_m_mmas, num_n_mmas, MMA_M, MMA_N, MMA_K, consumer_group_size], b: MMASmemDescriptor) -> Self
get_a
def get_a(self) -> Self.a_t
Returns:
Self.a_t
get_b
def get_b(self) -> Self.b_t
Returns:
Self.b_t