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

test_matmul_sm90_swapAB_comparison_v2

def test_matmul_sm90_swapAB_comparison_v2[MType: CoordLike, NType: CoordLike, KType: CoordLike, //, a_type: DType, b_type: DType, c_type: DType, BM: Int, BN: Int, BK: Int, MMA_M: Int, MMA_N: Int, MMA_K: Int, num_pipeline_stages: Int, num_consumer: Int, k_group_size: Int = Int(1), num_k_partitions: Int = Int(1), partitioned_multicast: Bool = False, BM_SWAPAB: Int = BM, BN_SWAPAB: Int = BN, BK_SWAPAB: Int = BK, MMA_M_SWAPAB: Int = MMA_M, MMA_N_SWAPAB: Int = MMA_N, MMA_K_SWAPAB: Int = MMA_K, num_pipeline_stages_swapAB: Int = num_pipeline_stages, num_consumer_swapAB: Int = num_consumer, k_group_size_swapAB: Int = k_group_size, num_k_partitions_swapAB: Int = num_k_partitions, partitioned_multicast_swapAB: Bool = partitioned_multicast, use_vendor_reference: Bool = False, default_epilogue: Bool = False, elementwise_compute_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]] = None](ctx: DeviceContext, m: MType, n: NType, k: KType)

Compare matmul results between normal execution and swapAB execution.

This version accepts config parameters directly as compile-time values and builds configs internally.

Both compute: C[M,N] = A[M,K] @ B[N,K]^T swapAB internally swaps A/B and transposes C on store, but result should match.

Parameters:

MType (CoordLike): The type of the M dimension.
NType (CoordLike): The type of the N dimension.
KType (CoordLike): The type of the K dimension.
a_type (DType): Data type of matrix A.
b_type (DType): Data type of matrix B.
c_type (DType): Data type of output matrix C.
BM (Int): Block tile M dimension for normal kernel.
BN (Int): Block tile N dimension for normal kernel.
BK (Int): Block tile K dimension for normal kernel.
MMA_M (Int): MMA M dimension for normal kernel.
MMA_N (Int): MMA N dimension for normal kernel.
MMA_K (Int): MMA K dimension for normal kernel.
num_pipeline_stages (Int): Number of pipeline stages for normal kernel.
num_consumer (Int): Number of consumers for normal kernel.
k_group_size (Int): K group size for normal kernel.
num_k_partitions (Int): Number of K partitions for normal kernel.
partitioned_multicast (Bool): Partitioned multicast for normal kernel.
BM_SWAPAB (Int): Block tile M dimension for swapAB kernel.
BN_SWAPAB (Int): Block tile N dimension for swapAB kernel.
BK_SWAPAB (Int): Block tile K dimension for swapAB kernel.
MMA_M_SWAPAB (Int): MMA M dimension for swapAB kernel.
MMA_N_SWAPAB (Int): MMA N dimension for swapAB kernel.
MMA_K_SWAPAB (Int): MMA K dimension for swapAB kernel.
num_pipeline_stages_swapAB (Int): Number of pipeline stages for swapAB kernel.
num_consumer_swapAB (Int): Number of consumers for swapAB kernel.
k_group_size_swapAB (Int): K group size for swapAB kernel.
num_k_partitions_swapAB (Int): Number of K partitions for swapAB kernel.
partitioned_multicast_swapAB (Bool): Partitioned multicast for swapAB kernel.
use_vendor_reference (Bool): If True, use vendor matmul (cuBLAS) as reference instead of normal kernel.
default_epilogue (Bool): If True, use default epilogue function that stores directly to output tensor.
elementwise_compute_lambda_fn (Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> SIMD[dtype, width]]): Optional compute lambda function to apply to each element before storing.

Args:

ctx (DeviceContext): The device context.
m (MType): The M dimension (can be static or dynamic).
n (NType): The N dimension (can be static or dynamic).
k (KType): The K dimension (can be static or dynamic).