multistage_gemm

def multistage_gemm[c_type: DType, a_type: DType, b_type: DType, //, *, transpose_b: Bool, config: MatmulConfig[a_type, b_type, c_type, transpose_b], elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type], ctx: DeviceContext)

TileTensor overload of multistage_gemm. Converts to LayoutTensor and dispatches to a GEMM kernel.

Parameters:

c_type (DType): DType of the output tile c elements (inferred).
a_type (DType): DType of the input tile a elements (inferred).
b_type (DType): DType of the input tile b elements (inferred).
transpose_b (Bool): Whether b is accessed transposed, so its row y is read as b[y, i] instead of b[i, y].
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Compile-time MatmulConfig selecting the block tile, warp tile, MMA shape, and pipeline stages for the kernel.
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 accumulated value before it is stored to c (defaults to None, which stores the raw accumulation).

Args:

c (TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output tile of shape (M, N) receiving the matmul result.
a (TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input tile of shape (M, K).
b (TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Input tile of shape (K, N), or (N, K) when transpose_b is set.
ctx (DeviceContext): Device context used to enqueue the kernel.

def multistage_gemm[c_type: DType, a_type: DType, b_type: DType, EpilogueFnType: ElementwiseEpilogueFn, //, *, transpose_b: Bool, config: MatmulConfig[a_type, b_type, c_type, transpose_b]](c: TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type], epilogue_fn: EpilogueFnType, ctx: DeviceContext)

Runs multistage_gemm, storing the output through an epilogue closure value.

Only the AMD CDNA transpose_b kernels support this overload.

Parameters:

c_type (DType): DType of the output tile c elements (inferred).
a_type (DType): DType of the input tile a elements (inferred).
b_type (DType): DType of the input tile b elements (inferred).
EpilogueFnType (ElementwiseEpilogueFn): Type of epilogue_fn (inferred).
transpose_b (Bool): Whether b is stored as (N, K). Must be True.
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Compile-time MatmulConfig selecting the block tile, warp tile, MMA shape, and pipeline stages for the kernel.

Args:

c (TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output tile of shape (M, N). Only its shape is read.
a (TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input tile of shape (M, K).
b (TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Input tile of shape (N, K).
epilogue_fn (EpilogueFnType): Stores each output element.
ctx (DeviceContext): Device context used to enqueue the kernel.

def multistage_gemm[c_type: DType, a_type: DType, b_type: DType, //, *, transpose_b: Bool, config: MatmulConfig[a_type, b_type, c_type, transpose_b], elementwise_lambda_fn: Optional[def[dtype: DType, width: SIMDLength, *, alignment: Int = Int(1)](IndexList[Int(2)], SIMD[dtype, width]) capturing thin -> None] = None](c: TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type], runtime_config: MatmulConfig[a_type, b_type, c_type, transpose_b], ctx: DeviceContext)

TileTensor overload of multistage_gemm with runtime config. Constrains c to mut=True because split_k_reduce requires a mutable output tensor.

Parameters:

c_type (DType): DType of the output tile c elements (inferred).
a_type (DType): DType of the input tile a elements (inferred).
b_type (DType): DType of the input tile b elements (inferred).
transpose_b (Bool): Whether b is accessed transposed, so its row y is read as b[y, i] instead of b[i, y].
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Compile-time MatmulConfig selecting the block tile, warp tile, MMA shape, and pipeline stages for the kernel.
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 accumulated value before it is stored to c (defaults to None, which stores the raw accumulation).

Args:

c (TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output tile of shape (M, N) receiving the matmul result; must be mutable because split_k_reduce writes the reduced sum back into it.
a (TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input tile of shape (M, K).
b (TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Input tile of shape (K, N), or (N, K) when transpose_b is set.
runtime_config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Runtime MatmulConfig carrying the number of K partitions used for the split-K reduction path; when num_k_partitions is greater than 1 the kernel writes partial sums to a workspace and reduces them into c.
ctx (DeviceContext): Device context used to enqueue the kernel and allocate the split-K workspace.

def multistage_gemm[c_type: DType, a_type: DType, b_type: DType, EpilogueFnType: ElementwiseEpilogueFn, //, *, transpose_b: Bool, config: MatmulConfig[a_type, b_type, c_type, transpose_b]](c: TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type], a: TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type], b: TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type], runtime_config: MatmulConfig[a_type, b_type, c_type, transpose_b], epilogue_fn: EpilogueFnType, ctx: DeviceContext)

Runs multistage_gemm with runtime config, storing the output through an epilogue closure value.

Only the AMD CDNA transpose_b kernels support this overload.

Parameters:

c_type (DType): DType of the output tile c elements (inferred).
a_type (DType): DType of the input tile a elements (inferred).
b_type (DType): DType of the input tile b elements (inferred).
EpilogueFnType (ElementwiseEpilogueFn): Type of epilogue_fn (inferred).
transpose_b (Bool): Whether b is stored as (N, K). Must be True.
config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Compile-time MatmulConfig selecting the block tile, warp tile, MMA shape, and pipeline stages for the kernel.

Args:

c (TileTensor[c_type, Engine=c.Engine, address_space=c.address_space, linear_idx_type=c.linear_idx_type]): Output tile of shape (M, N). Only its shape is read, unless the split-K path needs it.
a (TileTensor[a_type, Engine=a.Engine, address_space=a.address_space, linear_idx_type=a.linear_idx_type]): Input tile of shape (M, K).
b (TileTensor[b_type, Engine=b.Engine, address_space=b.address_space, linear_idx_type=b.linear_idx_type]): Input tile of shape (N, K).
runtime_config (MatmulConfig[a_type, b_type, c_type, transpose_b]): Runtime MatmulConfig carrying the number of K partitions used for the split-K reduction path.
epilogue_fn (EpilogueFnType): Stores each output element.
ctx (DeviceContext): Device context used to enqueue the kernel and allocate the split-K workspace.