leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

v4.4 release update v2. (#2999)

Browse Source
This commit is contained in:
Junkai-Wu
2026-02-04 09:48:31 +08:00
committed by GitHub
parent 1cfbb53a23
commit 6b3e607b85
91 changed files with 13242 additions and 1488 deletions
Download Patch File Download Diff File Expand all files Collapse all files

284
examples/python/CuTeDSL/blackwell/dense_gemm_persistent_dynamic.py
View File

@@ -34,8 +34,8 @@ import cuda.bindings.driver as cuda
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
from cutlass.cute.runtime import from_dlpack
import cutlass.utils as utils
from cutlass.utils import is_fp8_dtype, create_cute_tensor_for_fp8
import cutlass.pipeline as pipeline
from cutlass.pipeline import pipeline_init_arrive, pipeline_init_wait
from cutlass.cute.nvgpu import cpasync, tcgen05
@@ -1046,45 +1046,70 @@ class PersistentDenseGemmKernel:
# (MMA, MMA_M, MMA_N, STAGE)
tCtAcc_base = cute.make_tensor(tmem_ptr, tCtAcc_fake.layout)
#
# Persistent tile scheduling loop for epilogue
#
acc_consumer_state = pipeline.make_pipeline_state(
pipeline.PipelineUserType.Consumer, self.num_acc_stage
)
if cutlass.const_expr(self.use_tma_store):
assert tma_atom_c is not None and sC is not None
utils.gemm.sm100.epilogue_tma_store(
self,
tidx,
warp_idx,
acc_pipeline,
tiled_mma,
tma_atom_c,
tCtAcc_base,
sC,
tCgC,
epi_tile,
tile_sched,
epilogue_op,
clc_pipeline,
clc_consumer_state,
c_producer_group = pipeline.CooperativeGroup(
pipeline.Agent.Thread,
32 * len(self.epilogue_warp_id),
)
else:
utils.gemm.sm100.epilogue(
self,
tidx,
acc_pipeline,
tiled_mma,
tCtAcc_base,
tCgC,
epi_tile,
tile_sched,
epilogue_op,
tmem_dealloc_barrier,
None,
None,
clc_pipeline,
clc_consumer_state,
c_pipeline = pipeline.PipelineTmaStore.create(
num_stages=self.num_c_stage, producer_group=c_producer_group
)
while work_tile.is_valid_tile:
# Get tile coord from tile scheduler
cur_tile_coord = work_tile.tile_idx
mma_tile_coord_mnl = (
cur_tile_coord[0] // cute.size(tiled_mma.thr_id.shape),
cur_tile_coord[1],
cur_tile_coord[2],
)
num_tiles_executed = tile_sched.num_tiles_executed
if cutlass.const_expr(self.use_tma_store):
acc_consumer_state = utils.gemm.sm100.epilogue_tma_store(
self,
tidx,
warp_idx,
tma_atom_c,
tCtAcc_base,
sC,
tCgC,
epi_tile,
num_tiles_executed,
epilogue_op,
mma_tile_coord_mnl,
acc_consumer_state,
acc_pipeline,
c_pipeline,
)
else:
acc_consumer_state = utils.gemm.sm100.epilogue(
self,
tidx,
tCtAcc_base,
tCgC,
epi_tile,
epilogue_op,
mma_tile_coord_mnl,
acc_consumer_state,
acc_pipeline,
)
#
# Advance to next tile
#
clc_pipeline.consumer_wait(clc_consumer_state)
work_tile = tile_sched.get_current_work()
clc_pipeline.consumer_release(clc_consumer_state)
clc_consumer_state.advance()
if cutlass.const_expr(self.use_tma_store):
# Wait for C store complete
c_pipeline.producer_tail()
else:
# Synchronize before TMEM dealloc (done by the caller)
tmem_dealloc_barrier.arrive_and_wait()
#
# Dealloc the tensor memory buffer
#
@@ -1150,8 +1175,11 @@ class PersistentDenseGemmKernel:
return num_tmem_alloc_cols
def check_supported_dtypes(
self, ab_dtype: Type[cutlass.Numeric], c_dtype: Type[cutlass.Numeric]
) -> bool:
self,
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
):
"""
Check if the dtypes are valid
@@ -1173,8 +1201,10 @@ class PersistentDenseGemmKernel:
cutlass.Float8E4M3FN,
cutlass.Float8E5M2,
}
if ab_dtype not in valid_ab_dtypes:
raise testing.CantImplementError(f"Unsupported AB dtype: {ab_dtype}")
if a_dtype not in valid_ab_dtypes or b_dtype not in valid_ab_dtypes:
raise testing.CantImplementError(
f"Unsupported AB dtype: {a_dtype} and {b_dtype}"
)
if self.acc_dtype not in {cutlass.Float32, cutlass.Float16, cutlass.Int32}:
raise testing.CantImplementError(
@@ -1198,8 +1228,13 @@ class PersistentDenseGemmKernel:
cutlass.Int32: {cutlass.Uint8, cutlass.Int8},
}
# Check compatibility between accumulator type and AB type
if ab_dtype not in acc_ab_compatibility[self.acc_dtype]:
return False
if (
a_dtype not in acc_ab_compatibility[self.acc_dtype]
or b_dtype not in acc_ab_compatibility[self.acc_dtype]
):
raise testing.CantImplementError(
f"Unsupported AB dtype: {a_dtype} and {b_dtype} for accumulator dtype: {self.acc_dtype}"
)
# Define compatibility mapping between accumulator type and C type
acc_c_compatibility = {
@@ -1228,11 +1263,11 @@ class PersistentDenseGemmKernel:
}
# Check compatibility between accumulator type and C type
if c_dtype not in acc_c_compatibility[self.acc_dtype]:
return False
raise testing.CantImplementError(
f"Unsupported C dtype: {c_dtype} for accumulator dtype: {self.acc_dtype}"
)
return True
def check_mma_tiler_and_cluster_shape(self) -> bool:
def check_mma_tiler_and_cluster_shape(self):
"""Check if the mma tiler and cluster shape are valid.
:raises testing.CantImplementError: If the mma tiler and cluster shape are invalid
@@ -1273,12 +1308,13 @@ class PersistentDenseGemmKernel:
n: int,
k: int,
l: int,
ab_dtype: Type[cutlass.Numeric],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
) -> bool:
):
"""
Check if the tensor alignment is valid
@@ -1290,8 +1326,10 @@ class PersistentDenseGemmKernel:
:type k: int
:param l: The number of columns in the C tensor
:type l: int
:param ab_dtype: The data type of the A and B operands
:type ab_dtype: Type[cutlass.Numeric]
:param a_dtype: The data type of the A operand
:type a_dtype: Type[cutlass.Numeric]
:param b_dtype: The data type of the B operand
:type b_dtype: Type[cutlass.Numeric]
:param c_dtype: The data type of the output tensor
:type c_dtype: Type[cutlass.Numeric]
:param a_major: The major axis of the A tensor
@@ -1301,8 +1339,7 @@ class PersistentDenseGemmKernel:
:param c_major: The major axis of the C tensor
:type c_major: str
:return: True if the problem shape is valid, False otherwise
:rtype: bool
:raises testing.CantImplementError: If the tensor alignment is invalid
"""
# TODO: move to utils
@@ -1313,15 +1350,15 @@ class PersistentDenseGemmKernel:
return num_major_elements % num_contiguous_elements == 0
if (
not check_contiguous_16B_alignment(ab_dtype, a_major == "m", (m, k, l))
or not check_contiguous_16B_alignment(ab_dtype, b_major == "n", (n, k, l))
not check_contiguous_16B_alignment(a_dtype, a_major == "m", (m, k, l))
or not check_contiguous_16B_alignment(b_dtype, b_major == "n", (n, k, l))
or not check_contiguous_16B_alignment(c_dtype, c_major == "m", (m, n, l))
):
raise testing.CantImplementError(
f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {ab_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}"
f"Invalid tensor alignment: {m}, {n}, {k}, {l}, {a_dtype}, {b_dtype}, {c_dtype}, {a_major}, {b_major}, {c_major}"
)
def check_epilog_store_option(self, m: int, n: int) -> bool:
def check_epilog_store_option(self, m: int, n: int):
"""
Check if the epilogue store option is valid
@@ -1346,7 +1383,8 @@ class PersistentDenseGemmKernel:
def can_implement(
self,
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
@@ -1357,8 +1395,10 @@ class PersistentDenseGemmKernel:
:param mnkl: Problem size as a tuple (M, N, K, L).
:type mnkl: Tuple[int, int, int, int]
:param ab_dtype: Data type for input tensors A and B.
:type ab_dtype: Type[cutlass.Numeric]
:param a_dtype: Data type for input tensors A.
:type a_dtype: Type[cutlass.Numeric]
:param b_dtype: Data type for input tensors B.
:type b_dtype: Type[cutlass.Numeric]
:param c_dtype: Data type for output tensor C.
:type c_dtype: Type[cutlass.Numeric]
:param a_major: Major dimension of the A tensor layout ("m" or "k").
@@ -1373,14 +1413,14 @@ class PersistentDenseGemmKernel:
try:
# Skip unsupported types
self.check_supported_dtypes(ab_dtype, c_dtype)
self.check_supported_dtypes(a_dtype, b_dtype, c_dtype)
# Skip invalid mma tile shape and cluster shape
self.check_mma_tiler_and_cluster_shape()
m, n, k, l = mnkl
self.check_tensor_alignment(
m, n, k, l, ab_dtype, c_dtype, a_major, b_major, c_major
m, n, k, l, a_dtype, b_dtype, c_dtype, a_major, b_major, c_major
)
self.check_epilog_store_option(m, n)
except testing.CantImplementError:
@@ -1432,43 +1472,72 @@ def bmm(
@lru_cache(maxsize=1)
def prepare_tensors(
mnkl: Tuple[int, int, int, int],
ab_dtype: Type[cutlass.Numeric],
a_dtype: Type[cutlass.Numeric],
b_dtype: Type[cutlass.Numeric],
c_dtype: Type[cutlass.Numeric],
a_major: str,
b_major: str,
c_major: str,
init_random: bool = True,
normal_mean: float = 0.0,
normal_std: float = 1.0,
):
"""Prepare tensors for GEMM.
Returns:
Tuple of (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage):
- *_f32: Float32 tensors with the logical data (for reference and fp8 conversion)
- *_storage: Storage tensors for DLPack (uint8 for fp8, otherwise the target dtype)
"""
import torch
from cutlass.torch import dtype as torch_dtype
m, n, k, l = mnkl
if a_major == "k":
a = torch.empty((l, m, k), dtype=torch.float32, device="cuda")
a_f32 = torch.empty((l, m, k), dtype=torch.float32, device="cuda")
elif a_major == "m":
a = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
a_f32 = torch.empty((l, k, m), dtype=torch.float32, device="cuda").permute(
0, 2, 1
)
if b_major == "n":
b = torch.empty((l, k, n), dtype=torch.float32, device="cuda")
b_f32 = torch.empty((l, k, n), dtype=torch.float32, device="cuda")
elif b_major == "k":
b = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(0, 2, 1)
b_f32 = torch.empty((l, n, k), dtype=torch.float32, device="cuda").permute(
0, 2, 1
)
if c_major == "n":
c = torch.empty((l, m, n), dtype=torch.float32, device="cuda")
c_f32 = torch.empty((l, m, n), dtype=torch.float32, device="cuda")
elif c_major == "m":
c = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(0, 2, 1)
c_f32 = torch.empty((l, n, m), dtype=torch.float32, device="cuda").permute(
0, 2, 1
)
if init_random:
a.random_(-2, 3)
b.random_(-2, 3)
c.random_(-2, 3)
# Uniform random initialization in range [-2, 3)
a_f32.random_(-2, 3)
b_f32.random_(-2, 3)
c_f32.random_(-2, 3)
else:
# Normal (Gaussian) initialization with user-specified mean and std
a_f32.normal_(mean=normal_mean, std=normal_std)
b_f32.normal_(mean=normal_mean, std=normal_std)
c_f32.normal_(mean=normal_mean, std=normal_std)
return (
a.to(dtype=torch_dtype(ab_dtype)),
b.to(dtype=torch_dtype(ab_dtype)),
c.to(dtype=torch_dtype(c_dtype)),
)
# For float8 types, use uint8 as storage type to avoid dlpack limitation
# (dlpack doesn't support float8 types)
# For other types, convert to the target dtype
a_storage_dtype = torch.uint8 if is_fp8_dtype(a_dtype) else torch_dtype(a_dtype)
b_storage_dtype = torch.uint8 if is_fp8_dtype(b_dtype) else torch_dtype(b_dtype)
c_storage_dtype = torch.uint8 if is_fp8_dtype(c_dtype) else torch_dtype(c_dtype)
a_storage = a_f32.to(dtype=a_storage_dtype)
b_storage = b_f32.to(dtype=b_storage_dtype)
c_storage = c_f32.to(dtype=c_storage_dtype)
return (a_f32, b_f32, c_f32, a_storage, b_storage, c_storage)
@lru_cache(maxsize=1)
@@ -1499,7 +1568,7 @@ def compile_bmm(
)
# Check if configuration can be implemented
can_implement = gemm.can_implement(
mnkl, a.element_type, c.element_type, a_major, b_major, c_major
mnkl, a.element_type, b.element_type, c.element_type, a_major, b_major, c_major
)
if not can_implement:
raise testing.CantImplementError(
@@ -1594,21 +1663,30 @@ def run(
)
# Run and verify BMM with torch
a, b, c = prepare_tensors(mnkl, ab_dtype, c_dtype, a_major, b_major, c_major)
a_f32, b_f32, c_f32, a_storage, b_storage, c_storage = prepare_tensors(
mnkl, ab_dtype, ab_dtype, c_dtype, a_major, b_major, c_major
)
leading_dim_a = 2 if a_major == "k" else 1
leading_dim_b = 1 if b_major == "k" else 2
leading_dim_c = 2 if c_major == "n" else 1
a_ = from_dlpack(
a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_a)
b_ = from_dlpack(
b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_b)
c_ = from_dlpack(
c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_c)
# Create CuTe tensors, passing float32 source for fp8 conversion
a_ = create_cute_tensor_for_fp8(
a_storage, ab_dtype, leading_dim_a, source_f32_tensor=a_f32
)
b_ = create_cute_tensor_for_fp8(
b_storage, ab_dtype, leading_dim_b, source_f32_tensor=b_f32
)
c_ = create_cute_tensor_for_fp8(
c_storage, c_dtype, leading_dim_c, source_f32_tensor=c_f32
)
print("Compile Blackwell Persistent Dense GEMM with:")
print(f"ab_dtype: {ab_dtype}, c_dtype: {c_dtype}, acc_dtype: {acc_dtype}")
print(f"a_major: {a_major}, b_major: {b_major}, c_major: {c_major}")
print(f"mma_tiler_mn: {mma_tiler_mn}, cluster_shape_mn: {cluster_shape_mn}")
print(f"use_2cta_instrs: {use_2cta_instrs}, use_tma_store: {use_tma_store}")
compiled_fn = compile_bmm(
mnkl,
@@ -1640,13 +1718,15 @@ def run(
compiled_fn(a_, b_, c_, current_stream)
# Manually quantize to be comparable
# Use float32 source data for reference calculation
ref = (
torch.bmm(a.to(dtype=torch.float32), b.to(dtype=torch.float32))
torch.bmm(a_f32, b_f32)
.to(dtype=torch_dtype(c_dtype))
.to(dtype=torch.float32)
)
# Read back the result from CuTe tensor (c_storage was updated in-place)
torch.testing.assert_close(
c.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03
c_storage.to(dtype=torch.float32), ref, atol=tolerance, rtol=1e-03
)
if not benchmark:
@@ -1654,32 +1734,33 @@ def run(
def generate_tensors():
init_normal = ab_dtype not in [cutlass.Int8, cutlass.Uint8]
a, b, c = prepare_tensors(
a_f32, b_f32, c_f32, a_st, b_st, c_st = prepare_tensors(
mnkl,
ab_dtype,
ab_dtype,
c_dtype,
a_major,
b_major,
c_major,
init_random=not init_normal,
)
a_ = from_dlpack(
a, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_a)
b_ = from_dlpack(
b, assumed_align=16, force_tf32=ab_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_b)
c_ = from_dlpack(
c, assumed_align=16, force_tf32=c_dtype == cutlass.TFloat32
).mark_layout_dynamic(leading_dim=leading_dim_c)
a_ = create_cute_tensor_for_fp8(
a_st, ab_dtype, leading_dim_a, source_f32_tensor=a_f32
)
b_ = create_cute_tensor_for_fp8(
b_st, ab_dtype, leading_dim_b, source_f32_tensor=b_f32
)
c_ = create_cute_tensor_for_fp8(
c_st, c_dtype, leading_dim_c, source_f32_tensor=c_f32
)
return testing.JitArguments(a_, b_, c_, current_stream)
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
a.numel() * a.element_size()
+ b.numel() * b.element_size()
+ c.numel() * c.element_size()
a_storage.numel() * a_storage.element_size()
+ b_storage.numel() * b_storage.element_size()
+ c_storage.numel() * c_storage.element_size()
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
@@ -1735,6 +1816,7 @@ def prepare_parser():
action="store_true",
help="Enable 2CTA MMA instructions feature",
)
parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k")
parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k")
parser.add_argument("--c_major", choices=["n", "m"], type=str, default="n")

2977
examples/python/CuTeDSL/blackwell/sm103_dense_blockscaled_gemm_persistent.py Normal file
View File

File diff suppressed because it is too large Load Diff

155
examples/python/CuTeDSL/experimental/ampere/memcpy_simt_universal_copy.py Normal file
View File

@@ -0,0 +1,155 @@
# Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
import torch
import pytest
from cutlass import cute
from cutlass.cute import experimental as cute_ext
from cutlass.cute.runtime import from_dlpack
import cutlass.utils as utils
@cute.experimental.kernel
def memcpy_simt_universal_copy_kernel(
mA: cute.Tensor, mD: cute.Tensor, addend: cute.Float16
):
tile_mn = cute.core._pack_shape((128, 64))
gA = cute.zipped_divide(mA, tile_mn)
gD = cute.zipped_divide(mD, tile_mn)
cta_m, cta_n, cta_l = cute.arch.block_idx()
tid_x, _, _ = cute.arch.thread_idx()
gA_tile = gA[(None, None), (cta_m, cta_n, cta_l)]
gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)]
buffer = cute_ext.allocate(
cute.Float16,
cute.AddressSpace.rmem,
cute.make_layout(((8, 1), (1, 8)), stride=((1, 8), (1, 8))),
alignment=16,
)
tCgA = cute_ext.partition(
gA_tile,
tid_x,
layout_tv=cute.make_layout(((16, 8), (8, 1)), stride=((8, 128), (1, 1024))),
tiler=cute.core._pack_tile((128, 8)),
)
tCgD = cute_ext.partition(
gD_tile,
tid_x,
layout_tv=cute.make_layout(((16, 8), (8, 1)), stride=((8, 128), (1, 1024))),
tiler=cute.core._pack_tile((128, 8)),
)
# cute_ext.copy() automatically computes predicates based on the shape of
# the tensor passed to the @cute.experimental.kernel argument
cute_ext.copy(
tCgA,
buffer,
copy_atom=cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
tCgD.element_type,
num_bits_per_copy=128,
),
)
# Update the RMEM tensor in place using elementwise addition.
buffer.store(buffer.load() + addend)
# cute_ext.copy() automatically computes predicates based on the shape of
# the tensor passed to the @cute.experimental.kernel argument
cute_ext.copy(
buffer,
tCgD,
copy_atom=cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
tCgD.element_type,
num_bits_per_copy=128,
),
)
@cute.experimental.jit
def memcpy_simt_universal_copy(
src: cute.Tensor, dst: cute.Tensor, addend: cute.Float16
):
tile_mn = cute.core._pack_shape((128, 64))
div = cute.tiled_divide(src, tile_mn)
grid = (div.shape[1], div.shape[2], div.shape[3])
memcpy_simt_universal_copy_kernel(src, dst, addend).launch(
grid=grid,
block=(128, 1, 1),
smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_80")),
)
def run_simt_universal_memcpy(M, N, L):
src = torch.randn(L, N, M).permute(2, 1, 0).to(torch.float16).cuda()
dst = torch.randn(L, N, M).permute(2, 1, 0).to(torch.float16).cuda()
mA = (
from_dlpack(src, assumed_align=16)
.mark_layout_dynamic(leading_dim=0)
.mark_compact_shape_dynamic(
mode=0, stride_order=src.dim_order(), divisibility=8
)
)
mD = (
from_dlpack(dst, assumed_align=16)
.mark_layout_dynamic(leading_dim=0)
.mark_compact_shape_dynamic(
mode=0, stride_order=dst.dim_order(), divisibility=8
)
)
addend = 5.0
memcpy_simt_universal_copy(
mA,
mD,
cute.Float16(addend),
no_cache=True,
)
torch.testing.assert_close(src.cpu() + addend, dst.cpu())
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="Example memory copy example using CuTe auto predication features."
)
parser.add_argument("--mnl", default=[136, 7, 9], nargs="+", type=int)
args = parser.parse_args()
M, N, L = tuple(args.mnl)
run_simt_universal_memcpy(M, N, L)
print("PASS")

1027
examples/python/CuTeDSL/experimental/blackwell/dense_block_scaled_gemm.py Normal file
View File

File diff suppressed because it is too large Load Diff

1410
examples/python/CuTeDSL/experimental/blackwell/dense_gemm.py Normal file
View File

File diff suppressed because it is too large Load Diff

522
examples/python/CuTeDSL/experimental/blackwell/dense_gemm_2sm.py Normal file
View File

@@ -0,0 +1,522 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
"""
2SM Dense GEMM example using cute_ext decorators.
"""
import torch
import math
import cutlass
from cutlass import cute
from cutlass.cute import experimental as cute_ext
from cutlass.cute.runtime import from_dlpack
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.utils as utils
from cutlass.base_dsl.typing import Numeric
from typing import Type
def create_gemm_tensors_torch(
M,
N,
K,
majors: tuple[
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
],
dtypes: tuple[torch.dtype, torch.dtype, torch.dtype],
):
A = None
B = None
D = None
if majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.MN:
A = torch.empty(K, M).random_(-4, 4).permute(1, 0).to(dtypes[0]).cuda()
elif majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.K:
A = torch.empty(M, K).random_(-4, 4).permute(0, 1).to(dtypes[0]).cuda()
if majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.MN:
B = torch.empty(K, N).random_(-4, 4).permute(1, 0).to(dtypes[1]).cuda()
elif majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.K:
B = torch.empty(N, K).random_(-4, 4).permute(0, 1).to(dtypes[1]).cuda()
if majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.MN:
D = torch.empty(N, M).random_(-4, 4).permute(1, 0).to(dtypes[2]).cuda()
elif majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.K:
D = torch.empty(M, N).random_(-4, 4).permute(0, 1).to(dtypes[2]).cuda()
return A, B, D
def get_gemm_tensors(
M,
N,
K,
majors: tuple[
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
cute.nvgpu.tcgen05.OperandMajorMode,
],
dtypes: tuple[torch.dtype, torch.dtype, torch.dtype],
):
A, B, D = create_gemm_tensors_torch(M, N, K, majors, dtypes)
A_cute = from_dlpack(A, assumed_align=16).mark_layout_dynamic(
leading_dim=1 if majors[0] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0
)
B_cute = from_dlpack(B, assumed_align=16).mark_layout_dynamic(
leading_dim=1 if majors[1] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0
)
D_cute = from_dlpack(D, assumed_align=16).mark_layout_dynamic(
leading_dim=1 if majors[2] == cute.nvgpu.tcgen05.OperandMajorMode.K else 0
)
return A, B, D, A_cute, B_cute, D_cute
def sm100_4x4x1_kernel_builder(
use_tma_multicast: bool,
use_2cta_instrs: bool,
acc_dtype: Type[Numeric],
M: int,
N: int,
):
CLUSTER_SHAPE = (2, 1, 1)
GRID_SHAPE = (
math.ceil(M / 128),
math.ceil(N / 256),
1,
) # TODO (xpbowler): remove hard-code
NUM_WARPS_PER_CTA = 6
TMA_STORE_PIPE_DEPTH = 4
MAINLOOP_STAGE_DEPTH = 4 # pipeline depth of TMA->MMA
# pipeline depth of mainloop->epilogue. only useful if using persistent CTA
EPILOGUE_STAGE_DEPTH = 1
# m256n256k16 2SM MMA / m128n256k16 1SM MMA
mma_inst_shape_mnk = (256, 256, 16) if use_2cta_instrs else (128, 256, 16)
@cute_ext.kernel
def kernel(
mA: cute.Tensor,
mB: cute.Tensor,
mD: cute.Tensor,
):
d_layout = utils.LayoutEnum.from_tensor(mD)
d_dtype = mD.element_type
ab_dtype = mA.element_type
mma_inst_shape_m, mma_inst_shape_n, mma_inst_shape_k = mma_inst_shape_mnk
if cutlass.const_expr(use_2cta_instrs):
cta_group = cute.nvgpu.tcgen05.CtaGroup.TWO
else:
cta_group = cute.nvgpu.tcgen05.CtaGroup.ONE
tiled_mma = sm100_utils.make_trivial_tiled_mma(
ab_dtype,
utils.LayoutEnum.from_tensor(mA).mma_major_mode(),
utils.LayoutEnum.from_tensor(mB).mma_major_mode(),
acc_dtype,
cta_group,
(mma_inst_shape_m, mma_inst_shape_n),
)
mma_inst_tile_k = (
4 # 4 MMAs per MMA tile K. For 16b types, tcgen05.mma has K=16.
)
mma_inst_tile_m = mma_inst_tile_n = 1 # 1 MMAs per MMA tile M/N
bM = mma_inst_shape_m * mma_inst_tile_m
bN = mma_inst_shape_n * mma_inst_tile_n
bK = mma_inst_shape_k * mma_inst_tile_k
mnk_tiler = (bM, bN, bK)
cta_m, cta_n, _ = cute.arch.block_idx()
tid_x, _, _ = cute.arch.thread_idx()
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
cluster_layout_vmnk = cute.tiled_divide(
cute.make_layout(CLUSTER_SHAPE),
cute.core._pack_shape((cute.size(tiled_mma.thr_id.shape),)),
)
cluster_layout_v_size = cute.size(cluster_layout_vmnk.shape[0])
mma_coord_vmnk = (
cta_m % cluster_layout_v_size,
cta_m // cluster_layout_v_size,
cta_n,
)
gA = cute.zipped_divide(mA, (bM, bK)) # ((bM, bK), (M/bM, K/bK))
gA_tma = cute.zipped_divide(
mA, (bM // cluster_layout_v_size, bK)
) # ((bM/2, bK), (2*M/bM, K/bK))
tAgA = gA_tma[(None, None), (cta_m, None)] # ((bM/2, bK), (1, K/bK))
gB_tma = cute.zipped_divide(
mB, (bN // cluster_layout_v_size, bK)
) # ((bN/2, bK), (2*M/bM, K/bK))
# ((bN/2, bK), (1, K/bK))
tBgB = gB_tma[
(None, None),
(cluster_layout_v_size * cta_n + cta_m % cluster_layout_v_size, None),
]
gD_tma = cute.zipped_divide(
mD, (bM // cluster_layout_v_size, bN)
) # ((bM/2, bN), (2*M/bM, N/bN))
tDgD = gD_tma[(None, None), (cta_m, cta_n)] # ((bM/2, bN), (1, 1))
a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
mnk_tiler,
ab_dtype,
MAINLOOP_STAGE_DEPTH,
)
b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
mnk_tiler,
ab_dtype,
MAINLOOP_STAGE_DEPTH,
)
cta_tile_shape_mnk = cute.shape_div(mnk_tiler, (cluster_layout_v_size, 1, 1))
epi_tile = sm100_utils.compute_epilogue_tile_shape(
cta_tile_shape_mnk,
use_2cta_instrs,
d_layout,
d_dtype,
)
sc_smem_layout_staged = sm100_utils.make_smem_layout_epi(
d_dtype,
d_layout,
epi_tile,
TMA_STORE_PIPE_DEPTH,
)
acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2])
tmem_layout = tiled_mma.make_fragment_C(
cute.append(acc_shape, EPILOGUE_STAGE_DEPTH)
).layout
bufferA = cute_ext.allocate(
ab_dtype,
cute.AddressSpace.smem,
a_smem_layout_staged,
alignment=1024,
)
bufferB = cute_ext.allocate(
ab_dtype,
cute.AddressSpace.smem,
b_smem_layout_staged,
alignment=1024,
)
bufferAcc = cute_ext.allocate(
acc_dtype,
cute.AddressSpace.tmem,
tmem_layout,
alignment=16,
is2cta=use_2cta_instrs,
)
bufferC = cute_ext.allocate(
d_dtype,
cute.AddressSpace.smem,
sc_smem_layout_staged,
alignment=1024,
)
copy_atom_t2r = sm100_utils.get_tmem_load_op(
cta_tile_shape_mnk,
d_layout,
d_dtype,
acc_dtype,
epi_tile,
use_2cta_instrs,
)
# Take only one stage of the TMEM buffer for the epilogue
accumulators = cute.zipped_divide(bufferAcc, ((epi_tile), 1))
acc_epi_div = accumulators[((None, None), 0), 0]
# Create the TMEM copy atom based on the size of transfer within one iteration of epilogue
tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div)
# Calculate the per thread destination size per iteration for output of TMEM and input of SMEM
thr_copy_t2r = tiled_copy_t2r.get_slice(tid_x)
gC_mnl_epi = cute.flat_divide(tDgD, epi_tile)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
acc_d_rmem_layout = cute.make_fragment_like(
tTR_gC[(None, None, None, 0, 0)].layout
)
bufferRAcc = cute_ext.allocate(
acc_dtype,
cute.AddressSpace.rmem,
acc_d_rmem_layout,
alignment=32,
)
bufferRD = cute_ext.allocate(
d_dtype,
cute.AddressSpace.rmem,
acc_d_rmem_layout,
alignment=32,
)
tma_mcast_proj_A = 2
tma_mcast_proj_B = 1
mma_operation_type = tma_operation_type = None
acc_pipe = mainloop_pipe = None
if cutlass.const_expr(use_2cta_instrs):
mma_operation_type = cute_ext.OperationTypeEnum.SM100_MMA_2SM_SS
if cutlass.const_expr(use_tma_multicast):
tma_operation_type = (
cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM_MULTICAST
)
else:
tma_operation_type = cute_ext.OperationTypeEnum.SM100_TMA_LOAD_2SM
else:
mma_operation_type = cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS
if cutlass.const_expr(use_tma_multicast):
tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD_MULTICAST
else:
tma_operation_type = cute_ext.OperationTypeEnum.SM90_TMA_LOAD
# MMA <-> TMEM load pipeline
# if 2CTA MMA, warpgroup from both peer and leader CTA consumer.release
acc_pipe_consumer_arv_count = 256 if use_2cta_instrs else 128
acc_pipe = cute_ext.UMMAtoAsyncPipeline.create(
num_stages=EPILOGUE_STAGE_DEPTH,
mma_operation_type=mma_operation_type,
consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R,
consumer_arv_count=acc_pipe_consumer_arv_count,
cluster_layout_vmnk=cluster_layout_vmnk,
)
if cutlass.const_expr(use_tma_multicast):
# TMA load <-> MMA pipeline
mainloop_pipe = cute_ext.TMAToUMMAPipeline.create_with_mask(
num_stages=MAINLOOP_STAGE_DEPTH,
tma_operation_type=tma_operation_type,
mma_operation_type=mma_operation_type,
cluster_layout_vmnk=cluster_layout_vmnk,
)
else:
mainloop_pipe = cute_ext.TMAToUMMAPipeline.create(
num_stages=MAINLOOP_STAGE_DEPTH,
mma_operation_type=mma_operation_type,
tma_operation_type=tma_operation_type,
cluster_layout_vmnk=cluster_layout_vmnk,
)
tma_store_warp_id = 0
mma_warp_id = 4
tma_load_warp_id = 5
is_tma_thr = warp_idx == tma_load_warp_id
is_mma_thr = warp_idx == mma_warp_id
is_epi_thr = warp_idx < 4
is_leader_cta = mma_coord_vmnk[0] == 0
# SMEM -> GMEM
tma_store_pipe = cute_ext.TMAStorePipeline(
stages=TMA_STORE_PIPE_DEPTH,
arv_count=128,
barrier_id=1,
tma_warp_id=tma_store_warp_id,
)
k_tile_count = cute.size(gA, mode=[1, 1])
if is_tma_thr:
for k_tile in cutlass.range(0, k_tile_count, 1, unroll=1):
gA_k = tAgA[None, None, k_tile]
gB_k = tBgB[None, None, k_tile]
producer_stage_token, idx = (
mainloop_pipe.producer_acquire_and_get_stage()
)
mbar = cute_ext.get_mbarrier(producer_stage_token)
bufferA_sliced = bufferA[None, None, None, idx]
bufferB_sliced = bufferB[None, None, None, idx]
a_cta_v_map = cute_ext.get_cta_v_map_ab(mA, mnk_tiler, tiled_mma, "A")
b_cta_v_map = cute_ext.get_cta_v_map_ab(mB, mnk_tiler, tiled_mma, "B")
if cutlass.const_expr(use_tma_multicast):
cute_ext.tma_load_multicast(
gA_k,
bufferA_sliced,
mbar,
vmnk_layout=cluster_layout_vmnk,
cta_v_map=a_cta_v_map,
tma_operation_type=tma_operation_type,
multicast_mode=tma_mcast_proj_A,
)
cute_ext.tma_load_multicast(
gB_k,
bufferB_sliced,
mbar,
vmnk_layout=cluster_layout_vmnk,
cta_v_map=b_cta_v_map,
tma_operation_type=tma_operation_type,
multicast_mode=tma_mcast_proj_B,
)
else:
cute_ext.tma_load(
gA_k,
bufferA_sliced,
mbar,
cta_v_map=a_cta_v_map,
tma_operation_type=tma_operation_type,
)
cute_ext.tma_load(
gB_k,
bufferB_sliced,
mbar,
cta_v_map=b_cta_v_map,
tma_operation_type=tma_operation_type,
)
if is_leader_cta:
mainloop_pipe.producer_commit()
mainloop_pipe.producer_state = cute_ext.pipeline_advance_iterator(
mainloop_pipe.raw_pipeline, mainloop_pipe.producer_state
)
if is_mma_thr and is_leader_cta:
producer_stage_token, idx = acc_pipe.producer_acquire_and_get_stage()
accumulators_sliced = bufferAcc[None, None, None, idx]
mma_atom = cute.make_mma_atom(tiled_mma.op)
mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False)
for k_tile in cutlass.range(0, k_tile_count, 1, unroll=1):
_, mainloop_idx = mainloop_pipe.consumer_wait_and_get_stage()
bufferA_sliced_stage = cute.core.slice_(
bufferA, (None, None, None, mainloop_idx)
)
bufferB_sliced_stage = cute.core.slice_(
bufferB, (None, None, None, mainloop_idx)
)
for k_block in cutlass.range(mma_inst_tile_k, unroll_full=True):
cute_ext.dot(
mma_atom,
cute.append_ones(
bufferA_sliced_stage[None, None, k_block], up_to_rank=3
),
cute.append_ones(
bufferB_sliced_stage[None, None, k_block], up_to_rank=3
),
accumulators_sliced,
)
mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True)
mainloop_pipe.consumer_release_and_advance()
acc_pipe.producer_commit_and_advance()
if is_epi_thr:
_, idx = acc_pipe.consumer_wait_and_get_stage()
accumulators_sliced = bufferAcc[(None, None), 0, 0, idx]
acc_epi_div_tiled = cute.flat_divide(accumulators_sliced, epi_tile)
tiled_copy_r2s = cute.make_tiled_copy_D(
cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), d_dtype),
tiled_copy_t2r,
)
c_cta_v_map = cute_ext.get_cta_v_map_c(mD, epi_tile)
subtile_cnt = cute.size(acc_epi_div_tiled.shape, mode=[3])
for mn in range(subtile_cnt):
# TMEM -> RMEM
cute_ext.partition_and_copy(
tiled_copy_t2r.get_slice(tid_x),
acc_epi_div_tiled[None, None, 0, mn],
bufferRAcc,
)
# RMEM -> RMEM
bufferRD.store(bufferRAcc.load().to(d_dtype))
tma_store_pipe.acquire_sync()
store_idx = tma_store_pipe.get_index()
# RMEM -> SMEM
cute_ext.partition_and_copy(
tiled_copy_r2s.get_slice(tid_x),
bufferRD,
bufferC[None, None, store_idx],
)
tma_store_pipe.commit_sync()
if warp_idx == tma_store_warp_id:
cute_ext.tma_store(
bufferC[None, None, store_idx],
gC_mnl_epi[None, None, 0, mn],
cta_v_map=c_cta_v_map,
)
tma_store_pipe.release_advance()
tma_store_pipe.tail()
acc_pipe.consumer_release_and_advance()
# Return a callable that launches the kernel with proper grid/block/cluster
@cute_ext.jit
def launch_kernel(mA: cute.Tensor, mB: cute.Tensor, mD: cute.Tensor):
kernel(mA, mB, mD).launch(
grid=GRID_SHAPE,
block=(32 * NUM_WARPS_PER_CTA, 1, 1),
cluster=CLUSTER_SHAPE,
smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")),
)
return launch_kernel
if __name__ == "__main__":
M = 256
N = 256
K = 64
use_tma_multicast = True
use_2cta_instrs = True
acc_dtype = cutlass.Float32
majors = (
cute.nvgpu.tcgen05.OperandMajorMode.K,
cute.nvgpu.tcgen05.OperandMajorMode.K,
cute.nvgpu.tcgen05.OperandMajorMode.K,
)
dtypes = (torch.float16, torch.float16, torch.float16)
A_torch, B_torch, D_torch, A_cute, B_cute, D_cute = get_gemm_tensors(
M, N, K, majors, dtypes
)
kernel_launcher = sm100_4x4x1_kernel_builder(
use_tma_multicast, use_2cta_instrs, acc_dtype, M, N
)
compiled_kernel = cute_ext.compile(kernel_launcher, A_cute, B_cute, D_cute)
compiled_kernel(A_cute, B_cute, D_cute)
# Reference check (may fail on simulator/unsupported GPU)
try:
ref = torch.mm(A_torch.float(), B_torch.float().T)
torch.testing.assert_close(D_torch.float(), ref, atol=1e-2, rtol=1e-2)
print("PASS")
except RuntimeError as e:
if "no kernel image is available" in str(e):
print("SKIP: Reference check skipped - GPU not supported by PyTorch")
else:
raise

1742
examples/python/CuTeDSL/experimental/blackwell/dense_gemm_cute_pipeline.py Executable file
View File

File diff suppressed because it is too large Load Diff

825
examples/python/CuTeDSL/experimental/blackwell/dense_gemm_ptr_array.py Executable file
View File

@@ -0,0 +1,825 @@
# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: BSD-3-Clause
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are met:
# 1. Redistributions of source code must retain the above copyright notice, this
# list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright notice,
# this list of conditions and the following disclaimer in the documentation
# and/or other materials provided with the distribution.
# 3. Neither the name of the copyright holder nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
# FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
# DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
# CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
# OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
import argparse
import torch
from typing import Type, Tuple, List
import cutlass
from cutlass.cute import experimental as cute_ext
from cutlass.base_dsl.typing import Numeric
from cutlass import cute as cute
from cutlass import utils
from cutlass import torch as cutlass_torch
import cutlass.utils.blackwell_helpers as sm100_utils
import cutlass.cute.testing as testing
class DenseGemmPtrArrayKernel:
def __init__(
self,
mn_tiler: tuple[int, int],
mma_dtype: tuple[Type[Numeric], Type[Numeric], Type[Numeric]],
tmem_output_dtype: Type[Numeric],
batch_count: int, # Number of batches, each batch will have its own pointer for the matrix
A_shape: tuple, # Shape of the matrix A
A_stride: tuple, # Stride of the matrix A
B_shape: tuple, # Shape of the matrix B
B_stride: tuple, # Stride of the matrix B
D_shape: tuple, # Shape of the matrix D
D_stride: tuple, # Stride of the matrix D
epilogue_op=lambda x: x,
):
self.mn_tiler = mn_tiler
self.ab_dtype, self.acc_dtype, self.d_dtype = mma_dtype
self.tmem_output_dtype = tmem_output_dtype
self.use_2cta_instrs = False
self.TMA_STORE_STAGE = 4
self.epilogue_op = epilogue_op
self.batch_count = batch_count
self.A_shape = A_shape
self.A_stride = A_stride
self.B_shape = B_shape
self.B_stride = B_stride
self.D_shape = D_shape
self.D_stride = D_stride
"""
Helper function to convert an int64 to a cute.ptr of a certain type.
The cute.ptr is always located in Gmem.
This is used to load the pointers for A/B/D from the Ptr array.
"""
@cute.experimental.jit
def _get_pointer(self, address_as_int, cute_type):
cute_ptr = cute.make_ptr(
cute_type,
address_as_int,
mem_space=cute.AddressSpace.gmem,
assumed_align=16,
)
return cute_ptr
@cute.experimental.jit
def __call__(
self, mA_tensor: cute.Tensor, mB_tensor: cute.Tensor, mD_tensor: cute.Tensor
):
# Get the pointer to the first batch of D
d_ptr = self._get_pointer(mD_tensor[0], self.d_dtype)
d_ptr_base_tensor = cute.make_tensor(
d_ptr, layout=cute.make_layout(self.D_shape, stride=self.D_stride)
)
tile_mn = cute.core._pack_shape((*self.mn_tiler, 1))
div = cute.tiled_divide(d_ptr_base_tensor, tile_mn)
grid = (div.shape[1], div.shape[2], div.shape[3])
self.kernel(mA_tensor, mB_tensor, mD_tensor).launch(
grid=grid,
block=(192, 1, 1),
cluster=(1, 1, 1),
smem=cute.Int64(utils.get_smem_capacity_in_bytes("sm_100")),
)
@cute.experimental.kernel
def kernel(
self,
mA_tensor: cute.Tensor,
mB_tensor: cute.Tensor,
mD_tensor: cute.Tensor,
):
# Get pointers for the first batch to perform shape and stage calculations
A_0_ptr = self._get_pointer(mA_tensor[0], self.ab_dtype)
B_0_ptr = self._get_pointer(mB_tensor[0], self.ab_dtype)
D_0_ptr = self._get_pointer(mD_tensor[0], self.d_dtype)
mA = cute.make_tensor(
A_0_ptr, layout=cute.make_layout(self.A_shape, stride=self.A_stride)
)
mB = cute.make_tensor(
B_0_ptr, layout=cute.make_layout(self.B_shape, stride=self.B_stride)
)
mD = cute.make_tensor(
D_0_ptr, layout=cute.make_layout(self.D_shape, stride=self.D_stride)
)
tiled_mma = sm100_utils.make_trivial_tiled_mma(
self.ab_dtype,
utils.LayoutEnum.from_tensor(mA).mma_major_mode(),
utils.LayoutEnum.from_tensor(mB).mma_major_mode(),
self.acc_dtype,
cute.nvgpu.tcgen05.CtaGroup.ONE,
self.mn_tiler,
)
mma_inst_shape_k = cute.size(tiled_mma.shape_mnk, mode=[2])
mma_inst_tile_k = 4
mnk_tiler = (
self.mn_tiler[0],
self.mn_tiler[1],
mma_inst_shape_k * mma_inst_tile_k,
)
d_layout = utils.LayoutEnum.from_tensor(mD)
d_dtype = mD.element_type
tiler_mk = (mnk_tiler[0], mnk_tiler[2])
tiler_nk = (mnk_tiler[1], mnk_tiler[2])
tiler_mn = (mnk_tiler[0], mnk_tiler[1])
gA = cute.zipped_divide(mA, tiler_mk)
gB = cute.zipped_divide(mB, tiler_nk)
gD = cute.zipped_divide(mD, tiler_mn)
mainloop_stage = 2
acc_stage = 2
cta_m, cta_n, cta_l = cute.arch.block_idx()
tid_x, _, _ = cute.arch.thread_idx()
gA_tile = gA[(None, None), (cta_m, None, cta_l)]
gB_tile = gB[(None, None), (cta_n, None, cta_l)]
gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)]
# Compute A/B/C shared memory layout
a_smem_layout_staged = sm100_utils.make_smem_layout_a(
tiled_mma,
mnk_tiler,
self.ab_dtype,
mainloop_stage,
)
b_smem_layout_staged = sm100_utils.make_smem_layout_b(
tiled_mma,
mnk_tiler,
self.ab_dtype,
mainloop_stage,
)
cta_tile_shape_mnk = cute.shape_div(
mnk_tiler, (cute.size(tiled_mma.thr_id.shape), 1, 1)
)
epi_tile = sm100_utils.compute_epilogue_tile_shape(
cta_tile_shape_mnk,
self.use_2cta_instrs,
d_layout,
d_dtype,
)
sc_smem_layout_staged = sm100_utils.make_smem_layout_epi(
d_dtype,
d_layout,
epi_tile,
self.TMA_STORE_STAGE,
)
# UMMA ACC TMEM Layout
acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2])
tmem_layout = tiled_mma.make_fragment_C(
cute.append(acc_shape, acc_stage)
).layout
# Allocate UMMA Buffers
bufferA = cute_ext.allocate(
self.ab_dtype,
cute.AddressSpace.smem,
a_smem_layout_staged,
alignment=1024,
)
bufferB = cute_ext.allocate(
self.ab_dtype,
cute.AddressSpace.smem,
b_smem_layout_staged,
alignment=1024,
)
bufferAcc = cute_ext.allocate(
self.acc_dtype,
cute.AddressSpace.tmem,
tmem_layout,
alignment=16,
)
# Allocate SMEM buffer for C
bufferC = cute_ext.allocate(
d_dtype,
cute.AddressSpace.smem,
sc_smem_layout_staged,
alignment=1024,
)
# Create the TMEM load atom
copy_atom_t2r = sm100_utils.get_tmem_load_op(
cta_tile_shape_mnk,
d_layout,
self.tmem_output_dtype,
self.acc_dtype,
epi_tile,
self.use_2cta_instrs,
)
# Take only one stage of the TMEM buffer
accumulators = cute.zipped_divide(bufferAcc, ((epi_tile), 1))
acc_epi_div = accumulators[((None, None), 0), 0]
# Create the TMEM copy atom based on the size of transfer within one iteration of epilogue
tiled_copy_t2r = cute.nvgpu.tcgen05.make_tmem_copy(copy_atom_t2r, acc_epi_div)
# Calculate the per thread destination size per iteration for output of TMEM and input of SMEM
thr_copy_t2r = tiled_copy_t2r.get_slice(tid_x)
gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi)
acc_d_rmem_layout = cute.make_fragment_like(
tTR_gC[(None, None, None, 0, 0)].layout
)
# Allocate RMEM buffers
bufferRAcc = cute_ext.allocate(
self.acc_dtype,
cute.AddressSpace.rmem,
acc_d_rmem_layout,
alignment=32,
)
bufferRD = cute_ext.allocate(
d_dtype,
cute.AddressSpace.rmem,
acc_d_rmem_layout,
alignment=32,
)
# TMA -> UMMA
mainloop_pipe = cute_ext.TMAToUMMAPipeline.create(
num_stages=mainloop_stage,
mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS,
)
# UMMA -> TMEM
acc_pipe = cute_ext.UMMAtoAsyncPipeline.create(
num_stages=acc_stage,
mma_operation_type=cute_ext.OperationTypeEnum.SM100_MMA_1SM_SS,
consumer=cute_ext.OperationTypeEnum.SM100_COPY_T2R,
consumer_arv_count=128,
)
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
# warp assignment: [0]-tma_store, [0-3]-epi, [4]-mma, [5]-tma_load
tma_store_warp_id = 0
mma_warp_id = 4
tma_load_warp_id = 5
is_tma_thr = warp_idx == tma_load_warp_id
is_mma_thr = warp_idx == mma_warp_id
is_epi_thr = warp_idx < 4
# SMEM -> GMEM
tma_store_pipe = cute_ext.TMAStorePipeline(
stages=self.TMA_STORE_STAGE,
arv_count=128,
barrier_id=1,
tma_warp_id=tma_store_warp_id,
)
k_tile_size = cute.size(gA, mode=[1, 1])
# Outer loop over batches and perform GEMM for each batch as usual
# This is a dynamic for loop that lowers to an scf.for
# Note that the tensor loading is done in the if `thread` warp specialized
# sections. This is essential to ensure proper synchronization of tma loads
# and tma updates across batches.
for batch_idx in range(0, self.batch_count):
# Load pointers for the current batch
ptr_A = self._get_pointer(mA_tensor[batch_idx], self.ab_dtype)
ptr_B = self._get_pointer(mB_tensor[batch_idx], self.ab_dtype)
ptr_D = self._get_pointer(mD_tensor[batch_idx], self.d_dtype)
gALayout = cute.zipped_divide(mA, tiler_mk)
k_tile_size = cute.size(gALayout, mode=[1, 1])
if is_tma_thr:
mA = cute.make_tensor(
ptr_A, layout=cute.make_layout(self.A_shape, stride=self.A_stride)
)
mB = cute.make_tensor(
ptr_B, layout=cute.make_layout(self.B_shape, stride=self.B_stride)
)
gA = cute.zipped_divide(mA, tiler_mk)
gB = cute.zipped_divide(mB, tiler_nk)
gA_tile = gA[(None, None), (cta_m, None, cta_l)]
gB_tile = gB[(None, None), (cta_n, None, cta_l)]
for k in cutlass.range(0, k_tile_size, 1, unroll=1):
gA_k = gA_tile[None, None, k]
gB_k = gB_tile[None, None, k]
# Scoped state management - pipeline object manages state internally
(
producer_stage_token,
idx,
) = mainloop_pipe.producer_acquire_and_get_stage()
mbar = cute_ext.get_mbarrier(producer_stage_token)
## producer_body begin ##
bufferA_sliced = bufferA[None, None, None, idx]
bufferB_sliced = bufferB[None, None, None, idx]
a_cta_v_map = cute_ext.get_cta_v_map_ab(
mA, mnk_tiler, tiled_mma, "A"
)
b_cta_v_map = cute_ext.get_cta_v_map_ab(
mB, mnk_tiler, tiled_mma, "B"
)
cute_ext.tma_load(
gA_k,
bufferA_sliced,
mbar,
cta_v_map=a_cta_v_map,
)
cute_ext.tma_load(
gB_k,
bufferB_sliced,
mbar,
cta_v_map=b_cta_v_map,
)
## producer_body end ##
mainloop_pipe.producer_commit_and_advance()
# MMA section remains same as a regular GEMM
if is_mma_thr:
producer_stage_token, idx = acc_pipe.producer_acquire_and_get_stage()
## acc_producer_body begin ##
accumulators_sliced = bufferAcc[None, None, None, idx]
mma_atom = cute.make_mma_atom(tiled_mma.op)
mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, False)
for k_tile in cutlass.range(0, k_tile_size, 1, unroll=1):
# Scoped state management - pipeline object manages consumer state internally
(
_,
mainloop_idx,
) = mainloop_pipe.consumer_wait_and_get_stage()
## tma_consumer_body begin ##
bufferA_sliced_stage = cute.core.slice_(
bufferA, (None, None, None, mainloop_idx)
)
bufferB_sliced_stage = cute.core.slice_(
bufferB, (None, None, None, mainloop_idx)
)
for k_block in cutlass.range(mma_inst_tile_k, unroll_full=True):
bufferA_sliced = bufferA_sliced_stage[None, None, k_block]
bufferB_sliced = bufferB_sliced_stage[None, None, k_block]
cute_ext.dot(
mma_atom,
cute.append_ones(bufferA_sliced, up_to_rank=3),
cute.append_ones(bufferB_sliced, up_to_rank=3),
accumulators_sliced,
)
mma_atom.set(cute.nvgpu.tcgen05.Field.ACCUMULATE, True)
## tma_consumer_body end ##
mainloop_pipe.consumer_release_and_advance()
## acc_producer_body end ##
acc_pipe.producer_commit_and_advance()
if is_epi_thr:
# Load the D tensor in the warp specialized section
mD = cute.make_tensor(
ptr_D, layout=cute.make_layout(self.D_shape, stride=self.D_stride)
)
gD = cute.zipped_divide(mD, tiler_mn)
gD_tile = gD[(None, None), (cta_m, cta_n, cta_l)]
gC_mnl_epi = cute.flat_divide(gD_tile, epi_tile)
_, idx = acc_pipe.consumer_wait_and_get_stage()
## acc_consume_body begin ##
accumulators_sliced = bufferAcc[(None, None), 0, 0, idx]
acc_epi_div_tiled = cute.flat_divide(accumulators_sliced, epi_tile)
subtile_cnt = cute.size(acc_epi_div_tiled.shape, mode=[3])
for mn in range(subtile_cnt):
# TMEM -> RMEM
cute_ext.partition_and_copy(
tiled_copy_t2r.get_slice(tid_x),
acc_epi_div_tiled[None, None, 0, mn],
bufferRAcc,
)
# RMEM -> RMEM
bufferRD.store(self.epilogue_op(bufferRAcc.load().to(self.d_dtype)))
# Acquire pipeline stage and synchronize before RMEM->SMEM copy
tma_store_pipe.acquire_sync()
idx = tma_store_pipe.get_index()
# RMEM -> SMEM
tiled_copy_r2s = cute.make_tiled_copy_D(
cute.make_copy_atom(cute.nvgpu.CopyUniversalOp(), self.d_dtype),
tiled_copy_t2r,
)
cute_ext.partition_and_copy(
tiled_copy_r2s.get_slice(tid_x),
bufferRD,
bufferC[None, None, idx],
)
# Fence SMEM writes and synchronize before TMA store
tma_store_pipe.commit_sync()
# SMEM -> GMEM (only designated TMA store warp performs TMA store)
if warp_idx == tma_store_warp_id:
c_cta_v_map = cute_ext.get_cta_v_map_c(mD, epi_tile)
cute_ext.tma_store(
bufferC[None, None, idx],
gC_mnl_epi[None, None, 0, mn],
cta_v_map=c_cta_v_map,
)
# Release pipeline stage and advance
tma_store_pipe.release_advance()
tma_store_pipe.tail()
acc_pipe.consumer_release_and_advance()
def create_tensors(l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype):
torch.manual_seed(1111)
a_torch_cpu = cutlass_torch.matrix(l, m, k, a_major == "m", ab_dtype)
b_torch_cpu = cutlass_torch.matrix(l, n, k, b_major == "n", ab_dtype)
d_torch_cpu = cutlass_torch.matrix(l, m, n, d_major == "m", d_dtype)
a_tensor, a_torch_gpu = cutlass_torch.cute_tensor_like(
a_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
)
b_tensor, b_torch_gpu = cutlass_torch.cute_tensor_like(
b_torch_cpu, ab_dtype, is_dynamic_layout=True, assumed_align=16
)
d_tensor, d_torch_gpu = cutlass_torch.cute_tensor_like(
d_torch_cpu, d_dtype, is_dynamic_layout=True, assumed_align=16
)
return (
a_tensor,
b_tensor,
d_tensor,
a_torch_cpu,
b_torch_cpu,
d_torch_cpu,
a_torch_gpu,
b_torch_gpu,
d_torch_gpu,
)
# Helper creates a cute.Tensor from a List of device pointers
def make_tensor_of_ptrs(torch_tensor_array: List):
tensor_of_ptrs_torch = torch.tensor(
[t.data_ptr() for t in torch_tensor_array],
dtype=torch.int64,
device="cuda",
requires_grad=False,
)
tensor_of_ptrs_cute, backing_torch_tensor = cutlass_torch.cute_tensor_like(
tensor_of_ptrs_torch,
cutlass.Int64,
is_dynamic_layout=False,
assumed_align=16,
)
return tensor_of_ptrs_cute, backing_torch_tensor
def create_tensors_for_ptr_array(
l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype
):
# Store torch gpu pointers
As_torch_gpu = []
Bs_torch_gpu = []
Ds_torch_gpu = []
# Store cute tensors
A_cutes = []
B_cutes = []
D_cutes = []
for batch_idx in range(l):
torch.manual_seed(111 + batch_idx)
(
A_tensor,
B_tensor,
D_tensor,
A_torch_cpu,
B_torch_cpu,
D_torch_cpu,
A_torch_gpu,
B_torch_gpu,
D_torch_gpu,
) = create_tensors(
1, # outer loop creates a new tensor for each batch
m,
n,
k,
a_major,
b_major,
d_major,
ab_dtype,
d_dtype,
)
A_cutes.append(A_tensor)
B_cutes.append(B_tensor)
D_cutes.append(D_tensor)
As_torch_gpu.append(A_torch_gpu)
Bs_torch_gpu.append(B_torch_gpu)
Ds_torch_gpu.append(D_torch_gpu)
# Create cute tensors of pointers
a_tensor, a_backing_torch_tensor = make_tensor_of_ptrs(As_torch_gpu)
b_tensor, b_backing_torch_tensor = make_tensor_of_ptrs(Bs_torch_gpu)
d_tensor, d_backing_torch_tensor = make_tensor_of_ptrs(Ds_torch_gpu)
return (
a_tensor,
b_tensor,
d_tensor,
a_backing_torch_tensor,
b_backing_torch_tensor,
d_backing_torch_tensor,
A_cutes,
B_cutes,
D_cutes,
As_torch_gpu,
Bs_torch_gpu,
Ds_torch_gpu,
)
def compare(a_torch_cpu, b_torch_cpu, d_torch_gpu, d_dtype, tolerance):
ref = torch.einsum("mkl,nkl->mnl", a_torch_cpu, b_torch_cpu)
_, ref_torch_gpu = cutlass_torch.cute_tensor_like(
ref, d_dtype, is_dynamic_layout=True, assumed_align=16
)
ref_result = ref_torch_gpu.cpu()
torch.testing.assert_close(
d_torch_gpu.cpu(), ref_result, atol=tolerance, rtol=1e-05
)
def run(
mnkl: Tuple[int, int, int, int],
mma_tiler_mn: Tuple[int, int],
cluster_shape_mn: Tuple[int, int],
ab_dtype: Type[Numeric],
c_dtype: Type[Numeric],
acc_dtype: Type[Numeric],
a_major: str,
b_major: str,
c_major: str,
warmup_iterations: int = 0,
iterations: int = 1,
use_cold_l2: bool = False,
tolerance: float = 1e-02,
skip_ref_check: bool = False,
**kwargs,
):
"""Execute a Pointer array batched dense GEMM operation on Blackwell architecture with performance benchmarking.
The main difference between this and a regular bathced GEMM is that the inputs to the kernel are arrays of pointers.
Every batch of each operand (A/B/D) has its own pointer. Thus, the size of the array of pointers is the batch size.
These pointers NEED NOT be stored contiguously in memory.
This example also demonstrates how cute_ext.tma_load/cute_ext.tma_store performs automatic device side TMA updates.
Note that the dimensions of the operand for each batch are the same across all batches. That is, all batches of A have the same shape and stride, same for B and D.
This function prepares input tensors, configures and launches the GEMM kernel,
optionally performs reference validation, and benchmarks the execution performance.
:param mnkl: Problem size (M, N, K, L)
:type mnkl: Tuple[int, int, int, int]
:param mma_tiler_mn: MMA tiling size.
:type mma_tiler_mn: Tuple[int, int]
:param cluster_shape_mn: Cluster shape.
:type cluster_shape_mn: Tuple[int, int]
:param ab_dtype: Data type for input tensors A and B
:type ab_dtype: Type[Numeric]
:param d_dtype: Data type for output tensor D
:type d_dtype: Type[Numeric]
"""
print("Running Blackwell Dense GEMM test with:")
print(f"mnkl: {mnkl}")
print(f"AB dtype: {ab_dtype}, D dtype: {c_dtype}, Acc dtype: {acc_dtype}")
print(f"Matrix majors - A: {a_major}, B: {b_major}, D: {c_major}")
print(f"Mma Tiler (M, N): {mma_tiler_mn}, Cluster Shape (M, N): {cluster_shape_mn}")
print(f"Tolerance: {tolerance}")
print(f"Warmup iterations: {warmup_iterations}")
print(f"Iterations: {iterations}")
print(f"Skip reference checking: {skip_ref_check}")
print(f"Use cold L2: {'True' if use_cold_l2 else 'False'}")
m, n, k, l = mnkl
ab_dtype = ab_dtype
d_major = c_major
d_dtype = c_dtype
# a_tensor, b_tensor, d_tensor are cute Tensors where each element is an Int64 pointer to global memory
# A_cutes, B_cutes, D_cutes are lists of cute Tensors for each batch of A/B/D
(
a_tensor,
b_tensor,
d_tensor,
a_backing_torch_tensor,
b_backing_torch_tensor,
d_backing_torch_tensor,
A_cutes,
B_cutes,
D_cutes,
As_torch_gpu,
Bs_torch_gpu,
Ds_torch_gpu,
) = create_tensors_for_ptr_array(
l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype
)
ptr_array_dense_gemm = DenseGemmPtrArrayKernel(
mn_tiler=mma_tiler_mn,
mma_dtype=(ab_dtype, acc_dtype, d_dtype),
tmem_output_dtype=d_dtype,
batch_count=l,
A_shape=A_cutes[0].shape,
A_stride=A_cutes[0].stride,
B_shape=B_cutes[0].shape,
B_stride=B_cutes[0].stride,
D_shape=D_cutes[0].shape,
D_stride=D_cutes[0].stride,
)
compiled_dense_gemm = cute_ext.compile(
ptr_array_dense_gemm, a_tensor, b_tensor, d_tensor
)
compiled_dense_gemm(a_tensor, b_tensor, d_tensor)
if not skip_ref_check:
for batch_idx in range(l):
compare(
As_torch_gpu[batch_idx].cpu(),
Bs_torch_gpu[batch_idx].cpu(),
Ds_torch_gpu[batch_idx],
d_dtype,
tolerance,
)
print("check reference: PASS")
def generate_tensors():
(
a_tensor,
b_tensor,
d_tensor,
a_backing_torch_tensor,
b_backing_torch_tensor,
d_backing_torch_tensor,
A_cutes,
B_cutes,
D_cutes,
As_torch_gpu,
Bs_torch_gpu,
Ds_torch_gpu,
) = create_tensors_for_ptr_array(
l, m, n, k, a_major, b_major, d_major, ab_dtype, d_dtype
)
args = testing.JitArguments(a_tensor, b_tensor, d_tensor)
args.add_to_scope([A_cutes, B_cutes, D_cutes])
return args
workspace_count = 1
if use_cold_l2:
one_workspace_bytes = (
sum(
As_torch_gpu[batch_idx].numel() * As_torch_gpu[batch_idx].element_size()
for batch_idx in range(l)
)
+ sum(
Bs_torch_gpu[batch_idx].numel() * Bs_torch_gpu[batch_idx].element_size()
for batch_idx in range(l)
)
+ sum(
Ds_torch_gpu[batch_idx].numel() * Ds_torch_gpu[batch_idx].element_size()
for batch_idx in range(l)
)
)
workspace_count = testing.get_workspace_count(
one_workspace_bytes, warmup_iterations, iterations
)
exec_time = testing.benchmark(
compiled_dense_gemm,
workspace_generator=generate_tensors,
workspace_count=workspace_count,
warmup_iterations=warmup_iterations,
iterations=iterations,
)
return exec_time
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> Tuple[int, ...]:
try:
return tuple(int(x.strip()) for x in s.split(","))
except ValueError:
raise argparse.ArgumentTypeError(
"Invalid format. Expected comma-separated integers."
)
parser = argparse.ArgumentParser(description="Example of Dense GEMM on Blackwell.")
parser.add_argument(
"--mnkl",
type=parse_comma_separated_ints,
default=(256, 256, 512, 1),
help="mnkl dimensions (comma-separated)",
)
parser.add_argument(
"--mma_tiler_mn",
type=parse_comma_separated_ints,
default=(128, 128),
help="Mma tile shape (comma-separated)",
)
parser.add_argument(
"--cluster_shape_mn",
type=parse_comma_separated_ints,
default=(1, 1),
help="Cluster shape (comma-separated)",
)
parser.add_argument("--ab_dtype", type=cutlass.dtype, default=cutlass.Float32)
parser.add_argument("--d_dtype", type=cutlass.dtype, default=cutlass.Float32)
parser.add_argument("--acc_dtype", type=cutlass.dtype, default=cutlass.Float32)
parser.add_argument("--a_major", choices=["k", "m"], type=str, default="k")
parser.add_argument("--b_major", choices=["k", "n"], type=str, default="k")
parser.add_argument("--d_major", choices=["n", "m"], type=str, default="n")
parser.add_argument(
"--warmup_iterations", type=int, default=0, help="Warmup iterations"
)
parser.add_argument(
"--iterations", type=int, default=1, help="Number of iterations"
)
parser.add_argument("--use_cold_l2", action="store_true", help="Use cold L2")
parser.add_argument(
"--tolerance", type=float, default=1e-02, help="Tolerance for validation"
)
parser.add_argument(
"--skip_ref_check", action="store_true", help="Skip reference checking"
)
args = parser.parse_args()
if len(args.mnkl) != 4:
parser.error("--mnkl must contain exactly 4 values")
if len(args.mma_tiler_mn) != 2:
parser.error("--mma_tiler_mn must contain exactly 2 values")
exec_time = run(
args.mnkl,
args.mma_tiler_mn,
args.cluster_shape_mn,
args.ab_dtype,
args.d_dtype,
args.acc_dtype,
args.a_major,
args.b_major,
args.d_major,
args.warmup_iterations,
args.iterations,
args.use_cold_l2,
args.tolerance,
args.skip_ref_check,
)
print(f"Execution time: {exec_time} microseconds per iteration")