leavelet/cutlass
1
0
Fork 0
Code Issues Pull Requests Actions 26 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-21 14:26:30 -04:00
committed by GitHub
co-authored by Larry Wu
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
examples/python/CuTeDSL/ampere/elementwise_add.py
+13 -12
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,