For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).

Mojo function

conv2d_kernel_rdna

def conv2d_kernel_rdna[out_type: DType, in_type: DType, filter_type: DType, out_layout: TensorLayout, filter_nk_layout: TensorLayout, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, s_type: DType = get_accum_type[out_type](), BLOCK_K: Int = Int(32), BLOCK_M: Int = Int(128), BLOCK_N: Int = Int(128), WARPS_M: Int = Int(8), WARPS_N: Int = Int(2), WARP_TILE_M: Int = Int(1), WARP_TILE_N: Int = Int(4)](output: TileTensor[out_type, out_layout, MutAnyOrigin], input_ptr: Pointer[Scalar[in_type], ImmutAnyOrigin, _safe=False], filter_nk: TileTensor[filter_type, filter_nk_layout, MutAnyOrigin], M: Int, N: Int, K: Int, HW_out: Int, W_out: Int, H_in: Int, W_in: Int, C_in: Int, R: Int, S: Int, pad_h: Int, pad_w: Int)

Conv2D implicit GEMM kernel for RDNA 3+ GPUs.

Identical to the RDNA WMMA matmul kernel except the A-tile is loaded directly from the NHWC input with on-the-fly im2col coordinate computation, eliminating the intermediate im2col buffer.

The B-tile (filter) must be pre-transposed to [N, K] = [C_out, RSC_in] layout for vectorized loading with transpose_b=True semantics.

Output [M, N] in row-major maps directly to NHWC output.

Constraints (enforced by the dispatch layer):

stride = (1, 1) and dilation = (1, 1)
K % BLOCK_K == 0 and C_in % BLOCK_K == 0

Parameters:

out_type (DType): DType of the output tensor.
in_type (DType): DType of the input activation tensor.
filter_type (DType): DType of the filter tensor.
out_layout (TensorLayout): TensorLayout of the output TileTensor.
filter_nk_layout (TensorLayout): TensorLayout of the pre-transposed [N, K] filter TileTensor.
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional fused epilogue lambda applied per output element (defaults to None).
s_type (DType): DType used for the accumulator (defaults to the accumulator type for out_type).
BLOCK_K (Int): Tile size along the K reduction dimension (defaults to 32).
BLOCK_M (Int): Tile size along the M output-row dimension (defaults to 128).
BLOCK_N (Int): Tile size along the N output-column dimension (defaults to 128).
WARPS_M (Int): Number of warps tiling the M dimension (defaults to 8).
WARPS_N (Int): Number of warps tiling the N dimension (defaults to 2).
WARP_TILE_M (Int): MMA tiles per warp along M (defaults to 1).
WARP_TILE_N (Int): MMA tiles per warp along N (defaults to 4).

Args:

output (TileTensor[out_type, out_layout, MutAnyOrigin]): Output TileTensor of shape [M, N] in row-major, mapping to NHWC output.
input_ptr (Pointer[Scalar[in_type], ImmutAnyOrigin, _safe=False]): Pointer to the NHWC input activation tensor of in_type.
filter_nk (TileTensor[filter_type, filter_nk_layout, MutAnyOrigin]): Filter TileTensor pre-transposed to [N, K] = [C_out, RSC_in] layout.
M (Int): Number of output rows; equals batch * H_out * W_out.
N (Int): Number of output columns; equals C_out.
K (Int): Reduction dimension length; equals R * S * C_in.
HW_out (Int): Product of output spatial dimensions H_out * W_out.
W_out (Int): Output spatial width.
H_in (Int): Input spatial height.
W_in (Int): Input spatial width.
C_in (Int): Number of input channels.
R (Int): Filter height.
S (Int): Filter width.
pad_h (Int): Zero padding applied to the input height.
pad_w (Int): Zero padding applied to the input width.