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]: ...

62
python/CuTeDSL/cutlass/pipeline/__init__.py Normal file
View File

@@ -0,0 +1,62 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
from .helpers import (
Agent,
CooperativeGroup,
PipelineOp,
SyncObject,
MbarrierArray,
NamedBarrier,
TmaStoreFence,
PipelineUserType,
PipelineState,
make_pipeline_state,
pipeline_init_wait,
arrive,
arrive_unaligned,
wait,
wait_unaligned,
arrive_and_wait,
sync,
)
from .sm90 import (
PipelineAsync,
PipelineTmaAsync,
PipelineTmaMultiConsumersAsync,
PipelineTmaStore,
)
from .sm100 import (
PipelineTmaUmma,
PipelineAsyncUmma,
PipelineUmmaAsync,
)
__all__ = [
"Agent",
"CooperativeGroup",
"PipelineOp",
"SyncObject",
"MbarrierArray",
"NamedBarrier",
"TmaStoreFence",
"PipelineUserType",
"PipelineState",
"PipelineAsync",
"PipelineTmaAsync",
"PipelineTmaUmma",
"PipelineTmaMultiConsumersAsync",
"PipelineAsyncUmma",
"PipelineUmmaAsync",
"PipelineTmaStore",
]

645
python/CuTeDSL/cutlass/pipeline/helpers.py Normal file
View File

@@ -0,0 +1,645 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Union
import warnings
import cutlass.cute as cute
from cutlass.cutlass_dsl import Boolean, Int32, Int64, if_generate
from cutlass._mlir.dialects import llvm
import cutlass._mlir.dialects.cute as _cute_ir
##############################################################################
# Agent class
##############################################################################
class Agent(enum.Enum):
"""
Agent indicates what is participating in the pipeline synchronization.
"""
# Arbitrary grouping of N threads
Thread = enum.auto()
# Same as AsyncThread, but includes all threads in the block
ThreadBlock = enum.auto()
# Same as AsyncThread, but includes all threads in the cluster
ThreadBlockCluster = enum.auto()
class CooperativeGroup:
"""
CooperativeGroup contains size and alignment restrictions for an Agent.
"""
def __init__(self, agent: Agent, size: int = 1, alignment: int = 1):
if agent is Agent.Thread:
assert size > 0
if size == 32:
assert (
size == alignment
), "Error: Alignment does not match number of threads in a warp."
elif size == 128:
assert (
size == alignment
), "Error: Alignment does not match number of threads in a warpgroup."
elif agent is Agent.ThreadBlock:
raise NotImplementedError("Error: Not yet supported.")
elif agent is Agent.ThreadBlockCluster:
raise NotImplementedError("Error: Not yet supported.")
else:
# Should never reach this state
size = 0
if size <= 0:
raise ValueError(
"Error: The number of threads in a CooperativeGroup must be more than 0."
)
# Size indicates how many threads are participating in this CooperativeGroup
self.size = size
# Agent indicates the type of thread group
self.agent = agent
class PipelineOp(enum.Enum):
"""
PipelineOp assigns an operation to an agent corresponding to a specific hardware feature.
"""
# async-threads
AsyncThread = enum.auto()
# Blackwell (SM100a) MMA instruction
TCGen05Mma = enum.auto()
# Tensor Memory Accelerator load
TmaLoad = enum.auto()
# TMA Store consuming smem produced by AsyncThread
TmaStore = enum.auto()
# Composite of multiple PipelineOps
Composite = enum.auto()
def _get_pipeline_op(type_str):
return PipelineOp(type_str)
##############################################################################
# SyncObject class
##############################################################################
class SyncObject(ABC):
"""Abstract base class for hardware synchronization primitives.
This class defines the interface for different types of hardware synchronization
mechanisms including shared memory barriers, named barriers, and fences.
"""
@abstractmethod
def arrive(self) -> None:
pass
@abstractmethod
def wait(self) -> None:
pass
@abstractmethod
def arrive_and_wait(self) -> None:
pass
@abstractmethod
def arrive_and_drop(self) -> None:
pass
@abstractmethod
def get_barrier(self) -> Union[cute.Pointer, int, None]:
pass
@abstractmethod
def max(self) -> Union[int, None]:
pass
class MbarrierArray(SyncObject):
"""
MbarrierArray implements an abstraction for an array of smem barriers.
"""
def __init__(
self,
barrier_storage: cute.Pointer,
num_stages: int,
agent: tuple[PipelineOp, CooperativeGroup],
tx_count: int = 0,
) -> None:
self.barrier_storage = barrier_storage
self.tx_count = tx_count
self.num_stages = num_stages
self.op_type, self.cg = agent
self.arrive_count = self.cg.size
if self.num_stages <= 0:
raise ValueError("Error: Mbarrier stage count must be greater than 0.")
if self.arrive_count <= 0:
raise ValueError("Error: Mbarrier arrive count must be greater than 0.")
if self.op_type is PipelineOp.TmaLoad and self.tx_count < 0:
raise ValueError(
"Error: Mbarrier tx count must not be less than 0 for TMA ops."
)
# Store mbarrier base pointer
self.mbarrier_base = self.barrier_storage
# Mbarrier initialization in constructor
self.mbarrier_init()
def recast_to_new_op_type(self, new_op_type: PipelineOp) -> "MbarrierArray":
"""
Creates a copy of MbarrierArray with a different op_type without re-initializing barriers
"""
# Create new instance without initialization
new_mbarrier_array = object.__new__(MbarrierArray)
# Copy all attributes directly
new_mbarrier_array.barrier_storage = self.barrier_storage
new_mbarrier_array.op_type = new_op_type
new_mbarrier_array.cg = self.cg
new_mbarrier_array.num_stages = self.num_stages
new_mbarrier_array.tx_count = self.tx_count
new_mbarrier_array.arrive_count = self.arrive_count
new_mbarrier_array.mbarrier_base = self.mbarrier_base
return new_mbarrier_array
# Mbarrier initialization
def mbarrier_init(self) -> None:
"""
Initializes an array of mbarriers using warp 0.
"""
def then_body():
for index in range(self.num_stages):
cute.arch.mbarrier_init(self.get_barrier(index), self.arrive_count)
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
if_generate(warp_idx == 0, then_body)
def arrive(
self,
index: int,
dst: int,
cta_group: Optional[cute.nvgpu.tcgen05.CtaGroup] = None,
) -> None:
"""Select the arrive corresponding to this MbarrierArray's PipelineOp.
:param index: Index of the mbarrier in the array to arrive on
:type index: int
:param dst: Destination parameter for selective arrival, which can be either a mask or destination cta rank.
When None, both ``TCGen05Mma`` and ``AsyncThread`` will arrive on their local mbarrier.
- For ``TCGen05Mma``, ``dst`` serves as a multicast mask (e.g., 0b1011 allows arrive signal to be multicast to CTAs
in the cluster with rank = 0, 1, and 3).
- For ``AsyncThread``, ``dst`` serves as a destination cta rank (e.g., 3 means threads will arrive on
the mbarrier with rank = 3 in the cluster).
:type dst: int | None
:param cta_group: CTA group for ``TCGen05Mma``, defaults to None for other op types
:type cta_group: ``cute.nvgpu.tcgen05.CtaGroup``, optional
"""
if self.op_type is PipelineOp.AsyncThread:
self.arrive_mbarrier(index, dst)
elif self.op_type is PipelineOp.TCGen05Mma:
assert (
cta_group is not None
), "Error: CTA group must be provided for TCGen05Mma."
self.arrive_tcgen05mma(index, dst, cta_group)
elif self.op_type in [PipelineOp.TmaLoad]:
self.arrive_and_expect_tx(index, self.tx_count)
else:
assert (
False
), f"Error: MbarrierArray is not supported for PipelineOp: {_get_pipeline_op(self.op_type)}."
def arrive_mbarrier(self, index: int, dst_rank: Optional[int] = None) -> None:
if dst_rank is None:
cute.arch.mbarrier_arrive(self.get_barrier(index))
else:
cute.arch.mbarrier_arrive(self.get_barrier(index), dst_rank)
def arrive_tcgen05mma(
self, index: int, mask: Optional[int], cta_group: cute.nvgpu.tcgen05.CtaGroup
) -> None:
if mask is None:
with cute.arch.elect_one():
cute.nvgpu.tcgen05.commit(self.get_barrier(index))
else:
with cute.arch.elect_one():
cute.nvgpu.tcgen05.commit(self.get_barrier(index), mask, cta_group)
def arrive_and_expect_tx(self, index: int, tx_count: int) -> None:
with cute.arch.elect_one():
cute.arch.mbarrier_arrive_and_expect_tx(self.get_barrier(index), tx_count)
def try_wait(self, index: int, phase: int) -> Boolean:
return cute.arch.mbarrier_try_wait(self.get_barrier(index), phase)
def wait(self, index: int, phase: int) -> None:
cute.arch.mbarrier_wait(self.get_barrier(index), phase)
def arrive_and_wait(
self,
index: int,
phase: int,
dst: int,
cta_group: Optional[cute.nvgpu.tcgen05.CtaGroup] = None,
) -> None:
arrive(index, dst, cta_group)
wait(index, phase)
def arrive_and_drop(self) -> None:
raise NotImplementedError("Error: Not yet supported.")
def get_barrier(self, index: int) -> cute.Pointer:
return self.mbarrier_base + index
def max(self) -> int:
# Transaction barriers have a maximum arrive count of 511 (2^9 - 1).
# Non-transaction barriers have a maximum arrive count of 1,048,575 (2^20 - 1).
return 511
def __extract_mlir_values__(self):
return [self.barrier_storage]
def __new_from_mlir_values__(self, values):
return MbarrierArray(
values[0], self.num_stages, (self.op_type, self.cg), self.tx_count
)
@dataclass(frozen=True)
class NamedBarrier(SyncObject):
"""
NamedBarrier is an abstraction for named barriers managed by hardware.
There are 16 named barriers available, with barrier_ids 0-15.
See the `PTX documentation <https://https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-bar>`__.
"""
barrier_id: int
num_threads: int
def __post_init__(self) -> None:
if self.barrier_id < 0 or self.barrier_id >= 16:
raise ValueError("Error: NamedBarrier ID must be between 0 and 16.")
if self.barrier_id == 0:
warnings.warn(
"NamedBarrier ID 0 is by other driver APIs (i.e. sync_threads()) and should not be used."
)
def arrive(self) -> None:
"""
The aligned flavor of arrive is used when all threads in the CTA will execute the
same instruction. See PTX documentation.
"""
cute.arch.barrier_arrive(
barrier_id=self.barrier_id, number_of_threads=self.num_threads
)
def arrive_unaligned(self) -> None:
"""
The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA.
"""
llvm.inline_asm(
None,
[Int32(self.barrier_id).ir_value(), Int32(self.num_threads).ir_value()],
"barrier.arrive $0, $1;",
"r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
def wait(self) -> None:
"""
NamedBarriers do not have a standalone wait like mbarriers, only an arrive_and_wait.
If synchronizing two warps in a producer/consumer pairing, the arrive count would be
32 using mbarriers but 64 using NamedBarriers. Only threads from either the producer
or consumer are counted for mbarriers, while all threads participating in the sync
are counted for NamedBarriers.
"""
warnings.warn(
"NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()."
)
self.arrive_and_wait()
def wait_unaligned(self) -> None:
warnings.warn(
"NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()."
)
llvm.inline_asm(
None,
[Int32(self.barrier_id).ir_value(), Int32(self.num_threads).ir_value()],
"barrier.sync $0, $1;",
"r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
def arrive_and_wait(self) -> None:
cute.arch.barrier(
barrier_id=self.barrier_id, number_of_threads=self.num_threads
)
def arrive_and_drop(self) -> None:
raise NotImplementedError("Error: Not supported.")
def sync(self) -> None:
cute.arch.barrier(barrier_id=self.barrier_id)
def get_barrier(self) -> int:
return self.barrier_id
def max(self) -> int:
# Transaction barriers have a maximum arrive count of 4095 (2^12 - 1).
return 4095
class TmaStoreFence(SyncObject):
"""
TmaStoreFence is used for a multi-stage epilogue buffer.
"""
def __init__(self, num_stages: int = 0) -> None:
if num_stages <= 0:
raise ValueError("Mbarrier stage count must be greater than 0.")
self.num_stages = num_stages
def arrive(self) -> None:
cute.arch.cp_async_bulk_commit_group()
def wait(self) -> None:
cute.arch.cp_async_bulk_wait_group(self.num_stages - 1, read=True)
def arrive_and_wait(self) -> None:
self.arrive()
self.wait()
def arrive_and_drop(self) -> None:
raise NotImplementedError("Error: Not supported.")
# TmaStoreFence doesn't have mbarriers
def get_barrier(self) -> None:
assert (
False
), "Error: TmaStoreFence doesn't use mbarriers and cannot return a barrier."
def max(self) -> None:
raise NotImplementedError("Error: Not supported.")
def tail(self) -> None:
cute.arch.cp_async_bulk_wait_group(0, read=True)
##############################################################################
# PipelineState class
##############################################################################
class PipelineUserType(enum.Enum):
Producer = enum.auto()
Consumer = enum.auto()
class PipelineState:
"""
Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
"""
def __init__(self, stages: int, count, index, phase):
self._stages = stages
self._count = count
self._index = index
self._phase = phase
def clone(self) -> "PipelineState":
return PipelineState(self.stages, self._count, self.index, self.phase)
@property
def index(self) -> Int32:
return self._index
@property
def count(self) -> Int32:
return self._count
@property
def stages(self) -> int:
return self._stages
@property
def phase(self) -> Int32:
return self._phase
def reset_count(self):
self._count = Int32(0)
def advance(self):
self._index += 1
self._count += 1
def then_body(index, phase):
new_index = Int32(0)
new_phase = phase ^ 1
return new_index, new_phase
def else_body(index, phase):
return index, phase
self._index, self._phase = if_generate(
self._index == self.stages,
then_body,
else_body,
[self.index, self.phase],
[Int32, Int32],
)
def reverse(self):
self._index -= 1
self._count -= 1
def then_body(index, phase):
new_index = Int32(self.stages - 1)
new_phase = phase ^ 1
return new_index, new_phase
def else_body(index, phase):
return index, phase
self._index, self._phase = if_generate(
self._index == -1,
then_body,
else_body,
[self.index, self.phase],
[Int32, Int32],
)
def __get_mlir_types__(self):
return [self._count.type, self._index.type, self._phase.type]
def __extract_mlir_values__(self):
count = self._count
index = self._index
phase = self._phase
return [count.ir_value(), index.ir_value(), phase.ir_value()]
# This can be overridden by derived classes
def __new_from_mlir_values__(self, values):
return PipelineState(
self.stages, Int32(values[0]), Int32(values[1]), Int32(values[2])
)
def make_pipeline_state(type: PipelineUserType, stages: int):
"""
Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1.
"""
if type is PipelineUserType.Producer:
return PipelineState(
stages,
Int32(0),
Int32(0),
Int32(1),
)
elif type is PipelineUserType.Consumer:
return PipelineState(
stages,
Int32(0),
Int32(0),
Int32(0),
)
else:
assert (
False
), "Error: invalid PipelineUserType specified for make_pipeline_state."
##############################################################################
# Helper functions
##############################################################################
def pipeline_init_wait(cta_layout_vmnk: Optional[cute.Layout] = None):
"""
Fences the mbarrier init and syncs the threadblock or cluster
"""
cute.arch.mbarrier_init_fence()
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
# If not using clusters, sync the threadblock
_sync(Agent.ThreadBlock)
else:
# If using clusters, sync the cluster
_sync(Agent.ThreadBlockCluster)
def _sync(group: Agent):
"""
Syncs all threads within an agent.
"""
if group is Agent.Thread:
raise NotImplementedError("Error: Not supported.")
elif group is Agent.ThreadBlock:
cute.arch.sync_threads()
elif group is Agent.ThreadBlockCluster:
cute.arch.cluster_arrive()
cute.arch.cluster_wait()
else:
assert (
False
), "Error: No explicit sync instruction exists. Please use barriers (named / mbarrier) instead."
def _mbarrier_i64_to_ptr(val: Int64) -> cute.Pointer:
"""
Converts a smem pointer of type Int64 to cute.Pointer with 8B alignment
"""
return cute.make_ptr(
Int64,
val.ir_value(),
mem_space=_cute_ir.AddressSpace.smem,
assumed_align=8,
)
# NamedBarrier free functions
def arrive(barrier_id: int, num_threads: int):
"""
The aligned flavor of arrive is used when all threads in the CTA will execute the
same instruction. See PTX documentation.
"""
cute.arch.barrier_arrive(barrier_id=barrier_id, number_of_threads=num_threads)
def arrive_unaligned(barrier_id: int, num_threads: int):
"""
The unaligned flavor of arrive can be used with an arbitrary number of threads in the CTA.
"""
llvm.inline_asm(
None,
[Int32(barrier_id).ir_value(), Int32(num_threads).ir_value()],
"barrier.arrive $0, $1;",
"r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
def wait(barrier_id: int, num_threads: int):
"""
NamedBarriers do not have a standalone wait like mbarriers, only an arrive_and_wait.
If synchronizing two warps in a producer/consumer pairing, the arrive count would be
32 using mbarriers but 64 using NamedBarriers. Only threads from either the producer
or consumer are counted for mbarriers, while all threads participating in the sync
are counted for NamedBarriers.
"""
warnings.warn(
"NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()."
)
arrive_and_wait()
def wait_unaligned(barrier_id: int, num_threads: int):
warnings.warn(
"NamedBarrier wait also arrives on the barrier. Routing call to NamedBarrier.arrive_and_wait()."
)
llvm.inline_asm(
None,
[Int32(barrier_id).ir_value(), Int32(num_threads).ir_value()],
"barrier.sync $0, $1;",
"r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
def arrive_and_wait(barrier_id: int, num_threads: int):
cute.arch.barrier(barrier_id=barrier_id, number_of_threads=num_threads)
def sync(barrier_id: int = 0):
cute.arch.barrier(barrier_id=barrier_id)

452
python/CuTeDSL/cutlass/pipeline/sm100.py Normal file
View File

@@ -0,0 +1,452 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Union
import warnings
import cutlass.cute as cute
from cutlass.cutlass_dsl import Boolean, if_generate
from cutlass.pipeline import (
CooperativeGroup,
PipelineOp,
PipelineState,
pipeline_init_wait,
PipelineAsync,
)
##############################################################################
# Pipeline classes
##############################################################################
@dataclass(frozen=True)
class PipelineTmaUmma(PipelineAsync):
"""
PipelineTmaUmma is used for TMA producers and UMMA consumers (e.g. Blackwell mainloops).
"""
is_leader_cta: bool
cta_group: cute.nvgpu.tcgen05.CtaGroup
@staticmethod
def _compute_mcast_arrival_mask(cta_layout_vmnk: cute.Layout):
"""
Computes a mask for signaling arrivals to multicasting threadblocks.
"""
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
)
tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
)
block_in_cluster_coord_vmnk_peer = (
cta_in_cluster_coord_vmnk[0] ^ 1,
*cta_in_cluster_coord_vmnk[1:],
)
tma_mcast_mask_a_peer = cute.nvgpu.cpasync.create_tma_multicast_mask(
cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=2
)
tma_mcast_mask_b_peer = cute.nvgpu.cpasync.create_tma_multicast_mask(
cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=1
)
return (
tma_mcast_mask_a
| tma_mcast_mask_b
| tma_mcast_mask_a_peer
| tma_mcast_mask_b_peer
)
@staticmethod
def _compute_is_leader_cta(cta_layout_vmnk: cute.Layout):
"""
Computes leader threadblocks for 2CTA kernels. For 1CTA, all threadblocks are leaders.
"""
bidx, bidy, _ = cute.arch.block_idx()
mma_coord_vmnk = (
bidx % cute.size(cta_layout_vmnk, mode=[0]),
bidx // cute.size(cta_layout_vmnk, mode=[0]),
bidy,
None,
)
return mma_coord_vmnk[0] == 0
@staticmethod
def create(
*,
num_stages: int,
producer_group: CooperativeGroup,
consumer_group: CooperativeGroup,
tx_count: int,
barrier_storage: cute.Pointer = None,
cta_layout_vmnk: Optional[cute.Layout] = None,
):
"""
This helper function computes any necessary attributes and returns an instance of PipelineTmaUmma.
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
:type barrier_storage: cute.Pointer
:param num_stages: Number of buffer stages for this pipeline
:type num_stages: Int32
:param producer_group: CooperativeGroup for the producer agent
:type producer_group: CooperativeGroup
:param consumer_group: CooperativeGroup for the consumer agent
:type consumer_group: CooperativeGroup
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
:type tx_count: int
:param cta_layout_vmnk: Layout of the cluster shape
:type cta_layout_vmnk: cute.Layout | None
"""
if not isinstance(barrier_storage, cute.Pointer):
raise ValueError(
f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
)
producer_type = PipelineOp.TmaLoad
consumer_type = PipelineOp.TCGen05Mma
producer = (producer_type, producer_group)
consumer = (consumer_type, consumer_group)
sync_object_full = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8), num_stages, producer, tx_count
)
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
# No mcast mask if not using clusters
producer_mask = None
# All threadblocks are leaders if not using clusters
is_leader_cta = True
else:
producer_mask = PipelineTmaUmma._compute_mcast_arrival_mask(cta_layout_vmnk)
is_leader_cta = PipelineTmaUmma._compute_is_leader_cta(cta_layout_vmnk)
cta_group = (
cute.nvgpu.tcgen05.CtaGroup.ONE
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1
else cute.nvgpu.tcgen05.CtaGroup.TWO
)
consumer_mask = producer_mask
pipeline_init_wait(cta_layout_vmnk)
return PipelineTmaUmma(
sync_object_full,
sync_object_empty,
num_stages,
producer_mask,
consumer_mask,
is_leader_cta,
cta_group,
)
def consumer_release(self, state: PipelineState):
"""
UMMA consumer release buffer empty, cta_group needs to be provided.
"""
self.sync_object_empty.arrive(state.index, self.consumer_mask, self.cta_group)
def producer_acquire(
self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
):
"""
TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
"""
if_generate(
try_acquire_token is None or try_acquire_token == 0,
lambda: self.sync_object_empty.wait(state.index, state.phase),
)
if_generate(
self.is_leader_cta,
lambda: self.sync_object_full.arrive(state.index, self.producer_mask),
)
def producer_commit(self, state: PipelineState):
"""
TMA producer commit is a noop since TMA instruction itself updates the transaction count.
"""
pass
@dataclass(frozen=True)
class PipelineAsyncUmma(PipelineAsync):
"""
PipelineAsyncUmma is used for AsyncThread producers and UMMA consumers (e.g. Blackwell input fusion pipelines).
"""
cta_group: cute.nvgpu.tcgen05.CtaGroup
@staticmethod
def _compute_leading_cta_rank(cta_v_size):
"""
Computes the leading CTA rank.
"""
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
return cta_rank_in_cluster // cta_v_size * cta_v_size
@staticmethod
def _compute_is_leader_cta(cta_layout_vmnk: cute.Layout):
"""
Computes leader threadblocks for 2CTA kernels. For 1CTA, all threadblocks are leaders.
"""
bidx, bidy, _ = cute.arch.block_idx()
mma_coord_vmnk = (
bidx % cute.size(cta_layout_vmnk, mode=[0]),
bidx // cute.size(cta_layout_vmnk, mode=[0]),
bidy,
None,
)
return mma_coord_vmnk[0] == 0
@staticmethod
def _compute_peer_cta_mask(cta_layout_vmnk: cute.Layout):
"""
Computes a mask for signaling arrivals to multicasting threadblocks.
"""
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
mask_self = cute.nvgpu.cpasync.create_tma_multicast_mask(
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=0
)
block_in_cluster_coord_vmnk_peer = (
cta_in_cluster_coord_vmnk[0] ^ 1,
*cta_in_cluster_coord_vmnk[1:],
)
mask_peer = cute.nvgpu.cpasync.create_tma_multicast_mask(
cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=0
)
return mask_self | mask_peer
@staticmethod
def create(
*,
num_stages: int,
producer_group: CooperativeGroup,
consumer_group: CooperativeGroup,
barrier_storage: cute.Pointer = None,
cta_layout_vmnk: Optional[cute.Layout] = None,
):
"""
This helper function computes any necessary attributes and returns an instance of PipelineAsyncUmma.
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
:type barrier_storage: cute.Pointer
:param num_stages: Number of buffer stages for this pipeline
:type num_stages: Int32
:param producer_group: CooperativeGroup for the producer agent
:type producer_group: CooperativeGroup
:param consumer_group: CooperativeGroup for the consumer agent
:type consumer_group: CooperativeGroup
:param cta_layout_vmnk: Layout of the cluster shape
:type cta_layout_vmnk: cute.Layout | None
"""
if not isinstance(barrier_storage, cute.Pointer):
raise ValueError(
f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
)
producer_type = PipelineOp.AsyncThread
consumer_type = PipelineOp.TCGen05Mma
producer = (producer_type, producer_group)
consumer = (consumer_type, consumer_group)
sync_object_full = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8),
num_stages,
producer,
)
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)
cta_v_size = (
cute.size(cta_layout_vmnk, mode=[0]) if cta_layout_vmnk is not None else 1
)
cta_group = (
cute.nvgpu.tcgen05.CtaGroup.ONE
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1
else cute.nvgpu.tcgen05.CtaGroup.TWO
)
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1:
# No mcast mask if we're not using 2CTA tcgen05 MMA
producer_mask = None
consumer_mask = None
else:
# If we're using 2CTA UMMAs, producer will arrive the mbar on leading CTA
# We need to get the target cta_rank
producer_mask = PipelineAsyncUmma._compute_leading_cta_rank(cta_v_size)
# consumer needs to get the mask to signal
consumer_mask = PipelineAsyncUmma._compute_peer_cta_mask(cta_layout_vmnk)
pipeline_init_wait(cta_layout_vmnk)
return PipelineAsyncUmma(
sync_object_full,
sync_object_empty,
num_stages,
producer_mask,
consumer_mask,
cta_group,
)
def consumer_release(self, state: PipelineState):
"""
UMMA consumer release buffer empty, cta_group needs to be provided.
"""
self.sync_object_empty.arrive(state.index, self.consumer_mask, self.cta_group)
@dataclass(frozen=True)
class PipelineUmmaAsync(PipelineAsync):
"""
PipelineUmmaAsync is used for UMMA producers and AsyncThread consumers (e.g. Blackwell accumulator pipelines).
"""
cta_group: cute.nvgpu.tcgen05.CtaGroup
@staticmethod
def _compute_tmem_sync_mask(cta_layout_vmnk: cute.Layout):
"""
Computes a mask to signal completion of tmem buffers for 2CTA kernels.
"""
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
return cute.make_layout_image_mask(
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mode=0
)
@staticmethod
def _compute_peer_cta_rank():
"""
Computes a mask to signal release of tmem buffers for 2CTA kernels.
"""
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
return cta_rank_in_cluster // 2 * 2
@staticmethod
def create(
*,
num_stages: int,
producer_group: CooperativeGroup,
consumer_group: CooperativeGroup,
barrier_storage: cute.Pointer = None,
cta_layout_vmnk: Optional[cute.Layout] = None,
):
"""
This helper function computes any necessary attributes and returns an instance of PipelineUmmaAsync.
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
:type barrier_storage: cute.Pointer
:param num_stages: Number of buffer stages for this pipeline
:type num_stages: Int32
:param producer_group: CooperativeGroup for the producer agent
:type producer_group: CooperativeGroup
:param consumer_group: CooperativeGroup for the consumer agent
:type consumer_group: CooperativeGroup
:param cta_layout_vmnk: Layout of the cluster shape
:type cta_layout_vmnk: cute.Layout | None
"""
if not isinstance(barrier_storage, cute.Pointer):
raise ValueError(
f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
)
producer_type = PipelineOp.TCGen05Mma
consumer_type = PipelineOp.AsyncThread
producer = (producer_type, producer_group)
consumer = (consumer_type, consumer_group)
sync_object_full = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8), num_stages, producer
)
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
# Set mask to None if not using clusters (i.e. 1CTA kernels)
producer_mask = None
else:
producer_mask = PipelineUmmaAsync._compute_tmem_sync_mask(cta_layout_vmnk)
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1:
# Set mask to None if not using 2CTA intructions
consumer_mask = None
else:
consumer_mask = PipelineUmmaAsync._compute_peer_cta_rank()
cta_group = (
cute.nvgpu.tcgen05.CtaGroup.ONE
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1
else cute.nvgpu.tcgen05.CtaGroup.TWO
)
pipeline_init_wait(cta_layout_vmnk)
return PipelineUmmaAsync(
sync_object_full,
sync_object_empty,
num_stages,
producer_mask,
consumer_mask,
cta_group,
)
def producer_commit(self, state: PipelineState):
"""
UMMA producer commit buffer full, cta_group needs to be provided.
"""
self.sync_object_full.arrive(state.index, self.producer_mask, self.cta_group)
def producer_tail(self, state: PipelineState):
"""
Make sure the last used buffer empty signal is visible to producer.
Producer tail is usually executed by producer before exit, to avoid dangling
mbarrier arrive signals after kernel exit.
:param state: The pipeline state that points to next useful buffer
:type state: PipelineState
"""
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
is_leader_cta = cta_rank_in_cluster % 2 == 0
def then_body():
# Assume state contains that next useful buffer
# So we only need to advance to num_stages - 1 times to last used buffer
for i in range(self.num_stages - 1):
state.advance()
self.producer_acquire(state)
if_generate(is_leader_cta, then_body)

803
python/CuTeDSL/cutlass/pipeline/sm90.py Normal file
View File

@@ -0,0 +1,803 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# Use of this software is governed by the terms and conditions of the
# NVIDIA End User License Agreement (EULA), available at:
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
#
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
import enum
from abc import ABC, abstractmethod
from dataclasses import dataclass
from typing import Optional, Union
import warnings
import cutlass.cute as cute
from cutlass.cutlass_dsl import Boolean, Int32, if_generate
from cutlass.pipeline import (
CooperativeGroup,
PipelineOp,
SyncObject,
MbarrierArray,
TmaStoreFence,
PipelineUserType,
PipelineState,
make_pipeline_state,
pipeline_init_wait,
)
##############################################################################
# Pipeline classes
##############################################################################
@dataclass(frozen=True)
class PipelineAsync:
"""PipelineAsync is a generic pipeline class where both the producer and consumer are
AsyncThreads. It also serves as a base class for specialized pipeline classes.
This class implements a producer-consumer pipeline pattern where both sides operate
asynchronously. The pipeline maintains synchronization state using barrier objects
to coordinate between producer and consumer threads.
The pipeline state transitions of one pipeline entry(mbarrier) can be represented as:
.. table:: Pipeline State Transitions
:widths: auto
+-----------+-----------+-----------+-----------+-----------+-----------+
| Barrier | State | p.acquire | p.commit | c.wait | c.release |
+===========+===========+===========+===========+===========+===========+
| empty_bar | empty | <Return> | n/a | n/a | - |
+-----------+-----------+-----------+-----------+-----------+-----------+
| empty_bar | wait | <Block> | n/a | n/a | -> empty |
+-----------+-----------+-----------+-----------+-----------+-----------+
| full_bar | wait | n/a | -> full | <Block > | n/a |
+-----------+-----------+-----------+-----------+-----------+-----------+
| full_bar | full | n/a | - | <Return> | n/a |
+-----------+-----------+-----------+-----------+-----------+-----------+
Where:
- p: producer
- c: consumer
- <Block>: This action is blocked until transition to a state allow it to proceed by other side
- e.g. ``p.acquire()`` is blocked until ``empty_bar`` transition to ``empty`` state by ``c.release()``
.. code-block:: text
Array of mbarriers as circular buffer:
Advance Direction
<-------------------
Producer Consumer
| ^
V |
+-----------------+
--|X|X|W|D|D|D|D|R|X|<-.
/ +-----------------+ \\
| |
`------------------------'
Where:
- X: Empty buffer (initial state)
- W: Producer writing (producer is waiting for buffer to be empty)
- D: Data ready (producer has written data to buffer)
- R: Consumer reading (consumer is consuming data from buffer)
**Example:**
.. code-block:: python
# Create pipeline with 5 stages
pipeline = PipelineAsync.create(
num_stages=5, # number of pipeline stages
producer_group=producer_warp,
consumer_group=consumer_warp
barrier_storage=smem_ptr, # smem pointer for array of mbarriers in shared memory
)
# Producer side
producer = pipeline.make_pipeline_producer(producer_warp)
for i in range(num_iterations):
producer.acquire() # Wait for buffer to be empty
# Write data to pipeline buffer
producer.commit() # Signal buffer is full
producer.advance() # Move index to next stage
# Consumer side
consumer = pipeline.make_pipeline_consumer(consumer_warp)
for i in range(num_iterations):
consumer.wait() # Wait for buffer to be full
# Read data from pipeline buffer
consumer.release() # Signal buffer is empty
consumer.advance() # Move index to next stage
"""
sync_object_full: SyncObject
sync_object_empty: SyncObject
num_stages: int
producer_mask: Optional[Int32]
consumer_mask: Optional[Int32]
@staticmethod
def _make_sync_object(
barrier_storage: cute.Pointer,
num_stages: int,
agent: tuple[PipelineOp, CooperativeGroup],
tx_count: int = 0,
) -> SyncObject:
"""
Returns a SyncObject corresponding to an agent's PipelineOp.
"""
if agent[0] in [
PipelineOp.AsyncThread,
PipelineOp.TmaLoad,
PipelineOp.TCGen05Mma,
PipelineOp.Composite,
]:
return MbarrierArray(
barrier_storage=barrier_storage,
num_stages=num_stages,
agent=agent,
tx_count=tx_count,
)
elif agent[0] is PipelineOp.TmaStore:
# Path taken for AsyncTmaStore
return TmaStoreFence(num_stages=num_stages)
else:
assert False, "Error: Invalid PipelineOp specified."
@staticmethod
def create(
*,
num_stages: int,
producer_group: CooperativeGroup,
consumer_group: CooperativeGroup,
barrier_storage: cute.Pointer = None,
producer_mask: Int32 = None,
consumer_mask: Int32 = None,
):
"""Creates and initializes a new PipelineAsync instance.
This helper function computes necessary attributes and returns an instance of PipelineAsync
with the specified configuration for producer and consumer synchronization.
:param barrier_storage: Pointer to the shared memory address for this pipeline's mbarriers
:type barrier_storage: cute.Pointer
:param num_stages: Number of buffer stages for this pipeline
:type num_stages: int
:param producer_group: `CooperativeGroup` for the producer agent
:type producer_group: CooperativeGroup
:param consumer_group: `CooperativeGroup` for the consumer agent
:type consumer_group: CooperativeGroup
:param producer_mask: Mask for signaling arrives for the producer agent, defaults to ``None``
:type producer_mask: Int32, optional
:param consumer_mask: Mask for signaling arrives for the consumer agent, defaults to ``None``
:type consumer_mask: Int32, optional
:return: A new PipelineAsync instance
:rtype: PipelineAsync
:raises ValueError: If barrier_storage is not a cute.Pointer instance
"""
if not isinstance(barrier_storage, cute.Pointer):
raise ValueError(
f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
)
producer_type = PipelineOp.AsyncThread
consumer_type = PipelineOp.AsyncThread
producer = (producer_type, producer_group)
consumer = (consumer_type, consumer_group)
sync_object_full = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8), num_stages, producer
)
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)
pipeline_init_wait()
return PipelineAsync(
sync_object_full,
sync_object_empty,
num_stages,
producer_mask,
consumer_mask,
)
def producer_acquire(
self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
):
if_generate(
try_acquire_token is None or try_acquire_token == 0,
lambda: self.sync_object_empty.wait(state.index, state.phase),
)
def producer_try_acquire(self, state: PipelineState):
return self.sync_object_empty.try_wait(state.index, state.phase)
def producer_commit(self, state: PipelineState):
self.sync_object_full.arrive(state.index, self.producer_mask)
def consumer_wait(
self, state: PipelineState, try_wait_token: Optional[Boolean] = None
):
if_generate(
try_wait_token is None or try_wait_token == 0,
lambda: self.sync_object_full.wait(state.index, state.phase),
)
def consumer_try_wait(self, state: PipelineState):
return self.sync_object_full.try_wait(state.index, state.phase)
def consumer_release(self, state: PipelineState):
self.sync_object_empty.arrive(state.index, self.consumer_mask)
def producer_get_barrier(self, state: PipelineState) -> cute.Pointer:
return self.sync_object_full.get_barrier(state.index)
def producer_tail(self, state: PipelineState):
"""
Make sure the last used buffer empty signal is visible to producer.
Producer tail is usually executed by producer before exit, to avoid dangling
mbarrier arrive signals after kernel exit.
:param state: The pipeline state that points to next useful buffer
:type state: PipelineState
"""
# Assume state contains that next useful buffer
# So we only need to advance to num_stages - 1 times to last used buffer
for i in range(self.num_stages - 1):
state.advance()
self.producer_acquire(state)
# Util methods to manage produer and consumer
def make_pipeline_producer(self, group: CooperativeGroup):
state = make_pipeline_state(PipelineUserType.Producer, self.num_stages)
return PipelineProducer(self, state, group)
def make_pipeline_consumer(self, group: CooperativeGroup):
state = make_pipeline_state(PipelineUserType.Consumer, self.num_stages)
return PipelineConsumer(self, state, group)
@dataclass(frozen=True)
class PipelineTmaAsync(PipelineAsync):
"""
PipelineTmaAsync is used for TMA producers and AsyncThread consumers (e.g. Hopper mainloops).
"""
is_signalling_thread: Boolean
@staticmethod
@cute.jit
def init_empty_barrier_arrive_signal(cta_layout_vmnk: cute.Layout, tidx: Int32):
"""
Initialize the empty barrier arrive signal
This function returns the destination cta rank and a boolean indicating if the signalling thread is the same as the current thread
"""
# Logic to optimally schedule Empty Arrives
cluster_shape_vmnk = cta_layout_vmnk.shape
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
tidx = tidx % 32
is_signalling_thread = tidx < cute.size(cluster_shape_vmnk)
dst_rank = tidx % cute.size(cluster_shape_vmnk)
dst_cta_coord = cta_layout_vmnk.get_hier_coord(dst_rank)
cur_cta_coord = cta_layout_vmnk.get_hier_coord(cta_rank_in_cluster)
is_same_row = (
dst_cta_coord[0] == cur_cta_coord[0]
and dst_cta_coord[1] == cur_cta_coord[1]
and dst_cta_coord[3] == cur_cta_coord[3]
)
is_same_col = (
dst_cta_coord[0] == cur_cta_coord[0]
and dst_cta_coord[2] == cur_cta_coord[2]
and dst_cta_coord[3] == cur_cta_coord[3]
)
is_same_row_or_col = is_same_row or is_same_col
is_signalling_thread_final = is_signalling_thread and is_same_row_or_col
return dst_rank, is_signalling_thread_final
@staticmethod
def create(
*,
num_stages: int,
producer_group: CooperativeGroup,
consumer_group: CooperativeGroup,
tx_count: int,
barrier_storage: cute.Pointer = None,
cta_layout_vmnk: Optional[cute.Layout] = None,
tidx: Optional[Int32] = None,
):
"""
This helper function computes any necessary attributes and returns an instance of PipelineTmaAsync.
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
:type barrier_storage: cute.Pointer
:param num_stages: Number of buffer stages for this pipeline
:type num_stages: Int32
:param producer_group: CooperativeGroup for the producer agent
:type producer_group: CooperativeGroup
:param consumer_group: CooperativeGroup for the consumer agent
:type consumer_group: CooperativeGroup
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
:type tx_count: int
:param cta_layout_vmnk: Layout of the cluster shape
:type cta_layout_vmnk: cute.Layout | None
:param tidx: thread index to consumer async threads
:type tidx: Int32 | None
"""
if not isinstance(barrier_storage, cute.Pointer):
raise ValueError(
f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
)
producer_type = PipelineOp.TmaLoad
consumer_type = PipelineOp.AsyncThread
producer = (producer_type, producer_group)
consumer = (consumer_type, consumer_group)
sync_object_full = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8), num_stages, producer, tx_count
)
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)
if tidx is None:
tidx, _, _ = cute.arch.thread_idx()
if cta_layout_vmnk is None:
cta_layout_vmnk = cute.make_layout((1, 1, 1, 1))
(
dst_rank,
is_signalling_thread,
) = PipelineTmaAsync.init_empty_barrier_arrive_signal(cta_layout_vmnk, tidx)
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
dst_rank = None
else:
dst_rank = dst_rank
producer_mask = None
pipeline_init_wait(cta_layout_vmnk)
return PipelineTmaAsync(
sync_object_full,
sync_object_empty,
num_stages,
producer_mask,
dst_rank,
is_signalling_thread,
)
def producer_acquire(
self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
):
"""
TMA producer commit conditionally waits on buffer empty and sets the transaction barrier.
"""
if_generate(
try_acquire_token is None or try_acquire_token == 0,
lambda: self.sync_object_empty.wait(state.index, state.phase),
)
self.sync_object_full.arrive(state.index, self.producer_mask)
def producer_commit(self, state: PipelineState):
"""
TMA producer commit is a noop since TMA instruction itself updates the transaction count.
"""
pass
def consumer_release(self, state: PipelineState):
"""
TMA consumer release conditionally signals the empty buffer to the producer.
"""
if_generate(
self.is_signalling_thread,
lambda: self.sync_object_empty.arrive(state.index, self.consumer_mask),
)
@dataclass(frozen=True)
class PipelineTmaMultiConsumersAsync(PipelineAsync):
"""
PipelineTmaMultiConsumersAsync is used for TMA producers and UMMA+Async consumers.
"""
is_leader_cta: bool
sync_object_empty_umma: SyncObject
sync_object_empty_async: SyncObject
cta_group: cute.nvgpu.tcgen05.CtaGroup
@staticmethod
def create(
*,
num_stages: int,
producer_group: CooperativeGroup,
consumer_group_umma: CooperativeGroup,
consumer_group_async: CooperativeGroup,
tx_count: int,
barrier_storage: cute.Pointer = None,
cta_layout_vmnk: Optional[cute.Layout] = None,
):
"""
This helper function computes any necessary attributes and returns an instance of PipelineTmaMultiConsumersAsync.
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
:type barrier_storage: cute.Pointer
:param num_stages: Number of buffer stages for this pipeline
:type num_stages: Int32
:param producer_group: CooperativeGroup for the producer agent
:type producer_group: CooperativeGroup
:param consumer_group_umma: CooperativeGroup for the UMMA consumer agent
:type consumer_group_umma: CooperativeGroup
:param consumer_group_async: CooperativeGroup for the AsyncThread consumer agent
:type consumer_group_async: CooperativeGroup
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
:type tx_count: int
:param cta_layout_vmnk: Layout of the cluster shape
:type cta_layout_vmnk: cute.Layout | None
"""
if not isinstance(barrier_storage, cute.Pointer):
raise ValueError(
f"Expected barrier_storage to be a cute.Pointer, but got {type(barrier_storage)}"
)
producer_type = PipelineOp.TmaLoad
consumer_type = PipelineOp.Composite
consumer_type_umma = PipelineOp.TCGen05Mma
consumer_type_async = PipelineOp.AsyncThread
if consumer_group_umma.agent != consumer_group_async.agent:
raise ValueError(
"UMMA and AsyncThread consumer groups must be the same agent"
)
if cta_layout_vmnk is not None and cute.size(cta_layout_vmnk) != 1:
raise ValueError(
f"PipelineTmaMultiConsumersAsync is not verified for cta_layout_vmnk != 1, cta_layout_vmnk:{cta_layout_vmnk}"
)
consumer_group = CooperativeGroup(
consumer_group_umma.agent,
consumer_group_umma.size + consumer_group_async.size,
)
producer = (producer_type, producer_group)
consumer = (consumer_type, consumer_group)
sync_object_full = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8), num_stages, producer, tx_count
)
sync_object_empty = PipelineAsync._make_sync_object(
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
)
sync_object_empty_umma = sync_object_empty.recast_to_new_op_type(
consumer_type_umma
)
sync_object_empty_async = sync_object_empty.recast_to_new_op_type(
consumer_type_async
)
# No mcast mask if not using clusters
producer_mask = None
consumer_mask = None
# All threadblocks are leaders if not using clusters
is_leader_cta = True
cta_group = (
cute.nvgpu.tcgen05.CtaGroup.ONE
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1
else cute.nvgpu.tcgen05.CtaGroup.TWO
)
pipeline_init_wait(cta_layout_vmnk)
return PipelineTmaMultiConsumersAsync(
sync_object_full,
sync_object_empty,
num_stages,
producer_mask,
consumer_mask,
is_leader_cta,
sync_object_empty_umma,
sync_object_empty_async,
cta_group,
)
def producer_acquire(
self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
):
"""
TMA producer acquire waits on buffer empty and sets the transaction barrier for leader threadblocks.
"""
if_generate(
try_acquire_token is None or try_acquire_token == 0,
lambda: self.sync_object_empty.wait(state.index, state.phase),
)
if_generate(
self.is_leader_cta,
lambda: self.sync_object_full.arrive(state.index, self.producer_mask),
)
def producer_commit(self, state: PipelineState):
"""
TMA producer commit is a noop since TMA instruction itself updates the transaction count.
"""
pass
def consumer_release(self, state: PipelineState, op_type: PipelineOp):
if op_type == PipelineOp.TCGen05Mma:
self.sync_object_empty_umma.arrive(
state.index, self.consumer_mask, self.cta_group
)
elif op_type == PipelineOp.AsyncThread:
self.sync_object_empty_async.arrive(state.index, self.consumer_mask)
else:
raise ValueError(f"Invalid PipelineOp specified. op_type:{op_type}")
@dataclass(frozen=True)
class PipelineTmaStore(PipelineAsync):
"""
PipelineTmaStore is used for synchronizing TMA stores in the epilogue. It does not use mbarriers.
"""
@staticmethod
def create(
*,
num_stages: int,
producer_group: CooperativeGroup,
):
"""
This helper function computes any necessary attributes and returns an instance of PipelineTmaStore.
:param num_stages: Number of buffer stages for this pipeline
:type num_stages: Int32
:param producer_group: CooperativeGroup for the producer agent
:type producer_group: CooperativeGroup
"""
producer_type = PipelineOp.TmaStore
producer = (producer_type, producer_group)
sync_object_full = PipelineAsync._make_sync_object(None, num_stages, producer)
return PipelineTmaStore(sync_object_full, None, num_stages, None, None)
def producer_acquire(self):
self.sync_object_full.wait()
def producer_commit(self):
self.sync_object_full.arrive()
def consumer_wait(self):
assert False, "Error: PipelineTmaStore does not have a consumer agent."
def consumer_release(self):
assert False, "Error: PipelineTmaStore does not have a consumer agent."
def producer_tail(self):
self.sync_object_full.tail()
#################################################################
# Utilities to help user of pipeline to simplify the workflow
#################################################################
class PipelineProducer:
"""A class representing a producer in an asynchronous pipeline.
The Producer class manages the producer side of an asynchronous pipeline, handling
synchronization and state management for producing data. It provides methods for
acquiring, committing, and advancing through pipeline stages.
:ivar _pipeline: The asynchronous pipeline this producer belongs to
:type _pipeline: PipelineAsync
:ivar _state: The current state of the producer in the pipeline
:type _state: PipelineState
:ivar _group: The cooperative group this producer operates in
:type _group: CooperativeGroup
**Examples:**
.. code-block:: python
pipeline = PipelineAsync.create(...)
producer = pipeline.create_producer(producer_group, stages)
for i in range(iterations):
producer.acquire() # Wait for buffer to be empty
# Produce data
producer.commit() # Signal data is ready
producer.advance() # Move to next stage
"""
_pipeline: PipelineAsync
_state: PipelineState
_group: CooperativeGroup
def __init__(self, pipeline, state, group: CooperativeGroup):
"""Initialize a new Producer instance.
:param pipeline: The pipeline this producer belongs to
:type pipeline: PipelineAsync
:param state: Initial pipeline state
:type state: PipelineState
:param group: The cooperative group for synchronization
:type group: CooperativeGroup
"""
self._pipeline = pipeline
self._state = state
self._group = group
@property
def index(self):
"""Get the index of the current pipeline stage."""
return self._state.index
def get_barrier(self):
"""Get the barrier pointer for the current pipeline stage.
:return: Pointer to the barrier for the current stage
:rtype: cute.Pointer
"""
return self._pipeline.producer_get_barrier(self._state)
def acquire(self):
"""Wait for the current buffer to be empty before producing data.
This is a blocking operation.
"""
self._pipeline.producer_acquire(self._state)
def try_acquire(self):
"""Try to acquire the current buffer without blocking.
:return: True if acquisition was successful, False otherwise
:rtype: bool
"""
self._pipeline.producer_try_acquire(self._state)
def commit(self):
"""Signal that data production is complete for the current stage.
This allows consumers to start processing the data.
"""
self._pipeline.producer_commit(self._state)
def tail(self):
"""Ensure all used buffers are properly synchronized before producer exit.
This should be called before the producer finishes to avoid dangling signals.
"""
self._pipeline.producer_tail(self._state)
def advance(self):
"""Move to the next pipeline stage."""
self._state.advance()
def __extract_mlir_values__(self):
"""Extract MLIR values from the current state.
:return: List of MLIR values representing the current state
:rtype: list
"""
# TODO: need to handle pipeline as well
return self._state.__extract_mlir_values__()
def __new_from_mlir_values__(self, values):
"""Create a new Producer instance from MLIR values.
:param values: MLIR values to initialize the state
:type values: Any
:return: New Producer instance with state initialized from values
:rtype: Producer
"""
return PipelineProducer(
self._pipeline, self._state.__new_from_mlir_values__(values), self._group
)
class PipelineConsumer:
"""A class representing a consumer in an asynchronous pipeline.
The Consumer class manages the consumer side of an asynchronous pipeline, handling
synchronization and state management for consuming data. It provides methods for
waiting, releasing, and advancing through pipeline stages.
:ivar _pipeline: The asynchronous pipeline this consumer belongs to
:type _pipeline: PipelineAsync
:ivar _state: The current state of the consumer in the pipeline
:type _state: PipelineState
:ivar _group: The cooperative group this consumer operates in
:type _group: CooperativeGroup
**Examples:**
.. code-block:: python
pipeline = PipelineAsync.create(...)
consumer = pipeline.create_consumer(consumer_group, stages)
for i in range(iterations):
consumer.wait() # Wait for data to be ready
# Consume data
consumer.release() # Signal buffer is empty
consumer.advance() # Move to next stage
"""
_pipeline: PipelineAsync
_state: PipelineState
_group: CooperativeGroup
def __init__(self, pipeline, state: PipelineState, group: CooperativeGroup):
"""Initialize a new Consumer instance.
:param pipeline: The pipeline this consumer belongs to
:type pipeline: PipelineAsync
:param state: Initial pipeline state
:type state: PipelineState
:param group: The cooperative group for synchronization
:type group: CooperativeGroup
"""
self._pipeline = pipeline
self._group = group
self._state = state
@property
def index(self):
"""Get the index of the current pipeline stage."""
return self._state.index
def wait(self):
"""Wait for data to be ready in the current buffer.
This is a blocking operation.
"""
self._pipeline.consumer_wait(self._state)
def try_wait(self):
"""Try to check if data is ready without blocking.
:return: True if data is ready, False otherwise
:rtype: bool
"""
self._pipeline.consumer_try_wait(self._state)
def release(self):
"""Signal that data consumption is complete for the current stage.
This allows producers to start producing new data.
"""
self._pipeline.consumer_release(self._state)
def advance(self):
"""Move to the next pipeline stage."""
self._state.advance()
def __extract_mlir_values__(self):
"""Extract MLIR values from the current state.
:return: List of MLIR values representing the current state
:rtype: list
"""
return self._state.__extract_mlir_values__()
def __new_from_mlir_values__(self, values):
"""Create a new Consumer instance from MLIR values.
:param values: MLIR values to initialize the state
:type values: Any
:return: New Consumer instance with state initialized from values
:rtype: Consumer
"""
# TODO: need to call pipeline.__new_from_mlir_values__ recursively
return PipelineConsumer(
self._pipeline, self._state.__new_from_mlir_values__(values), self._group
)

123
python/CuTeDSL/cutlass/torch.py
View File

@@ -29,6 +29,7 @@ from cutlass.cute.typing import (
from cutlass.cute.runtime import from_dlpack
import cutlass.cute as cute
import torch
from cuda import cuda
def dtype(ty: Type[Numeric]):
@@ -94,12 +95,13 @@ def create_and_permute_torch_tensor(
init_config: Optional[
Union[RandomInitConfig, ScalarInitConfig, GaussianInitConfig]
] = None,
device: Optional[torch.device] = None,
) -> "torch.Tensor":
"""
Create a torch tensor with specified shape and dtype. Optionally permute it and initialize it with specified init type and config
"""
init_dtype = torch.int32 if init_type == TensorInitType.RANDOM else torch.float32
init_torch_tensor = torch.empty(*shape, dtype=init_dtype)
init_torch_tensor = torch.empty(*shape, dtype=init_dtype, device=device)
if init_type == TensorInitType.SKIP:
assert init_config is None
f32_torch_tensor = init_torch_tensor
@@ -167,3 +169,122 @@ def convert_cute_tensor(
# Copy and convert from f32 cute tensor to dtype cute tensor
cute.testing.convert(fp32_cute_tensor, cute_tensor)
return cute_tensor
def default_stream() -> cuda.CUstream:
"""
Get default CUstream from torch stream
"""
torch_stream = torch.cuda.default_stream()
stream = cuda.CUstream(torch_stream.cuda_stream)
return stream
def current_stream() -> cuda.CUstream:
"""
Get current CUstream from torch stream
"""
torch_stream = torch.cuda.current_stream()
stream = cuda.CUstream(torch_stream.cuda_stream)
return stream
def matrix(
l: int,
mode0: int,
mode1: int,
is_mode0_major: bool,
cutlass_dtype: Type[Numeric],
init_type: TensorInitType = TensorInitType.RANDOM,
init_config: Optional[
Union[RandomInitConfig, ScalarInitConfig, GaussianInitConfig]
] = None,
device: Optional[torch.device] = None,
) -> torch.Tensor:
"""
Create a torch tensor for matrix
:param l: length of the matrix
:param mode0: mode0 of the matrix
:param mode1: mode1 of the matrix
:param is_mode0_major: whether the matrix is mode0 major
:param cutlass_dtype: cutlass dtype of the matrix
:param init_type: type of initialization
:param init_config: configuration for initialization
:param device: target torch device
"""
shape = (l, mode1, mode0) if is_mode0_major else (l, mode0, mode1)
permute_order = (2, 1, 0) if is_mode0_major else (1, 2, 0)
if cutlass_dtype.is_float and cutlass_dtype.width <= 8:
torch_dtype = torch.int8
else:
torch_dtype = dtype(cutlass_dtype)
if init_type == TensorInitType.RANDOM and init_config is None:
if torch_dtype.is_signed:
min_val = -2
max_val = 2
else:
min_val = 0
max_val = 4
init_config = RandomInitConfig(min_val=min_val, max_val=max_val)
# Create dtype torch tensor
torch_tensor = create_and_permute_torch_tensor(
shape,
torch_dtype,
permute_order=permute_order,
init_type=init_type,
init_config=init_config,
device=device,
)
return torch_tensor
def cute_tensor_like(
data_ref: torch.Tensor,
cutlass_dtype: Type[Numeric],
is_dynamic_layout: bool,
assumed_align: Optional[int] = None,
) -> tuple[Tensor, torch.Tensor]:
"""
Create a cute tensor use a torch tensor as the data source
:param data_ref: torch tensor as the data source
:param cutlass_dtype: cutlass dtype of the cute tensor
:param is_dynamic_layout: whether the cute tensor uses dynamic layout
:param assumed_align: assumed alignment of the cute tensor
"""
# allocate device buffer for cute tensor
if cutlass_dtype.is_float and cutlass_dtype.width <= 8:
torch_dtype = torch.int8
else:
torch_dtype = dtype(cutlass_dtype)
torch_tensor = torch.empty_like(data_ref, dtype=torch_dtype, device="cuda")
# create cute tensor using the device buffer
cute_tensor = from_dlpack(torch_tensor, assumed_align=assumed_align)
cute_tensor.element_type = cutlass_dtype
if is_dynamic_layout:
for i, stride in enumerate(torch_tensor.stride()):
if stride == 1:
leading_dim = i
break
cute_tensor = cute_tensor.mark_layout_dynamic(leading_dim=leading_dim)
# initialize the cute tensor data
if cutlass_dtype.is_float and cutlass_dtype.width <= 8:
cute_tensor = convert_cute_tensor(
data_ref.to(dtype=torch.float32),
cute_tensor,
cutlass_dtype,
is_dynamic_layout,
)
else:
torch_tensor.copy_(data_ref.to(dtype=torch_dtype))
return cute_tensor, torch_tensor

16
python/CuTeDSL/cutlass/utils/__init__.py
View File

@@ -15,20 +15,6 @@ from .static_persistent_tile_scheduler import (
StaticPersistentTileScheduler,
)
from .pipeline import (
Agent,
CooperativeGroup,
PipelineUserType,
PipelineState,
make_pipeline_state,
PipelineAsync,
PipelineTmaAsync,
PipelineTmaUmma,
PipelineUmmaAsync,
PipelineTmaStore,
pipeline_init_wait,
)
from .hardware_info import (
HardwareInfo,
)
@@ -65,6 +51,8 @@ from .smem_allocator import SmemAllocator
from .layout import LayoutEnum
__all__ = [
"SmemAllocator",
"LayoutEnum",
"WorkTileInfo",
"PersistentTileSchedulerParams",
"StaticPersistentTileScheduler",

143
python/CuTeDSL/cutlass/utils/blackwell_helpers.py
View File

@@ -51,8 +51,13 @@ from cutlass.cute.nvgpu.tcgen05 import (
is_tmem_load,
get_tmem_copy_properties,
)
from cutlass.cute.nvgpu.cpasync import (
CopyBulkTensorTileG2SMulticastOp,
CopyBulkTensorTileG2SOp,
)
from cutlass.utils.layout import LayoutEnum
@dsl_user_op
def compute_epilogue_tile_shape(
cta_tile_shape: cute.Shape,
@@ -716,6 +721,7 @@ def make_smem_layout_b(
return b_smem_layout_staged
@dsl_user_op
def get_smem_layout_atom_epi(
layout: LayoutEnum,
@@ -827,6 +833,7 @@ SMEM_CAPACITY = {
"sm120": SmemCapacity.SM120_SMEM_CAPACITY_BYTES.value,
}
@dsl_user_op
def make_trivial_tiled_mma(
ab_dtype: Type[Numeric],
@@ -908,3 +915,139 @@ def make_trivial_tiled_mma(
raise TypeError(f"unsupported ab_dtype, got {ab_dtype}")
return cute.make_tiled_mma(cute.make_mma_atom(mma_op))
@dsl_user_op
def cluster_shape_to_tma_atom_A(
cluster_shape_mnk: cute.Shape, atom_thr_id: cute.Layout, *, loc=None, ip=None
) -> Union[CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileG2SOp]:
"""
Select the appropriate TMA copy atom for A based on the number of SMs and the multicast flag.
:param cluster_shape_mnk: The shape of the cluster
:type cluster_shape_mnk: cute.Shape
:param atom_thr_id: The thread ID of the atom
:type atom_thr_id: cute.Layout
:return: The appropriate TMA copy atom kind
:rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp
:raise ValueError: If the atom_sm_cnt is invalid
:raise ValueError: If the cluster shape is not divisible by the atom SM count
"""
atom_sm_cnt = cute.size(atom_thr_id, loc=loc, ip=ip)
mcast = not (cute.size(cluster_shape_mnk, mode=[1], loc=loc, ip=ip) == 1)
cluster_size = cute.size(cluster_shape_mnk, loc=loc, ip=ip)
if not isinstance(cluster_size, int) or not isinstance(atom_sm_cnt, int):
raise ValueError(
f"Dynamic cluster shape or atom SM count is not supported: {cluster_shape_mnk} and {atom_thr_id}"
)
if cute.size(cluster_shape_mnk, mode=[0], loc=loc, ip=ip) % atom_sm_cnt != 0:
raise ValueError(
f"Cluster shape not divisible by MMA size: {cluster_shape_mnk} and {atom_thr_id}"
)
if atom_sm_cnt == 2 and mcast:
return CopyBulkTensorTileG2SMulticastOp(CtaGroup.TWO)
elif atom_sm_cnt == 2 and not mcast:
return CopyBulkTensorTileG2SOp(CtaGroup.TWO)
elif atom_sm_cnt == 1 and mcast:
return CopyBulkTensorTileG2SMulticastOp(CtaGroup.ONE)
elif atom_sm_cnt == 1 and not mcast:
return CopyBulkTensorTileG2SOp(CtaGroup.ONE)
raise ValueError(
f"Unsupported Configuration for SM100 TMA: {cluster_shape_mnk} and {atom_thr_id}"
)
@dsl_user_op
def cluster_shape_to_tma_atom_B(
cluster_shape_mnk: cute.Shape, atom_thr_id: cute.Layout, *, loc=None, ip=None
) -> Union[CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileG2SOp]:
"""
Select the appropriate TMA copy atom for Bbased on the number of SMs and the multicast flag.
:param cluster_shape_mnk: The shape of the cluster
:type cluster_shape_mnk: cute.Shape
:param atom_thr_id: The thread ID of the atom
:type atom_thr_id: cute.Layout
:return: The appropriate TMA copy atom kind
:rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp
:raise ValueError: If the atom_sm_cnt is invalid
:raise ValueError: If the cluster shape is not divisible by the atom SM count
"""
atom_sm_cnt = cute.size(atom_thr_id, loc=loc, ip=ip)
mcast = not (cute.size(cluster_shape_mnk, mode=[0], loc=loc, ip=ip) == atom_sm_cnt)
cluster_size = cute.size(cluster_shape_mnk, loc=loc, ip=ip)
if not isinstance(cluster_size, int) or not isinstance(atom_sm_cnt, int):
raise ValueError(
f"Dynamic cluster shape or atom SM count is not supported: {cluster_shape_mnk} and {atom_thr_id}"
)
if cute.size(cluster_shape_mnk, mode=[0], loc=loc, ip=ip) % atom_sm_cnt != 0:
raise ValueError(
f"Cluster shape not divisible by MMA size: {cluster_shape_mnk} and {atom_thr_id}"
)
if atom_sm_cnt == 2 and mcast:
return CopyBulkTensorTileG2SMulticastOp(CtaGroup.TWO)
elif atom_sm_cnt == 2 and not mcast:
return CopyBulkTensorTileG2SOp(CtaGroup.TWO)
elif atom_sm_cnt == 1 and mcast:
return CopyBulkTensorTileG2SMulticastOp(CtaGroup.ONE)
elif atom_sm_cnt == 1 and not mcast:
return CopyBulkTensorTileG2SOp(CtaGroup.ONE)
raise ValueError(
f"Unsupported Configuration for SM100 TMA: {cluster_shape_mnk} and {atom_thr_id}"
)
@dsl_user_op
def cluster_shape_to_tma_atom_SFB(
cluster_shape_mnk: cute.Shape, atom_thr_id: cute.Layout, *, loc=None, ip=None
) -> Union[CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileG2SOp]:
"""
Select the appropriate TMA copy atom for SFB based on the number of SMs and the multicast flag.
:param cluster_shape_mnk: The shape of the cluster
:type cluster_shape_mnk: cute.Shape
:param atom_thr_id: The thread ID of the atom
:type atom_thr_id: cute.Layout
:return: The appropriate TMA copy atom kind
:rtype: cpasync.CopyBulkTensorTileG2SMulticastOp or cpasync.CopyBulkTensorTileG2SOp
:raise ValueError: If the atom_sm_cnt is invalid
:raise ValueError: If the cluster shape is not divisible by the atom SM count
"""
atom_sm_cnt = cute.size(atom_thr_id, loc=loc, ip=ip)
mcast = not (cute.size(cluster_shape_mnk, mode=[0], loc=loc, ip=ip) == 1)
cluster_size = cute.size(cluster_shape_mnk, loc=loc, ip=ip)
if not isinstance(cluster_size, int) or not isinstance(atom_sm_cnt, int):
raise ValueError(
f"Dynamic cluster shape or atom SM count is not supported: {cluster_shape_mnk} and {atom_thr_id}"
)
if cute.size(cluster_shape_mnk, mode=[0], loc=loc, ip=ip) % atom_sm_cnt != 0:
raise ValueError(
f"Cluster shape not divisible by MMA size: {cluster_shape_mnk} and {atom_thr_id}"
)
if atom_sm_cnt == 2:
return CopyBulkTensorTileG2SMulticastOp(CtaGroup.TWO)
elif atom_sm_cnt == 1 and mcast:
return CopyBulkTensorTileG2SMulticastOp(CtaGroup.ONE)
elif atom_sm_cnt == 1 and not mcast:
return CopyBulkTensorTileG2SOp(CtaGroup.ONE)
raise ValueError(
f"Unsupported Configuration for SM100 TMA: {cluster_shape_mnk} and {atom_thr_id}"
)

8
python/CuTeDSL/cutlass/utils/hopper_helpers.py
View File

@@ -96,6 +96,10 @@ def make_trivial_tiled_mma(
acc_dtype: Type[Numeric],
atom_layout_mnk: Tuple[int, int, int],
tiler_mn: Tuple[int, int],
a_source: OperandSource = OperandSource.SMEM,
*,
loc=None,
ip=None,
) -> cute.TiledMma:
"""Make a tiled MMA atom with given data type, leading dimension, cta group and mma tile shape.
By default, the MMA atom is created with SMEM operand source for A.
@@ -131,7 +135,7 @@ def make_trivial_tiled_mma(
a_dtype,
acc_dtype,
(*tiler_mn, 16),
OperandSource.SMEM,
a_source,
a_leading_mode,
b_leading_mode,
)
@@ -144,7 +148,7 @@ def make_trivial_tiled_mma(
b_dtype,
acc_dtype,
(*tiler_mn, 32),
OperandSource.SMEM,
a_source,
a_leading_mode,
b_leading_mode,
)

1023
python/CuTeDSL/cutlass/utils/pipeline.py
View File

File diff suppressed because it is too large Load Diff

137
python/CuTeDSL/cutlass/utils/smem_allocator.py
View File

@@ -18,41 +18,25 @@ from cutlass.cute.arch import get_dyn_smem
class SmemAllocator:
"""
A class for managing shared memory allocation on GPU.
"""A class for managing shared memory allocation on GPU.
This class manages a chunk of shared memory and provide APIs for sub-allocation
This class manages a chunk of shared memory and provides APIs for sub-allocation
inside the chunk.
Attributes
----------
_base : cute.Pointer as i8 typed dynamic value
The current base address of the shared memory.
:ivar _base: The current base address of the shared memory as an i8 typed dynamic value.
:type _base: cute.Pointer
:ivar _allocated_bytes: The total number of bytes allocated in shared memory.
:type _allocated_bytes: int
_allocated_bytes:
The bytes allocated in shared memory.
Methods
-------
allocate(num_bytes, alignment)
Allocates num_bytes in the shared memory with the given byte alignment.
allocate_value(value_ty, num_elems)
Allocates num_elems of value_ty values in the shared memory.
allocate_tensor(value_ty, layout, alignment)
Allocates a tensor in the shared memory with given layout and byte alignment.
Notes
-----
This class is responsible for managing the allocation of tensors in shared memory.
.. note::
This class is responsible for managing the allocation of tensors in shared memory.
The base pointer is aligned to 1024 bytes upon initialization.
"""
def __init__(self):
"""
Initializes the SmemAllocator instance with dynamic smem base ptr,
which is i8 type and aligned to 1024.
"""Initialize the SmemAllocator instance.
Creates a dynamic shared memory base pointer of type i8, aligned to 1024 bytes.
"""
self._base = get_dyn_smem(Int8, alignment=1024)
self._allocated_bytes = 0
@@ -64,30 +48,19 @@ class SmemAllocator:
def allocate(self, size_or_type: cute.struct, byte_alignment: int): ...
def allocate(self, size_or_type, byte_alignment: int = 1) -> int:
"""Allocate a block of memory with specified size and alignment.
This method adjusts the base pointer to ensure proper alignment and updates
the internal state to track allocated memory.
:param size_or_type: The number of bytes to allocate or a struct class
:type size_or_type: Union[int, cute.struct]
:param byte_alignment: The byte alignment requirement, defaults to 1 (no alignment)
:type byte_alignment: int, optional
:return: Pointer to the start of the allocated memory block or struct instance
:rtype: cute.Pointer
:raises ValueError: If size is negative or alignment is less than 1
"""
Allocates a block of memory with the specified size and byte alignment.
This method adjusts the base cute.Pointer to ensure that the allocated memory
is aligned according to the specified byte alignment. It updates the internal
state to reflect the new base cute.Pointer and the total allocated bytes.
Parameters
----------
size_or_type : int or struct
The number of bytes to allocate or struct class.
byte_alignment : int
The byte alignment requirement for the allocation. Defaults to 1 (no alignment).
Returns
----------
A cute.Pointer to the start of the allocated memory block or struct instance.
Raises
----------
ValueError
If num_bytes is negative or if byte_alignmemt is less than 1.
"""
if isinstance(size_or_type, cute.struct):
alignment = max(byte_alignment, size_or_type.__alignof__())
base_ptr = self.allocate(size_or_type.__sizeof__(), alignment)
@@ -110,27 +83,16 @@ class SmemAllocator:
return ptr
def allocate_array(self, element_type: Type[Numeric], num_elems: int = 1):
"""
Allocates num_elems values of element_type in shared memory.
"""Allocate an array of elements in shared memory.
This method calls allocate() to return a byte ptr, pointing to start of shared
memory. Then calls cute.recast_ptr() to recast this byte cute.Pointer to element_type.
Parameters
----------
element_type : Type[Numeric]
The type of the values in the tensor.
num_elems : int, optional
The number of elements for each allocation. Defaults to 1.
Returns
----------
A value_type cute.Pointer to the start of the allocated memory block.
Raises
----------
ValueError
If num_elems is less than 1.
:param element_type: The type of elements to allocate
:type element_type: Type[Numeric]
:param num_elems: Number of elements to allocate, defaults to 1
:type num_elems: int, optional
:return: Pointer to the start of the allocated array
:rtype: cute.Pointer
:raises ValueError: If num_elems is less than 1
:raises TypeError: If element_type is not a Numeric type
"""
if num_elems < 1:
raise ValueError("num_elems must be at least 1")
@@ -152,28 +114,21 @@ class SmemAllocator:
byte_alignment: int = 1,
swizzle: cute.Swizzle = None,
):
"""
Allocates a tensor in the shared memory with value type, layout and byte alignment.
"""Allocate a tensor in shared memory.
Parameters
----------
element_type : Type[Numeric]
The type of the values in the tensor.
layout : int | DynamicInt | cute.Layout | cute.ComposedLayout
The layout of the tensor.
byte_alignment : int, optional
The byte alignment requirement for the allocation. Defaults to 1 (no alignment).
swizzle : cute.Swizzle
A swizzle for the iterator (for position-dependent swizzling).
Returns
-------
tensor : cute.Tensor
The allocated tensor with specified value type, layout and byte alignment.
Notes
-----
The base address is updated to point to the next available memory location.
:param element_type: The type of elements in the tensor
:type element_type: Type[Numeric]
:param layout: The layout specification for the tensor
:type layout: Union[int, cute.Layout, cute.ComposedLayout]
:param byte_alignment: The byte alignment requirement, defaults to 1
:type byte_alignment: int, optional
:param swizzle: Swizzle for position-dependent swizzling, defaults to None
:type swizzle: cute.Swizzle, optional
:return: The allocated tensor with specified properties
:rtype: cute.Tensor
:raises TypeError: If element_type is not a Numeric type or if swizzle conflicts with layout
:raises ValueError: If allocation is not byte-aligned
:raises NotImplementedError: If dynamic layout is specified
"""
if not isinstance(element_type, NumericMeta):
raise TypeError(