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

v4.3 update. (#2709)

Browse Source 
* v4.3 update.

* Update the cute_dsl_api changelog's doc link

* Update version to 4.3.0

* Update the example link

* Update doc to encourage user to install DSL from requirements.txt

---------

Co-authored-by: Larry Wu <larwu@nvidia.com>
This commit is contained in:
Junkai-Wu
2025-10-22 02:26:30 +08:00
committed by GitHub
parent e6e2cc29f5
commit b1d6e2c9b3
244 changed files with 59272 additions and 10455 deletions
Download Patch File Download Diff File Expand all files Collapse all files

12
examples/python/CuTeDSL/ampere/call_bypass_dlpack.py
View File

@@ -89,7 +89,7 @@ def tensor_op_gemm_wrapper(
k: cutlass.Int32,
l: cutlass.Int32,
):
print(f"\n[DSL INFO] Input Parameters:")
print("\n[DSL INFO] Input Parameters:")
print(f"[DSL INFO] mnkl: {(m, n, k, l)}")
# Assume alignment of shape to call tensorop_gemm example
@@ -111,7 +111,7 @@ def tensor_op_gemm_wrapper(
tensor_op_gemm = TensorOpGemm(
a_ptr.value_type, c_ptr.value_type, cutlass.Float32, (2, 2, 1)
)
print(f"\n[DSL INFO] Created TensorOpGemm instance")
print("\n[DSL INFO] Created TensorOpGemm instance")
print(f"[DSL INFO] Input dtype: {a_ptr.value_type}")
print(f"[DSL INFO] Output dtype: {c_ptr.value_type}")
print(f"[DSL INFO] Accumulation dtype: {cutlass.Float32}")
@@ -119,11 +119,11 @@ def tensor_op_gemm_wrapper(
# No need to compile inside jit function
tensor_op_gemm(mA, mB, mC)
print(f"\n[DSL INFO] Executed TensorOpGemm")
print("\n[DSL INFO] Executed TensorOpGemm")
def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
print(f"\nRunning TensorOpGemm test with:")
print("\nRunning TensorOpGemm test with:")
print(f"Tensor dimensions: {mnkl}")
# (M,K,L)
@@ -139,7 +139,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
mnkl[3], mnkl[0], mnkl[1], dtype=torch.float16, device="cuda"
).permute(1, 2, 0)
print(f"Input tensor shapes:")
print("Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
@@ -158,7 +158,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
ref = torch.einsum("mkl,nkl->mnl", a, b)
torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
print(f"\n[DSL INFO] Results verified successfully!")
print("\n[DSL INFO] Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3, :3]}")

14
examples/python/CuTeDSL/ampere/call_from_jit.py
View File

@@ -169,7 +169,7 @@ def tensor_op_gemm_wrapper(
acc_dtype: Type[cutlass.Numeric],
atom_layout_mnk: cutlass.Constexpr[tuple[int, int, int]],
):
print(f"\n[DSL INFO] Input Parameters:")
print("\n[DSL INFO] Input Parameters:")
print(f"[DSL INFO] mnkl: {mnkl}")
print(f"[DSL INFO] buffer_a: {buffer_a}")
print(f"[DSL INFO] buffer_b: {buffer_b}")
@@ -181,7 +181,7 @@ def tensor_op_gemm_wrapper(
mB = buffer_b.to_tensor(cute.select(mnkl, mode=[3, 1, 2]))
mC = buffer_c.to_tensor(cute.select(mnkl, mode=[3, 0, 1]))
print(f"\n[DSL INFO] Created Tensors:")
print("\n[DSL INFO] Created Tensors:")
print(f"[DSL INFO] mA = {mA}")
print(f"[DSL INFO] mB = {mB}")
print(f"[DSL INFO] mC = {mC}")
@@ -192,7 +192,7 @@ def tensor_op_gemm_wrapper(
acc_dtype,
atom_layout_mnk,
)
print(f"\n[DSL INFO] Created TensorOpGemm instance")
print("\n[DSL INFO] Created TensorOpGemm instance")
print(f"[DSL INFO] Input dtype: {buffer_a.ptr.value_type}")
print(f"[DSL INFO] Output dtype: {buffer_c.ptr.value_type}")
print(f"[DSL INFO] Accumulation dtype: {acc_dtype}")
@@ -200,11 +200,11 @@ def tensor_op_gemm_wrapper(
# No need to compile inside jit function
tensor_op_gemm(mA, mB, mC)
print(f"\n[DSL INFO] Executed TensorOpGemm")
print("\n[DSL INFO] Executed TensorOpGemm")
def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
print(f"\nRunning TensorOpGemm test with:")
print("\nRunning TensorOpGemm test with:")
print(f"Tensor dimensions: {mnkl}")
ab_dtype = cutlass.Float16
@@ -220,7 +220,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
mnkl[3], mnkl[0], mnkl[1], dtype=torch_dtype(c_dtype), device="cuda"
)
print(f"Input tensor shapes:")
print("Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
@@ -251,7 +251,7 @@ def run_tensor_op_gemm_wrapper(mnkl: Tuple[int, int, int, int]):
ref = torch.einsum("lmk,lnk->lmn", a, b)
torch.testing.assert_close(c, ref, atol=1e-05, rtol=1e-05)
print(f"\n[DSL INFO] Results verified successfully!")
print("\n[DSL INFO] Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3, :3]}")

25
examples/python/CuTeDSL/ampere/elementwise_add.py
View File

@@ -28,11 +28,10 @@
import argparse
import torch
import time
from typing import Type
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
@@ -154,7 +153,7 @@ def elementwise_add_kernel(
blkCrd = cC[blk_coord] # (TileM, TileN)
# Note: these prints only run at compile/jit time
print(f"[DSL INFO] Sliced Tensors per thread block:")
print("[DSL INFO] Sliced Tensors per thread block:")
print(f"[DSL INFO] blkA = {blkA.type}")
print(f"[DSL INFO] blkB = {blkB.type}")
print(f"[DSL INFO] blkC = {blkC.type}")
@@ -182,9 +181,9 @@ def elementwise_add_kernel(
frgC = cute.make_fragment_like(thrC)
thrCrd = thr_copy_C.partition_S(blkCrd)
frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
print(f"[DSL INFO] Sliced Tensors per thread:")
print("[DSL INFO] Sliced Tensors per thread:")
print(f"[DSL INFO] thrA = {thrA.type}")
print(f"[DSL INFO] thrB = {thrB.type}")
print(f"[DSL INFO] thrC = {thrC.type}")
@@ -233,18 +232,18 @@ def elementwise_add(mA, mB, mC, copy_bits: cutlass.Constexpr = 128):
val_layout = cute.make_ordered_layout((4, vector_size), order=(1, 0))
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
print(f"[DSL INFO] Input Tensors:")
print("[DSL INFO] Input Tensors:")
print(f"[DSL INFO] mA = {mA.type}")
print(f"[DSL INFO] mB = {mB.type}")
print(f"[DSL INFO] Tiling Parameters:")
print("[DSL INFO] Tiling Parameters:")
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
print(f"[DSL INFO] tv_layout = {tv_layout}")
gA = cute.zipped_divide(mA, tiler_mn) # ((TileM,TileN),(RestM,RestN))
gB = cute.zipped_divide(mB, tiler_mn) # ((TileM,TileN),(RestM,RestN))
gC = cute.zipped_divide(mC, tiler_mn) # ((TileM,TileN),(RestM,RestN))
print(f"[DSL INFO] Tiled Tensors:")
print("[DSL INFO] Tiled Tensors:")
print(f"[DSL INFO] gA = {gA.type}")
print(f"[DSL INFO] gB = {gB.type}")
print(f"[DSL INFO] gC = {gC.type}")
@@ -271,7 +270,7 @@ def run_elementwise_add(
warmup_iterations=2,
iterations=200,
):
print(f"\nRunning Elementwise Add test with:")
print("\nRunning Elementwise Add test with:")
print(f"Tensor dimensions: [{M}, {N}]")
print(f"Input and Output Data type: {dtype}")
@@ -285,7 +284,7 @@ def run_elementwise_add(
c = torch.zeros_like(a)
print(f"Input tensor shapes:")
print("Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
@@ -307,7 +306,9 @@ def run_elementwise_add(
print("Compiling kernel with cute.compile ...")
start_time = time.time()
compiled_func = cute.compile(elementwise_add, a_tensor, b_tensor, c_tensor)
compiled_func = cute.compile(
elementwise_add, a_tensor, b_tensor, c_tensor, options="--generate-line-info"
)
compilation_time = time.time() - start_time
print(f"Compilation time: {compilation_time:.4f} seconds")
@@ -386,7 +387,7 @@ if __name__ == "__main__":
args = parser.parse_args()
if not torch.cuda.is_available():
raise RuntimeError(f"Ampere GPU is required to run this example!")
raise RuntimeError("Ampere GPU is required to run this example!")
run_elementwise_add(
args.M,

277
examples/python/CuTeDSL/ampere/elementwise_apply.py
View File

@@ -30,17 +30,18 @@
import argparse
import operator
import time
from typing import Type, List
from functools import partial
from typing import List, Type
import cuda.bindings.driver as cuda
import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.torch as cutlass_torch
import torch
from cutlass.cute.runtime import from_dlpack
import cutlass
"""
An Elementwise Apply Example using CuTe DSL.
@@ -78,103 +79,83 @@ while maintaining high performance through efficient memory access patterns.
@cute.kernel
def elementwise_apply_kernel(
op: cutlass.Constexpr,
inputs: List[cute.Tensor],
gC: cute.Tensor,
mInputs: List[cute.Tensor],
mC: cute.Tensor,
cC: cute.Tensor, # coordinate tensor
shape: cute.Shape,
tv_layout: cute.Layout, # (tid, vid) -> logic coord
):
tidx, _, _ = cute.arch.thread_idx()
bidx, _, _ = cute.arch.block_idx()
bidx, bidy, _ = cute.arch.block_idx()
###############################################################################
# Slice to local tile of thread block
###############################################################################
blk_crd = ((None, None), (bidx, bidy))
# slice for CTAs
cta_coord = ((None, None), bidx)
# logical coord -> address
# Leverage the meta-programming capability of the DSL to slice the tensors for each input
# All for loops below on input tensors would be fully unrolled automatically at compile time
ctaInputs = [t[cta_coord] for t in inputs] # (TileM, TileN)
ctaC = gC[cta_coord] # (TileM, TileN)
ctaCrd = cC[cta_coord] # (TileM, TileN)
# logical coord -> memory address
gInputs = [t[blk_crd] for t in mInputs] # (TileM, TileN)
gC = mC[blk_crd] # (TileM, TileN)
gCrd = cC[blk_crd] # (TileM, TileN)
print(f"[DSL INFO] Sliced Tensors per thread block:")
for i in cutlass.range_constexpr(len(ctaInputs)):
print(f"[DSL INFO] ctaInputs{i} = {ctaInputs[i].type}")
print(f"[DSL INFO] ctaC = {ctaC.type}")
print(f"[DSL INFO] ctaCrd = {ctaCrd.type}")
print("[DSL INFO] Sliced Tensors per thread block:")
for i in cutlass.range_constexpr(len(gInputs)):
print(f"[DSL INFO] ctaInputs{i} = {gInputs[i].type}")
print(f"[DSL INFO] gC = {gC.type}")
print(f"[DSL INFO] gCrd = {gCrd.type}")
# compose with CTA TV layout
# (tid, vid) -> address
tidfrgInputs = [cute.composition(t, tv_layout) for t in ctaInputs]
tidfrgC = cute.composition(ctaC, tv_layout)
tidfrgCrd = cute.composition(ctaCrd, tv_layout)
# print(f"{tv_layout = }")
# print(f"{tidfrgAB[0] = }")
###############################################################################
# Compose with thread block TV layout to map thread & value indices to memory address
###############################################################################
# (tid, vid) -> memory address
tidfrgInputs = [cute.composition(t, tv_layout) for t in gInputs]
tidfrgC = cute.composition(gC, tv_layout)
tidfrgCrd = cute.composition(gCrd, tv_layout)
thr_coord = (tidx, (None, None))
# repeat None like vid to remove hierarchy of layout
thr_crd = (tidx, cute.repeat_like(None, tidfrgInputs[0][1]))
# slice for threads
###############################################################################
# Slice to local tile of thread
###############################################################################
# vid -> address
thrInputs = [t[thr_coord] for t in tidfrgInputs] # (V)
thrC = tidfrgC[thr_coord] # (V)
thrCrd = tidfrgCrd[thr_coord]
thrInputs = [t[thr_crd] for t in tidfrgInputs] # (V)
thrC = tidfrgC[thr_crd] # (V)
thrCrd = tidfrgCrd[thr_crd]
print(f"[DSL INFO] Sliced Tensors per thread:")
print("[DSL INFO] Sliced Tensors per thread:")
for i in cutlass.range_constexpr(len(thrInputs)):
print(f"[DSL INFO] thrInputs{i} = {thrInputs[i].type}")
print(f"[DSL INFO] thrC = {thrC.type}")
print(f"[DSL INFO] thrCrd = {thrCrd.type}")
# allocate fragments for gmem->rmem
frgInputs = [cute.make_fragment_like(t, t.element_type) for t in thrInputs]
frgC = cute.make_fragment_like(thrC, gC.element_type)
frgPred = cute.make_fragment(thrCrd.shape, cutlass.Boolean)
###############################################################################
# Compute predicate for out of boundary checks
###############################################################################
frgPred = cute.make_rmem_tensor(thrCrd.shape, cutlass.Boolean)
print(f"[DSL INFO] frgPred = {frgPred.type}")
for i in cutlass.range(cute.size(frgPred), unroll=1):
for i in cutlass.range_constexpr(cute.size(frgPred)):
frgPred[i] = cute.elem_less(thrCrd[i], shape)
# if tidx == 0 and bidx == 0:
# cute.print_tensor(frgPred)
##########################################################
# Move data to reg address space
# Load data and compute result
##########################################################
# declare the atoms which will be used later for memory copy
# Compile time validation: expect same element type for all input tensors so as to reuse the copy atom for load
assert all(t.element_type == inputs[0].element_type for t in inputs)
copy_atom_load = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
inputs[0].element_type,
num_bits_per_copy=inputs[0].element_type.width,
)
copy_atom_store = cute.make_copy_atom(
cute.nvgpu.CopyUniversalOp(),
gC.element_type,
num_bits_per_copy=gC.element_type.width,
)
for thrInput, frgInput in zip(thrInputs, frgInputs):
cute.copy(copy_atom_load, thrInput, frgInput, pred=frgPred)
# Load data before use. The compiler will optimize the copy and load
# operations to convert some memory ld/st into register uses.
result = op(*[frgInput.load() for frgInput in frgInputs])
# Save the results back to registers. Here we reuse b's registers.
frgC.store(result)
# Copy the results back to c
cute.copy(copy_atom_store, frgC, thrC, pred=frgPred)
result = op(*[thrInput.load() for thrInput in thrInputs])
thrC.store(result)
@cute.jit
def elementwise_apply(
op: cutlass.Constexpr,
a: cute.Tensor,
b: cute.Tensor,
result: cute.Tensor,
stream: cuda.CUstream,
op: cutlass.Constexpr, inputs, result: cute.Tensor, stream: cuda.CUstream
):
"""CUDA kernel applying binary operator on each element of two n-D input tensors in
CuTe Python and store to result tensor.
@@ -232,51 +213,71 @@ def elementwise_apply(
# Opt-3: SOL with 2D thread tile
# * mA layout: (4096, 4096):(4096, 1)
# * TV layout map to (16, 128) logical tile
# * TV layout map to (64, 256) logical tile
# * tidx maps to mode-1 and input layout is contiguous on mode-1 for coalesced load-store
thr_layout = cute.make_layout((4, 32), stride=(32, 1))
val_layout = cute.make_layout((4, 4), stride=(4, 1))
# Use 128bit(16B) load as canonicalized form of val_layout then recast to target element-type
coalesced_ldst_bytes = 16
# Compile time validation: expect same element type for all input tensors
assert all(t.element_type == inputs[0].element_type for t in inputs)
dtype = inputs[0].element_type
thr_layout = cute.make_ordered_layout((4, 64), order=(1, 0))
val_layout = cute.make_ordered_layout((16, coalesced_ldst_bytes), order=(1, 0))
val_layout = cute.recast_layout(dtype.width, 8, val_layout)
tiler_mn, tv_layout = cute.make_layout_tv(thr_layout, val_layout)
print(f"[DSL INFO] Input Tensors:")
print(f"[DSL INFO] a = {a.type}")
print(f"[DSL INFO] b = {b.type}")
print(f"[DSL INFO] result = {result.type}")
print("[DSL INFO] Input Tensors:")
for i, t in enumerate(inputs):
print(f"[DSL INFO] inputs{i} = {t}")
print(f"[DSL INFO] result = {result}")
print(f"[DSL INFO] Tiling Parameters:")
print("[DSL INFO] Tiling Parameters:")
print(f"[DSL INFO] tiler_mn = {tiler_mn} per thread block")
print(f"[DSL INFO] tv_layout = {tv_layout}")
gA = cute.zipped_divide(a, tiler_mn) # ((TileM, TileN), (RestM, RestN))
gB = cute.zipped_divide(b, tiler_mn) # ((TileM, TileN), (RestM, RestN))
gC = cute.zipped_divide(result, tiler_mn) # ((TileM, TileN), (RestM, RestN))
print("[DSL INFO] Tiled Tensors:")
mInputs = [cute.zipped_divide(input, tiler_mn) for input in inputs]
# ((TileM, TileN), (RestM, RestN))
mC = cute.zipped_divide(result, tiler_mn)
print(f"[DSL INFO] Tiled Tensors:")
print(f"[DSL INFO] gA = {gA.type}")
print(f"[DSL INFO] gB = {gB.type}")
print(f"[DSL INFO] gC = {gC.type}")
# (RestM, RestN) -> (RestN, RestM)
remap_block = cute.make_ordered_layout(
cute.select(mInputs[0].shape[1], mode=[1, 0]), order=(1, 0)
)
for i, t in enumerate(mInputs):
print(f"[DSL INFO] gInputs{i} = {mInputs[i]}")
mInputs[i] = cute.composition(t, (None, remap_block))
print(f"[DSL INFO] gInputs{i} (remapped) = {mInputs[i]}")
mC = cute.composition(mC, (None, remap_block))
print(f"[DSL INFO] gC = {mC}")
idC = cute.make_identity_tensor(result.shape)
cC = cute.zipped_divide(idC, tiler=tiler_mn)
print(f"[DSL INFO] coord tensor = {cC.type}")
print(f"[DSL INFO] coord tensor = {cC}")
# Launch the kernel asynchronously
# Async token(s) can also be specified as dependencies
elementwise_apply_kernel(
op,
[gA, gB], # Group input tensors into a list as a single argument
gC,
cC,
result.shape,
tv_layout,
).launch(
grid=[cute.size(gC, mode=[1]), 1, 1],
# Group input tensors into a list as a single argument
elementwise_apply_kernel(op, mInputs, mC, cC, result.shape, tv_layout).launch(
# Compute production at each mode of mC.shape[1] to get multi-dimensional grid size
grid=cute.product_each(mC.shape[1]),
block=[cute.size(tv_layout, mode=[0]), 1, 1],
stream=stream,
)
def run_elementwise_apply_and_verify(
@cutlass.dsl_user_op
def leaky_relu(x, alpha, *, loc=None, ip=None):
return cute.where(x > 0, x, alpha * x, loc=loc, ip=ip)
def leaky_relu_ref(x, alpha):
return torch.where(x > 0, x, alpha * x)
def run_and_verify(
op,
M,
N,
@@ -287,14 +288,23 @@ def run_elementwise_apply_and_verify(
iterations=100,
):
if not torch.cuda.is_available():
raise RuntimeError(f"Ampere GPU is required to run this example!")
raise RuntimeError("NVIDIA GPU is required to run this example!")
if op == "leaky_relu":
op = partial(leaky_relu, alpha=0.01)
ref_op = partial(leaky_relu_ref, alpha=0.01)
num_inputs = 1
else:
op = getattr(operator, op)
ref_op = op
num_inputs = 2
# Create non default CUDA stream from PyTorch
torch_stream = torch.cuda.Stream()
# Get the raw stream pointer as a CUstream
current_stream = cuda.CUstream(torch_stream.cuda_stream)
print(f"\nRunning Elementwise Apply test with:")
print("\nRunning Elementwise Apply test with:")
print(f"Tensor dimensions: [{M}, {N}]")
print(f"Input and Output Data type: {dtype}")
print(f"Warmup iterations: {warmup_iterations}")
@@ -303,85 +313,78 @@ def run_elementwise_apply_and_verify(
torch_dtype = cutlass_torch.dtype(dtype)
# Allocate tensors with random values.
a = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
b = torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
c = torch.zeros_like(a)
inputs = [
torch.randn(M, N, device=torch.device("cuda"), dtype=torch_dtype)
for _ in range(num_inputs)
]
c = torch.zeros_like(inputs[0])
print(f"Input tensor shapes:")
print(f"a: {a.shape}, dtype: {a.dtype}")
print(f"b: {b.shape}, dtype: {b.dtype}")
print("Input tensor shapes:")
for i in range(num_inputs):
print(f"inputs[{i}]: {inputs[i].shape}, dtype: {inputs[i].dtype}")
print(f"c: {c.shape}, dtype: {c.dtype}\n")
epsilon = 1.2
if op in (operator.truediv, operator.floordiv):
b = torch.where(b == 0, torch.tensor(epsilon), b)
inputs[1] = torch.where(inputs[1] == 0, torch.tensor(epsilon), inputs[1])
print("Executing elementwise apply kernel...")
inputs_ = [from_dlpack(t, assumed_align=16) for t in inputs]
c_ = from_dlpack(c, assumed_align=16).mark_layout_dynamic()
print("Compiling kernel with cute.compile ...")
start_time = time.time()
compiled_fn = cute.compile[cute.GenerateLineInfo(True)](
elementwise_apply, op, inputs_, c_, current_stream
)
compilation_time = time.time() - start_time
print(f"Compilation time: {compilation_time:.4f} seconds")
if not skip_ref_check:
elementwise_apply(
op,
from_dlpack(a),
from_dlpack(b),
from_dlpack(c).mark_layout_dynamic(),
current_stream,
)
print("Executing elementwise apply kernel...")
compiled_fn(inputs_, c_, current_stream)
print("Verifying results...")
torch.testing.assert_close(op(a, b), c)
torch.testing.assert_close(ref_op(*inputs), c)
print("Results verified successfully!")
print(f"First few elements of result: \n{c[:3, :3]}")
if not benchmark:
return
compiled_func = cute.compile(
elementwise_apply,
op,
from_dlpack(a),
from_dlpack(b),
from_dlpack(c).mark_layout_dynamic(),
current_stream,
)
# When compiled we inlined op in the kernel, so we do not pass it when benchmarking
print("Benchmarking elementwise apply kernel...")
avg_time_us = testing.benchmark(
compiled_func,
kernel_arguments=testing.JitArguments(
from_dlpack(a),
from_dlpack(b),
from_dlpack(c).mark_layout_dynamic(),
current_stream,
),
compiled_fn,
kernel_arguments=testing.JitArguments(inputs_, c_, current_stream),
warmup_iterations=warmup_iterations,
iterations=iterations,
use_cuda_graphs=True,
stream=current_stream,
)
avg_time = avg_time_us / 1e3
num_elements = sum(input.numel() for input in inputs) + c.numel()
# Print execution results
print(f"Kernel execution time: {avg_time:.4f} ms")
print(f"Kernel execution time: {avg_time_us / 1e3:.4f} ms")
print(
f"Achieved memory throughput: {(3 * a.numel() * dtype.width // 8) / (avg_time / 1000) / 1e9:.2f} GB/s"
f"Achieved memory throughput: {(num_elements * dtype.width // 8) / (avg_time_us * 1000):.2f} GB/s"
)
print(f"First few elements of result: \n{c[:3, :3]}")
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="example of elementwise apply to demonstrate building elementwise kernels"
description="Demonstration of building customizable elementwise CUDA kernels using the CuTe DSL"
)
parser.add_argument("--M", default=128, type=int)
parser.add_argument("--N", default=128, type=int)
parser.add_argument("--M", default=4096, type=int)
parser.add_argument("--N", default=4096, type=int)
parser.add_argument("--op", default="add", type=str)
parser.add_argument("--warmup_iterations", default=2, type=int)
parser.add_argument("--iterations", default=100, type=int)
parser.add_argument("--skip_ref_check", action="store_true")
parser.add_argument("--benchmark", action="store_true")
args = parser.parse_args()
run_elementwise_apply_and_verify(
getattr(operator, args.op),
run_and_verify(
args.op,
args.M,
args.N,
dtype=cutlass.Float32,

30
examples/python/CuTeDSL/ampere/flash_attention_v2.py
View File

@@ -28,7 +28,7 @@
import argparse
from types import SimpleNamespace
from typing import Type, Union, Callable
from typing import Type, Callable
import torch
import cuda.bindings.driver as cuda
@@ -38,6 +38,7 @@ import cutlass.cute as cute
from cutlass.cute.nvgpu import cpasync, warp
import cutlass.torch as cutlass_torch
from cutlass.cute.runtime import from_dlpack
import cutlass.pipeline as pipeline
import cutlass.utils as utils
"""
@@ -126,6 +127,10 @@ class FlashAttentionForwardAmpere:
self._num_threads = num_threads
self._is_causal = is_causal
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1, num_threads=num_threads
)
@staticmethod
def can_implement(
dtype, head_dim, m_block_size, n_block_size, num_threads, is_causal
@@ -450,7 +455,7 @@ class FlashAttentionForwardAmpere:
acc_shape_O = thr_mma.partition_shape_C(
(self._m_block_size, self._head_dim_padded)
)
acc_O = cute.make_fragment(acc_shape_O, cutlass.Float32)
acc_O = cute.make_rmem_tensor(acc_shape_O, cutlass.Float32)
acc_O.fill(0.0)
# ///////////////////////////////////////////////////////////////////////////////
@@ -506,7 +511,7 @@ class FlashAttentionForwardAmpere:
tKVcKV = gmem_thr_copy_QKV.partition_S(cKV)
# Allocate predicate tensors for m and n, here we only allocate the tile of k, and do special process for mn.
# This is to reduce register pressure and gets 2-3% performance gain compared with allocating the whole tile.
tQpQ = cute.make_fragment(
tQpQ = cute.make_rmem_tensor(
cute.make_layout(
(
tQsQ.shape[0][1],
@@ -517,7 +522,7 @@ class FlashAttentionForwardAmpere:
),
cutlass.Boolean,
)
tKVpKV = cute.make_fragment(
tKVpKV = cute.make_rmem_tensor(
cute.make_layout(
(
tKsK.shape[0][1],
@@ -571,11 +576,11 @@ class FlashAttentionForwardAmpere:
# Softmax intermediate result: row_max and row_sum
# ///////////////////////////////////////////////////////////////////////////////
# shape: (atom_v_m * rest_m)
row_max = cute.make_fragment(
row_max = cute.make_rmem_tensor(
(acc_O.shape[0][0] * acc_O.shape[1]), cutlass.Float32
)
# shape: (atom_v_m * rest_m)
row_sum = cute.make_fragment(
row_sum = cute.make_rmem_tensor(
(acc_O.shape[0][0] * acc_O.shape[1]), cutlass.Float32
)
row_max.fill(-cutlass.Float32.inf)
@@ -710,7 +715,7 @@ class FlashAttentionForwardAmpere:
tOgO = gmem_thr_copy_O.partition_D(gO)
tOrO = cute.make_fragment_like(tOgO, self._dtype)
# sync before all smem stores are done.
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# load acc O from smem to rmem for wider vectorization
cute.copy(
gmem_tiled_copy_O,
@@ -724,7 +729,7 @@ class FlashAttentionForwardAmpere:
(m_block, 0),
)
tOcO = gmem_thr_copy_O.partition_D(cO)
tOpO = cute.make_fragment(
tOpO = cute.make_rmem_tensor(
cute.make_layout(
(tOgO.shape[0][1], tOgO.shape[1], tOgO.shape[2]),
stride=(tOgO.shape[2], 0, 1),
@@ -778,12 +783,12 @@ class FlashAttentionForwardAmpere:
acc_shape_S = mma_params.thr_mma.partition_shape_C(
(self._m_block_size, self._n_block_size)
)
acc_S = cute.make_fragment(acc_shape_S, cutlass.Float32)
acc_S = cute.make_rmem_tensor(acc_shape_S, cutlass.Float32)
acc_S.fill(0.0)
# wait for smem tile QK before mma calculation for S
cute.arch.cp_async_wait_group(0)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# load smem tile V for O, special process for the first tile to avoid loading nan.
# The `if` here is a constexpr, won't be generated in the IR.
if is_first_n_block:
@@ -847,7 +852,7 @@ class FlashAttentionForwardAmpere:
# wait for smem tile V for O
cute.arch.cp_async_wait_group(0)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
if basic_params.n_block > 0:
cute.copy(
@@ -1170,7 +1175,7 @@ def run(
f"Unsupported testcase {dtype}, {head_dim}, {m_block_size}, {n_block_size}, {num_threads}, {is_causal}"
)
print(f"Running Ampere SM80 FlashAttentionForward test with:")
print("Running Ampere SM80 FlashAttentionForward test with:")
print(f" dtype: {dtype}")
print(f" batch_size: {batch_size}")
print(f" seqlen_q: {seqlen_q}")
@@ -1285,6 +1290,7 @@ def run(
return avg_time_us # Return execution time in microseconds
if __name__ == "__main__":
parser = argparse.ArgumentParser(
description="example of flash attention v2 with CuTe on GPU"

1149
examples/python/CuTeDSL/ampere/hstu_attention.py Normal file
View File

File diff suppressed because it is too large Load Diff

245
examples/python/CuTeDSL/ampere/inline_ptx.py Normal file
View File

@@ -0,0 +1,245 @@
# Copyright (c) 2025 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.
from functools import partial
from typing import Union
import torch
import cutlass.cute as cute
from cutlass.cute.runtime import from_dlpack
from cutlass._mlir.dialects import llvm
from cutlass.cute.typing import Boolean, Int32, Int, Constexpr
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass.cute.arch.nvvm_wrappers import FULL_MASK, WARP_SIZE
"""
A simple example to show how to wrap PTX instructions by using inline_asm op in llvm dialect.
Situations like:
1. Instructions that are not already exposed by CuTe DSL via `nvvm` module
2. Sequences of instructions that the compiler otherwise does not generate optimally
motivate developers to inline PTX themselves.
In this example, we inline the vote.sync.ballot.b32, vote.sync.any.pred, vote.sync.all.pred,
vote.sync.uni.pred, and use the corresponding ops in nvvm_wrappers.py for the test.
You can refer to the documentation of `inline_asm op in llvm dialect <https://mlir.llvm.org/docs/Dialects/LLVM/#llvminline_asm-llvminlineasmop>`_
and `vote.sync <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`_
for more details.
To run this example:
.. code-block:: bash
python examples/ampere/inline_ptx.py
The example will run the vote kernel with inline PTX and nvvm dialect separately.
The results from inline PTX and nvvm dialect will be verified correspondingly.
"""
@dsl_user_op
def ptx_vote_sync_op(
pred: Boolean, kind: str, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Union[Int32, Boolean]:
return_type = Boolean
return_type_str = "pred"
return return_type(
llvm.inline_asm(
T.bool(),
[
Boolean(pred).ir_value(loc=loc, ip=ip),
Int32(mask).ir_value(loc=loc, ip=ip),
],
f"""{{\n\t
.reg .pred ps;\n\t
.reg .pred pd;\n\t
setp.ne.b32 ps, $1, 0;\n\t
vote.sync.{kind}.{return_type_str} pd, ps, $2;\n\t
selp.b32 $0, 1, 0, pd;\n\t
}}""",
"=r,r,i",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
)
ptx_vote_any_sync = partial(ptx_vote_sync_op, kind="any")
ptx_vote_all_sync = partial(ptx_vote_sync_op, kind="all")
ptx_vote_uni_sync = partial(ptx_vote_sync_op, kind="uni")
@dsl_user_op
def ptx_vote_ballot_sync(
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Union[Int32, Boolean]:
return_type = Int32
return_type_str = "b32"
return return_type(
llvm.inline_asm(
T.i32(),
[
Boolean(pred).ir_value(loc=loc, ip=ip),
Int32(mask).ir_value(loc=loc, ip=ip),
],
f"""{{\n\t
.reg .pred p;\n\t
setp.ne.b32 p, $1, 0;\n\t
vote.sync.ballot.{return_type_str} $0, p, $2;\n\t
}}""",
"=r,r,i",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
)
@cute.kernel
def vote_kernel(
mBallot: cute.Tensor,
mAny: cute.Tensor,
mAll: cute.Tensor,
mUni: cute.Tensor,
use_inline_ptx: Constexpr[bool],
):
tidx, _, _ = cute.arch.thread_idx()
vote_ballot = (
ptx_vote_ballot_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_ballot_sync(tidx < 10)
)
vote_any = (
ptx_vote_any_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_any_sync(tidx < 10)
)
vote_all = (
ptx_vote_all_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_all_sync(tidx < 10)
)
vote_uni = (
ptx_vote_uni_sync(tidx < 10)
if use_inline_ptx
else cute.arch.vote_uni_sync(tidx < 10)
)
mBallot[tidx] = vote_ballot
mAny[tidx] = vote_any
mAll[tidx] = vote_all
mUni[tidx] = vote_uni
@cute.jit
def vote(
mBallot: cute.Tensor,
mAny: cute.Tensor,
mAll: cute.Tensor,
mUni: cute.Tensor,
use_inline_ptx: Constexpr[bool],
):
vote_kernel(
mBallot,
mAny,
mAll,
mUni,
use_inline_ptx,
).launch(
grid=[1, 1, 1],
block=[cute.size(WARP_SIZE, mode=[0]), 1, 1],
)
def run():
ballot_ptx = torch.randint(
0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32
)
any_ptx = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
all_ptx = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
uni_ptx = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
mBallotPTX = from_dlpack(ballot_ptx).mark_layout_dynamic()
mAnyPTX = from_dlpack(any_ptx).mark_layout_dynamic()
mAllPTX = from_dlpack(all_ptx).mark_layout_dynamic()
mUniPTX = from_dlpack(uni_ptx).mark_layout_dynamic()
# get the results from ptx
vote(mBallotPTX, mAnyPTX, mAllPTX, mUniPTX, use_inline_ptx=True)
ballot_nvvm = torch.randint(
0, 100, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.int32
)
any_nvvm = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
all_nvvm = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
uni_nvvm = torch.randint(
0, 2, (WARP_SIZE,), device=torch.device("cuda"), dtype=torch.bool
)
mBallotNVVM = from_dlpack(ballot_nvvm).mark_layout_dynamic()
mAnyNVVM = from_dlpack(any_nvvm).mark_layout_dynamic()
mAllNVVM = from_dlpack(all_nvvm).mark_layout_dynamic()
mUniNVVM = from_dlpack(uni_nvvm).mark_layout_dynamic()
# get the results from nvvm
vote(mBallotNVVM, mAnyNVVM, mAllNVVM, mUniNVVM, use_inline_ptx=False)
print("Verifying ballot results...")
torch.testing.assert_close(ballot_ptx, ballot_nvvm)
print("Verifying any results...")
torch.testing.assert_close(any_ptx, any_nvvm)
print(torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE))))
assert torch.all(any_ptx == any(i < 10 for i in range(WARP_SIZE)))
print("Verifying all results...")
torch.testing.assert_close(all_ptx, all_nvvm)
assert torch.all(all_ptx == all(i < 10 for i in range(WARP_SIZE)))
print("Verifying uni results...")
torch.testing.assert_close(uni_ptx, uni_nvvm)
assert torch.all(uni_ptx == (len(set(i < 10 for i in range(WARP_SIZE))) == 1))
print("Results verified successfully!")
if __name__ == "__main__":
run()

57
examples/python/CuTeDSL/ampere/sgemm.py
View File

@@ -36,7 +36,7 @@ import torch
import cutlass
import cutlass.cute as cute
import cutlass.cute.testing as testing
import cutlass.torch as cutlass_torch
import cutlass.pipeline as pipeline
import cutlass.utils as utils
from cutlass.cute.runtime import from_dlpack
@@ -103,6 +103,9 @@ class SGemm:
assert self._bM % 16 == 0, "multiple of 16 required for tile dimension M"
assert self._bN % 16 == 0, "multiple of 16 required for tile dimension N"
assert self._num_stages >= 3, "num_stages must be greater than or equal to 3"
self.cta_sync_barrier = pipeline.NamedBarrier(
barrier_id=1, num_threads=num_threads
)
@cute.jit
def __call__(
@@ -166,9 +169,8 @@ class SGemm:
mA.element_type,
num_bits_per_copy=mB.element_type.width,
)
if cutlass.const_expr(self.a_major_mode == utils.LayoutEnum.COL_MAJOR):
num_vectorized = 4 if (mA.layout.max_alignment % 16 == 0) else 1
num_vectorized = 4 if (mA.layout[0].max_alignment % 16 == 0) else 1
atom_async_copy_A = cute.make_copy_atom(
cute.nvgpu.cpasync.CopyG2SOp(),
mA.element_type,
@@ -182,7 +184,7 @@ class SGemm:
vA = cute.make_layout((num_vectorized, 1))
if cutlass.const_expr(self.b_major_mode == utils.LayoutEnum.COL_MAJOR):
num_vectorized = 4 if (mB.layout.max_alignment % 16 == 0) else 1
num_vectorized = 4 if (mB.layout[0].max_alignment % 16 == 0) else 1
atom_async_copy_B = cute.make_copy_atom(
cute.nvgpu.cpasync.CopyG2SOp(),
mA.element_type,
@@ -294,7 +296,7 @@ class SGemm:
# tile (instead of the last one) irregular in shape when k is irregular.
# We first handle the irregular tile to avoid checking for this
# condition within the mainloop.
residue_k = mA.shape[1] - cutlass.Int32(self._bK) * gA.shape[2]
residue_k = mA.shape[1] - self._bK * gA.shape[2]
gA = cute.domain_offset((0, residue_k, 0), gA)
gB = cute.domain_offset((0, residue_k, 0), gB)
@@ -342,7 +344,7 @@ class SGemm:
tAcA = thr_copy_A.partition_S(cA)
tBcB = thr_copy_B.partition_S(cB)
# Allocate predicate tensors for m and n
tApA = cute.make_fragment(
tApA = cute.make_rmem_tensor(
cute.make_layout(
(
tAsA.shape[0][1],
@@ -353,7 +355,7 @@ class SGemm:
),
cutlass.Boolean,
)
tBpB = cute.make_fragment(
tBpB = cute.make_rmem_tensor(
cute.make_layout(
(
tBsB.shape[0][1],
@@ -365,7 +367,7 @@ class SGemm:
cutlass.Boolean,
)
# Allocate predicate tensors for m, n and k for residue k-tile
tApA_residue_k = cute.make_fragment(
tApA_residue_k = cute.make_rmem_tensor(
cute.make_layout(
(
tAsA.shape[0][1],
@@ -380,7 +382,7 @@ class SGemm:
),
cutlass.Boolean,
)
tBpB_residue_k = cute.make_fragment(
tBpB_residue_k = cute.make_rmem_tensor(
cute.make_layout(
(
tBsB.shape[0][1],
@@ -508,7 +510,7 @@ class SGemm:
if k_block_max > 1:
# Wait until our first prefetched tile is loaded in
cute.arch.cp_async_wait_group(k_pipe_max - 2)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# Prefetch the first rmem from the first k-tile
cute.autovec_copy(tCsA_p[None, None, 0], tCrA[None, None, 0])
cute.autovec_copy(tCsB_p[None, None, 0], tCrB[None, None, 0])
@@ -545,7 +547,7 @@ class SGemm:
tCsA_p = tCsA[None, None, None, smem_pipe_read]
tCsB_p = tCsB[None, None, None, smem_pipe_read]
cute.arch.cp_async_wait_group(k_pipe_max - 2)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
# Load A, B from shared memory to registers for k_block + 1
k_block_next = (k_block + 1) % k_block_max # static
@@ -611,13 +613,13 @@ class SGemm:
# them without vectorization.
# ///////////////////////////////////////////////////////////////////////////////
cute.arch.cp_async_wait_group(0)
cute.arch.barrier()
self.cta_sync_barrier.arrive_and_wait()
tCrC.store(epilogue_op(tCrC.load()))
# predicate
cC = cute.make_identity_tensor(gC.shape)
tCpC = thr_mma.partition_C(cC)
predC = cute.make_fragment(tCrC.layout, cutlass.Boolean)
predC = cute.make_rmem_tensor(tCrC.layout, cutlass.Boolean)
residue_m = mC.shape[0] - cutlass.Int32(self._bM) * bidx
residue_n = mC.shape[1] - cutlass.Int32(self._bN) * bidy
for i in range(cute.size(tCrC.shape)):
@@ -664,7 +666,7 @@ def run(
:return: Execution time of the GEMM kernel in microseconds
:rtype: float
"""
print(f"Running Ampere SIMT GEMM example:")
print("Running Ampere SIMT GEMM example:")
print(f"mnk: {mnk}")
print(f"A major: {a_major}, B major: {b_major}, C major: {c_major}")
print(f"Static shape: {static_shape}")
@@ -697,14 +699,17 @@ def run(
divisibility_b = b.shape[1] if b_major == "k" else b.shape[0]
divisibility_c = c.shape[1] if c_major == "n" else c.shape[0]
a_tensor = (
from_dlpack(a, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if a_major == "k" else 0))
.mark_compact_shape_dynamic(
mode=(1 if a_major == "k" else 0),
divisibility=divisibility_a,
if static_shape:
a_tensor = (
from_dlpack(a, assumed_align=16)
.mark_layout_dynamic(leading_dim=(1 if a_major == "k" else 0))
.mark_compact_shape_dynamic(
mode=(1 if a_major == "k" else 0),
divisibility=divisibility_a,
)
)
)
else:
a_tensor = from_dlpack(a, assumed_align=16)
b_tensor = (
from_dlpack(b, assumed_align=16)
@@ -733,12 +738,8 @@ def run(
print("Compiling kernel with cute.compile ...")
start_time = time.time()
compiled_fn = cute.compile(
sgemm,
a_tensor,
b_tensor,
c_tensor,
stream=current_stream,
compiled_fn = cute.compile[cute.GenerateLineInfo](
sgemm, a_tensor, b_tensor, c_tensor, stream=current_stream
)
compilation_time = time.time() - start_time
print(f"Compilation time: {compilation_time:.4f} seconds")
@@ -833,7 +834,7 @@ if __name__ == "__main__":
parser.add_argument(
"--mnk", type=parse_comma_separated_ints, default=(256, 256, 64)
)
parser.add_argument("--a_major", choices=["k", "m"], default="k")
parser.add_argument("--a_major", choices=["k", "m"], default="m")
parser.add_argument("--b_major", choices=["k", "n"], default="k")
parser.add_argument("--c_major", choices=["n", "m"], default="n")
parser.add_argument("--warmup_iterations", default=2, type=int)

39
examples/python/CuTeDSL/ampere/smem_allocator.py
View File

@@ -69,7 +69,7 @@ class complex:
class SharedStorage:
# struct elements with natural alignment
a: cute.struct.MemRange[cutlass.Float32, 32] # array
b: cutlass.Int64 # saclar
b: cutlass.Int64 # scalar
c: complex # nested struct
# struct elements with strict alignment
x: cute.struct.Align[
@@ -90,10 +90,17 @@ def kernel(
dst_c: cute.Tensor,
):
# Note: SMEM_SIZE bytes (specified in kernel().launch(smem=...)) can be reserved for developer to utilize
# Note: alignment of inital allocator base ptr is 1024
# Note: alignment of initial allocator base ptr is 1024
allocator = cutlass.utils.SmemAllocator()
# base ptr of allocator points at: SMEM_ADDR_START (the starting address of available shared memory)
# -- Allocate a scalar
int_ptr = allocator.allocate(cutlass.Int32)
# base ptr of allocator now points at: SMEM_ADDR_AFTER_INT = SMEM_ADDR_START + aligned_size(int)
assert int_ptr.dtype == cutlass.Int32, "Expected Int32, but got {}".format(
int_ptr.dtype
)
# -- Allocate a struct --
# Note: when specified alignment, max(alignment, alignof(struct)) will be applied
# reserves the section of struct in smem, elements in the struct can be accessed by ptr
@@ -153,7 +160,7 @@ def kernel(
@cute.jit
def run_allocation_kernel(
def host(
const_a: cutlass.Constexpr,
dst_a: cute.Tensor,
const_b: cutlass.Constexpr,
@@ -161,22 +168,18 @@ def run_allocation_kernel(
const_c: cutlass.Constexpr,
dst_c: cute.Tensor,
):
# additional size for the example, 64(section) + 112(array) + 128(tensor) < 384
addtional_bytes = 384
# Note: launch shared memory size is: SMEM_SIZE = 512 + 384 = 896 bytes
# Note: Shared Memory size is automatically calculated now
kernel(const_a, dst_a, const_b, dst_b, const_c, dst_c).launch(
grid=(1, 1, 1),
block=(1, 1, 1),
smem=SharedStorage.size_in_bytes() + addtional_bytes,
grid=(1, 1, 1), block=(1, 1, 1)
)
def veify_allocation_kernel(const_a, const_b, const_c):
def run_and_verify(const_a, const_b, const_c):
dst_a = torch.zeros((8, 4), dtype=torch.float32, device="cuda")
dst_b = torch.zeros((8, 2), dtype=torch.float32, device="cuda")
dst_c = torch.zeros((16, 2), dtype=torch.float32, device="cuda")
run_allocation_kernel(
host(
const_a,
from_dlpack(dst_a),
const_b,
@@ -185,9 +188,15 @@ def veify_allocation_kernel(const_a, const_b, const_c):
from_dlpack(dst_c),
)
np.testing.assert_equal(const_a, dst_a.detach().cpu().numpy()[0])
np.testing.assert_equal(const_b, dst_b.detach().cpu().numpy()[0])
np.testing.assert_equal(const_c, dst_c.detach().cpu().numpy()[0])
assert const_a == dst_a.cpu()[0, 0], (
f"Expected {const_a}, but got {dst_a.cpu()[0, 0]}"
)
assert const_b == dst_b.cpu()[0, 0], (
f"Expected {const_b}, but got {dst_b.cpu()[0, 0]}"
)
assert const_c == dst_c.cpu()[0, 0], (
f"Expected {const_c}, but got {dst_c.cpu()[0, 0]}"
)
if __name__ == "__main__":
@@ -197,4 +206,4 @@ if __name__ == "__main__":
const_a = 0.5
const_b = 1.0
const_c = 2.0
veify_allocation_kernel(const_a, const_b, const_c)
run_and_verify(const_a, const_b, const_c)

23
examples/python/CuTeDSL/ampere/tensorop_gemm.py
View File

@@ -28,10 +28,8 @@
import argparse
import math
import time
from typing import Tuple, Type
import cuda.bindings.driver as cuda
import torch
import cutlass
@@ -121,12 +119,12 @@ class TensorOpGemm:
self.mma_inst_shape = (16, 8, 16)
mmaM, mmaN, mmaK = self.mma_inst_shape
assert (
self.bM % (atom_lay_M * mmaM) == 0
), "bM must be divisible by MMA instruction"
assert (
self.bN % (atom_lay_N * mmaN) == 0
), "bN must be divisible by MMA instruction"
assert self.bM % (atom_lay_M * mmaM) == 0, (
"bM must be divisible by MMA instruction"
)
assert self.bN % (atom_lay_N * mmaN) == 0, (
"bN must be divisible by MMA instruction"
)
assert atom_lay_K == 1, "this example does not support atom layout K > 1"
assert self.bK % mmaK == 0, "bK must be divisible by MMA instruction"
assert self.num_stages >= 3, "num_stages must be greater than or equal to 3"
@@ -428,7 +426,7 @@ class TensorOpGemm:
# at the granularity of a copy atom, so the predicate tensor does not
# need separate booleans for individual elements within a copy
# atom (for example, the elements of tAgA.shape[0][0].)
tApA = cute.make_fragment(
tApA = cute.make_rmem_tensor(
cute.make_layout(
(
tAgA.shape[0][1],
@@ -439,7 +437,7 @@ class TensorOpGemm:
),
cutlass.Boolean,
)
tBpB = cute.make_fragment(
tBpB = cute.make_rmem_tensor(
cute.make_layout(
(
tBsB.shape[0][1],
@@ -707,7 +705,7 @@ class TensorOpGemm:
cute.autovec_copy(tCsC_epilogue, tCrC_epilogue)
# Create predication tensor for m
tCpC = cute.make_fragment(
tCpC = cute.make_rmem_tensor(
cute.make_layout(
(
tCgC_epilogue.shape[0][1],
@@ -851,7 +849,7 @@ def run(
use_cold_l2: bool = False,
**kwargs,
):
print(f"Running Ampere tensor core GEMM example:")
print("Running Ampere tensor core GEMM example:")
print(f"mnkl: {mnkl}")
print(
f"A dtype: {ab_dtype}, B dtype: {ab_dtype}, C dtype: {c_dtype}, Acc dtype: {acc_dtype}"
@@ -944,6 +942,7 @@ def run(
return avg_time_us # Return execution time in microseconds
if __name__ == "__main__":
def parse_comma_separated_ints(s: str) -> Tuple[int, ...]: