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

dispatch_im2col_matmul_conv3d

def dispatch_im2col_matmul_conv3d[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(3)], dilation: IndexList[Int(3)], symmetric_padding: IndexList[Int(3)], num_groups: Int, ctx: DeviceContext) -> Bool

Try to dispatch a 3-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 1x1x1 (the vectorized naive kernel wins on tiny shapes), and K too small for the matmul fast path.

Parameters:

• ​input_type (DType): Element type of the input tensor; must be DType.bfloat16.
• ​filter_type (DType): Element type of the filter tensor.
• ​output_type (DType): Element type of the output tensor.
• ​filter_is_fcrs (Bool): True if the filter is in FCQRS layout, False if in QRSCF layout (inferred, 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 elementwise epilogue applied to each output element with 5-D coordinates (inferred, defaults to None).
• ​m_tile_byte_budget (Int): Upper bound on bytes allocated for one M-tile of the im2col matrix, controlling scratch memory usage (inferred, defaults to 256 MiB).

Args:

• ​input (TileTensor[input_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): Input tensor in NDHWC layout with shape [batch, D, H, W, C].
• ​filter (TileTensor[filter_type, Storage=filter.Storage, address_space=filter.address_space, linear_idx_type=filter.linear_idx_type]): Filter tensor with static shape, either FCQRS or QRSCF depending on filter_is_fcrs.
• ​output (TileTensor[output_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Output tensor in NDHWC layout with shape [batch, D_out, H_out, W_out, F].
• ​stride (IndexList[Int(3)]): Spatial strides for the depth, height, and width axes.
• ​dilation (IndexList[Int(3)]): Spatial dilation for the depth, height, and width axes; must be 1 for each.
• ​symmetric_padding (IndexList[Int(3)]): Symmetric padding for the depth, height, and width axes.
• ​num_groups (Int): Number of convolution groups; must be 1.
• ​ctx (DeviceContext): Device context for enqueueing GPU kernels and allocating scratch buffers.

Returns:

Bool: True if the conv was handled, False if the caller should fall back to another implementation.