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 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], filter_nk: TileTensor[filter_type, filter_nk_layout, MutAnyOrigin], M: Int32, N: Int32, K: Int32, HW_out: Int32, W_out: Int32, H_in: Int32, W_in: Int32, C_in: Int32, R: Int32, S: Int32, pad_h: Int32, pad_w: Int32)

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]): 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 (Int32): Number of output rows; equals batch * H_out * W_out.
• ​N (Int32): Number of output columns; equals C_out.
• ​K (Int32): Reduction dimension length; equals R * S * C_in.
• ​HW_out (Int32): Product of output spatial dimensions H_out * W_out.
• ​W_out (Int32): Output spatial width.
• ​H_in (Int32): Input spatial height.
• ​W_in (Int32): Input spatial width.
• ​C_in (Int32): Number of input channels.
• ​R (Int32): Filter height.
• ​S (Int32): Filter width.
• ​pad_h (Int32): Zero padding applied to the input height.
• ​pad_w (Int32): Zero padding applied to the input width.