Mojo module
matmul_amd
Aliases
RegTileType
alias RegTileType[_dtype: DType, layout: Layout] = LayoutTensor[_dtype, layout, MutableAnyOrigin, address_space=AddressSpace(5), alignment=align_of[SIMD[_dtype, simd_width_of[_dtype]()]]()]
Parameters
SMemTileType
alias SMemTileType[_dtype: DType, layout: Layout] = LayoutTensor[_dtype, layout, MutableAnyOrigin, address_space=AddressSpace(3), alignment=align_of[SIMD[_dtype, simd_width_of[_dtype]()]]()]
Parameters
SMemWarpTileType
alias SMemWarpTileType[_dtype: DType, layout: Layout, warp_rows: Int, warp_cols: Int] = LayoutTensor[_dtype, LayoutTensor._compute_tile_layout[True, _dtype, layout, MutableAnyOrigin, AddressSpace(3), Layout.__init__(IntTuple[__origin_of()](1), IntTuple[__origin_of()](1)), _get_layout_type(layout, AddressSpace(3)), _get_index_type(layout, AddressSpace(3)), False, align_of[SIMD[_dtype, simd_width_of[_dtype]()]](), warp_rows, warp_cols]()[0], MutableAnyOrigin, address_space=AddressSpace(3), layout_int_type=_get_layout_type(layout, AddressSpace(3)), linear_idx_type=_get_index_type(layout, AddressSpace(3)), masked=_tile_is_masked[layout, warp_rows, warp_cols](), alignment=align_of[SIMD[_dtype, simd_width_of[_dtype]()]]()]
Parameters
Structs
-
IteratorScatterGatherAmd: -
MmaOpAMD: -
MMATileBuffers: Manages memory for a single matrix (A or B) in GEMM computation. -
ScatterGatherAmd:
Functions
-
gemm_kernel_amd: AMD-optimized GEMM kernel for matrix multiplication C = A * B.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!