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_im2col_matmul_conv2d

def dispatch_im2col_matmul_conv2d[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, m_tile_byte_budget: Int = Int(268435456)](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

Try to dispatch a 2-D conv as explicit im2col + generic matmul.

Returns True if the conv was handled; False if the caller should fall back to another implementation (naive Mojo kernel, cuDNN, etc.).

Skips on: non-bf16 dtype, grouped conv, dilation != 1, kernel size 1x1 (the vectorized naive kernel wins on tiny shapes), and K too small for the matmul fast path.

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).
m_tile_byte_budget (Int): Byte budget cap for the im2col [M, K] scratch tile used to chunk the M axis (defaults to _DEFAULT_M_TILE_BYTE_BUDGET).

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