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_amd_4wave_conv3d

def dispatch_amd_4wave_conv3d[input_type: DType, filter_type: DType, output_type: DType, filter_is_fcqrs: Bool, elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None, block_m_override: Int = Int(0), block_n_override: Int = Int(0), block_k_override: Int = Int(0)](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 Conv3D to amd_4wave_conv on MI355X. Returns True if handled; False if the caller should fall through (typically to dispatch_im2col_matmul_conv3d).

Parameters:

input_type (DType): DType of the input tensor; must be float8_e4m3fn, bfloat16, or float16.
filter_type (DType): DType of the filter tensor; must equal input_type.
output_type (DType): DType of the output tensor; bfloat16 for FP8 input, otherwise tracks input_type.
filter_is_fcqrs (Bool): True if the filter is laid out as FCQRS [F, C, Q, R, S]; False if QRSCF [Q*R*S*C, F].
elementwise_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None]): Optional epilogue applied to the conv output (defaults to None).
block_m_override (Int): Override for the BM tile size, 0 uses the heuristic (defaults to 0).
block_n_override (Int): Override for the BN tile size, 0 uses the heuristic (defaults to 0).
block_k_override (Int): Override for the BK tile size, 0 uses the heuristic (defaults to 0).

Args:

input (TileTensor[input_type, Storage=input.Storage, address_space=input.address_space, linear_idx_type=input.linear_idx_type]): Input NDHWC tensor of shape [N, D, H, W, C_in]; all spatial dims must be static.
filter (TileTensor[filter_type, Storage=filter.Storage, address_space=filter.address_space, linear_idx_type=filter.linear_idx_type]): Filter tensor; FCQRS [F, C, Q, R, S] when filter_is_fcqrs else QRSCF [Q, R, S, C, F].
output (TileTensor[output_type, Storage=output.Storage, address_space=output.address_space, linear_idx_type=output.linear_idx_type]): Output NDHWC tensor of shape [N, D_out, H_out, W_out, C_out].
stride (IndexList[Int(3)]): Per-axis stride [stride_d, stride_h, stride_w]; all three must be equal and in {1, 2}.
dilation (IndexList[Int(3)]): Per-axis dilation; must be (1, 1, 1).
symmetric_padding (IndexList[Int(3)]): Per-axis symmetric padding [pad_d, pad_h, pad_w]; each in {0, 1, 2} and pad_h must equal pad_w.
num_groups (Int): Convolution group count; must be 1.
ctx (DeviceContext): DeviceContext used to enqueue the transpose and conv kernels.

Returns:

Bool