Skip to main content

Mojo function

dispatch_rdna_conv2d

dispatch_rdna_conv2d[input_type: DType, filter_type: DType, output_type: DType, filter_is_fcrs: Bool, maybe_epilogue_func: Optional[elementwise_simd_epilogue_type] = None](input: TileTensor[input_type, input.LayoutType, input.origin, address_space=input.address_space, linear_idx_type=input.linear_idx_type, element_size=input.element_size], filter: TileTensor[filter_type, filter.LayoutType, filter.origin, address_space=filter.address_space, linear_idx_type=filter.linear_idx_type, element_size=filter.element_size], output: TileTensor[output_type, output.LayoutType, output.origin, address_space=output.address_space, linear_idx_type=output.linear_idx_type, element_size=output.element_size], stride: IndexList[2], dilation: IndexList[2], symmetric_padding: IndexList[2], num_groups: Int, ctx: DeviceContext) -> Bool

Try to dispatch Conv2D on RDNA via implicit GEMM (im2col fused into WMMA).

Returns True if the convolution was handled, False if the caller should fall back to another implementation (e.g. MIOpen).

Uses the implicit GEMM kernel when C_in is aligned to BLOCK_K (covers all FLUX VAE shapes), falling back to explicit im2col + matmul otherwise.

Returns:

Bool

Was this page helpful?