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

enqueue_apple_conv2d

def enqueue_apple_conv2d[in_type: DType, c_type: DType = DType.float32, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type], input: TileTensor[in_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], filter_nk: TileTensor[in_type, Storage=filter_nk.Storage, address_space=filter_nk.address_space, linear_idx_type=filter_nk.linear_idx_type], conv: ConvIm2colParams, ctx: DeviceContext)

Enqueue the fused online-im2col conv2d (AppleM5MatMul.run_conv).

SM100/Apple M5 (compute_capability == 5). No [M, K] scratch is materialised: the A operand is gathered from input (4-D NHWC) on the fly. filter_nk is the filter pre-transposed to (C_out, K=R*S*C_in) row-major (the NK layout dispatch_im2col_matmul_conv2d already builds); the GEMM uses transpose_b=True. C is (M=N_batch*H_out*W_out, N=C_out) row-major (a flat view of the NHWC output). Grid mirrors enqueue_apple_matmul (single-pass; no split-K for conv yet).

Parameters:

in_type (DType): Element type of the NHWC input and filter_nk operands (bf16 for now).
c_type (DType): Output element type (fp16, bf16, fp32). Accumulation is fp32 (defaults to float32).
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; receives SIMD[c_type, width] at absolute (row, col) (AMD's contract) (defaults to None).

Args:

c (TileTensor[c_type, Storage=c.Storage, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output matrix (M, N) row-major with M = N_batch*H_out*W_out and N = C_out (a flat view of the NHWC output).
input (TileTensor[in_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): 4-D NHWC source tensor; its flat pointer drives the im2col gather.
filter_nk (TileTensor[in_type, Storage=filter_nk.Storage, address_space=filter_nk.address_space, linear_idx_type=filter_nk.linear_idx_type]): Filter pre-transposed to (C_out, K=R*S*C_in) row-major (the NK layout); used with transpose_b=True.
conv (ConvIm2colParams): Conv geometry for the im2col gather.
ctx (DeviceContext): Device context to enqueue the kernel on; must be Apple M5 (compute_capability == 5).

Raises:

If the attached GPU is not Apple M5 (compute_capability != 5).