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

v4.3.1 update. (#2817)

Browse Source
This commit is contained in:
Junkai-Wu
2025-11-27 22:49:30 +08:00
committed by GitHub
parent 2052fd3885
commit 1de3a576cc
44 changed files with 3316 additions and 510 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/CuTeDSL/cutlass/cute/__init__.py
View File

@@ -205,6 +205,7 @@ KeepCUBIN = _dsl.KeepCUBIN
KeepPTX = _dsl.KeepPTX
GPUArch = _dsl.GPUArch
LinkLibraries = _dsl.LinkLibraries
EnableTVMFFI = _dsl.EnableTVMFFI
# attach the TVM FFI ABI interface postprocessor to the DSL
from . import _tvm_ffi_args_spec_converter

240
python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py
View File

@@ -42,7 +42,7 @@ from .typing import (
)
import cuda.bindings.driver as cuda
from typing import List, Dict, Any, Optional
from typing import List, Dict, Any, Optional, Tuple, get_origin, get_args
import inspect
NumericToTVMFFIDtype = {
@@ -91,6 +91,11 @@ def _get_llvm_address_space_from_memspace(
return 1
return None
def _is_gpu_memspace(
memspace: _cute_ir.AddressSpace,
) -> bool:
return memspace != _cute_ir.AddressSpace.generic
class SymIntId:
def __init__(self, sym_int: SymInt):
@@ -103,118 +108,187 @@ class SymIntId:
return self.sym_int is other.sym_int
def _tvm_ffi_args_spec_converter(
function_name: str,
args_spec: inspect.FullArgSpec,
dynamic_args: List[Any],
dynamic_kwargs: Dict[str, Any],
):
"""Convert cute algebra args to tvm ffi spec params.
This function converts the cute arguments specs to tvm ffi spec params.
"""
exec_args = ExecutionArgs(args_spec, function_name)
rectified_args = exec_args.get_rectified_args(dynamic_args, dynamic_kwargs)
arg_names = exec_args.args_spec.args + exec_args.args_spec.kwonlyargs
class ConverterContext:
"""Context for managing variable allocation during TVM FFI args conversion."""
params = []
num_dyn_shape_vars = 0
num_dyn_stride_vars = 0
sym_int_id_mapping = {}
def __init__(self):
self.num_dyn_shape_vars = 0
self.num_dyn_stride_vars = 0
self.sym_int_id_mapping = {}
def alloc_shape_name():
nonlocal num_dyn_shape_vars
name = f"n{num_dyn_shape_vars}"
num_dyn_shape_vars += 1
def alloc_shape_name(self) -> str:
"""Allocate a new dynamic shape variable name."""
name = f"n{self.num_dyn_shape_vars}"
self.num_dyn_shape_vars += 1
return name
def alloc_stride_name():
nonlocal num_dyn_stride_vars
name = f"s{num_dyn_stride_vars}"
num_dyn_stride_vars += 1
def alloc_stride_name(self) -> str:
"""Allocate a new dynamic stride variable name."""
name = f"s{self.num_dyn_stride_vars}"
self.num_dyn_stride_vars += 1
return name
def alloc_or_reuse_symint_var(value, name_alloc_func):
nonlocal sym_int_id_mapping
def alloc_or_reuse_symint_var(self, value: SymInt, name_alloc_func):
"""Allocate or reuse a symbolic integer variable."""
sym_int_id = SymIntId(value)
if sym_int_id in sym_int_id_mapping:
return sym_int_id_mapping[sym_int_id]
if sym_int_id in self.sym_int_id_mapping:
return self.sym_int_id_mapping[sym_int_id]
name = name_alloc_func()
if value.width == 32:
dtype = NumericToTVMFFIDtype[Int32]
else:
dtype = NumericToTVMFFIDtype[Int64]
var = spec.Var(name, dtype, divisibility=value.divisibility)
sym_int_id_mapping[sym_int_id] = var
self.sym_int_id_mapping[sym_int_id] = var
return var
for arg, arg_name in zip(rectified_args, arg_names):
arg_type = args_spec.annotations.get(arg_name, None)
if isinstance(arg, Numeric) and arg.dtype in AcceptableNumericTypesForScalar:
params.append(spec.Var(arg_name, NumericToTVMFFIDtype[arg.dtype]))
elif is_cute_algebra_type(arg_type):
shape = []
for i in range(len(arg)):
if isinstance(arg[i], int):
shape.append(arg[i])
elif isinstance(arg[i], SymInt):
shape.append(alloc_or_reuse_symint_var(arg[i], alloc_shape_name))
else:
shape.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[arg[i].dtype]))
params.append(spec.Shape(arg_name, shape))
elif isinstance(arg, Tensor):
shapes = []
for i, dyn_mask in enumerate(arg.dynamic_shapes_mask):
if not dyn_mask:
shapes.append(arg.shape[i])
elif isinstance(arg.shape[i], SymInt):
shapes.append(alloc_or_reuse_symint_var(arg.shape[i], alloc_shape_name))
else:
shapes.append(spec.Var(alloc_shape_name(), NumericToTVMFFIDtype[Int32]))
strides = []
for i, dyn_mask in enumerate(arg.dynamic_strides_mask):
if not dyn_mask:
strides.append(arg.stride[i])
elif isinstance(arg.stride[i], SymInt):
strides.append(alloc_or_reuse_symint_var(arg.stride[i], alloc_stride_name))
else:
if hasattr(arg, "_use_32bit_stride") and arg._use_32bit_stride:
dtype = NumericToTVMFFIDtype[Int32]
else:
dtype = NumericToTVMFFIDtype[Int64]
strides.append(spec.Var(alloc_stride_name(), dtype))
def _convert_single_arg(
arg,
arg_name: str,
arg_type,
ctx: ConverterContext
) -> spec.Param:
"""Convert a single argument to a spec.Param.
Parameters
----------
arg : Any
The argument value to convert.
arg_name : str
The name of the argument.
arg_type : type
The type annotation of the argument.
ctx : ConverterContext
The converter context for managing variable allocation.
Returns
-------
spec.Param
The converted parameter specification.
"""
if arg is None:
return spec.ConstNone(arg_name)
elif (isinstance(arg, Numeric) and arg.dtype in AcceptableNumericTypesForScalar):
return spec.Var(arg_name, NumericToTVMFFIDtype[arg.dtype])
elif arg_type in AcceptableNumericTypesForScalar:
return spec.Var(arg_name, NumericToTVMFFIDtype[arg_type])
elif is_cute_algebra_type(arg_type):
shape = []
for i in range(len(arg)):
if isinstance(arg[i], int):
shape.append(arg[i])
elif isinstance(arg[i], SymInt):
shape.append(ctx.alloc_or_reuse_symint_var(arg[i], ctx.alloc_shape_name))
else:
shape.append(spec.Var(ctx.alloc_shape_name(), NumericToTVMFFIDtype[arg[i].dtype]))
return spec.Shape(arg_name, shape)
elif isinstance(arg, Tensor):
shapes = []
for i, dyn_mask in enumerate(arg.dynamic_shapes_mask):
if not dyn_mask:
shapes.append(arg.shape[i])
elif isinstance(arg.shape[i], SymInt):
shapes.append(ctx.alloc_or_reuse_symint_var(arg.shape[i], ctx.alloc_shape_name))
else:
shapes.append(spec.Var(ctx.alloc_shape_name(), NumericToTVMFFIDtype[Int32]))
strides = []
for i, dyn_mask in enumerate(arg.dynamic_strides_mask):
if not dyn_mask:
strides.append(arg.stride[i])
elif isinstance(arg.stride[i], SymInt):
strides.append(ctx.alloc_or_reuse_symint_var(arg.stride[i], ctx.alloc_stride_name))
else:
if hasattr(arg, "_use_32bit_stride") and arg._use_32bit_stride:
dtype = NumericToTVMFFIDtype[Int32]
else:
dtype = NumericToTVMFFIDtype[Int64]
strides.append(spec.Var(ctx.alloc_stride_name(), dtype))
if hasattr(arg, "_tvm_ffi_tensor"):
tvm_ffi_tensor = arg._tvm_ffi_tensor
dtype = tvm_ffi_tensor.dtype
tvm_ffi_cute_tensor = spec.Tensor(
arg_name,
shapes,
arg._tvm_ffi_tensor.dtype,
strides=strides,
data_alignment=arg._assumed_align,
device_type=tvm_ffi_tensor.device.type
)
else:
# for FakeTensor, strictly follow the shape and stride from the cute tensor
device_type = "cuda" if _is_gpu_memspace(arg.memspace) else "cpu"
tvm_ffi_cute_tensor = spec.Tensor(
arg_name,
shapes,
NumericToTVMFFIDtype[arg.element_type],
strides=strides,
data_alignment=arg._assumed_align,
device_type=device_type,
)
if arg.element_type == Float4E2M1FN:
tvm_ffi_cute_tensor = spec.create_map_tensor_dtype_f4x2_to_f4_spec(
tvm_ffi_cute_tensor
)
params.append(tvm_ffi_cute_tensor)
elif isinstance(arg, Pointer):
address_space = None
if hasattr(arg, "memspace"):
address_space = _get_llvm_address_space_from_memspace(arg.memspace)
params.append(spec.DataPointer(arg_name, address_space=address_space))
elif isinstance(arg, _FakeStream):
if arg.use_tvm_ffi_env_stream:
params.append(spec.EnvStream(arg_name))
else:
params.append(spec.Stream(arg_name))
elif isinstance(arg, cuda.CUstream):
params.append(spec.Stream(arg_name))
return tvm_ffi_cute_tensor
elif isinstance(arg, Pointer) or arg_type == Pointer:
address_space = None
if hasattr(arg, "memspace"):
address_space = _get_llvm_address_space_from_memspace(arg.memspace)
return spec.DataPointer(arg_name, address_space=address_space)
elif isinstance(arg, _FakeStream):
if arg.use_tvm_ffi_env_stream:
return spec.EnvStream(arg_name)
else:
raise DSLRuntimeError(f"Unsupported argument type: {type(arg)}")
# The following code can obtain signature of the function
# that maybe useful for future debugging and usecases.
# signature = spec.signature(function_name, params)
return spec.Stream(arg_name)
elif isinstance(arg, cuda.CUstream):
return spec.Stream(arg_name)
elif arg_type is not None and get_origin(arg_type) is tuple:
# Handle Tuple[X, Y, ...] type annotations
tuple_element_types = get_args(arg_type)
if not isinstance(arg, (tuple, list)):
raise DSLRuntimeError(f"Expected tuple for argument {arg_name}, got {type(arg)}")
if len(arg) != len(tuple_element_types):
raise DSLRuntimeError(
f"Tuple length mismatch for argument {arg_name}: "
f"expected {len(tuple_element_types)}, got {len(arg)}"
)
# Recursively convert each tuple element
tuple_params = []
for i, (elem, elem_type) in enumerate(zip(arg, tuple_element_types)):
elem_name = f"{arg_name}[{i}]"
elem_param = _convert_single_arg(elem, elem_name, elem_type, ctx)
tuple_params.append(elem_param)
return spec.TupleParam(arg_name, tuple_params)
else:
raise DSLRuntimeError(f"Unsupported argument type: {type(arg)}")
def _tvm_ffi_args_spec_converter(
function_name: str,
args_spec: inspect.FullArgSpec,
full_args: List[Any],
full_kwargs: Dict[str, Any],
):
"""Convert cute algebra args to tvm ffi spec params.
This function converts the cute arguments specs to tvm ffi spec params.
"""
exec_args = ExecutionArgs(args_spec, function_name)
rectified_args = exec_args.get_rectified_args_from_original_args(full_args, full_kwargs)
arg_names = exec_args.args_spec.args + exec_args.args_spec.kwonlyargs
params = []
ctx = ConverterContext()
for arg, arg_name in zip(rectified_args, arg_names):
arg_type = args_spec.annotations.get(arg_name, None)
param = _convert_single_arg(arg, arg_name, arg_type, ctx)
params.append(param)
return params

3
python/CuTeDSL/cutlass/cute/arch/mbar.py
View File

@@ -209,6 +209,7 @@ def mbarrier_conditional_try_wait(
def mbarrier_arrive(
mbar_ptr: Pointer,
peer_cta_rank_in_cluster: Optional[Int] = None,
arrive_count: Int = 1,
*,
loc=None,
ip=None,
@@ -239,7 +240,7 @@ def mbarrier_arrive(
nvvm.mbarrier_txn(
mbar_llvm_ptr,
Int32(1).ir_value(loc=loc, ip=ip),
Int32(arrive_count).ir_value(loc=loc, ip=ip),
kind=nvvm.MBarrierTxnKind.ARRIVE,
space=space,
loc=loc,

26
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py
View File

@@ -201,13 +201,13 @@ class MmaOp(Tcgen05MmaOp):
if self.cta_group == CtaGroup.ONE:
if m not in [64, 128]:
raise OpError(self, f"expects the M-mode to be 64 or 128, but got {m}")
if m == 64:
if (n < 8) or (n > 256) or (n % 8 != 0):
if self.b_dtype.width == 8 and self.b_major_mode == OperandMajorMode.MN:
if (n < 16) or (n > 256) or (n % 16 != 0):
raise OpError(
self,
f"expects the N-mode to satisfy 8 <= N <= 256 and N % 8 == 0, but got {n}",
f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}",
)
elif m == 128:
else:
if (n < 8) or (n > 256) or (n % 8 != 0):
raise OpError(
self,
@@ -216,11 +216,18 @@ class MmaOp(Tcgen05MmaOp):
else:
if m not in [128, 256]:
raise OpError(self, f"expects the M-mode to be 128 or 256, but got {m}")
if (n < 16) or (n > 256) or (n % 16 != 0):
raise OpError(
self,
f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}",
)
if self.b_dtype.width == 8 and self.b_major_mode == OperandMajorMode.MN:
if (n < 32) or (n > 256) or (n % 32 != 0):
raise OpError(
self,
f"expects the N-mode to satisfy 32 <= N <= 256 and N % 32 == 0, but got {n}",
)
else:
if (n < 16) or (n > 256) or (n % 16 != 0):
raise OpError(
self,
f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}",
)
def __str__(self) -> str:
return (
@@ -302,6 +309,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp):
admissible_archs = [
Arch.sm_100a,
Arch.sm_103a,
]
def __post_init__(self) -> None:

50
python/CuTeDSL/cutlass/cute/runtime.py
View File

@@ -10,7 +10,6 @@
# is strictly prohibited.
import ctypes
import os
import sys
from pathlib import Path
from functools import lru_cache
@@ -20,6 +19,7 @@ from typing import Union, Optional, Type, List
# MLIR modules imports
from cutlass._mlir import ir
from cutlass.base_dsl.env_manager import get_prefix_dsl_libs
import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cuda as _cuda_dialect
@@ -144,6 +144,7 @@ class _Tensor(Tensor):
self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor)
self._dlpack_data = self._tvm_ffi_tensor.__dlpack__()
self._dltensor_wrapper = None
self._assumed_align = assumed_align
self._is_dynamic = False
@@ -387,7 +388,16 @@ class _Tensor(Tensor):
return CoreTensor(values[0].value, self._dtype)
def __tvm_ffi_object__(self):
return self._tvm_ffi_tensor
try:
return self._tvm_ffi_tensor
except AttributeError:
raise DSLRuntimeError(
(
"runtime._Tensor is not a TVM-FFI tensor. "
"Enable TVM-FFI with `from_dlpack(..., enable_tvm_ffi=True)` "
"or `CUTE_DSL_ENABLE_TVM_FFI=1`."
)
)
def _get_cute_type_str(inp):
@@ -411,7 +421,8 @@ class _FakeCompactTensor(Tensor):
self._dtype = dtype
self._shape = shape
self._stride_order = stride_order or tuple(range(len(shape)))
self._memspace = memspace or AddressSpace.gmem
# cannot use memspace or AddressSpace.gmem because AddressSpace.generic is 0
self._memspace = memspace if memspace is not None else AddressSpace.gmem
self._assumed_align = assumed_align or -(-dtype.width // 8)
self._use_32bit_stride = use_32bit_stride
@@ -510,7 +521,8 @@ class _FakeTensor(Tensor):
self._dtype = dtype
self._shape = shape
self._stride = stride
self._memspace = memspace or AddressSpace.generic
# cannot use memspace or AddressSpace.generic because AddressSpace.generic is 0
self._memspace = memspace if memspace is not None else AddressSpace.gmem
self._assumed_align = assumed_align
if assumed_align is None:
# use the bytes width of the element dtype. The alignment is at least one byte align.
@@ -605,7 +617,7 @@ def make_fake_compact_tensor(
:param shape: Shape of the tensor.
:type shape: tuple[int, ...]
:param stride_order: Order in which strides (memory layout) are assigned to the tensor dimensions.
If None, the default layout is row-major. Otherwise, it should be a permutation of the dimension indices.
If None, the default layout is col-major. Otherwise, it should be a permutation of the dimension indices.
:type stride_order: tuple[int, ...], optional
:param memspace: Memory space where the fake tensor resides. Optional.
:type memspace: str, optional
@@ -644,6 +656,7 @@ def make_fake_compact_tensor(
use_32bit_stride=use_32bit_stride,
)
def make_fake_tensor(dtype, shape, stride, *, memspace=None, assumed_align=None):
"""
Create a fake tensor with the specified element type, shape, and stride.
@@ -859,21 +872,22 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]:
"""
def _get_cuda_dialect_runtime_path():
libs = os.environ.get("CUTE_DSL_LIBS")
if libs:
sep = ";" if sys.platform.startswith("win32") else ":"
for path in libs.split(sep):
if path.endswith("libcuda_dialect_runtime.so"):
return path
try:
# find package library from wheel package
pkg_base = Path(__file__).resolve().parent.parent
lib_path = pkg_base / "lib" / "libcuda_dialect_runtime.so"
if lib_path.is_file():
return str(lib_path)
except OSError:
libs = get_prefix_dsl_libs("CUTE_DSL")
if libs is None:
return None
# check if the separator is ; for windows
if sys.platform.startswith("win32") and ";" in libs:
libs = libs.split(";")
else:
libs = libs.split(":")
for path in libs:
if path.endswith("libcuda_dialect_runtime.so"):
return path
return None
libs = []
cuda_dialect_runtime_path = _get_cuda_dialect_runtime_path()
if cuda_dialect_runtime_path:

2
python/CuTeDSL/cutlass/cute/testing.py
View File

@@ -20,6 +20,7 @@ from typing import Type, Union, Callable, Optional, Dict, List, Any
import cuda.bindings.driver as cuda_driver
import cuda.bindings.runtime as cuda_runtime
import cutlass
import cutlass.base_dsl.jit_executor
from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op
@@ -233,6 +234,7 @@ def convert(src: cute.Tensor, dst: cute.Tensor):
src.shape[leading_mode] % elem_per_copy == 0
and dst.shape[leading_mode] % elem_per_copy == 0
)
_convert(src, dst, leading_mode, elem_per_copy)