Mojo struct
StMatrixConfig
@register_passable(trivial)
struct StMatrixConfig[c_type: DType, stageN: Int, c_swizzle: TensorMapSwizzle = TensorMapSwizzle.SWIZZLE_128B, transpose_c: Bool = False]
Configuration for st.matrix store operations.
Computes the various constants needed for st.matrix operations based on the output tile configuration.
Template Parameters: c_type: Output data type (e.g., bfloat16). stageN: Stage width in elements. c_swizzle: TMA swizzle mode. transpose_c: Whether output is transposed.
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
Movable,
UnknownDestructibility
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
stmtx_simd_width
comptime stmtx_simd_width = 4 if ((stageN % 16) == 0) else 2
stsmx_lane_size
comptime stsmx_lane_size = (16 // size_of[c_type]())
stsmx_row_size
comptime stsmx_row_size = (32 // size_of[c_type]()) if ((stageN % 16) == 0) else (16 // size_of[c_type]())
swizzle_width
comptime swizzle_width = (c_swizzle.bytes() // size_of[c_type]())
Methods
make_swizzle
static make_swizzle() -> Swizzle
Create the swizzle pattern for st.matrix operations.
Returns:
Swizzle: Swizzle instance for the configured swizzle mode.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!