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 module

tensor_core

Tensor Core Module for High-Performance Matrix Operations.

Provides abstractions for using GPU Tensor Cores to perform optimized matrix operations. It supports both NVIDIA and AMD GPU architectures with hardware-specific optimizations.

Key Components:

• TensorCore: Core struct that encapsulates tensor core operations with support for various data types and matrix shapes. It handles loading matrix fragments, performing matrix multiply-accumulate operations, and storing results.

• Matrix Fragment Management: Functions for loading and storing matrix fragments to/from shared memory with hardware-specific optimizations.

• Matrix Multiply-Accumulate (MMA): Optimized implementations of matrix multiplication operations using tensor cores.

Supported Operations:

• Matrix loading with various layouts and swizzling patterns
• Matrix multiply-accumulate (D = A * B + C)
• Matrix storing with hardware-specific optimizations

Supported Data Types:

• NVIDIA: float32, bfloat16, float16, float8_e4m3fn, float8_e5m2
• AMD: float32, bfloat16, float16

Supported Matrix Shapes:

• NVIDIA: 16x8x8, 16x8x4, 16x8x16, 8x8x4, 16x8x32
• AMD: 16x16x4, 16x16x16, 32x32x8

comptime values

shape_16x16x16

comptime shape_16x16x16 = IndexList(Int(16), Int(16), Int(16), __list_literal__=NoneType(None))

AMDGPU tensor core shape 16x16x16.

shape_16x16x32

comptime shape_16x16x32 = IndexList(Int(16), Int(16), Int(32), __list_literal__=NoneType(None))

AMDGPU tensor core shape 16x16x32.

shape_16x16x4

comptime shape_16x16x4 = IndexList(Int(16), Int(16), Int(4), __list_literal__=NoneType(None))

AMDGPU tensor core shape 16x16x4.

shape_16x8x16

comptime shape_16x8x16 = IndexList(Int(16), Int(8), Int(16), __list_literal__=NoneType(None))

Tensor core shape 16x8x16.

shape_16x8x32

comptime shape_16x8x32 = IndexList(Int(16), Int(8), Int(32), __list_literal__=NoneType(None))

Tensor core shape 16x8x32.

shape_16x8x4

comptime shape_16x8x4 = IndexList(Int(16), Int(8), Int(4), __list_literal__=NoneType(None))

Tensor core shape 16x8x4.

shape_16x8x8

comptime shape_16x8x8 = IndexList(Int(16), Int(8), Int(8), __list_literal__=NoneType(None))

Tensor core shape 16x8x8.

shape_32x32x16

comptime shape_32x32x16 = IndexList(Int(32), Int(32), Int(16), __list_literal__=NoneType(None))

AMDGPU tensor core shape 32x32x16.

shape_32x32x64

comptime shape_32x32x64 = IndexList(Int(32), Int(32), Int(64), __list_literal__=NoneType(None))

AMDGPU tensor core shape 32x32x64.

shape_32x32x8

comptime shape_32x32x8 = IndexList(Int(32), Int(32), Int(8), __list_literal__=NoneType(None))

AMDGPU tensor core shape 32x32x8.

shape_8x8x4

comptime shape_8x8x4 = IndexList(Int(8), Int(8), Int(4), __list_literal__=NoneType(None))

Tensor core shape 8x8x4.

shape_null

comptime shape_null = IndexList(Int(0), Int(0), Int(0), __list_literal__=NoneType(None))

Null tensor core shape (0x0x0).

Structs

• ​TensorCore: TensorCore provides an abstraction for GPU tensor core hardware to perform optimized matrix operations.
• ​TiledTensorCore: TiledTensorCore provides a wrapper around TensorCore to support multiple MMAs along the K dimension.

Functions

• ​get_fragment_size: Calculates the fragment size per thread for a given MMA shape.
• ​get_mma_shape: Returns the appropriate matrix multiply-accumulate (MMA) shape for tensor core operations.
• ​load_b_nt: Loads the b operand tile for AMD tensor core MFMA from (N, K) storage.
• ​load_b_tr: Loads the b operand tile for AMD tensor core MFMA instructions using transposed memory access.
• ​num_matrix_reg: Calculates the number of matrix registers required per thread.