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
UMMAInsDescriptor
struct UMMAInsDescriptor[mma_kind: UMMAKind]
Descriptor for UMMA instructions.
This struct represents a descriptor that encodes information about UMMA instructions. The descriptor contains the following bit fields:
- Sparsity (2 bits): The sparsity of the input matrices. Currently defaults to dense matrices.
- Saturate for integer types (1 bits): Whether to saturate the result for integer types. Currently not supported.
- Matrix D type (2 bits): Data type of matrix D.
- Matrix A type (3 bits): Data type of matrix A.
- Matrix B type (3 bits): Data type of matrix B.
- Negate A matrix (1 bit): Whether to negate matrix A. Currently defaults to False.
- Negate B matrix (1 bit): Whether to negate matrix B. Currently defaults to False.
- Transpose A (1 bit): Whether to transpose matrix A.
- Transpose B (1 bit): Whether to transpose matrix B.
- N, Dimension of Matrix B (6 bits): Number of columns in matrix B. 3 LSBs are unused.
- M, Dimension of Matrix A (6 bits): Number of rows in matrix A. 3 LSBs are unused.
Parameters
- mma_kind (
UMMAKind): The kind of UMMA instruction.
Fields
- desc (
UInt32): The 32-bit descriptor value that encodes UMMA instruction information. This field stores the complete descriptor with all bit fields packed into a single 32-bit integer:- Bits 0-1: Sparsity selector(2 bits)
- Bits 2: Sparsity enable(1 bit)
- Bits 3: Saturate for integer types (1 bit)
- Bits 4-5: Matrix D type (2 bits)
- Bits 6: Reserved (1 bit)
- Bits 7-9: Matrix A type (3 bits)
- Bits 10-12: Matrix B type (3 bits)
- Bits 13: Negate A matrix (1 bit)
- Bits 14: Negate B matrix (1 bit)
- Bits 15: Transpose A (1 bit)
- Bits 16: Transpose B (1 bit)
- Bits 17-22: N, Dimension of Matrix B (6 bits)
- Bits 23: Reserved (1 bit)
- Bits 24-28: M, Dimension of Matrix A (5 bits)
- Bits 29: Reserved (1 bit)
- Bits 30-31: Maximum shift while attempting B matrix (2 bits)
Implemented traits
AnyType,
Copyable,
ImplicitlyCopyable,
ImplicitlyDeletable,
Movable,
RegisterPassable,
TrivialRegisterPassable
Methods
__init__
def __init__(value: UInt32) -> Self
Initialize descriptor with raw 32-bit value.
This constructor allows creating a descriptor directly from a 32-bit integer that already contains the properly formatted bit fields for the descriptor.
Args:
- value (
UInt32): A 32-bit integer containing the complete descriptor bit layout.
create
static def create[d_type: DType, a_type: DType, b_type: DType, output_shape: IndexList[Int(2), element_type=DType.uint32], /, *, transpose_a: Bool = False, transpose_b: Bool = True]() -> Self
Create a descriptor for UMMA instructions.
This function creates a descriptor for UMMA instructions based on the provided parameters.
Parameters:
- d_type (
DType): The data type of matrix D. - a_type (
DType): The data type of matrix A. - b_type (
DType): The data type of matrix B. - output_shape (
IndexList[Int(2), element_type=DType.uint32]): The shape of the output matrix. - transpose_a (
Bool): Whether to transpose matrix A. - transpose_b (
Bool): Whether to transpose matrix B.
Returns:
Self: A 32-bit integer containing the complete descriptor bit layout.
static def create[d_type: DType, a_type: DType, b_type: DType, scale_type: DType, output_shape: IndexList[Int(2), element_type=DType.uint32], /, *, transpose_a: Bool = False, transpose_b: Bool = True]() -> Self
Create a descriptor for UMMA MXF8F6F4 instructions.
This function creates a descriptor for UMMA MXF8F6F4 instructions based on the provided parameters.
Parameters:
- d_type (
DType): The data type of matrix D. - a_type (
DType): The data type of matrix A. - b_type (
DType): The data type of matrix B. - scale_type (
DType): The data type of the scale factors (only applicable to MXF8F6F4). - output_shape (
IndexList[Int(2), element_type=DType.uint32]): The shape of the output matrix. - transpose_a (
Bool): Whether to transpose matrix A. - transpose_b (
Bool): Whether to transpose matrix B.
Returns:
Self: A 32-bit integer containing the complete descriptor bit layout.
update_desc_with_sf_id
static def update_desc_with_sf_id[sf_id: UInt32](inst_desc) -> Self
Update the descriptor with the scale factor ID.
Parameters:
- sf_id (
UInt32): The scale factor ID.
Args:
- inst_desc (
Self): The descriptor to update.
Returns:
Self: The updated descriptor.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!