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

v4.1 release

Browse Source
This commit is contained in:
Junkai-Wu
2025-07-03 20:07:53 +08:00
committed by GitHub
parent b995f93317
commit a1aaf2300a
155 changed files with 18407 additions and 6068 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -68,7 +68,9 @@ from .core import (
select,
front,
is_major,
leading_dim,
find,
find_if,
coalesce,
group_modes,
cosize,
@@ -221,7 +223,9 @@ __all__ = [
"select",
"front",
"is_major",
"leading_dim",
"find",
"find_if",
"coalesce",
"group_modes",
"cosize",

8
python/CuTeDSL/cutlass/cute/arch/__init__.py
View File

@@ -25,12 +25,13 @@ __all__ = [
#
# mbar.py
#
"mbarrier_init_arrive_cnt",
"mbarrier_init",
"mbarrier_init_fence",
"mbarrier_init_tx_bytes",
"mbarrier_arrive_and_expect_tx",
"mbarrier_expect_tx",
"mbarrier_wait",
"mbarrier_try_wait",
"conditional_mbarrier_try_wait",
"mbarrier_conditional_try_wait",
"mbarrier_arrive",
#
# nvvm_wrappers.py
@@ -51,6 +52,7 @@ __all__ = [
"shuffle_sync_down",
"shuffle_sync_bfly",
"barrier",
"barrier_arrive",
"sync_threads",
"sync_warp",
"fence_acq_rel_cta",

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

@@ -69,7 +69,16 @@ def elect_one(*, loc=None, ip=None) -> IfOpRegion:
pass
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(arch, ["sm_90", "sm_90a", "sm_100a"], "arch")
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
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)

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

@@ -8,6 +8,7 @@
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
from typing import Optional
from cutlass.cutlass_dsl import CuTeDSL, T, if_generate, dsl_user_op
@@ -26,7 +27,7 @@ from ...impl_utils import check_value_in
@dsl_user_op
def mbarrier_init_arrive_cnt(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None:
def mbarrier_init(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None:
"""
Initializes a mbarrier with the specified thread arrival count.
@@ -46,16 +47,25 @@ def mbarrier_init_fence(*, loc=None, ip=None) -> None:
A fence operation that applies to the mbarrier initializations.
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(arch, ["sm_90", "sm_90a", "sm_100a"], "arch")
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
nvvm.fence_mbarrier_init(loc=loc, ip=ip)
@dsl_user_op
def mbarrier_init_tx_bytes(
def mbarrier_arrive_and_expect_tx(
mbar_ptr: Pointer, bytes: Int, peer_cta_rank_in_cluster=None, *, loc=None, ip=None
) -> None:
"""
Initializes a mbarrier with the specified number of transaction bytes.
Arrives on a mbarrier and expects a specified number of transaction bytes.
:param mbar_ptr: A pointer to the mbarrier in SMEM
:type mbar_ptr: Pointer
@@ -66,7 +76,16 @@ def mbarrier_init_tx_bytes(
SMEM.
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(arch, ["sm_90", "sm_90a", "sm_100a"], "arch")
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
@@ -91,6 +110,56 @@ def mbarrier_init_tx_bytes(
)
@dsl_user_op
def mbarrier_expect_tx(
mbar_ptr: Pointer, bytes: Int, peer_cta_rank_in_cluster=None, *, loc=None, ip=None
) -> None:
"""
Expects a specified number of transaction bytes without an arrive.
:param mbar_ptr: A pointer to the mbarrier in SMEM
:type mbar_ptr: Pointer
:param bytes: The number of transaction bytes
:type bytes: Int
:param peer_cta_rank_in_cluster: An optional CTA rank in cluster. If provided, the pointer to
the mbarrier is converted to a remote address in the peer CTA's
SMEM.
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
mbar_llvm_ptr = nvvm.mapa(
mbar_llvm_ptr.type,
mbar_llvm_ptr,
Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
space = nvvm.MBarrierSpaceKind.CLUSTER
else:
space = nvvm.MBarrierSpaceKind.CTA
nvvm.mbarrier_txn(
mbar_llvm_ptr,
Int32(bytes).ir_value(loc=loc, ip=ip),
kind=nvvm.MBarrierTxnKind.EXPECT_TX,
space=space,
loc=loc,
ip=ip,
)
@dsl_user_op
def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None:
"""
@@ -102,7 +171,16 @@ def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None:
:type phase: Int
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(arch, ["sm_90", "sm_90a", "sm_100a"], "arch")
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
timeout_ns = 10000000
# This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX
@@ -129,7 +207,16 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo
:rtype: Boolean
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(arch, ["sm_90", "sm_90a", "sm_100a"], "arch")
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
return Boolean(
nvvm.mbarrier_wait_parity(
@@ -144,7 +231,7 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo
@dsl_user_op
def conditional_mbarrier_try_wait(
def mbarrier_conditional_try_wait(
cond, mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None
) -> Boolean:
"""
@@ -159,7 +246,16 @@ def conditional_mbarrier_try_wait(
:rtype: Boolean
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(arch, ["sm_90", "sm_90a", "sm_100a"], "arch")
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
return if_generate(
cond,
lambda: mbarrier_try_wait(mbar_ptr, phase, loc=loc, ip=ip),
@@ -171,7 +267,11 @@ def conditional_mbarrier_try_wait(
@dsl_user_op
def mbarrier_arrive(
mbar_ptr: Pointer, peer_cta_rank_in_cluster: Int = None, *, loc=None, ip=None
mbar_ptr: Pointer,
peer_cta_rank_in_cluster: Optional[Int] = None,
*,
loc=None,
ip=None,
) -> None:
"""
Arrives on an mbarrier.
@@ -185,7 +285,16 @@ def mbarrier_arrive(
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(arch, ["sm_90", "sm_90a", "sm_100a"], "arch")
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
mbar_llvm_ptr.type,

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

@@ -225,6 +225,25 @@ def barrier(*, barrier_id=None, number_of_threads=None, loc=None, ip=None) -> No
barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip
)
@dsl_user_op
def barrier_arrive(
*, barrier_id=None, number_of_threads=None, loc=None, ip=None
) -> None:
if barrier_id is not None:
barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip)
if number_of_threads is None:
raise ValueError(
"barrier_arrive needs pass number_of_threads to arrive the barrier",
)
number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip)
nvvm.barrier_arrive(
barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip
)
@dsl_user_op
def sync_threads(*, loc=None, ip=None) -> None:
"""
@@ -545,3 +564,20 @@ def exp2(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
asm_dialect=llvm.AsmDialect.AD_ATT,
)
)
# TODO: add `fastmath` flag for this op
@dsl_user_op
def exp(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
LOG2_E = 1.4426950408889634
return exp2(a * LOG2_E, loc=loc, ip=ip)
# TODO: add `fastmath` flag for this op
@dsl_user_op
def exp_packed_f32x2(
a: Tuple[Float32, Float32], *, loc=None, ip=None
) -> Tuple[Float32, Float32]:
LOG2_E = Float32(1.4426950408889634)
b = mul_packed_f32x2(a, (LOG2_E, LOG2_E), loc=loc, ip=ip)
return exp2(b[0], loc=loc, ip=ip), exp2(b[1], loc=loc, ip=ip)

685
python/CuTeDSL/cutlass/cute/core.py
View File

File diff suppressed because it is too large Load Diff

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

@@ -26,7 +26,7 @@ __all__ = [
#
# helpers.py
#
"make_tma_tile_atom",
"make_tiled_tma_atom",
"tma_partition",
"create_tma_multicast_mask",
"prefetch_descriptor",

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

@@ -127,7 +127,12 @@ class CopyBulkTensorTileG2SOp(CopyOp):
cta_group: CtaGroup = CtaGroup.ONE
admissible_archs = ["sm_90", "sm_90a", "sm_100a"]
admissible_archs = [
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
]
def __post_init__(self) -> None:
if not isinstance(self.cta_group, CtaGroup):
@@ -159,7 +164,7 @@ class CopyBulkTensorTileG2SOp(CopyOp):
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "CopyBulkTensorTileG2SNonExecTrait":
raise NotImplementedError(
"Use cpasync.make_tma_tile_atom to obtain a copy Atom for TMA"
"Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA"
)
def _to_ir(self) -> _cute_nvgpu_ir.TiledTmaLoadEnum:
@@ -224,7 +229,12 @@ class CopyBulkTensorTileG2SMulticastOp(CopyOp):
cta_group: CtaGroup = CtaGroup.ONE
admissible_archs = ["sm_90", "sm_90a", "sm_100a"]
admissible_archs = [
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
]
def __post_init__(self):
if not isinstance(self.cta_group, CtaGroup):
@@ -256,7 +266,7 @@ class CopyBulkTensorTileG2SMulticastOp(CopyOp):
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "CopyBulkTensorTileG2SMulticastNonExecTrait":
raise NotImplementedError(
"Use cpasync.make_tma_tile_atom to obtain a copy Atom for TMA"
"Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA"
)
def _to_ir(self) -> _cute_nvgpu_ir.TiledTmaLoadEnum:
@@ -326,7 +336,12 @@ class CopyBulkTensorTileS2GOp(CopyOp):
This Operation uses TMA in the ``.tile`` mode.
"""
admissible_archs = ["sm_90", "sm_90a", "sm_100a"]
admissible_archs = [
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
]
def __post_init__(self):
# Arch verification
@@ -345,7 +360,7 @@ class CopyBulkTensorTileS2GOp(CopyOp):
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "CopyBulkTensorTileS2GTrait":
raise NotImplementedError(
"Use cpasync.make_tma_tile_atom to obtain a copy Atom for TMA"
"Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA"
)

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

@@ -29,14 +29,14 @@ from .copy import (
@dsl_user_op
def make_tma_tile_atom(
def make_tiled_tma_atom(
op: Union[
CopyBulkTensorTileG2SOp,
CopyBulkTensorTileG2SMulticastOp,
CopyBulkTensorTileS2GOp,
],
gmem_tensor: Tensor,
smem_layout: Layout,
smem_layout: Union[Layout, core.ComposedLayout],
cta_tiler: Tiler,
num_multicast: int = 1,
*,
@@ -45,7 +45,7 @@ def make_tma_tile_atom(
ip=None,
) -> Tuple[core.CopyAtom, Tensor]:
"""
Makes a TMA Copy Atom in the ``.tile`` mode to copy tiles of a GMEM tensor to/from and SMEM
Makes a TMA Copy Atom in the ``.tile`` mode to copy tiles of a GMEM tensor to/from SMEM
buffer with the given Layout.
Given
@@ -71,7 +71,7 @@ def make_tma_tile_atom(
:param gmem_tensor: The GMEM tensor involved in the Copy
:type gmem_tensor: Tensor
:param smem_layout: The SMEM layout to construct the Copy Atom for
:type smem_layout: Layout
:type smem_layout: Union[Layout, core.ComposedLayout]
:param cta_tiler: The CTA Tiler to use
:type cta_tiler: Tiler
:param num_multicast: The multicast factor
@@ -94,6 +94,12 @@ def make_tma_tile_atom(
ip=ip,
)
# Wrap smem_layout in a composed layout to make it a TMA-friendly layout
if isinstance(smem_layout, Layout):
smem_layout = core.make_composed_layout(
core.make_swizzle(0, 4, 3), 0, smem_layout
)
if isinstance(op, CopyBulkTensorTileG2SOp):
if num_multicast != 1:
raise ValueError(

102
python/CuTeDSL/cutlass/cute/nvgpu/helpers.py
View File

@@ -34,7 +34,7 @@ from .cpasync.copy import (
@dsl_user_op
def make_tma_tile_atom_A(
def make_tiled_tma_atom_A(
op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
gmem_tensor: Tensor,
smem_layout: Layout,
@@ -46,6 +46,51 @@ def make_tma_tile_atom_A(
loc=None,
ip=None,
) -> Tuple[core.CopyAtom, Tensor]:
"""
Makes a TMA Copy atom mapping to ``.tile`` mode for ``cp.async.bulk.tensor`` PTX operation
accounting for the MK projections of the TiledMMA for A tensor loads.
Given
- a GMEM tensor
- a SMEM layout
- a MMA Tiler
- a TiledMma
- a Cluster-level shape
this function figures out the bulk tensor asynchronous copy instruction to use with the maximum
"TMA vector length" to copy tiles of the GMEM tensor to an SMEM buffer with the provided
layout and consistent with the provided Tiler & tiled_mma (considering the M-mode & K-mode).
The Cluster-level shape is used to determine the multicast factor across the N-mode for A tensor loads.
This function returns two results:
1. the Copy Atom
2. the so-called TMA tensor used to map logical coordinates of the GMEM tensor to coordinates
that the TMA unit can consume. TMA tensors have so-called basis stride elements so that the
associated layout can output coordinates. Otherwise, TMA tensors can be partitioned
similarly to any other CuTe tensors using the algebra.
:param op: The Copy Operation to construct an Atom for
:type op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp]
:param gmem_tensor: The GMEM tensor to be loaded by this copy atom
:type gmem_tensor: Tensor
:param smem_layout: Shared memory layout to load the tensor into (PDSL)
:type smem_layout: Layout
:param mma_tiler_mnk: The MMA Tiler shape (TILE_M, TILE_N, TILE_K) in MNK dimensions
:type mma_tiler_mnk: Shape
:param tiled_mma: The TiledMMA that will consume the load as operands
:type tiled_mma: core.TiledMma
:param cluster_shape_vmnk: The Cluster-level shape in VMNK dimensions
:type cluster_shape_vmnk: Shape
:param internal_type: An optional parameter for the internal data type to when element
type does not match the copy type
:type internal_type: Type[Numeric]
:return: A copy atom for this operation and the associated TMA coord tensor
:rtype: Tuple[core.CopyAtom, Tensor]
"""
if internal_type is not None:
if not isinstance(internal_type, NumericMeta):
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
@@ -54,7 +99,7 @@ def make_tma_tile_atom_A(
op,
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
"op",
"make_tma_tile_atom_A",
"make_tiled_tma_atom_A",
)
ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
@@ -94,7 +139,7 @@ def make_tma_tile_atom_A(
@dsl_user_op
def make_tma_tile_atom_B(
def make_tiled_tma_atom_B(
op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
gmem_tensor: Tensor,
smem_layout: Layout,
@@ -106,6 +151,51 @@ def make_tma_tile_atom_B(
loc=None,
ip=None,
) -> Tuple[core.CopyAtom, Tensor]:
"""
Makes a TMA Copy atom mapping to ``.tile`` mode for ``cp.async.bulk.tensor`` PTX operation
accounting for the NK projections of the TiledMMA for B tensor loads.
Given
- a GMEM tensor
- a SMEM layout
- a MMA Tiler
- a TiledMma
- a Cluster-level shape
this function figures out the bulk tensor asynchronous copy instruction to use with the maximum
"TMA vector length" to copy tiles of the GMEM tensor to an SMEM buffer with the provided
layout and consistent with the provided Tiler & tiled_mma (considering the N-mode & K-mode).
The Cluster-level shape is used to determine the multicast factor across the M-mode for B tensor loads.
This function returns two results:
1. the Copy Atom
2. the so-called TMA tensor used to map logical coordinates of the GMEM tensor to coordinates
that the TMA unit can consume. TMA tensors have so-called basis stride elements so that the
associated layout can output coordinates. Otherwise, TMA tensors can be partitioned
similarly to any other CuTe tensors using the algebra.
:param op: The Copy Operation to construct an Atom for
:type op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp]
:param gmem_tensor: The GMEM tensor to be loaded by this copy atom
:type gmem_tensor: Tensor
:param smem_layout: Shared memory layout to load the tensor into (PDSL)
:type smem_layout: Layout
:param mma_tiler_mnk: The MMA Tiler shape (TILE_M, TILE_N, TILE_K) in MNK dimensions
:type mma_tiler_mnk: Shape
:param tiled_mma: The TiledMMA that will consume the load as operands
:type tiled_mma: core.TiledMma
:param cluster_shape_vmnk: The Cluster-level shape in VMNK dimensions
:type cluster_shape_vmnk: Shape
:param internal_type: An optional parameter for the internal data type to when element
type does not match the copy type
:type internal_type: Type[Numeric]
:return: A Copy Atom for this Operation and the associated TMA tensor
:rtype: Tuple[core.CopyAtom, Tensor]
"""
if internal_type is not None:
if not isinstance(internal_type, NumericMeta):
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
@@ -114,7 +204,7 @@ def make_tma_tile_atom_B(
op,
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
"op",
"make_tma_tile_atom_B",
"make_tiled_tma_atom_B",
)
ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
@@ -154,6 +244,6 @@ def make_tma_tile_atom_B(
__all__ = [
"make_tma_tile_atom_A",
"make_tma_tile_atom_B",
"make_tiled_tma_atom_A",
"make_tiled_tma_atom_B",
]

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

@@ -98,7 +98,10 @@ class _LdBase(CopyOp):
repeat: Repetition = Repetition.x1
pack: Pack = Pack.NONE
admissible_archs = ["sm_100a"]
admissible_archs = [
"sm_100a",
"sm_100f",
]
def __post_init__(self) -> None:
# Arch verification
@@ -284,7 +287,10 @@ class _StBase(CopyOp):
repeat: Repetition
unpack: Unpack = Unpack.NONE
admissible_archs = ["sm_100a"]
admissible_archs = [
"sm_100a",
"sm_100f",
]
def __post_init__(self) -> None:
# Arch verification

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

@@ -136,7 +136,10 @@ class MmaOp(MmaOp):
a_major_mode: OperandMajorMode
b_major_mode: OperandMajorMode
admissible_archs = ["sm_100a"]
admissible_archs = [
"sm_100a",
"sm_100f",
]
def __post_init__(self) -> None:
# Verify arch

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

@@ -339,10 +339,10 @@ class MmaF8Op(MmaOp):
"expects the 'b_dtype' Op parameter to be one of Float8E5M2 or Float8E4M3FN",
)
# Accumulator data type verification
if self.acc_dtype != Float32:
if self.acc_dtype not in [Float16, Float32]:
raise OpError(
self,
"expects the 'acc_dtype' Op parameter to be Float32",
"expects the 'acc_dtype' Op parameter to be one of Float16 or Float32",
)
# Verify the instruction shape
instruction_k = 32

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

@@ -20,6 +20,7 @@ from typing import Union
from cutlass._mlir import ir
import cutlass._mlir.dialects.cute as _cute_ir
from cutlass.base_dsl.dsl import is_dynamic_expression
from cutlass.cutlass_dsl import TensorFormat, JitArgAdapterRegistry
# Local modules imports
@@ -45,7 +46,8 @@ from .typing import (
BFloat16,
Float8E5M2,
)
from .core import find, _Tensor as CoreTensor
from . import core
from .core import _Tensor as CoreTensor
class _Pointer(Pointer):
@@ -131,6 +133,9 @@ class _Pointer(Pointer):
def memspace(self):
return self._addr_space
def align(self, min_align: int, *, loc=None, ip=None) -> Pointer:
raise NotImplementedError("align is not supported in runtime")
def verify(self, expected_py_type):
if expected_py_type is Pointer:
return True
@@ -361,7 +366,7 @@ class _Tensor(Tensor):
* If nested leading dimensions are found, returns a tuple of indices
* If no leading dimension is found, returns None
"""
return find(1, self.stride, exclude_when=(1, self.shape))
return core.leading_dim(self.shape, self.stride)
def fill(self, value: Numeric):
raise TypeError(f"fill function is not supported in runtime")
@@ -479,12 +484,8 @@ class TensorAdapter:
Convert a DLPack protocol supported tensor/array to a cute tensor.
"""
# Need reference these capsules to avoid being garbage collected
tensor_capsules = []
def __init__(self, arg):
self._arg = from_dlpack(arg).mark_layout_dynamic()
self.tensor_capsules.append(self._arg)
def __new_from_mlir_values__(self, values):
return self._arg.__new_from_mlir_values__(values)

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

@@ -9,29 +9,26 @@
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
import random
import numpy as np
import functools
import hashlib
from cutlass.cutlass_dsl import (
const,
T,
CuTeDSL,
BaseDSL,
t,
Constexpr,
detect_gpu_arch,
)
import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.ir as ir
from cutlass._mlir.dialects import nvvm, cf, vector, builtin
from cutlass.cute import core
from cutlass.cute import nvgpu
from typing import Type
import inspect
import logging
import os
from enum import Enum
from inspect import isclass
from itertools import product
from time import time
from typing import Any, Callable, Dict, List, Optional, Type, Union
import cuda.bindings.driver as cuda_driver
import cuda.bindings.runtime as cuda_runtime
import numpy as np
import cutlass._mlir.ir as ir
import cutlass.base_dsl.jit_executor
import cutlass.cute as cute
from cutlass._mlir.dialects import builtin, cf, nvvm, vector
from cutlass.cute import core, nvgpu
from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, t
def assert_(cond, msg=None):
@@ -248,9 +245,10 @@ def sample_pytest(rand_cfg=None):
import functools
import os
import random
import pytest
import sys
import pytest
seed, sample_ratio = rand_cfg
random.seed(seed)
@@ -270,3 +268,311 @@ def sample_pytest(rand_cfg=None):
return wrapper
return decorator
#########################################
# Benchmarking utilities
#########################################
class JitArguments:
"""
A type to hold both args and kwargs for passing to a kernel while benchmarking.
"""
def __init__(self, *args, **kwargs):
self.args = args
self.kwargs = kwargs
def _cuda_success(
err: Union[tuple, cuda_runtime.cudaError_t, cuda_driver.CUresult], message: str
):
"""
Helper function to check CUDA API errors.
"""
if isinstance(err, tuple):
_cuda_success(err[0], message)
elif isinstance(err, cuda_runtime.cudaError_t):
error_message = cuda_runtime.cudaGetErrorString(err)[1].decode("utf-8")
if err != cuda_runtime.cudaError_t.cudaSuccess:
raise RuntimeError(f"{message} : {error_message}")
elif isinstance(err, cuda_driver.CUresult):
if err != cuda_driver.CUresult.CUDA_SUCCESS:
error_message = cuda_driver.cuGetErrorString(err)[1].decode("utf-8")
raise RuntimeError(f"{message} : {error_message}")
else:
raise TypeError(
f"{err} is an unexpected type : it should be a cudaError_t or CUresult"
)
def _does_kernel_use_stream(
kernel: Callable, stream: cuda_driver.CUstream, *args, **kwargs
):
"""
This function checks if the kernel uses the provided non-default stream.
It does this by capturing the stream and then checking if any kernels were launched.
:param kernel: The kernel to check
:type kernel: Callable
:param stream: The stream to check
:type stream: cuda_driver.CUstream
:return: True if the kernel uses the stream, False otherwise
:rtype: bool
"""
assert int(stream) != int(
cuda_driver.CUstream_flags.CU_STREAM_DEFAULT
), "Stream must be a non-default stream"
err = cuda_runtime.cudaStreamBeginCapture(
stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal
)
_cuda_success(err, "Error on stream capture")
kernel(*args, **kwargs)
err, graph = cuda_runtime.cudaStreamEndCapture(stream)
_cuda_success(err, "Error on stream capture")
# Get number of nodes in warmup graph to check it matches what is expected
err, _, num_nodes = cuda_runtime.cudaGraphGetNodes(graph)
_cuda_success(err, "Error on querying graph")
return num_nodes > 0
def benchmark(
callable: Callable,
*,
warmup_iterations: int = 10,
profiling_iterations: int = 100,
stream: Optional[cuda_driver.CUstream] = None,
kernel_arguments: Optional[JitArguments] = None,
workspace_generator: Optional[Callable[[], JitArguments]] = None,
workspace_count: int = 1,
use_cuda_graphs: bool = False,
) -> float:
"""Benchmarks a callable function with the specified parameters.
For example,
.. code-block:: python
from cutlass.cute.testing import benchmark
@cute.jit
def user_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor, stream: cuda_driver.CUstream):
# contents of the function
pass
time_us = benchmark(user_function, kernel_arguments=JitArguments(a, b, c, stream)
warmup_iterations=10, profiling_iterations=100
stream=stream)
To prevent skewing results by repeately accessing the L2 cache, use the workspace_count and workspace_generator
parameters to cycle through a number of different workspaces.
.. code-block:: python
from cutlass.cute.testing import benchmark
@cute.jit
def user_function(a: cute.Tensor, b: cute.Tensor, c: cute.Tensor):
# contents of the function
pass
def workspace_generator():
# create a, b, and c
return JitArguments(a, b, c)
time_us = benchmark(user_function,
workspace_generator=workspace_generator,
workspace_count=10,
warmup_iterations=10000,
profiling_iterations=1000)
To benchmark you may always configure the function being profiled (callable), the warmup iterations, and
the number of profiling iterations.
Whenever the kernel being benchmarked runs in a non-default stream, the stream must be provided through the stream parameter.
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
:type callable: Callable
:param warmup_iterations: Number of warmup iterations, defaults to 10
:type warmup_iterations: int, optional
:param profiling_iterations: Number of benchmark iterations, defaults to 100
:type profiling_iterations: int, optional
:param stream: Stream kernel is launched in, defaults to CUDA stream default
:type stream: CUstream, None
:param kernel_arguments: Kernel arguments to launch callable with, defaults to None
:type kernel_arguments: JitArguments, None
:param workspace_generator: Function that returns kernel arguments, defaults to None
:type workspace_generator: Callable
:param workspace_count: Number of workspaces (arguments) to loop through, looping through enough workspaces will keep the L2 cache cold
:type workspace_count: int, optional
:param use_cuda_graphs: Whether to use cuda graphs, defaults to False
:type use_cuda_graphs: bool, optional
:return: The benchmark time in microseconds
:rtype: float
"""
if stream is None:
stream = cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT)
if workspace_count < 1:
raise ValueError("workspace_count must be at least 1")
time_us = float("nan")
if workspace_generator == None:
# If no workspace generator is provided, we need a single workspace
if workspace_count != 1:
raise ValueError("Need a single workspace if not providing a generator")
# If no workspace generator is provided, we need a kernel_argument
if kernel_arguments == None:
raise ValueError(
"Please pass a kernel argument if not providing a generator"
)
workspace_generator = lambda: kernel_arguments
workspaces = [workspace_generator() for _ in range(workspace_count)]
for workspace in workspaces:
if type(workspace) != JitArguments:
raise TypeError(
"workspace_generator and/or kernel_arguments should use JitArguments type"
)
def _loop_and_call_kernel(iterations: int, workspace_index: int = 0):
for _ in range(iterations):
current_workspace = workspaces[workspace_index]
callable(*current_workspace.args, **current_workspace.kwargs)
workspace_index = (workspace_index + 1) % workspace_count
return workspace_index
# Create CUDA events for timing
err, start_event = cuda_driver.cuEventCreate(
cuda_driver.CUevent_flags.CU_EVENT_DEFAULT
)
_cuda_success(err, "Error on creating event")
err, end_event = cuda_driver.cuEventCreate(
cuda_driver.CUevent_flags.CU_EVENT_DEFAULT
)
_cuda_success(err, "Error on creating event")
elapsed_time = float("nan")
if use_cuda_graphs:
# Check if the callable is a JitExecutor
if not isinstance(callable, cutlass.base_dsl.jit_executor.JitExecutor):
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(
"Measuring with CUDA Graphs requires executing in a non-default stream"
)
workspace_index = 0
# Capture warmup graph
err = cuda_runtime.cudaStreamBeginCapture(
stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal
)
_cuda_success(err, "Error on stream capture")
workspace_index = _loop_and_call_kernel(warmup_iterations)
err, gwarm = cuda_runtime.cudaStreamEndCapture(stream)
_cuda_success(err, "Error on stream capture")
# Get number of nodes in warmup graph to check it matches what is expected
err, _, num_nodes = cuda_runtime.cudaGraphGetNodes(gwarm)
_cuda_success(err, "Error on querying graph")
# Assertion is >= since we may launch multiple kernels in one host function
if num_nodes < warmup_iterations:
raise ValueError(
f"CUDA stream passed to benchmark does not match the stream the kernel was launched in"
)
# Capture profiling graph
err = cuda_runtime.cudaStreamBeginCapture(
stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal
)
_cuda_success(err, "Error on stream capture")
_loop_and_call_kernel(profiling_iterations, workspace_index)
err, gprofile = cuda_runtime.cudaStreamEndCapture(stream)
_cuda_success(err, "Error on stream capture")
# Instantiate graphs
err, gwarm = cuda_runtime.cudaGraphInstantiate(gwarm, 0)
_cuda_success(err, "Error on graph instantiation")
err, gprofile = cuda_runtime.cudaGraphInstantiate(gprofile, 0)
_cuda_success(err, "Error on graph instantiation")
# Launch warmup graph
err = cuda_runtime.cudaGraphLaunch(gwarm, stream)
_cuda_success(err, "Error on graph launch")
# Record start time
err = cuda_driver.cuEventRecord(start_event, stream)
_cuda_success(err, "Error on recording event")
# Launch profiling graph
err = cuda_runtime.cudaGraphLaunch(gprofile, stream)
_cuda_success(err, "Error on graph launch")
# Record end time
err = cuda_driver.cuEventRecord(end_event, stream)
_cuda_success(err, "Error on recording event")
err = cuda_driver.cuEventSynchronize(end_event)
_cuda_success(err, "Error on synchronizing event")
# Get elapsed time
err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event)
_cuda_success(err, "Error on querying event")
# Destroy graphs
err = cuda_runtime.cudaGraphExecDestroy(gwarm)
_cuda_success(err, "Error on destroying graph")
err = cuda_runtime.cudaGraphExecDestroy(gprofile)
_cuda_success(err, "Error on destroying graph")
else:
if int(stream) != int(
cuda_driver.CUstream_flags.CU_STREAM_DEFAULT
) and not _does_kernel_use_stream(
callable, stream, *workspaces[0].args, **workspaces[0].kwargs
):
raise ValueError(
"CUDA stream passed to benchmark does not match the stream the kernel was launched in"
)
# Not using graphs
# Warmup
workspace_index = _loop_and_call_kernel(warmup_iterations)
# Record start event
err = cuda_driver.cuEventRecord(start_event, stream)
_cuda_success(err, "Error on recording event")
_loop_and_call_kernel(profiling_iterations, workspace_index)
# Record end event
err = cuda_driver.cuEventRecord(end_event, stream)
_cuda_success(err, "Error on recording event")
# Synchronize end event
err = cuda_driver.cuEventSynchronize(end_event)
_cuda_success(err, "Error on synchronizing event")
err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event)
_cuda_success(err, "Error on querying event")
# Destroy events
err = cuda_driver.cuEventDestroy(start_event)
_cuda_success(err, "Error on destroying event")
err = cuda_driver.cuEventDestroy(end_event)
_cuda_success(err, "Error on destroying event")
return elapsed_time / profiling_iterations * 1e3

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

@@ -68,6 +68,8 @@ class Pointer(ABC):
@property
def dtype(self) -> Type[Numeric]: ...
def align(self, min_align: int) -> "Pointer": ...
def __get_mlir_types__(self) -> List[ir.Type]: ...
def __extract_mlir_values__(self) -> List[ir.Value]: ...