Mojo module
matmul
comptime values
SCHED_MASK_DS_READ
comptime SCHED_MASK_DS_READ = 0
SCHED_MASK_DS_WRITE
comptime SCHED_MASK_DS_WRITE = 1
SCHED_MASK_MFMA
comptime SCHED_MASK_MFMA = 3
SCHED_MASK_VMEM_READ
comptime SCHED_MASK_VMEM_READ = 2
SMemWarpTileType
comptime SMemWarpTileType[_dtype: DType, layout: Layout, warp_rows: Int, warp_cols: Int] = LayoutTensor[_dtype, LayoutTensor._compute_tile_layout[warp_rows, warp_cols]()[0], MutAnyOrigin, address_space=AddressSpace.SHARED, layout_int_type=_get_layout_type(layout, AddressSpace.SHARED), linear_idx_type=_get_index_type(layout, AddressSpace.SHARED), masked=_tile_is_masked[layout, warp_rows, warp_cols]()]
Type alias for warp-level shared memory tiles with specified dimensions.
Parameters
Structs
-
MmaOpAMD: -
MMATileBuffers: Manages memory for a single matrix (A or B) in GEMM computation.
Functions
-
gemm_kernel_amd: AMD-optimized GEMM kernel for matrix multiplication C = A * B. -
write_output_fragments: Write output fragments from registers to global memory with optional elementwise operations.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!