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

v4.3.4 update. (#2892)

Browse Source
This commit is contained in:
Junkai-Wu
2025-12-22 00:49:12 +08:00
committed by GitHub
parent 331e2f451c
commit 7f5fe3edf1
31 changed files with 839 additions and 240 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -210,7 +210,7 @@ EnableTVMFFI = _dsl.EnableTVMFFI
# attach the TVM FFI ABI interface postprocessor to the DSL
from . import _tvm_ffi_args_spec_converter
_tvm_ffi_args_spec_converter.attach_args_spec_converter()
_tvm_ffi_args_spec_converter.attach_args_spec_converter(_dsl.CuTeDSL._get_dsl())
# Explicitly export all symbols for documentation generation
__all__ = [

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

@@ -395,8 +395,6 @@ def _tvm_ffi_args_spec_converter(
return params, kwargs_wrapper_spec
def attach_args_spec_converter():
"""Attach TVM FFI ABI interface postprocessor to the DSL."""
from .. import cutlass_dsl as _dsl
_dsl.CuTeDSL._get_dsl()._tvm_ffi_args_spec_converter = _tvm_ffi_args_spec_converter
def attach_args_spec_converter(dsl):
"""Attach TVM FFI ABI interface postprocessor to the DSL instance."""
dsl._tvm_ffi_args_spec_converter = _tvm_ffi_args_spec_converter

4
python/CuTeDSL/cutlass/cute/arch/elect.py
View File

@@ -10,7 +10,7 @@
# is strictly prohibited.
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import CuTeDSL, T, dsl_user_op
from cutlass.cutlass_dsl import BaseDSL, T, dsl_user_op
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir.dialects import nvvm, scf
@@ -69,7 +69,7 @@ def elect_one(*, loc=None, ip=None) -> IfOpRegion:
# Only one thread in the warp executes the code in this context
pass
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
is_thread_leader = nvvm.elect_sync(T.bool())
if_op = scf.IfOp(is_thread_leader, loc=loc, ip=ip)
return IfOpRegion(if_op.then_block, loc=loc, ip=ip)

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

@@ -11,7 +11,7 @@
from typing import Optional
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import CuTeDSL, T, if_generate, dsl_user_op
from cutlass.cutlass_dsl import BaseDSL, T, if_generate, dsl_user_op
from cutlass._mlir.dialects import nvvm
@@ -44,7 +44,7 @@ def mbarrier_init_fence(*, loc=None, ip=None) -> None:
"""
A fence operation that applies to the mbarrier initializations.
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
nvvm.fence_mbarrier_init(loc=loc, ip=ip)
@@ -63,7 +63,7 @@ def mbarrier_arrive_and_expect_tx(
the mbarrier is converted to a remote address in the peer CTA's
SMEM.
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
@@ -103,7 +103,7 @@ def mbarrier_expect_tx(
the mbarrier is converted to a remote address in the peer CTA's
SMEM.
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
@@ -138,7 +138,7 @@ def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None:
:param phase: The phase to wait for (either 0 or 1)
:type phase: Int
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
timeout_ns = 10000000
# This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX
@@ -164,7 +164,7 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo
:return: A boolean value indicating whether the wait operation was successful
:rtype: Boolean
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
return Boolean(
nvvm.mbarrier_wait_parity(
@@ -193,7 +193,7 @@ def mbarrier_conditional_try_wait(
:return: A boolean value indicating whether the wait operation was successful
:rtype: Boolean
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
return if_generate(
cond,
lambda: mbarrier_try_wait(mbar_ptr, phase, loc=loc, ip=ip),
@@ -225,7 +225,7 @@ def mbarrier_arrive(
"""
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
mbar_llvm_ptr.type,
@@ -259,7 +259,7 @@ def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> N
:param mbar_ptr: A pointer to the mbarrier in SMEM
:type mbar_ptr: Pointer
"""
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
nvvm.cp_async_mbarrier_arrive_shared(

7
python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py
View File

@@ -12,7 +12,8 @@
from cutlass.base_dsl.arch import Arch
from cutlass.base_dsl.common import DSLRuntimeError
from cutlass.cutlass_dsl import CuTeDSL, dsl_user_op
from cutlass.cutlass_dsl import BaseDSL, dsl_user_op
from cutlass._mlir import ir
from cutlass._mlir.dialects import builtin, arith, llvm, vector
@@ -53,7 +54,7 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
:return: The output 1D vector of bfloat16 with the same length as the input vector.
:rtype: 1D vector of bfloat16
"""
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if not arch in cvt_i8_bf16_intrinsic.supported_archs:
raise DSLRuntimeError(f"cvt_i8_bf16_intrinsic is not supported on {arch}")
src_pos = 0
@@ -130,7 +131,7 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
:return: The output 1D vector of bfloat16 with the same length as the input vector.
:rtype: 1D vector of bfloat16
"""
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if not arch in cvt_i4_bf16_intrinsic.supported_archs:
raise DSLRuntimeError(f"cvt_i4_bf16_intrinsic is not supported on {arch}")
src_pos = 0

40
python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py
View File

@@ -1305,6 +1305,46 @@ def exp_packed_f32x2(
return exp2(b[0], loc=loc, ip=ip), exp2(b[1], loc=loc, ip=ip)
@dsl_user_op
def griddepcontrol_wait(*, loc=None, ip=None) -> None:
"""
This instruction is used to wait for the previous kernel's grid ending
(all blocks of the previous kernel have finished and memflushed), i.e.,
the instruction after this instruction will not be issued until the previous
grid has finished.
"""
llvm.inline_asm(
res=None,
operands_=[],
asm_string="griddepcontrol.wait;",
constraints="",
has_side_effects=True,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def griddepcontrol_launch_dependents(*, loc=None, ip=None) -> None:
"""
Issuing the launch_dependents instruction hints a dependent kernel to launch earlier.
launch_dependents doesn't impact the functionality but the performance:
Launching a dependent kernel too early can compete with current kernels,
while launching too late can lead to a long latency.
"""
llvm.inline_asm(
res=None,
operands_=[],
asm_string="griddepcontrol.launch_dependents;",
constraints="",
has_side_effects=True,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@dsl_user_op
def cvt_f4e2m1_f16(src, *, loc=None, ip=None):

14
python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py
View File

@@ -15,7 +15,7 @@ from typing import Optional, Type
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import CuTeDSL
from cutlass.cutlass_dsl import BaseDSL
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
@@ -146,7 +146,7 @@ class CopyBulkTensorTileG2SOp(TmaCopyOp):
self, "expects the 'cta_group' parameter to be a CtaGroup instance"
)
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
arch: Arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
@@ -263,7 +263,7 @@ class CopyBulkTensorTileG2SMulticastOp(TmaCopyOp):
self, "expects the 'cta_group' parameter to be a CtaGroup instance"
)
# Arch verification
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
@@ -386,7 +386,7 @@ class CopyBulkTensorTileS2GOp(TmaCopyOp):
def __post_init__(self):
# Arch verification
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
@@ -561,7 +561,7 @@ class CopyBulkG2SOp(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
arch: Arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
@@ -646,7 +646,7 @@ class CopyBulkG2SMulticastOp(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
arch: Arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
@@ -740,7 +740,7 @@ class CopyBulkS2GOp(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
arch: Arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,

8
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py
View File

@@ -15,7 +15,7 @@ from typing import Type
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import CuTeDSL
from cutlass.cutlass_dsl import BaseDSL
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
@@ -113,7 +113,7 @@ class _LdBase(CopyOp):
:raises OpError: If pack parameter is not a Pack instance
"""
# Arch verification
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if arch not in self.admissible_archs:
raise OpError(
self,
@@ -416,7 +416,7 @@ class _StBase(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if arch not in self.admissible_archs:
raise OpError(
self,
@@ -625,7 +625,7 @@ class _S2TCopyBase(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if not arch.is_family_of(Arch.sm_100f):
raise OpError(
self,

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

@@ -15,7 +15,7 @@ from typing import Type, Any
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import CuTeDSL, T
from cutlass.cutlass_dsl import BaseDSL, T
import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
@@ -162,7 +162,7 @@ class MmaOp(Tcgen05MmaOp):
def __post_init__(self) -> None:
# Verify arch
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if arch not in self.admissible_archs:
raise OpError(
self,
@@ -314,7 +314,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp):
def __post_init__(self) -> None:
# Verify arch
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if arch not in self.admissible_archs:
raise OpError(
self,
@@ -471,7 +471,7 @@ class SparseMmaOp(Tcgen05MmaOp):
def __post_init__(self) -> None:
# Verify arch
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if arch not in self.admissible_archs:
raise OpError(
self,

4
python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py
View File

@@ -15,7 +15,7 @@ from typing import Type, Any
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import CuTeDSL, T
from cutlass.cutlass_dsl import BaseDSL, T
import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
@@ -130,7 +130,7 @@ class MmaOp(WarpGroupMmaOp):
def __post_init__(self) -> None:
# Verify arch
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if not arch == Arch.sm_90a:
raise OpError(
self,

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

@@ -925,7 +925,17 @@ def load_module(file_path: str, *, enable_tvm_ffi: bool = True):
if enable_tvm_ffi:
import tvm_ffi
return tvm_ffi.load_module(file_path)
try:
# keep_module_alive=False means the module will be unloaded
# after the returned module goes out of scope, this is useful
# for frequent loading and unloading of modules. The only requirement
# is that the module do not return object that have deleter in the module
# and the returned object lives longer than the module.
# DSL functions to not have such issue so it is desirable to set this to False.
return tvm_ffi.load_module(file_path, keep_module_alive=False)
except TypeError:
# compatible with tvm-ffi < 0.1.6
return tvm_ffi.load_module(file_path)
else:
raise DSLRuntimeError(
"Unimplemented, please load the module with enable_tvm_ffi=True."

4
python/CuTeDSL/cutlass/cute/tensor.py
View File

@@ -20,7 +20,7 @@ from cutlass.cutlass_dsl import (
T,
cutlass_arith,
_binary_op_type_promote,
CuTeDSL,
BaseDSL,
)
from cutlass._mlir import ir
import cutlass._mlir.dialects.cute as _cute_ir
@@ -1776,7 +1776,7 @@ class TensorSSA(cutlass_arith.ArithValue):
fast_cvt_func = cvt_i8_bf16_intrinsic
elif src_dtype == Int4 and dtype == BFloat16:
fast_cvt_func = cvt_i4_bf16_intrinsic
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if fast_cvt_func is not None and arch in fast_cvt_func.supported_archs:
res_vect = fast_cvt_func(src, size(self.shape), loc=loc, ip=ip)
else:

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

@@ -407,7 +407,7 @@ def benchmark(
To use CUDA graphs, the callable must be a compiled @cute.jit annotated function.
When using CUDA graphs, the kernel must be launched in a non-default stream.
:param callable: The function to benchmark
:param callable: The function to benchmark. For jit function, it must be compiled functions.
:type callable: Callable
:param warmup_iterations: Number of warmup iterations, defaults to 10
:type warmup_iterations: int, optional
@@ -475,15 +475,6 @@ def benchmark(
elapsed_time = float("nan")
if use_cuda_graphs:
# Check if the callable is a JitCompiledFunction or JitExecutor
# These are functions that can be called to launch kernels
compiled_types = (
cutlass.base_dsl.jit_executor.JitCompiledFunction,
cutlass.base_dsl.jit_executor.JitExecutor,
)
if not isinstance(callable, compiled_types):
raise TypeError("Function must be precompiled to be used with CUDA Graphs")
# Check if the stream is a non-default stream
if int(stream) == int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT):
raise ValueError(