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

dispatch_fused_im2col_conv2d_apple

def dispatch_fused_im2col_conv2d_apple[input_type: DType, filter_type: DType, output_type: DType, //, filter_is_fcrs: Bool = False, maybe_epilogue_func: Optional[def[dtype: DType, rank: Int, width: SIMDLength, alignment: Int = Int(1)](IndexList[rank], SIMD[dtype, width]) capturing thin -> None] = None](input: TileTensor[input_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type], filter: TileTensor[filter_type, Storage=filter.Storage, address_space=filter.address_space, linear_idx_type=filter.linear_idx_type], output: TileTensor[output_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type], stride: IndexList[Int(2)], dilation: IndexList[Int(2)], symmetric_padding: IndexList[Int(2)], num_groups: Int, ctx: DeviceContext) -> Bool

Apple M5 fused online-im2col conv2d (no [M, K] scratch materialised).

Apple silicon (compute_capability == 5). Mirrors the MI355 conv pattern: the filter is transposed to the (C_out, K=R*S*C_in) NK matrix once (the same kernel the materialised path uses), then enqueue_apple_conv2d runs the simdgroup-tiled GEMM with the A operand gathered from the NHWC input on the fly -- so the im2col matrix never touches global memory. This wins across both compute- and memory-bound regimes, removing the materialised path's memory-bound loss (and the conv_gpu Apple memory-bound naive guard).

Returns True if handled; False to fall back. Self-gates: bf16, groups=1, dilation=1, kernel > 1x1, K=RSC_in >= 16. Unlike the materialised dispatcher, there is NO N=C_out >= 16 gate: the fused path has no [M, K] scratch round-trip, so small C_out (down to 1) takes the MMA path and beats the naive thread-per-pixel conv (which is also broken on Metal for the C_out=3 VAE->RGB shape). The MMA handles tiny N via the existing edge-tile mask (b_valid_cols zero-fill in the B load + acol < n in the epilogue), all comptime-fixed by SG_N=32 and independent of the runtime N.

Parameters:

input_type (DType): Element DType of the input activation tensor (inferred).
filter_type (DType): Element DType of the filter tensor (inferred).
output_type (DType): Element DType of the output tensor (inferred).
filter_is_fcrs (Bool): True if the filter is laid out as [F, C, R, S]; False for [R, S, C, F] (defaults to False).
maybe_epilogue_func (Optional[def[dtype: DType, rank: Int, width: SIMDLength, alignment: Int = Int(1)](IndexList[rank], SIMD[dtype, width]) capturing thin -> None]): Optional SIMD elementwise epilogue applied to each output element in 4D (batch, h, w, channel) coordinates (defaults to None).

Args:

input (TileTensor[input_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): 4D NHWC input activation tensor of shape [batch, H, W, C_in].
filter (TileTensor[filter_type, Storage=filter.Storage, address_space=filter.address_space, linear_idx_type=filter.linear_idx_type]): 4D filter tensor; [R, S, C_in, C_out] or [C_out, C_in, R, S] depending on filter_is_fcrs.
output (TileTensor[output_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): 4D NHWC output tensor of shape [batch, H_out, W_out, C_out] to write into.
stride (IndexList[Int(2)]): Spatial stride as [stride_h, stride_w].
dilation (IndexList[Int(2)]): Spatial dilation as [dilation_h, dilation_w]; must be 1 to dispatch.
symmetric_padding (IndexList[Int(2)]): Symmetric padding as [pad_h, pad_w].
num_groups (Int): Group count; must be 1 to dispatch.
ctx (DeviceContext): Device context used to enqueue kernels and allocate scratch buffers.

Returns:

Bool