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

Merge branch 'main' into tvm-ffi

Browse Source
This commit is contained in:
Yuan Xiaolan
2026-02-14 13:35:28 +08:00
committed by GitHub
parent de161925a5 d4bbf728ca
commit 395ab575f6
201 changed files with 55213 additions and 4810 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/CuTeDSL/cutlass/cute/__init__.py
+53 -131
View File
@@ -181,19 +181,21 @@ from .atom import (
make_tiled_copy_C_atom,
make_cotiled_copy,
copy_atom_call,
mma_atom_call,
)
from .algorithm import gemm, copy, basic_copy, basic_copy_if, autovec_copy, prefetch
from . import typing as typing_module
from . import core
from . import arch
from . import export
from . import nvgpu
from . import testing
from . import runtime
from . import math
# Export all math ops without "math."
from .math import *
@@ -212,7 +214,6 @@ GenerateLineInfo = _dsl.GenerateLineInfo
KeepCUBIN = _dsl.KeepCUBIN
KeepPTX = _dsl.KeepPTX
GPUArch = _dsl.GPUArch
LinkLibraries = _dsl.LinkLibraries
EnableTVMFFI = _dsl.EnableTVMFFI
# attach the TVM FFI ABI interface postprocessor to the DSL
@@ -222,16 +223,52 @@ _tvm_ffi_args_spec_converter.attach_args_spec_converter(_dsl.CuTeDSL._get_dsl())
# Explicitly export all symbols for documentation generation
__all__ = [
# Core types
*core.__all__,
# ==================== cutlass._mlir.dialects.cute ====================
"AddressSpace",
"CacheEvictionPriority",
# ==================== .typing ====================
"Tensor",
"Layout",
"ComposedLayout",
"Swizzle",
"E",
"ScaledBasis",
"SymInt",
"is_integer",
"is_int_tuple",
# ==================== .core ====================
*core.__all__,
# ==================== .tuple ====================
"transform_leaf",
"find_if",
"find",
"flatten_to_tuple",
"unflatten",
"product",
"product_like",
"product_each",
"elem_less",
"tuple_cat",
"transform_apply",
"filter_tuple",
# ==================== .tensor ====================
"TensorSSA",
"ReductionOp",
"make_tensor",
"make_identity_tensor",
"make_fragment",
"make_fragment_like",
"make_rmem_tensor_like",
"make_rmem_tensor",
"recast_tensor",
"domain_offset",
"print_tensor",
"full",
"full_like",
"empty_like",
"ones_like",
"zeros_like",
"where",
"any_",
"all_",
# ==================== .atom ====================
"Atom",
"MmaAtom",
"CopyAtom",
@@ -239,106 +276,6 @@ __all__ = [
"TiledMma",
"ThrMma",
"ThrCopy",
"TensorSSA",
"ReductionOp",
"SymInt",
# Basic utility functions
"assume",
"is_integer",
"is_int_tuple",
"is_static",
"has_underscore",
"shape",
"printf",
"print_tensor",
"pretty_str",
# Layout functions
"make_layout",
"recast_layout",
"make_identity_layout",
"make_ordered_layout",
"make_layout_like",
"make_composed_layout",
"make_layout_tv",
"make_layout_image_mask",
"get_nonswizzle_portion",
"get_swizzle_portion",
# Tensor functions
"make_ptr",
"make_tensor",
"make_identity_tensor",
"make_fragment",
"make_fragment_like",
"make_rmem_tensor",
"make_rmem_tensor_like",
"recast_ptr",
"recast_tensor",
# Tensor manipulation
"get",
"select",
"front",
"is_major",
"leading_dim",
"find",
"find_if",
"transform_leaf",
"basis_value",
"basis_get",
"coalesce",
"group_modes",
"cosize",
"size_in_bytes",
# Tuple operations
"flatten_to_tuple",
"flatten",
"unflatten",
"product",
"product_like",
"product_each",
"prepend",
"append",
"prepend_ones",
"append_ones",
"elem_less",
"tuple_cat",
"transform_apply",
"filter_tuple",
# Math operations
"ceil_div",
"round_up",
# Layout operations
"slice_and_offset",
"crd2idx",
"domain_offset",
"filter_zeros",
"filter",
"tile_to_shape",
"shape_div",
"dice",
# Layout algebra
"composition",
"complement",
"right_inverse",
"left_inverse",
"max_common_layout",
"max_common_vector",
"is_congruent",
"is_weakly_congruent",
# Product operations
"logical_product",
"zipped_product",
"tiled_product",
"flat_product",
"raked_product",
"blocked_product",
# Division operations
"flat_divide",
"logical_divide",
"zipped_divide",
"tiled_divide",
"local_partition",
"local_tile",
# MMA and Copy atom operations
"make_atom",
"make_mma_atom",
"make_tiled_mma",
@@ -353,39 +290,24 @@ __all__ = [
"make_tiled_copy_C_atom",
"make_cotiled_copy",
"copy_atom_call",
# Algorithm operations
"mma_atom_call",
# ==================== .algorithm ====================
"gemm",
"copy",
"basic_copy",
"basic_copy_if",
"autovec_copy",
"copy",
"prefetch",
"gemm",
# Tensor creation
"full",
"full_like",
"empty_like",
"ones_like",
"zeros_like",
"where",
"any_",
"all_",
"repeat_as_tuple",
"repeat",
"repeat_like",
# User defined struct
"struct",
# FastDivmod operations
"FastDivmodDivisor",
"fast_divmod_create_divisor",
# Modules
# ==================== .extension ====================
# ==================== .math ====================
*math.__all__,
# ==================== submodules ====================
"arch",
"export",
"nvgpu",
"testing",
"runtime",
# Math utils
*math.__all__,
# Decorators and code generation
# ==================== DSL (cutlass_dsl) ====================
"jit",
"kernel",
"register_jit_arg_adapter",
python/CuTeDSL/cutlass/cute/algorithm.py
+101 -37
View File
@@ -10,7 +10,7 @@
# is strictly prohibited.
import math
from typing import Optional, Dict, Any, List, Tuple
from typing import Optional, Dict, Any, List, Tuple, Union
from cutlass._mlir import ir
from cutlass.cutlass_dsl import for_generate, yield_out, if_generate, dsl_user_op
@@ -29,15 +29,35 @@ from .core import (
append_ones,
group_modes,
)
from .atom import MmaAtom, CopyAtom, make_atom
from .atom import (
MmaAtom,
CopyAtom,
make_atom,
_normalize_variadic_tensor_operand,
copy_atom_call,
)
from .nvgpu.common import CacheEvictionPriority
def _normalize_gemm_operand_list(
x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str
) -> List["Tensor"]:
if isinstance(x, Tensor):
return [x]
if isinstance(x, (list, tuple)):
if len(x) == 0:
raise ValueError(f"`{name}` must contain at least one Tensor")
if not all(isinstance(t, Tensor) for t in x):
raise TypeError(f"All elements of `{name}` must be Tensor")
return list(x) # type: ignore
raise TypeError(f"`{name}` must be a Tensor or a sequence of Tensors")
@dsl_user_op
def gemm(
atom: MmaAtom,
d: Tensor,
a: Tensor,
b: Tensor,
a: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
b: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
c: Tensor,
*,
loc=None,
@@ -62,14 +82,17 @@ def gemm(
- Dispatch [4]: (V,M) x (V,N) => (V,M,N) => (V,M,1) x (V,N,1) => (V,M,N)
- Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N)
Operand flexibility:
- `a` and `b` can be a single Tensor (regular GEMM) or a sequence `[operand, scale_factor]` for block-scaled GEMM.
:param atom: MMA atom
:type atom: MmaAtom
:param d: Destination tensor
:type d: Tensor
:param a: First source tensor
:type a: Tensor
:param b: Second source tensor
:type b: Tensor
:param a: First source tensor or sequence for advanced modes (e.g., `[a, sfa]`)
:type a: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param b: Second source tensor or sequence for advanced modes (e.g., `[b, sfb]`)
:type b: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param c: Third source tensor
:type c: Tensor
:param loc: Source location for MLIR, defaults to None
@@ -82,8 +105,13 @@ def gemm(
:rtype: None
"""
a_rank = rank(a.shape)
b_rank = rank(b.shape)
# Normalize A/B to lists for variadic IR operands, while keeping old API working.
a_list = _normalize_gemm_operand_list(a, "a")
b_list = _normalize_gemm_operand_list(b, "b")
# Rank validations based on the primary A/B tensors (guaranteed non-empty)
a_rank = rank(a_list[0].shape)
b_rank = rank(b_list[0].shape)
c_rank = rank(c.shape)
d_rank = rank(d.shape)
@@ -104,7 +132,9 @@ def gemm(
raise ValueError("`c` must have rank 3 when `a` has rank 3")
value = atom._unpack(loc=loc, ip=ip, **kwargs)
return _cute_ir.gemm(value, d.value, a.value, b.value, c.value, loc=loc, ip=ip)
a_vals = [t.value for t in a_list]
b_vals = [t.value for t in b_list]
return _cute_ir.gemm(value, d.value, a_vals, b_vals, c.value, loc=loc, ip=ip)
@dsl_user_op
@@ -132,7 +162,7 @@ def basic_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
src.element_type.mlir_type, src.element_type.width
)
simt_copy = make_atom(simt_copy_ty, loc=loc, ip=ip)
return _cute_ir.copy(simt_copy, src.value, dst.value, loc=loc, ip=ip)
return _cute_ir.copy(simt_copy, [src.value], [dst.value], loc=loc, ip=ip)
s = size(dst, loc=loc, ip=ip)
# Always generate an scf.for Op when one of the tensors is dynamic
@@ -186,7 +216,14 @@ def _basic_copy_if_static(
@dsl_user_op
def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
def autovec_copy(
src: Tensor,
dst: Tensor,
*,
l1c_evict_priority: CacheEvictionPriority = CacheEvictionPriority.EVICT_NORMAL,
loc=None,
ip=None,
) -> None:
"""
Auto-vectorization SIMT copy policy.
@@ -239,11 +276,15 @@ def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
# Dispatch to copy with atom
simt_type = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
src.element_type.mlir_type, num_bits_per_copy
src.element_type.mlir_type,
num_bits_per_copy,
0,
0,
l1c_evict_priority._to_ir(),
)
simt_copy = make_atom(simt_type, loc=loc, ip=ip)
return _cute_ir.copy(
simt_copy, tiled_src.value, tiled_dst.value, loc=loc, ip=ip
simt_copy, [tiled_src.value], [tiled_dst.value], loc=loc, ip=ip
)
# Failed to vectorize, use a basic copy
@@ -258,19 +299,21 @@ def _parse_auto_multicast_args(
This function consumes the following key from kwargs if present:
- 'auto_multicast': dict
dict: { 'multicast_layout': str, 'use_2cta': bool }
dict: { 'multicast_layout': str, 'use_2cta': bool, 'from_block_api': bool }
Returns:
List of (attr_name, ir.Attribute) pairs to be attached to the op.
Recognized attributes:
- ('multicast_layout', #cute.layout<...>) when a layout string is provided
- ('use_2cta', unit) when use_2cta is True
- ('from_block_api', unit) when from_block_api is True
"""
attr_pairs: List[Tuple[str, ir.Attribute]] = []
# Pop known keys to avoid leaking to trait unpack
auto_multicast = kwargs.pop("auto_multicast", None)
from_block_api: bool = False
use_2cta: bool = False
layout_str: Optional[str] = None
@@ -281,6 +324,7 @@ def _parse_auto_multicast_args(
)
layout_str = auto_multicast.get("multicast_layout", None)
use_2cta = bool(auto_multicast.get("use_2cta", False))
from_block_api = bool(auto_multicast.get("from_block_api", False))
if layout_str is not None:
if not isinstance(layout_str, str):
@@ -293,7 +337,8 @@ def _parse_auto_multicast_args(
ir.Attribute.parse(f'#cute.layout<"{layout_str}">'),
)
)
if from_block_api:
attr_pairs.append(("from_block_api", ir.UnitAttr.get()))
if use_2cta:
attr_pairs.append(("use_2cta", ir.UnitAttr.get()))
@@ -303,8 +348,8 @@ def _parse_auto_multicast_args(
@dsl_user_op
def copy(
atom: CopyAtom,
src: Tensor,
dst: Tensor,
src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
*,
pred: Optional[Tensor] = None,
loc=None,
@@ -315,10 +360,10 @@ def copy(
:param atom: Copy atom specifying the transfer operation
:type atom: CopyAtom
:param src: Source tensor with layout profile ``(V, Rest...)``
:type src: Tensor
:param dst: Destination tensor with layout profile ``(V, Rest...)``
:type dst: Tensor
:param src: Source tensor or list of tensors with layout profile ``(V, Rest...)``
:type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param dst: Destination tensor or list of tensors with layout profile ``(V, Rest...)``
:type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param pred: Optional predication tensor for conditional transfers, defaults to None
:type pred: Optional[Tensor], optional
:param loc: Source location information, defaults to None
@@ -346,6 +391,12 @@ def copy(
Source and destination tensors must be partitioned in accordance with the Copy Atom specifications.
Post-partitioning, both tensors will exhibit a ``(V, Rest...)`` layout profile.
The operands `src` and `dst` are variadic, each containing a variable number of tensors:
- For regular copy, `src` and `dst` contain single source and destination tensors respectively.
- For copy with auxiliary operands, `src` and `dst` contain the primary tensors followed by
their respective auxiliary tensors.
**Precondition:** The size of mode 1 must be equal for both source and destination tensors:
``size(src, mode=[1]) == size(dst, mode=[1])``
@@ -371,41 +422,54 @@ def copy(
for future releases.
"""
if isinstance(src.type, _cute_ir.MemRefType) and isinstance(
dst.type, _cute_ir.MemRefType
# Normalize src/dst to lists for variadic IR operands
src_list = _normalize_variadic_tensor_operand(src, "src")
dst_list = _normalize_variadic_tensor_operand(dst, "dst")
# Validate primary tensors (first element)
src_primary = src_list[0]
dst_primary = dst_list[0]
if isinstance(src_primary.type, _cute_ir.MemRefType) and isinstance(
dst_primary.type, _cute_ir.MemRefType
):
if src.element_type.width != dst.element_type.width:
if src_primary.element_type.width != dst_primary.element_type.width:
raise TypeError(
"`copy` currently only supports equal source and destination "
"element type bit width"
)
if rank(src) != rank(dst):
if rank(src_primary) != rank(dst_primary):
raise ValueError(
"Expected source and destination tensors to have the same rank, "
f"but got {rank(src)} and {rank(dst)}"
f"but got {rank(src_primary)} and {rank(dst_primary)}"
)
# Canonicalize to at least rank-2 tensors
src = group_modes(append_ones(src, up_to_rank=2), 1)
dst = group_modes(append_ones(dst, up_to_rank=2), 1)
# Canonicalize all tensors to at least rank-2
src_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in src_list]
dst_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in dst_list]
if pred is not None:
pred = group_modes(append_ones(pred, up_to_rank=2), 1)
if is_static(src.shape[1]) and is_static(dst.shape[1]):
if size(src, mode=[1]) != size(dst, mode=[1]):
# Recompute primary references after canonicalization
src_primary = src_list[0]
dst_primary = dst_list[0]
if is_static(src_primary.shape[1]) and is_static(dst_primary.shape[1]):
if size(src_primary, mode=[1]) != size(dst_primary, mode=[1]):
raise ValueError(
"Expected source and destination tensors to have the same size in mode-1, "
f"but got {size(src, mode=[1])} and {size(dst, mode=[1])}"
f"but got {size(src_primary, mode=[1])} and {size(dst_primary, mode=[1])}"
)
multicast_attr_pairs = _parse_auto_multicast_args(kwargs)
value = atom._unpack(loc=loc, ip=ip, **kwargs)
if isinstance(pred, Tensor):
pred = pred.value
pred_value = pred.value if isinstance(pred, Tensor) else pred
op = _cute_ir.copy(value, src.value, dst.value, pred=pred, loc=loc, ip=ip)
src_vals = [t.value for t in src_list]
dst_vals = [t.value for t in dst_list]
op = _cute_ir.copy(value, src_vals, dst_vals, pred=pred_value, loc=loc, ip=ip)
for name, attr in multicast_attr_pairs:
op.attributes[name] = attr
python/CuTeDSL/cutlass/cute/arch/__init__.py
+10 -5
View File
@@ -11,7 +11,6 @@
from .elect import *
from .mbar import *
from .numeric_conversion import *
from .nvvm_wrappers import *
from .smem import *
from .tmem import *
@@ -74,6 +73,8 @@ __all__ = [
"vote_any_sync",
"vote_all_sync",
"vote_uni_sync",
"warp_redux_sync",
"atomic_max_float32",
"atomic_add",
"atomic_and",
"atomic_or",
@@ -95,15 +96,19 @@ __all__ = [
"fma_packed_f32x2",
"mul_packed_f32x2",
"add_packed_f32x2",
"sub_packed_f32x2",
"fmax",
"rcp_approx",
"exp2",
"cvt_i8x4_to_f32x4",
"cvt_i8x2_to_f32x2",
"cvt_i8_bf16",
"cvt_i8x2_to_bf16x2",
"cvt_i8x4_to_bf16x4",
"cvt_f32x2_bf16x2",
"warp_redux_sync",
# Constants
"WARP_SIZE",
# Forward from auto-generated nvvm python
"ProxyKind",
"SharedSpace",
"RoundingModeKind",
#
# smem.py
#
python/CuTeDSL/cutlass/cute/arch/clc.py
+19 -11
View File
@@ -10,8 +10,11 @@
# is strictly prohibited.
from typing import Tuple
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir.dialects import nvvm, vector
from cutlass._mlir import ir
from cutlass._mlir.dialects import nvvm, llvm, vector, arith
from ..typing import Int32, Pointer, Int128
@@ -20,6 +23,7 @@ from ..typing import Int32, Pointer, Int128
def issue_clc_query(
mbar_ptr: Pointer,
clc_response_ptr: Pointer,
multicast: bool = True,
loc=None,
ip=None,
) -> None:
@@ -36,12 +40,20 @@ def issue_clc_query(
"""
mbar_llvm_ptr = mbar_ptr.llvm_ptr
clc_response_llvm_ptr = clc_response_ptr.llvm_ptr
nvvm.clusterlaunchcontrol_try_cancel_multicast(
clc_response_llvm_ptr,
mbar_llvm_ptr,
loc=loc,
ip=ip,
)
if multicast:
nvvm.clusterlaunchcontrol_try_cancel_multicast(
clc_response_llvm_ptr,
mbar_llvm_ptr,
loc=loc,
ip=ip,
)
else:
nvvm.clusterlaunchcontrol_try_cancel(
clc_response_llvm_ptr,
mbar_llvm_ptr,
loc=loc,
ip=ip,
)
@dsl_user_op
@@ -78,7 +90,6 @@ def clc_response(
)
# Query if the cluster was canceled
pred = nvvm.clusterlaunchcontrol_query_cancel_is_canceled(
T.bool(),
clc_result_i128,
loc=loc,
ip=ip,
@@ -87,7 +98,6 @@ def clc_response(
# Get first CTA ID x component
m_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_x(
T.i32(),
clc_result_i128,
loc=loc,
ip=ip,
@@ -95,7 +105,6 @@ def clc_response(
# Get first CTA ID y component
n_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_y(
T.i32(),
clc_result_i128,
loc=loc,
ip=ip,
@@ -103,7 +112,6 @@ def clc_response(
# Get first CTA ID z component
l_idx_i32 = nvvm.clusterlaunchcontrol_query_cancel_get_first_ctaid_z(
T.i32(),
clc_result_i128,
loc=loc,
ip=ip,
python/CuTeDSL/cutlass/cute/arch/elect.py
+1 -2
View File
@@ -9,7 +9,6 @@
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL, T, dsl_user_op
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
@@ -72,6 +71,6 @@ def elect_one(*, loc=None, ip=None) -> IfOpRegion:
from cutlass.base_dsl.arch import Arch
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
is_thread_leader = nvvm.elect_sync(T.bool())
is_thread_leader = nvvm.elect_sync()
if_op = scf.IfOp(is_thread_leader, loc=loc, ip=ip)
return IfOpRegion(if_op.then_block, loc=loc, ip=ip)
python/CuTeDSL/cutlass/cute/arch/mbar.py
+28 -18
View File
@@ -13,7 +13,7 @@ from typing import Optional
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL, T, if_generate, dsl_user_op
from cutlass._mlir.dialects import nvvm
from cutlass._mlir.dialects import nvvm, llvm
from ..typing import Pointer, Int, Boolean, Int32, AddressSpace
@@ -35,10 +35,7 @@ def mbarrier_init(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None:
:type cnt: Int
"""
nvvm.mbarrier_init_shared(
mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
Int32(cnt).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
mbar_ptr.llvm_ptr, Int32(cnt).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
)
@@ -68,15 +65,18 @@ def mbarrier_arrive_and_expect_tx(
"""
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
mbar_llvm_ptr.type,
mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem)
mbar_llvm_ptr = nvvm.mapa(
mbar_cluster_type,
mbar_llvm_ptr,
Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
mbar_shared_type = llvm.PointerType.get(AddressSpace.smem)
mbar_llvm_ptr = llvm.addrspacecast(mbar_shared_type, mbar_llvm_ptr)
space = nvvm.MBarrierSpaceKind.CLUSTER
else:
space = nvvm.MBarrierSpaceKind.CTA
@@ -108,15 +108,18 @@ def mbarrier_expect_tx(
"""
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem)
mbar_llvm_ptr = nvvm.mapa(
mbar_llvm_ptr.type,
mbar_cluster_type,
mbar_llvm_ptr,
Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
mbar_shared_type = llvm.PointerType.get(AddressSpace.smem)
mbar_llvm_ptr = llvm.addrspacecast(mbar_shared_type, mbar_llvm_ptr)
space = nvvm.MBarrierSpaceKind.CLUSTER
else:
space = nvvm.MBarrierSpaceKind.CTA
@@ -147,7 +150,7 @@ def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None:
# This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX
# The timeout in ns only applies to the latter and this call is truly blocking
nvvm.mbarrier_try_wait_parity_shared(
mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
mbar_ptr.llvm_ptr,
Int32(phase).ir_value(loc=loc, ip=ip),
Int32(timeout_ns).ir_value(loc=loc, ip=ip),
loc=loc,
@@ -171,8 +174,7 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo
return Boolean(
nvvm.mbarrier_wait_parity(
T.bool(),
mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
mbar_ptr.llvm_ptr,
Int32(phase).ir_value(loc=loc, ip=ip),
nvvm.MBarrierWaitKind.TRY,
loc=loc,
@@ -226,17 +228,20 @@ def mbarrier_arrive(
the mbarrier is converted to a remote address in the peer CTA's
SMEM.
"""
mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
if peer_cta_rank_in_cluster is not None:
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
mbar_llvm_ptr.type,
mbar_cluster_type = llvm.PointerType.get(AddressSpace.dsmem)
mbar_llvm_ptr = nvvm.mapa(
mbar_cluster_type,
mbar_llvm_ptr,
Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
mbar_shared_type = llvm.PointerType.get(AddressSpace.smem)
mbar_llvm_ptr = llvm.addrspacecast(mbar_shared_type, mbar_llvm_ptr)
space = nvvm.MBarrierSpaceKind.CLUSTER
else:
space = nvvm.MBarrierSpaceKind.CTA
@@ -264,5 +269,10 @@ def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> N
"""
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
nvvm.cp_async_mbarrier_arrive_shared(mbar_llvm_ptr, noinc=True, loc=loc, ip=ip)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
nvvm.cp_async_mbarrier_arrive_shared(
mbar_llvm_ptr,
noinc=True,
loc=loc,
ip=ip,
)
python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py
+75 -123
View File
@@ -9,16 +9,17 @@
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
from cutlass.base_dsl.arch import Arch
from cutlass.base_dsl.common import DSLRuntimeError
from cutlass.cutlass_dsl import BaseDSL, dsl_user_op
from cutlass._mlir import ir
from cutlass._mlir.dialects import builtin, arith, llvm, vector
from cutlass._mlir.dialects import arith, llvm, vector
from .nvvm_wrappers import (
cvt_i8_bf16,
cvt_i8x2_to_bf16x2,
cvt_i8x4_to_bf16x4,
cvt_f32x2_bf16x2,
cvt_i8x4_to_f32x4,
cvt_i8x2_to_f32x2,
@@ -26,22 +27,11 @@ from .nvvm_wrappers import (
cvt_i4x4_to_bf16x4,
cvt_i4x2_to_bf16x2,
cvt_i4_bf16,
cvt_f4e2m1x8_to_f16x8,
cvt_f4e2m1x4_to_f16x4,
cvt_f4e2m1x2_to_f16x2,
cvt_f4e2m1_f16,
cvt_f32_bf16,
sext_unpacked_i4x4_to_i8x4,
)
from ..typing import Int4, Int8, Float32, BFloat16, Int32
from ..typing import (
Int4,
Int8,
Int32,
Float16,
Float32,
BFloat16,
Float32,
)
@dsl_user_op
def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
@@ -64,6 +54,7 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
vec_f32x2_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc)
vec_dst_type = ir.VectorType.get([length], BFloat16.mlir_type, loc=loc)
vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)
arch = BaseDSL._get_dsl().get_arch_enum()
# try to use vectorized version
if length >= 4:
num_vec4 = length // 4
@@ -71,45 +62,66 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
vec_i8x4 = vector.extract_strided_slice(
vec_i8x4_type, vec_i8, [src_pos], [4], [1], loc=loc, ip=ip
)
vec_f32x4 = cvt_i8x4_to_f32x4(vec_i8x4, loc=loc, ip=ip)
vec_f32x2_lo = vector.extract_strided_slice(
vec_f32x2_type, vec_f32x4, [0], [2], [1], loc=loc, ip=ip
)
vec_f32x2_hi = vector.extract_strided_slice(
vec_f32x2_type, vec_f32x4, [2], [2], [1], loc=loc, ip=ip
)
vec_bf16x2_lo = cvt_f32x2_bf16x2(vec_f32x2_lo, loc=loc, ip=ip)
vec_bf16x2_hi = cvt_f32x2_bf16x2(vec_f32x2_hi, loc=loc, ip=ip)
vec_dst = vector.insert_strided_slice(
vec_bf16x2_lo, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
vec_dst = vector.insert_strided_slice(
vec_bf16x2_hi, vec_dst, [src_pos + 2], [1], loc=loc, ip=ip
)
if arch in cvt_i8_bf16_intrinsic.s26_bf16_supported_archs:
vec_bf16x4 = cvt_i8x4_to_bf16x4(vec_i8x4, loc=loc, ip=ip)
vec_dst = vector.insert_strided_slice(
vec_bf16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
else:
vec_f32x4 = cvt_i8x4_to_f32x4(vec_i8x4, loc=loc, ip=ip)
vec_f32x2_lo = vector.extract_strided_slice(
vec_f32x2_type, vec_f32x4, [0], [2], [1], loc=loc, ip=ip
)
vec_f32x2_hi = vector.extract_strided_slice(
vec_f32x2_type, vec_f32x4, [2], [2], [1], loc=loc, ip=ip
)
vec_bf16x2_lo = cvt_f32x2_bf16x2(vec_f32x2_lo, loc=loc, ip=ip)
vec_bf16x2_hi = cvt_f32x2_bf16x2(vec_f32x2_hi, loc=loc, ip=ip)
vec_dst = vector.insert_strided_slice(
vec_bf16x2_lo, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
vec_dst = vector.insert_strided_slice(
vec_bf16x2_hi, vec_dst, [src_pos + 2], [1], loc=loc, ip=ip
)
src_pos += 4
length -= 4
if length >= 2:
vec_i8x2 = vector.extract_strided_slice(
vec_i8x2_type, vec_i8, [src_pos], [2], [1], loc=loc, ip=ip
)
vec_f32x2 = cvt_i8x2_to_f32x2(vec_i8x2, loc=loc, ip=ip)
vec_bf16x2 = cvt_f32x2_bf16x2(vec_f32x2, loc=loc, ip=ip)
if arch in cvt_i8_bf16_intrinsic.s26_bf16_supported_archs:
vec_bf16x2 = cvt_i8x2_to_bf16x2(vec_i8x2, loc=loc, ip=ip)
else:
vec_f32x2 = cvt_i8x2_to_f32x2(vec_i8x2, loc=loc, ip=ip)
vec_bf16x2 = cvt_f32x2_bf16x2(vec_f32x2, loc=loc, ip=ip)
vec_dst = vector.insert_strided_slice(
vec_bf16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
src_pos += 2
length -= 2
if length >= 1:
val_bf16 = cvt_i8_bf16(
vector.extractelement(
if arch in cvt_i8_bf16_intrinsic.s26_bf16_supported_archs:
val_bf16 = cvt_i8_bf16(
vector.extractelement(
vec_i8,
position=arith.constant(Int32.mlir_type, src_pos),
loc=loc,
ip=ip,
),
loc=loc,
ip=ip,
)
else:
src_i8 = vector.extractelement(
vec_i8,
position=arith.constant(Int32.mlir_type, src_pos),
loc=loc,
ip=ip,
),
loc=loc,
ip=ip,
)
)
src_i32 = llvm.sext(Int32.mlir_type, src_i8, loc=loc, ip=ip)
src_f32 = llvm.sitofp(Float32.mlir_type, src_i32, loc=loc, ip=ip)
val_bf16 = cvt_f32_bf16(src_f32, loc=loc, ip=ip)
vec_dst = vector.insertelement(
val_bf16,
vec_dst,
@@ -121,7 +133,7 @@ def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
@dsl_user_op
def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
def cvt_i4_bf16_intrinsic(vec_i4, length, *, with_shuffle=False, loc=None, ip=None):
"""
Fast conversion from int4 to bfloat16. It converts a vector of int4 to a vector of bfloat16.
@@ -129,6 +141,13 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
:type vec_i4: 1D vector of int4
:param length: The length of the input vector.
:type length: int
:param with_shuffle: Whether the input vec_i4 follows a specific shuffle pattern.
If True, for consecutive 8 int4 values with indices of (0, 1, 2, 3, 4, 5, 6, 7),
the input elements are shuffled to (0, 2, 1, 3, 4, 6, 5, 7). For tailing elements less than 8,
the shuffle pattern is (0, 2, 1, 3) for 4 elements. No shuffle is needed for less than 4 elements.
Shuffle could help to produce converted bf16 values in the natural order of (0, 1, 2 ,3 ,4 ,5 ,6 ,7)
without extra prmt instructions and thus better performance.
:type with_shuffle: bool
:return: The output 1D vector of bfloat16 with the same length as the input vector.
:rtype: 1D vector of bfloat16
"""
@@ -141,6 +160,7 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
vec_i4x2_type = ir.VectorType.get([2], Int4.mlir_type, loc=loc)
vec_dst_type = ir.VectorType.get([length], BFloat16.mlir_type, loc=loc)
vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)
# try to use vectorized version
if length >= 8:
num_vec8 = length // 8
@@ -148,7 +168,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
vec_i4x8 = vector.extract_strided_slice(
vec_i4x8_type, vec_i4, [src_pos], [8], [1], loc=loc, ip=ip
)
vec_bf16x8 = cvt_i4x8_to_bf16x8(vec_i4x8, loc=loc, ip=ip)
vec_bf16x8 = cvt_i4x8_to_bf16x8(
vec_i4x8, with_shuffle=with_shuffle, loc=loc, ip=ip
)
vec_dst = vector.insert_strided_slice(
vec_bf16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
@@ -158,7 +180,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
vec_i4x4 = vector.extract_strided_slice(
vec_i4x4_type, vec_i4, [src_pos], [4], [1], loc=loc, ip=ip
)
vec_bf16x4 = cvt_i4x4_to_bf16x4(vec_i4x4, loc=loc, ip=ip)
vec_bf16x4 = cvt_i4x4_to_bf16x4(
vec_i4x4, with_shuffle=with_shuffle, loc=loc, ip=ip
)
vec_dst = vector.insert_strided_slice(
vec_bf16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
@@ -168,7 +192,9 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
vec_i4x2 = vector.extract_strided_slice(
vec_i4x2_type, vec_i4, [src_pos], [2], [1], loc=loc, ip=ip
)
vec_bf16x2 = cvt_i4x2_to_bf16x2(vec_i4x2, loc=loc, ip=ip)
vec_bf16x2 = cvt_i4x2_to_bf16x2(
vec_i4x2, with_shuffle=with_shuffle, loc=loc, ip=ip
)
vec_dst = vector.insert_strided_slice(
vec_bf16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
@@ -195,84 +221,6 @@ def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
return vec_dst
@dsl_user_op
def cvt_f4e2m1_f16_intrinsic(vec_f4e2m1, length, *, loc=None, ip=None):
"""
Convert a vector of float4e2m1 to a vector of float16.
:param vec_f4e2m1: The input vector of float4e2m1.
:type vec_f4e2m1: 1D vector of float4e2m1
:param length: The length of the input vector.
:type length: int
:return: The output 1D vector of float16 with the same length as the input vector.
:rtype: 1D vector of float16
"""
src_pos = 0
vec_src_i4 = builtin.unrealized_conversion_cast(
[ir.VectorType.get([length], Int4.mlir_type, loc=loc)],
[vec_f4e2m1],
loc=loc,
ip=ip,
)
vec_i4x8_type = ir.VectorType.get([8], Int4.mlir_type, loc=loc)
vec_i4x4_type = ir.VectorType.get([4], Int4.mlir_type, loc=loc)
vec_i4x2_type = ir.VectorType.get([2], Int4.mlir_type, loc=loc)
vec_dst_type = ir.VectorType.get([length], Float16.mlir_type, loc=loc)
vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)
# try to use vectorized version
if length >= 8:
num_vec8 = length // 8
for _ in range(num_vec8):
vec_f4e2m1x8 = vector.extract_strided_slice(
vec_i4x8_type, vec_src_i4, [src_pos], [8], [1], loc=loc, ip=ip
)
vec_f16x8 = cvt_f4e2m1x8_to_f16x8(vec_f4e2m1x8, loc=loc, ip=ip)
vec_dst = vector.insert_strided_slice(
vec_f16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
src_pos += 8
length -= 8
if length >= 4:
vec_f4e2m1x4 = vector.extract_strided_slice(
vec_i4x4_type, vec_src_i4, [src_pos], [4], [1], loc=loc, ip=ip
)
vec_f16x4 = cvt_f4e2m1x4_to_f16x4(vec_f4e2m1x4, loc=loc, ip=ip)
vec_dst = vector.insert_strided_slice(
vec_f16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
src_pos += 4
length -= 4
if length >= 2:
vec_f4e2m1x2 = vector.extract_strided_slice(
vec_i4x2_type, vec_src_i4, [src_pos], [2], [1], loc=loc, ip=ip
)
vec_f16x2 = cvt_f4e2m1x2_to_f16x2(vec_f4e2m1x2, loc=loc, ip=ip)
vec_dst = vector.insert_strided_slice(
vec_f16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip
)
src_pos += 2
length -= 2
if length >= 1:
val_f16 = cvt_f4e2m1_f16(
vector.extractelement(
vec_src_i4,
position=arith.constant(Int32.mlir_type, src_pos),
loc=loc,
ip=ip,
),
loc=loc,
ip=ip,
)
vec_dst = vector.insertelement(
val_f16,
vec_dst,
position=arith.constant(Int32.mlir_type, src_pos),
loc=loc,
ip=ip,
)
return vec_dst
@dsl_user_op
def sext_unpacked_i4_i8_intrinsic(vec_unpacked_i4, length, *, loc=None, ip=None):
"""
@@ -295,9 +243,7 @@ def sext_unpacked_i4_i8_intrinsic(vec_unpacked_i4, length, *, loc=None, ip=None)
vec_unpacked_i4x4 = vector.extract_strided_slice(
vec_i8x4_type, vec_unpacked_i4, [pos], [4], [1], loc=loc, ip=ip
)
vec_i8x4 = sext_unpacked_i4x4_to_i8x4(
vec_unpacked_i4x4, loc=loc, ip=ip
)
vec_i8x4 = sext_unpacked_i4x4_to_i8x4(vec_unpacked_i4x4, loc=loc, ip=ip)
vec_i8 = vector.insert_strided_slice(
vec_i8x4, vec_i8, [pos], [1], loc=loc, ip=ip
)
@@ -312,6 +258,12 @@ cvt_i8_bf16_intrinsic.supported_archs = (
*Arch.HopperArchs(),
*Arch.BlackwellArchs(),
)
cvt_i8_bf16_intrinsic.s26_bf16_supported_archs = (
Arch.sm_100a,
Arch.sm_110a,
Arch.sm_120a,
Arch.sm_121a,
)
cvt_i4_bf16_intrinsic.supported_archs = (
Arch.sm_100a,
Arch.sm_110a,
python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py
+531 -125
View File
@@ -10,25 +10,16 @@
# is strictly prohibited.
from functools import partial
from typing import Optional, Tuple, Union, Callable, Literal
from typing import Any, Optional, Tuple, Union, Callable, Literal
from typing_extensions import deprecated
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass.cutlass_dsl import T, dsl_user_op, target_version
import cutlass.cutlass_dsl as cutlass_dsl
from cutlass._mlir import ir
from cutlass._mlir.dialects import arith, llvm, nvvm, vector
# Forward nvvm enums
from cutlass._mlir.dialects.nvvm import (
ProxyKind,
SharedSpace,
Tcgen05WaitKind,
SetMaxRegisterAction,
RoundingModeKind,
)
from ..core import size
from ..typing import (
@@ -95,25 +86,23 @@ def _enhance_enum_with_str_mapping(enum_class):
"""
Convert a string literal to the corresponding enum member.
:param s: String representation of the enum member, or an enum member itself (deprecated)
:param s: String representation of the enum member
:return: The enum member (or None if s is None)
:raises ValueError: If the string is not a valid enum member
:raises TypeError: If an enum is passed instead of a string
"""
import warnings
if s is None:
return None
# Check if user passed an enum (should be a string literal instead)
# This catches cases where user passes e.g., RoundingModeKind.RN instead of "rn"
from enum import Enum
# Check if s is already an enum member of the correct type
if isinstance(s, cls):
warnings.warn(
f"Passing enum member directly to {cls.__name__}.from_str() is deprecated. "
f"Please use string literals instead (e.g., '{str(s)}' instead of {cls.__name__}.{s.name}).",
DeprecationWarning,
stacklevel=2,
if isinstance(s, Enum):
raise TypeError(
f"Expected a string literal for {cls.__name__}, but got enum '{type(s).__name__}.{s.name}'. "
f"Please pass a string instead (e.g., '{str(s)}' instead of {type(s).__name__}.{s.name}). "
f"Valid string options are: {sorted(str_to_enum_map.keys())}"
)
return s
if s not in str_to_enum_map:
valid_options = sorted(str_to_enum_map.keys())
raise ValueError(
@@ -446,6 +435,7 @@ def warp_reduction(
offset = offset // 2
return val
warp_reduction_max = partial(
warp_reduction,
op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else cutlass_dsl.max(x, y),
@@ -460,34 +450,13 @@ def barrier(*, barrier_id=None, number_of_threads=None, loc=None, ip=None) -> No
"""
if barrier_id is not None:
barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip)
else:
barrier_id = Int32(0).ir_value(loc=loc, ip=ip)
if number_of_threads is not None:
number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip)
llvm.inline_asm(
None,
[barrier_id, number_of_threads],
"bar.sync $0, $1;",
"r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
else:
llvm.inline_asm(
None,
[barrier_id],
"bar.sync $0;",
"r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
nvvm.barrier(
barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip
)
@dsl_user_op
@@ -496,8 +465,6 @@ def barrier_arrive(
) -> None:
if barrier_id is not None:
barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip)
else:
barrier_id = Int32(0).ir_value(loc=loc, ip=ip)
if number_of_threads is None:
raise ValueError(
@@ -505,14 +472,8 @@ def barrier_arrive(
)
number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip)
llvm.inline_asm(
None,
[barrier_id, number_of_threads],
"bar.arrive $0, $1;",
"r,r",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
nvvm.barrier_arrive(
barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip
)
@@ -638,19 +599,73 @@ def cluster_arrive_relaxed(*, aligned=None, loc=None, ip=None) -> None:
@dsl_user_op
def fence_proxy(
kind: ProxyKind,
kind: Literal[
"alias", "async", "async.global", "async.shared", "tensormap", "generic"
],
*,
space: Optional[SharedSpace] = None,
space: Optional[Literal["cta", "cluster"]] = None,
use_intrinsic=None,
loc=None,
ip=None,
) -> None:
"""
Fence operation to ensure memory consistency between proxies.
:param kind: Proxy kind string literal:
- "alias" : Alias proxy
- "async" : Async proxy
- "async.global" : Async global proxy
- "async.shared" : Async shared proxy
- "tensormap" : Tensormap proxy
- "generic" : Generic proxy
:type kind: Literal["alias", "async", "async.global", "async.shared", "tensormap", "generic"]
:param space: Shared memory space scope string literal (optional):
- "cta" : CTA (Cooperative Thread Array) scope
- "cluster" : Cluster scope
:type space: Optional[Literal["cta", "cluster"]]
:param use_intrinsic: Whether to use intrinsic version
"""
from cutlass._mlir.dialects.nvvm import (
SharedSpace,
ProxyKind,
)
# Enhance enum with str mapping
SharedSpace = _enhance_enum_with_str_mapping(SharedSpace)
ProxyKind = _enhance_enum_with_str_mapping(ProxyKind)
kind = ProxyKind.from_str(kind)
space = SharedSpace.from_str(space)
nvvm.fence_proxy(
kind=kind, space=space, use_intrinsic=use_intrinsic, loc=loc, ip=ip
kind=kind,
space=space,
use_intrinsic=use_intrinsic,
loc=loc,
ip=ip,
)
@dsl_user_op
def vote_sync_op(
pred: Boolean, kind: nvvm.VoteSyncKind, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Union[Int32, Boolean]:
"""
Performs a vote operation across the warp.
"""
return_type = Int32 if kind == nvvm.VoteSyncKind.ballot else Boolean
return return_type(
nvvm.vote_sync(
T.i32() if kind == nvvm.VoteSyncKind.ballot else T.bool(),
Int32(mask).ir_value(loc=loc, ip=ip),
Boolean(pred).ir_value(loc=loc, ip=ip),
kind,
loc=loc,
ip=ip,
)
)
def vote_ballot_sync(
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Int32:
@@ -668,45 +683,7 @@ def vote_ballot_sync(
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`__.
"""
return Int32(
nvvm.vote_ballot_sync(
T.i32(),
Int32(mask).ir_value(loc=loc, ip=ip),
Boolean(pred).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
)
@dsl_user_op
def vote_sync_op(
pred: Boolean, kind: str, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Union[Int32, Boolean]:
return_type = Boolean
return_type_str = "pred"
return return_type(
llvm.inline_asm(
T.bool(),
[
Boolean(pred).ir_value(loc=loc, ip=ip),
Int32(mask).ir_value(loc=loc, ip=ip),
],
f"""{{\n\t
.reg .pred ps;\n\t
.reg .pred pd;\n\t
setp.ne.b32 ps, $1, 0;\n\t
vote.sync.{kind}.{return_type_str} pd, ps, $2;\n\t
selp.b32 $0, 1, 0, pd;\n\t
}}""",
"=r,r,i",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
)
return vote_sync_op(pred, nvvm.VoteSyncKind.ballot, mask, loc=loc, ip=ip)
@dsl_user_op
@@ -714,7 +691,7 @@ def vote_any_sync(
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
) -> Boolean:
"""True if source predicate is True for any non-exited threads in mask. Negate the source
predicate to compute .not_all.
predicate to compute .none.
:param pred: The predicate value for the current thread
:type pred: Boolean
@@ -727,7 +704,7 @@ def vote_any_sync(
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`__.
"""
return vote_sync_op(pred, "any", mask, loc=loc, ip=ip)
return vote_sync_op(pred, nvvm.VoteSyncKind.any, mask, loc=loc, ip=ip)
@dsl_user_op
@@ -748,7 +725,7 @@ def vote_all_sync(
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`__.
"""
return vote_sync_op(pred, "all", mask, loc=loc, ip=ip)
return vote_sync_op(pred, nvvm.VoteSyncKind.all, mask, loc=loc, ip=ip)
@dsl_user_op
@@ -767,7 +744,7 @@ def vote_uni_sync(
threads in mask
:rtype: Boolean
"""
return vote_sync_op(pred, "uni", mask, loc=loc, ip=ip)
return vote_sync_op(pred, nvvm.VoteSyncKind.uni, mask, loc=loc, ip=ip)
@dsl_user_op
@@ -821,8 +798,8 @@ def fence_view_async_tmem_op(
from cutlass._mlir.dialects.nvvm import Tcgen05WaitKind
# Enhance enum and convert string literal to enum type
Tcgen05WaitKind_enhanced = _enhance_enum_with_str_mapping(Tcgen05WaitKind)
kind = Tcgen05WaitKind_enhanced.from_str(kind)
Tcgen05WaitKind = _enhance_enum_with_str_mapping(Tcgen05WaitKind)
kind = Tcgen05WaitKind.from_str(kind)
nvvm.tcgen05_wait(kind=kind, loc=loc, ip=ip)
@@ -847,9 +824,8 @@ def fence_view_async_shared(
This function is usually used for async execution unit (like TMA, UMMA) after the load/store operations.
"""
nvvm.fence_proxy(
nvvm.ProxyKind.async_shared, space=nvvm.SharedSpace.shared_cta, loc=loc, ip=ip
)
# Use the fence_proxy wrapper function with string literals
fence_proxy(kind="async.shared", space="cta", loc=loc, ip=ip)
@dsl_user_op
@@ -859,6 +835,7 @@ def setmaxregister_increase(
loc=None,
ip=None,
):
from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction
return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip)
@@ -869,6 +846,7 @@ def setmaxregister_decrease(
loc=None,
ip=None,
):
from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction
return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip)
@@ -880,6 +858,7 @@ def warpgroup_reg_alloc(
loc=None,
ip=None,
) -> None:
from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction
nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip)
@@ -891,8 +870,10 @@ def warpgroup_reg_dealloc(
loc=None,
ip=None,
) -> None:
from cutlass._mlir.dialects.nvvm import SetMaxRegisterAction
nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip)
@dsl_user_op
def calc_packed_f32x2_op(
src_a: Tuple[Float32, Float32],
@@ -908,8 +889,8 @@ def calc_packed_f32x2_op(
from cutlass._mlir.dialects.nvvm import RoundingModeKind
# Enhance enum and convert string literal to enum type
RoundingModeKind_enhanced = _enhance_enum_with_str_mapping(RoundingModeKind)
rnd = RoundingModeKind_enhanced.from_str(rnd)
RoundingModeKind = _enhance_enum_with_str_mapping(RoundingModeKind)
rnd = RoundingModeKind.from_str(rnd)
vec_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc)
vec_src_a = vector.from_elements(
@@ -959,15 +940,17 @@ mul_packed_f32x2 = partial(
add_packed_f32x2 = partial(
calc_packed_f32x2_op, src_c=None, calc_func=nvvm.add_packed_f32x2
)
sub_packed_f32x2 = partial(
calc_packed_f32x2_op, src_c=None, calc_func=nvvm.sub_packed_f32x2
)
@dsl_user_op
def fmax(
a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None
) -> Float32:
return Float32(
nvvm.fmax(
T.f32(),
Float32(a).ir_value(loc=loc, ip=ip),
Float32(b).ir_value(loc=loc, ip=ip),
loc=loc,
@@ -975,12 +958,25 @@ def fmax(
)
)
@dsl_user_op
def fmin(
a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None
) -> Float32:
return Float32(
nvvm.fmin(
Float32(a).ir_value(loc=loc, ip=ip),
Float32(b).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
)
@dsl_user_op
def rcp_approx(a: Union[float, Float32], *, loc=None, ip=None):
return Float32(
nvvm.rcp_approx_ftz_f(
T.f32(), Float32(a).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
)
nvvm.rcp_approx_ftz_f(Float32(a).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
)
@@ -1024,6 +1020,68 @@ def cvt_i8_bf16(src_i8, *, loc=None, ip=None):
return val_bf16
@dsl_user_op
def cvt_i8x2_to_bf16x2(src_vec2, *, loc=None, ip=None):
# pack 2 int8 into 1 int16 value
src_i16 = llvm.bitcast(Int16.mlir_type, src_vec2, loc=loc, ip=ip)
val_i32 = llvm.inline_asm(
Int32.mlir_type,
[
src_i16,
],
"""{\n\t
.reg .b16 scale;\n\t
mov.b16 scale, 0x8585;\n\t
cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, $1, scale;\n\t
}""",
"=r,h",
)
vec_bf16x2_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc)
vec_bf16x2 = llvm.bitcast(vec_bf16x2_type, val_i32, loc=loc, ip=ip)
return vec_bf16x2
@dsl_user_op
def cvt_i8x4_to_bf16x4(src_vec4, *, loc=None, ip=None):
# pack 4 int8 into 1 int32 value
src_i32 = llvm.bitcast(Int32.mlir_type, src_vec4, loc=loc, ip=ip)
rst01 = llvm.inline_asm(
Int32.mlir_type,
[
src_i32,
],
"""{\n\t
.reg .b16 pair<2>;\n\t
.reg .b16 scale;\n\t
mov.b32 {pair0, pair1}, $1;\n\t
mov.b16 scale, 0x8585;\n\t
cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair0, scale;\n\t
}""",
"=r,r",
)
rst23 = llvm.inline_asm(
Int32.mlir_type,
[
src_i32,
],
"""{\n\t
.reg .b16 pair<2>;\n\t
.reg .b16 scale;\n\t
mov.b32 {pair0, pair1}, $1;\n\t
mov.b16 scale, 0x8585;\n\t
cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair1, scale;\n\t
}""",
"=r,r",
)
vec_type = ir.VectorType.get([2], Int32.mlir_type, loc=loc)
rst_i32 = vector.from_elements(vec_type, [rst01, rst23], loc=loc, ip=ip)
vec_bf16x4_type = ir.VectorType.get([4], BFloat16.mlir_type, loc=loc)
vec_bf16x4 = llvm.bitcast(vec_bf16x4_type, rst_i32, loc=loc, ip=ip)
return vec_bf16x4
# Convert vector of 2 float values to vector of 2 bfloat16 values with satfinite rounding
@dsl_user_op
def cvt_f32x2_bf16x2(src_vec2, *, loc=None, ip=None):
@@ -1263,12 +1321,116 @@ def prmt(src, src_reg_shifted, prmt_indices, *, loc=None, ip=None):
@dsl_user_op
def cvt_i4_bf16(src_i4, *, loc=None, ip=None):
# i4 -> i32 -> f32 -> bf
src_i32 = llvm.zext(Int32.mlir_type, src_i4, loc=loc, ip=ip)
src_i32 = llvm.sext(Int32.mlir_type, src_i4, loc=loc, ip=ip)
src_f32 = llvm.sitofp(Float32.mlir_type, src_i32, loc=loc, ip=ip)
bf16_val = cvt_f32_bf16(src_f32, loc=loc, ip=ip)
return bf16_val
# Convert multiple shuffled int4 values to bfloat16 values.
# The input elements are assumed to be already shuffled following a specific shuffle pattern.
# Specifically, for consecutive 8 int4 values with indices of (0, 1, 2, 3, 4, 5, 6, 7),
# they are shuffled to (0, 2, 1, 3, 4, 6, 5, 7). For tailing elements less than 8, the
# shuffle pattern is (0, 2, 1, 3) for 4 elements. No shuffle is needed for less than 4 elements.
# Shuffle could help to produce converted bf16 values in the natural order of (0, 1, 2 ,3 ,4 ,5 ,6 ,7)
# without extra prmt instructions and thus better performance.
# The number of elements to be converted must be be even as specified by num_elts.
# Int4 values are packed into int32 values with upper bits filled with 0 if there are less than 4 int4 values.
# Results bfloat16 values are also packed into int32 values.
@dsl_user_op
def cvt_i4_to_bf16_with_shuffle_impl(src_i32, num_elts, *, loc=None, ip=None):
from cutlass import CUDA_VERSION
if CUDA_VERSION.major < 13:
raise cutlass_dsl.DSLCudaVerNotImplemented(
feature="cvt_i4_to_bf16_with_shuffle_impl", required_version="13.1"
)
num_i32_elts = num_elts // 2
mask_odd = arith.constant(Int32.mlir_type, 0xF0F0F0F0, loc=loc, ip=ip)
mask_even = arith.constant(Int32.mlir_type, 0x0F0F0F0F, loc=loc, ip=ip)
src_odd = arith.andi(src_i32, mask_odd, loc=loc, ip=ip)
src_even = arith.andi(src_i32, mask_even, loc=loc, ip=ip)
c4 = arith.constant(Int32.mlir_type, 4, loc=loc, ip=ip)
src_even = arith.shli(src_even, c4, loc=loc, ip=ip)
rst13 = llvm.inline_asm(
Int32.mlir_type,
[
src_odd,
],
"""{\n\t
.reg .b16 pair<2>;\n\t
.reg .b16 scale;\n\t
mov.b32 {pair0, pair1}, $1;\n\t
mov.b16 scale, 0x8181;\n\t
cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair0, scale;\n\t
}""",
"=r,r",
)
rst57 = llvm.inline_asm(
Int32.mlir_type,
[
src_odd,
],
"""{\n\t
.reg .b16 pair<2>;\n\t
.reg .b16 scale;\n\t
mov.b32 {pair0, pair1}, $1;\n\t
mov.b16 scale, 0x8181;\n\t
cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair1, scale;\n\t
}""",
"=r,r",
)
rst02 = llvm.inline_asm(
Int32.mlir_type,
[
src_even,
],
"""{\n\t
.reg .b16 pair<2>;\n\t
.reg .b16 scale;\n\t
mov.b16 scale, 0x8181;\n\t
mov.b32 {pair0, pair1}, $1;\n\t
cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair0, scale;\n\t
}""",
"=r,r",
)
rst46 = llvm.inline_asm(
Int32.mlir_type,
[
src_even,
],
"""{\n\t
.reg .b16 pair<2>;\n\t
.reg .b16 scale;\n\t
mov.b16 scale, 0x8181;\n\t
mov.b32 {pair0, pair1}, $1;\n\t
cvt.rn.satfinite.scaled::n2::ue8m0.bf16x2.s2f6x2 $0, pair1, scale;\n\t
}""",
"=r,r",
)
vec_type = ir.VectorType.get([num_i32_elts], Int32.mlir_type, loc=loc)
if num_elts == 2:
prmt_index = arith.constant(Int32.mlir_type, 0x00005410, loc=loc, ip=ip)
rst = llvm.inline_asm(
Int32.mlir_type,
[
rst02,
rst13,
prmt_index,
],
"prmt.b32 $0, $1, $2, $3;",
"=r,r,r,r",
)
vec_rsts = vector.from_elements(vec_type, [rst], loc=loc, ip=ip)
elif num_elts == 4:
vec_rsts = vector.from_elements(vec_type, [rst02, rst13], loc=loc, ip=ip)
else:
vec_rsts = vector.from_elements(
vec_type, [rst02, rst13, rst46, rst57], loc=loc, ip=ip
)
return vec_rsts
# Convert multiple int4 values to bfloat16 values.
# The number of elements to be converted must be be even as specified by num_elts.
# Int4 values are packed into int32 values with upper bits filled with 0 if there are less than 4 int4 values.
@@ -1357,11 +1519,12 @@ def cvt_i4_to_bf16_impl(src_i32, num_elts, *, loc=None, ip=None):
# Convert 2 int4 values to 2 bfloat16 values
@dsl_user_op
def cvt_i4x2_to_bf16x2(src_vec2, *, loc=None, ip=None):
def cvt_i4x2_to_bf16x2(src_vec2, *, with_shuffle=False, loc=None, ip=None):
cvt_func = cvt_i4_to_bf16_with_shuffle_impl if with_shuffle else cvt_i4_to_bf16_impl
# pack 2 int4 into 1 int32 value and fill upper bits with 0
src_i8 = llvm.bitcast(Int8.mlir_type, src_vec2, loc=loc, ip=ip)
src_i32 = llvm.zext(Int32.mlir_type, src_i8, loc=loc, ip=ip)
rst_i32 = cvt_i4_to_bf16_impl(src_i32, 2, loc=loc, ip=ip)
rst_i32 = cvt_func(src_i32, 2, loc=loc, ip=ip)
vec_bf16x2_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc)
vec_bf16x2 = llvm.bitcast(vec_bf16x2_type, rst_i32, loc=loc, ip=ip)
return vec_bf16x2
@@ -1369,11 +1532,12 @@ def cvt_i4x2_to_bf16x2(src_vec2, *, loc=None, ip=None):
# Convert 4 int4 values to 4 bfloat16 values
@dsl_user_op
def cvt_i4x4_to_bf16x4(src_vec4, *, loc=None, ip=None):
def cvt_i4x4_to_bf16x4(src_vec4, *, with_shuffle=False, loc=None, ip=None):
cvt_func = cvt_i4_to_bf16_with_shuffle_impl if with_shuffle else cvt_i4_to_bf16_impl
# pack 4 int4 into 1 int32 value and fill upper bits with 0
src_i16 = llvm.bitcast(Int16.mlir_type, src_vec4, loc=loc, ip=ip)
src_i32 = llvm.zext(Int32.mlir_type, src_i16, loc=loc, ip=ip)
rst_i32 = cvt_i4_to_bf16_impl(src_i32, 4, loc=loc, ip=ip)
rst_i32 = cvt_func(src_i32, 4, loc=loc, ip=ip)
vec_bf16x4_type = ir.VectorType.get([4], BFloat16.mlir_type, loc=loc)
vec_bf16x4 = llvm.bitcast(vec_bf16x4_type, rst_i32, loc=loc, ip=ip)
return vec_bf16x4
@@ -1381,14 +1545,16 @@ def cvt_i4x4_to_bf16x4(src_vec4, *, loc=None, ip=None):
# Convert 8 int4 values to 8 bfloat16 values
@dsl_user_op
def cvt_i4x8_to_bf16x8(src_vec8, *, loc=None, ip=None):
def cvt_i4x8_to_bf16x8(src_vec8, *, with_shuffle=False, loc=None, ip=None):
cvt_func = cvt_i4_to_bf16_with_shuffle_impl if with_shuffle else cvt_i4_to_bf16_impl
# pack 8 int4 into 1 int32 value and fill upper bits with 0
src_i32 = llvm.bitcast(Int32.mlir_type, src_vec8, loc=loc, ip=ip)
rst_i32 = cvt_i4_to_bf16_impl(src_i32, 8, loc=loc, ip=ip)
rst_i32 = cvt_func(src_i32, 8, loc=loc, ip=ip)
vec_bf16x8_type = ir.VectorType.get([8], BFloat16.mlir_type, loc=loc)
vec_bf16x8 = llvm.bitcast(vec_bf16x8_type, rst_i32, loc=loc, ip=ip)
return vec_bf16x8
# Sign extend 4 int4 unpacked in 8b containers
@dsl_user_op
def sext_unpacked_i4x4_to_i8x4(src_vec4, *, loc=None, ip=None):
@@ -1485,6 +1651,196 @@ def griddepcontrol_launch_dependents(*, loc=None, ip=None) -> None:
@dsl_user_op
def _warp_redux_sync_nvvm(
value: Numeric,
kind: Literal[
"fmax",
"fmin",
"max",
"min",
"add",
"xor",
"or",
"and",
],
mask_and_clamp: Int = FULL_MASK,
abs: bool = False,
nan: bool = None,
*,
loc=None,
ip=None,
) -> Numeric:
from cutlass._mlir.dialects.nvvm import ReduxKind
# Enhance enum and convert string literal to enum type
ReduxKind = _enhance_enum_with_str_mapping(ReduxKind)
kind = ReduxKind.from_str(kind)
value_type = type(value)
value_ir = value.ir_value(loc=loc, ip=ip)
return value_type(
nvvm.redux_sync(
res=value_ir.type,
val=value_ir,
kind=kind,
mask_and_clamp=Int32(mask_and_clamp).ir_value(loc=loc, ip=ip),
abs=abs,
nan=nan,
loc=loc,
ip=ip,
)
)
@dsl_user_op
def _warp_redux_sync_ptx(
value: Numeric,
kind: Literal[
"fmax",
"fmin",
"max",
"min",
],
mask_and_clamp: Int = FULL_MASK,
abs: bool = None,
nan: bool = None,
*,
loc=None,
ip=None,
) -> Numeric:
value_type = type(value)
value_ir = value.ir_value(loc=loc, ip=ip)
mlir_type = value_type.mlir_type
mask_ir = Int32(mask_and_clamp).ir_value(loc=loc, ip=ip)
kind_ptx_str = kind
if kind == "fmax":
kind_ptx_str = "max"
elif kind == "fmin":
kind_ptx_str = "min"
modifiers = []
if nan is True:
modifiers.append("NaN")
if abs is True:
modifiers.append("abs")
modifier_str = "." + ".".join(modifiers) if modifiers else ""
ptx_instr = f"redux.sync.{kind_ptx_str}{modifier_str}.f32 $0, $1, $2;"
return value_type(
llvm.inline_asm(
mlir_type,
[value_ir, mask_ir],
f"{ptx_instr}",
f"=f,f,i",
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
)
)
@dsl_user_op
def warp_redux_sync(
value: Numeric,
kind: Literal[
"fmax",
"fmin",
"max",
"min",
"add",
"xor",
"or",
"and",
],
mask_and_clamp: Int = FULL_MASK,
*,
abs: bool = None,
nan: bool = None,
loc=None,
ip=None,
) -> Numeric:
"""
Perform warp-level reduction operation across threads.
Reduces values from participating threads in a warp according to the specified operation.
All threads in the mask receive the same result.
:param value: Input value to reduce
:type value: Numeric
:param kind: Reduction operation. Supported operations:
- Integer types (Int32/Uint32): "add", "and", "max", "min", "or", "xor"
- Float types (Float32): "fmax", "fmin" (or "max"/"min" which auto-convert to "fmax"/"fmin")
:type kind: Literal["add", "and", "max", "min", "or", "xor", "fmin", "fmax"]
:param mask_and_clamp: Warp participation mask (default: FULL_MASK = 0xFFFFFFFF)
:type mask_and_clamp: Int
:param abs: Apply absolute value before reduction (float types only)
:type abs: bool
:param nan: Enable NaN propagation for fmax/fmin operations (float types only)
:type nan: Optional[bool]
:return: Reduced value (same for all participating threads)
:rtype: Numeric
"""
# Convert value to Numeric type if needed
if not isinstance(value, Numeric):
value = as_numeric(value)
# Determine value type and choose appropriate implementation
value_type = type(value)
mlir_type = value_type.mlir_type
# Use inline PTX for float types, NVVM for integer types
if mlir_type == T.f32():
return _warp_redux_sync_ptx(
value, kind, mask_and_clamp, abs, nan, loc=loc, ip=ip
)
else:
return _warp_redux_sync_nvvm(
value, kind, mask_and_clamp, abs, nan, loc=loc, ip=ip
)
@dsl_user_op
def atomic_max_float32(
ptr,
value: Float32,
*,
positive_only: bool = True,
loc=None,
ip=None,
) -> Float32:
"""
Performs an atomic max operation on a float32 value in global memory.
This implementation works correctly for non-negative values (>= 0) using direct bitcast.
:param ptr: Pointer to the memory location
:param value: The float32 value to compare and potentially store (should be >= 0 for correct results)
:type value: Float32
:param positive_only: If True (default), assumes input values are non-negative.
This parameter is provided for API compatibility and future extensions.
:type positive_only: bool
:return: The old value at the memory location
:rtype: Float32
"""
from cutlass._mlir.dialects.nvvm import AtomicOpKind
value_int = llvm.bitcast(T.i32(), value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
old_value_int = nvvm.atomicrmw(
AtomicOpKind.MAX,
ptr,
value_int,
loc=loc,
ip=ip,
)
return Float32(llvm.bitcast(T.f32(), old_value_int, loc=loc, ip=ip))
def _normalize_ptr(addr, *, loc=None, ip=None) -> ir.Value:
"""
Helper function to normalize pointer types to MLIR ir.Value.
@@ -1549,7 +1905,7 @@ def _atomic(
:rtype: Union[Numeric, ir.Value]
"""
from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind
from cutlass.utils.version_info import CUDA_VERSION
from cutlass import CUDA_VERSION
# Enhance enums and convert string literals to enum types
AtomicOpKind = _enhance_enum_with_str_mapping(AtomicOpKind)
@@ -1848,7 +2204,7 @@ def atomic_cas(
:rtype: Numeric
"""
from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind
from cutlass.utils.version_info import CUDA_VERSION
from cutlass import CUDA_VERSION
# Enhance enums and convert string literals to enum types
MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind)
@@ -1882,6 +2238,17 @@ def atomic_cas(
loc=loc,
ip=ip,
)
elif CUDA_VERSION.major == 13 and CUDA_VERSION.minor == 1:
result = nvvm.atomicrmw(
op=AtomicOpKind.CAS,
ptr=ptr,
a=val_ir,
b=cmp_ir,
mem_order=sem,
syncscope=scope,
loc=loc,
ip=ip,
)
else:
result = nvvm.atomicrmw(
op=AtomicOpKind.CAS,
@@ -2190,3 +2557,42 @@ def cvt_f4e2m1x8_to_f16x8(src_vec8, *, loc=None, ip=None):
vec_f16x8_type = ir.VectorType.get([8], Float16.mlir_type, loc=loc)
vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_f32x4, loc=loc, ip=ip)
return vec_f16x8
@dsl_user_op
def mapa(ptr, cta_rank_in_cluster=0, *, loc=None, ip=None):
"""
Map a pointer to distributed shared memory across cluster.
Portable wrapper that uses the appropriate NVVM API based on CUDA version:
- CUDA 13.1+: Uses nvvm.mapa with dsmem address space
- CUDA 12.9: Uses nvvm.mapa_shared_cluster
Args:
ptr: Pointer to shared memory (llvm_ptr attribute will be used)
cta_rank_in_cluster: CTA rank within the cluster (default 0)
Returns:
Mapped LLVM pointer to shared memory
"""
if target_version(min_version="13.1"):
dsmem_ptr_ty = llvm.PointerType.get(7) # dsmem
smem_ptr_ty = llvm.PointerType.get(3) # smem
llvm_ptr = nvvm.mapa(
dsmem_ptr_ty,
ptr.llvm_ptr,
Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return llvm.addrspacecast(smem_ptr_ty, llvm_ptr, loc=loc, ip=ip)
else:
llvm_ptr = ptr.llvm_ptr
return nvvm.mapa_shared_cluster(
llvm_ptr.type,
llvm_ptr,
Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
python/CuTeDSL/cutlass/cute/arch/smem.py
+1 -1
View File
@@ -17,7 +17,7 @@ import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
from ..typing import Pointer, Numeric, NumericMeta
from ..typing import Pointer, Numeric, NumericMeta, Layout
@dsl_user_op
python/CuTeDSL/cutlass/cute/arch/tmem.py
+1 -4
View File
@@ -55,7 +55,6 @@ def get_max_tmem_alloc_cols(compute_capability: str) -> int:
return TMEM_MAX_ALLOC_COLUMNS_MAP[compute_capability]
def get_min_tmem_alloc_cols(compute_capability: str) -> int:
"""Get the minimum TMEM allocation columns for a given compute capability.
@@ -179,11 +178,9 @@ def dealloc_tmem(
:param num_columns: The number of columns in the TMEM allocation
:type num_columns: Int
:param is_two_cta: Optional boolean parameter for 2-CTA MMAs
:param arch: The architecture of the GPU.
:type arch: str
"""
tmem_max_alloc_cols = get_max_tmem_alloc_cols(arch)
tmem_min_alloc_cols = get_min_tmem_alloc_cols(arch)
tmem_max_alloc_cols = get_max_tmem_alloc_cols(arch)
if isinstance(num_columns, int):
if (
num_columns < tmem_min_alloc_cols
python/CuTeDSL/cutlass/cute/atom.py
+113 -32
View File
@@ -10,7 +10,7 @@
# is strictly prohibited.
from abc import ABC, ABCMeta, abstractmethod
from typing import Type, Union, Optional, Any, overload
from typing import Type, Union, Optional, Any, List, Tuple, overload
from .typing import Shape, Layout, Tile, Tensor, Numeric, Int32
from .core import (
@@ -285,6 +285,8 @@ class MmaAtom(Atom):
if self.op is not None:
self.op._verify_fragment_B(input, loc=loc, ip=ip)
input = input.value
if isinstance(input, tuple):
input = _pack_shape(input, loc=loc, ip=ip)
return _cute_ir.mma_make_fragment(
_cute_ir.MmaOperand.B, self._trait.value, input, loc=loc, ip=ip
)
@@ -1111,25 +1113,66 @@ def make_tiled_copy_C_atom(atom: CopyAtom, mma: TiledMma, *, loc=None, ip=None):
return _make_tiled_copy(atom, layout_tv, tiler_mn, loc=loc, ip=ip)
def _normalize_variadic_tensor_operand(
x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str
) -> List["Tensor"]:
"""Normalize a Tensor or sequence of Tensors to a list of Tensors.
Helper function for operations with variadic operands.
"""
if isinstance(x, Tensor):
return [x]
if isinstance(x, (list, tuple)):
if len(x) == 0:
raise ValueError(f"`{name}` must contain at least one Tensor")
if not all(isinstance(t, Tensor) for t in x):
raise TypeError(f"All elements of `{name}` must be Tensor")
return list(x) # type: ignore
raise TypeError(f"`{name}` must be a Tensor or a sequence of Tensors")
@dsl_user_op
def copy_atom_call(
atom: CopyAtom,
src: Tensor,
dst: Tensor,
src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
*,
pred: Optional[Tensor] = None,
loc=None,
ip=None,
**kwargs,
) -> None:
"""Executes a single copy atom operation between two tensors.
"""
Execute a single copy atom operation.
The copy_atom_call operation executes a copy atom with the given operands.
Source and destination tensors have layout profile ``(V)``.
The ``V-mode`` represents either:
- A singular mode directly consumable by the provided Copy Atom
- A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``,
For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``.
- Certain Atoms may require additional operation-specific keyword arguments.
- Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned
for future releases.
Both ``src`` and ``dst`` operands are variadic, containing a variable number of tensors:
- For regular copy, ``src`` and ``dst`` each contain a single tensor.
- For copy with auxiliary operands, they contain the main tensor followed by
auxiliary tensors. For example:
:param atom: Copy atom specifying the transfer operation
:type atom: CopyAtom
:param src: Source tensor with layout profile ``(V)``
:type src: Tensor
:param dst: Destination tensor with layout profile ``(V)``
:type dst: Tensor
:param src: Source tensor(s) with layout profile ``(V)``. Can be a single Tensor
or a list/tuple of Tensors for operations with auxiliary source operands.
:type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param dst: Destination tensor(s) with layout profile ``(V)``. Can be a single Tensor
or a list/tuple of Tensors for operations with auxiliary destination operands.
:type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param pred: Optional predication tensor for conditional transfers, defaults to None
:type pred: Optional[Tensor], optional
:param loc: Source location information, defaults to None
@@ -1142,51 +1185,89 @@ def copy_atom_call(
:return: None
:rtype: None
The copy_atom_call operation executes a single copy atom with the given operands.
Source and destination tensors with layout profile like ``(V)``.
The ``V-mode`` represents either:
- A singular mode directly consumable by the provided Copy Atom
- A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``,
For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``.
**Examples**:
.. code-block:: python
# Basic copy atom operation
# Regular copy atom operation
cute.copy_atom_call(copy_atom, src, dst)
# Predicated copy atom operation
cute.copy_atom_call(copy_atom, src, dst, pred=pred)
.. note::
- Certain Atoms may require additional operation-specific keyword arguments.
- Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned
for future releases.
"""
if isinstance(src.type, _cute_ir.MemRefType) and isinstance(
dst.type, _cute_ir.MemRefType
# Normalize src/dst to lists for variadic IR operands, while keeping old API working.
src_list = _normalize_variadic_tensor_operand(src, "src")
dst_list = _normalize_variadic_tensor_operand(dst, "dst")
# Validate first src/dst for element type width check
if isinstance(src_list[0].type, _cute_ir.MemRefType) and isinstance(
dst_list[0].type, _cute_ir.MemRefType
):
if src.element_type.width != dst.element_type.width:
if src_list[0].element_type.width != dst_list[0].element_type.width:
raise TypeError(
"`copy_atom_call` currently only supports equal source and destination "
"element type bit width"
)
if rank(src, mode=[0]) > 2 or rank(dst, mode=[0]) > 2:
if rank(src_list[0], mode=[0]) > 2 or rank(dst_list[0], mode=[0]) > 2:
raise NotImplementedError(
"V-mode (mode-0) with rank > 2 is not supported yet, "
f"but got rank(src, mode=[0]) = {rank(src, mode=[0])} and rank(dst, mode=[0]) = {rank(dst, mode=[0])}"
f"but got rank(src, mode=[0]) = {rank(src_list[0], mode=[0])} and rank(dst, mode=[0]) = {rank(dst_list[0], mode=[0])}"
)
value = atom._unpack(loc=loc, ip=ip, **kwargs)
if isinstance(pred, Tensor):
pred = pred.value
return _cute_ir.copy_atom_call(
value, src.value, dst.value, pred=pred, loc=loc, ip=ip
src_vals = [t.value for t in src_list]
dst_vals = [t.value for t in dst_list]
return _cute_ir.copy_atom_call(value, src_vals, dst_vals, pred=pred, loc=loc, ip=ip)
@dsl_user_op
def mma_atom_call(
atom: MmaAtom,
d: Tensor,
a: Tensor,
b: Tensor,
c: Tensor,
*,
loc=None,
ip=None,
**kwargs,
) -> None:
"""
Execute a single MMA atom operation.
The mma_atom_call operation executes an MMA atom with the given operands.
This performs a matrix multiplication and accumulation operation:
D = A * B + C
Note: The tensors 'd', 'a', 'b', and 'c' must only have a single fragment.
:param atom: The MMA atom to execute
:type atom: MmaAtom
:param d: Destination tensor (output accumulator)
:type d: Tensor
:param a: First source tensor (matrix A)
:type a: Tensor
:param b: Second source tensor (matrix B)
:type b: Tensor
:param c: Third source tensor (input accumulator C)
:type c: Tensor
:param loc: Source location for MLIR, defaults to None
:type loc: Optional[Location], optional
:param ip: Insertion point, defaults to None
:type ip: Optional[InsertionPoint], optional
Examples:
.. code-block:: python
# Call an MMA atom operation
cute.mma_atom_call(mma_atom, d_tensor, a_tensor, b_tensor, c_tensor)
"""
value = atom._unpack(loc=loc, ip=ip, **kwargs)
return _cute_ir.mma_atom_call(
value, d.value, a.value, b.value, c.value, loc=loc, ip=ip
)
python/CuTeDSL/cutlass/cute/core.py
+99 -73
View File
@@ -10,14 +10,14 @@
# is strictly prohibited.
from functools import partial, reduce
import inspect
from inspect import isclass
from typing import Any, Dict, List, Optional, Tuple, Type, Union, overload
from cutlass import const_expr
from typing_extensions import deprecated
from cutlass._mlir import ir
from cutlass._mlir.dialects import builtin, llvm, vector
from cutlass._mlir.dialects import builtin, llvm, vector, arith, nvvm
from cutlass._mlir.dialects import cute as _cute_ir
from cutlass._mlir.dialects.cute import (
Ratio as _Ratio,
@@ -125,6 +125,7 @@ __all__ = [
"shape",
"recast_ptr",
"make_ptr",
"get_remote_smem_ptr_in_cluster",
"composition",
"complement",
"right_inverse",
@@ -247,7 +248,7 @@ def _unpack_x_tuple(t: Union[ir.Type, ir.Value], *, loc=None, ip=None) -> XTuple
vals = []
else:
vals = get_leaves(t, loc=loc, ip=ip)
if not isinstance(vals, list):
if not isinstance(vals, ir.OpResultList):
vals = [vals]
else:
raise TypeError(f"expects static type or value, but got {t}")
@@ -383,9 +384,9 @@ class IntValue(cutlass_arith.ArithValue):
@property
def divisibility(self):
assert isinstance(
self.get_typed_value().type, _cute_ir.IntTupleType
), f"expected self.get_typed_value() to be int_tuple type, but got {self.get_typed_value().type}"
assert isinstance(self.get_typed_value().type, _cute_ir.IntTupleType), (
f"expected self.get_typed_value() to be int_tuple type, but got {self.get_typed_value().type}"
)
return self.get_typed_value().type.get_divisibility([0])
def __str__(self):
@@ -429,7 +430,9 @@ class IntValue(cutlass_arith.ArithValue):
@dsl_user_op
@_binary_op
def __add__(self, other, *, loc=None, ip=None):
return _cute_ir.tuple_add(self.get_typed_value(), other, loc=loc, ip=ip)
return _cute_ir.tuple_add(
self.get_typed_value(loc=loc, ip=ip), other, loc=loc, ip=ip
)
@dsl_user_op
@_binary_op
@@ -461,8 +464,10 @@ class IntValue(cutlass_arith.ArithValue):
@dsl_user_op
@_binary_op
def __radd__(self, other, *, loc=None, ip=None):
return _cute_ir.tuple_add(other, self.get_typed_value(), loc=loc, ip=ip)
def __radd__(self, other, *, loc=None, ip=None) -> "IntValue":
return _cute_ir.tuple_add(
other, self.get_typed_value(loc=loc, ip=ip), loc=loc, ip=ip
)
@dsl_user_op
@_binary_op
@@ -1207,10 +1212,6 @@ class _ComposedLayout(ComposedLayout):
@property
@dsl_user_op
def shape(self, *, loc=None, ip=None) -> Shape:
return self.shape_method(loc=loc, ip=ip)
@dsl_user_op
def shape_method(self, *, loc=None, ip=None) -> Shape:
return _unpack_x_tuple(
_cute_ir.get_shape(self.value, loc=loc, ip=ip), loc=loc, ip=ip
)
@@ -1262,9 +1263,9 @@ class _ComposedLayout(ComposedLayout):
# In this context, a _ComposedLayout instance is an encapsulated ir.Value which is automatically created
# by value caster for ComposedLayout typed values
assert len(values) == 1, f"Expected 1 value, but got {len(values)}"
assert isinstance(
values[0], (_ComposedLayout, ir.Value)
), f"Expected _ComposedLayout or ir.Value, but got {type(values[0])}"
assert isinstance(values[0], (_ComposedLayout, ir.Value)), (
f"Expected _ComposedLayout or ir.Value, but got {type(values[0])}"
)
return _ComposedLayout(
values[0] if isinstance(values[0], ir.Value) else values[0].value,
)
@@ -1313,9 +1314,9 @@ class _Pointer(Pointer):
# In this context, a _Pointer instance is an encapsulated ir.Value which is automatically created
# by value caster for cute.ptr typed values
assert len(values) == 1, f"Expected 1 value, but got {len(values)}"
assert isinstance(
values[0], (_Pointer, ir.Value)
), f"Expected _Pointer or ir.Value, but got {type(values[0])}"
assert isinstance(values[0], (_Pointer, ir.Value)), (
f"Expected _Pointer or ir.Value, but got {type(values[0])}"
)
return _Pointer(
values[0] if isinstance(values[0], ir.Value) else values[0].value
)
@@ -1359,29 +1360,12 @@ class _Pointer(Pointer):
"""
Get the LLVM pointer representation of this pointer.
:param loc: Source location for MLIR, defaults to None
:type loc: Optional[Location]
:param ip: Insertion point for MLIR, defaults to None
:type ip: Optional[InsertionPoint]
:return: The LLVM pointer representation
:rtype: ir.Value
"""
return self.to_llvm_ptr(loc=loc, ip=ip)
@dsl_user_op
@lru_cache_ir()
def to_llvm_ptr(self, *, loc=None, ip=None) -> ir.Value:
"""
Get the LLVM pointer representation of this pointer. (Used by internal API to propagate loc and ip)
:param loc: Source location for MLIR, defaults to None
:type loc: Optional[Location]
:param ip: Insertion point for MLIR, defaults to None
:type ip: Optional[InsertionPoint]
:return: The LLVM pointer representation
:rtype: ir.Value
"""
llvm_ptr_ty = llvm.PointerType.get(self.memspace.value)
llvm_ptr_ty = llvm.PointerType.get(
self.memspace.value if self.memspace != AddressSpace.rmem else 0
)
return builtin.unrealized_conversion_cast(
[llvm_ptr_ty], [self.value], loc=loc, ip=ip
)
@@ -1679,7 +1663,16 @@ def printf(*args, loc=None, ip=None) -> None:
elif isinstance(arg0, tuple):
# Assume it's a tile
return _pack_tile(arg0)
elif isinstance(arg0, (_Tensor, _Pointer, _ComposedLayout)):
elif isinstance(arg0, _Tensor):
arg0._check_can_load_store()
if isinstance(arg0.layout, ComposedLayout) and isinstance(
arg0.layout.inner, Swizzle
):
raise NotImplementedError(
"tensor with swizzled layout (PISL) is not supported in printf, please use swizzled pointer (PDSL) instead"
)
return arg0.value
elif isinstance(arg0, (_Pointer, _ComposedLayout)):
return arg0.value
else:
raise TypeError(f"unsupported argument type in printf, got {type(arg)}")
@@ -1751,6 +1744,7 @@ def make_swizzle(b, m, s, *, loc=None, ip=None):
return Swizzle(static(ty, loc=loc, ip=ip))
@dsl_user_op
def static(value, *, loc=None, ip=None):
return _cute_ir.static(value, loc=loc, ip=ip)
@@ -3409,39 +3403,90 @@ def make_ptr(
loc=None,
ip=None,
) -> Pointer:
# Perform checks
if dtype is None or not isinstance(dtype, NumericMeta):
raise TypeError(f"expects dtype to be a type of Numeric, but got {dtype}")
if not isinstance(mem_space, AddressSpace):
raise TypeError(f"expects mem_space to be an AddressSpace, but got {mem_space}")
if isinstance(value, ir.Value) and llvm.PointerType.isinstance(value.type):
value = llvm.ptrtoint(T.i64(), value)
if not is_integer(value):
raise TypeError(f"expects integer value, but got {type(value)}")
value = Int32(value) if mem_space == AddressSpace.tmem else Int64(value)
# TMEM addresses are 32b wide
is_tmem = mem_space == AddressSpace.tmem
value = Int32(value) if mem_space == AddressSpace.tmem else Int64(value)
# Set the alignment of the pointer
bytes_per_elt = max(1, dtype.width // 8)
if assumed_align is None:
assumed_align = bytes_per_elt
if bytes_per_elt % assumed_align != 0 and assumed_align % bytes_per_elt != 0:
raise ValueError(
f"{bytes_per_elt=} is not a multiple of {assumed_align=} and vice versa."
)
aligned_ty = _cute_ir.ConstrainedIntType.get(assumed_align, type(value).width)
aligned_intptr = _cute_ir.assume(
aligned_ty, value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip
)
# Construct the pointer Type
data_ty = T.i8() if dtype is None else dtype.mlir_type
ptr_ty = _cute_ir.PtrType.get(data_ty, mem_space, assumed_align)
return _cute_ir.inttoptr(ptr_ty, aligned_intptr, loc=loc, ip=ip)
@dsl_user_op
def get_remote_smem_ptr_in_cluster(
smem_ptr: Pointer,
cta_rank_in_cluster: Int,
*,
loc=None,
ip=None,
) -> Pointer:
"""
Get the remote shared memory CuTe pointer in a cluster.
:param smem_ptr: The current shared memory pointer
:type smem_ptr: Pointer
:param cta_rank_in_cluster: The peer CTA rank in cluster to get the remote pointer for
:type cta_rank_in_cluster: Int
:param loc: Source location for MLIR, defaults to None
:type loc: Optional[Location]
:param ip: Insertion point, defaults to None
:type ip: Optional[InsertionPoint]
:return: The remote shared memory CuTe pointer
:rtype: Pointer
"""
cur_llvm_ptr = smem_ptr.llvm_ptr
remote_llvm_ptr = nvvm.mapa(
llvm.PointerType.get(7), # LLVM dsmem address space
cur_llvm_ptr,
Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
remote_llvm_ptr_cast = llvm.addrspacecast(
llvm.PointerType.get(AddressSpace.smem), remote_llvm_ptr, loc=loc, ip=ip
)
remote_ptr = make_ptr(
smem_ptr.dtype,
remote_llvm_ptr_cast,
AddressSpace.smem,
assumed_align=smem_ptr.alignment,
loc=loc,
ip=ip,
)
if const_expr(smem_ptr.value.type.is_swizzled):
sw = Swizzle(static(smem_ptr.value.type.swizzle_type))
remote_ptr = recast_ptr(
remote_ptr, swizzle_=sw, dtype=smem_ptr.dtype, loc=loc, ip=ip
)
return remote_ptr
#
# Layout algebra
#
@@ -3868,9 +3913,7 @@ def local_tile(
return _cute_ir.local_tile(
input=input.value,
tile=tiler_val,
static_tile=None,
coord=coord_val,
static_coord=None,
proj=proj,
loc=loc,
ip=ip,
@@ -3907,9 +3950,9 @@ def make_layout_image_mask(
sliced_lay, offset = slice_and_offset(slicer, lay, loc=loc, ip=ip)
# Given that we replace only one mode with _, the rank of the slice should be 1
assert rank(sliced_lay) == 1
assert is_static(
sliced_lay
), "make_layout_image_mask requires the layout to be static"
assert is_static(sliced_lay), (
"make_layout_image_mask requires the layout to be static"
)
# Create the mask of the image
mcast_mask = Int16(0)
@@ -3952,6 +3995,7 @@ def leading_dim(shape: Shape, stride: Stride) -> Union[int, Tuple[int, ...], Non
return find_if(stride, pred_fn=pred_fn)
@dsl_user_op
def make_layout_tv(
thr_layout: Layout, val_layout: Layout, *, loc=None, ip=None
@@ -4468,9 +4512,9 @@ class struct:
"""
Return the round-up offset up to the next multiple of align.
"""
assert align > 0 and not (
align & (align - 1)
), "align should be a strictly positive power of 2."
assert align > 0 and not (align & (align - 1)), (
"align should be a strictly positive power of 2."
)
return (offset + (align - 1)) & ~(align - 1)
@@ -4607,29 +4651,11 @@ class FastDivmodDivisor:
new_obj = object.__new__(FastDivmodDivisor)
new_obj._divisor = values[0]
return new_obj
def __repr__(self):
return f"FastDivmodDivisor({self._divisor.type})"
# Set explicit signature for Sphinx documentation to avoid issues with @dsl_user_op decorator
FastDivmodDivisor.__init__.__signature__ = inspect.Signature(
[
inspect.Parameter("self", inspect.Parameter.POSITIONAL_OR_KEYWORD),
inspect.Parameter(
"divisor",
inspect.Parameter.POSITIONAL_OR_KEYWORD,
annotation=Integer,
),
inspect.Parameter(
"is_power_of_2",
inspect.Parameter.POSITIONAL_OR_KEYWORD,
default=None,
annotation=bool,
),
]
)
@dsl_user_op
def fast_divmod_create_divisor(
divisor: Integer, *, loc=None, ip=None
python/CuTeDSL/cutlass/cute/experimental/README.md
+57
View File
@@ -0,0 +1,57 @@
# CuTe Experimental APIs
> **Note:** APIs in this module are experimental and subject to change.
>
> This module serves as a staging area for new CuTe functionality that is still under active development. Performance, compile time, and interoperability with CuTe are works in progress. API signatures, behavior, and naming conventions may change without notice between releases.
>
> Once these APIs are stabilized, they will be migrated to the main `cute` submodules.
>
> Users are encouraged to experiment with these APIs but should be prepared to update their code as the interfaces evolve.
## Core APIs (`core.py`)
- `elect_sync` — Elects one thread within a warp
- `get_mbarrier` — Returns the mbarrier pointer for a given stage token
- `create_pipeline` — Creates a circular buffer of synchronization primitives indexed by stage count
- `create_pipeline_with_mask` — Creates a pipeline with an arrival mask for cluster-scoped synchronization
- `pipeline_advance_iterator` — Advances a pipeline iterator to the next stage
- `producer_acquire` / `producer_commit` — Producer-side pipeline synchronization
- `consumer_wait` / `consumer_release` / `consumer_tail` — Consumer-side pipeline synchronization
- `get_pipeline_produce_stage` / `get_pipeline_consume_stage` — Gets pipeline stage tokens
## Memory APIs (`memory.py`)
- `allocate` — Allocate a buffer with given type, layout, and address space
- `tma_load` — Copy tensor from global memory to shared memory using TMA
- `tma_load_multicast` — Copy tensor from global memory to shared memory using TMA with multicast
- `tma_store` — Copy tensor from shared memory to global memory using TMA
- `copy` — Copy tensor from src to dst using a given copy atom
## Algorithm APIs (`algorithm.py`)
- `simt_auto_vec_copy` — Copies a tensor between buffers with single thread (auto-vectorized)
- `partition` — Partition a buffer into a given layout and tiler
- `partition_and_copy` — Combines partitioning and copying in a single operation
## Math APIs (`math.py`)
- `dot` — Computes a dot product of two tensors using an MMA atom
- `dot_block_scaled` — Computes a block-scaled dot product with scale factors
## Pipeline Classes (`pipeline.py`)
- `GenericPipeline` — Generic pipeline for any producer/consumer combination
- `TMAToUMMAPipeline` — Pipeline for TMA load to UMMA consumption
- `TMAToAsyncPipeline` — Pipeline for TMA load to async consumer
- `AsyncToUMMAPipeline` — Pipeline for async producer to UMMA consumption
- `UMMAtoAsyncPipeline` — Pipeline for UMMA producer to async consumer
- `TMAStorePipeline` — Pipeline for SMEM producer to TMA store consumer
## Utilities (`utils.py`)
- `get_cta_v_map_ab` — Compute CTA-V map for A/B operands
- `get_cta_v_map_c` — Compute CTA-V map for C operand
- `make_tmem_layout_acc` — Derive TMEM accumulator buffer layout from a tiled MMA
- `make_tmem_layout_a` — Derive TMEM A-operand buffer layout from a tiled MMA
- `make_t2r_rmem_layout` — Derive per-thread RMEM buffer layout for the T2R epilogue copy
python/CuTeDSL/cutlass/cute/experimental/__init__.py
Executable → Regular
+12 -3
View File
@@ -9,6 +9,15 @@
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
raise NotImplementedError(
"CuTe Experimental module is only supported on Cuda toolkit 13.1 and above!"
)
from ... import cutlass_dsl as _dsl
jit = _dsl.CuteExperimentalDSL.jit
kernel = _dsl.CuteExperimentalDSL.kernel
compile = _dsl.CompileCallable()
from .algorithm import *
from .core import *
from .math import *
from .memory import *
from .pipeline import *
from .utils import *
python/CuTeDSL/cutlass/cute/experimental/algorithm.py
+150
View File
@@ -0,0 +1,150 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 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 cutlass import cute
from cutlass.cutlass_dsl import dsl_user_op
from cutlass._mlir.dialects import lir as cutlass_lir
from .memory import copy
@dsl_user_op
def simt_auto_vec_copy(
src: cute.Tensor, dst: cute.Tensor, async_op=False, loc=None, ip=None
):
"""
Copies a tensor between two cute.memref buffers with single thread.
:param src: Source tensor
:type src: cute.Tensor
:param dst: Destination tensor
:type dst: cute.Tensor
:param async_op: Whether to use asynchronous operation, defaults to False
:type async_op: bool, optional
"""
if async_op:
cutlass_lir.SimtAutoVecCopyOp(
src.value, dst.value, async_=True, cache="always", loc=loc, ip=ip
)
else:
cutlass_lir.SimtAutoVecCopyOp(src.value, dst.value, loc=loc, ip=ip)
@dsl_user_op
def partition(
buffer: cute.Tensor, agent_id: cute.Int32, *, layout_tv, tiler, loc=None, ip=None
) -> cute.Tensor:
"""
Partition a buffer into a given layout and tiler.
:param buffer: Buffer to partition
:type buffer: cute.Tensor
:param agent_id: Agent ID
:type agent_id: cute.Int32
:param layout_tv: Layout tensor
:type layout_tv: cute.Tensor
:param tiler: Tiler
:type tiler: cute.Tensor
"""
assert isinstance(agent_id, cute.Int32), (
f"Expected agent_id to be cute.Int32, got {type(agent_id)}"
)
partition_op = cutlass_lir.PartitionOp(
buffer.value,
agent_id.ir_value(),
layout_tv=layout_tv.type.attribute,
tiler=tiler.type.attribute,
loc=loc,
ip=ip,
)
return partition_op.result
@dsl_user_op
def partition_and_copy(
tiled_copy: cute.core.ThrCopy,
src: cute.Tensor,
dst: cute.Tensor,
*,
loc=None,
ip=None,
):
"""
Copies a tensor between two cute.memref buffer
:param tiled_copy: Tiled copy
:type tiled_copy: cute.core.ThrCopy
:param src: Source tensor
:type src: cute.Tensor
:param dst: Destination tensor
:type dst: cute.Tensor
"""
src_partitioned = src
dst_partitioned = dst
tid_x = tiled_copy.thr_idx
if src.memspace != cute.AddressSpace.rmem:
src_partitioned = partition(
src,
tid_x,
layout_tv=tiled_copy.layout_src_tv_tiled,
tiler=cute.core._pack_tile(tiled_copy.tiler_mn),
)
if dst.memspace != cute.AddressSpace.rmem:
dst_partitioned = partition(
dst,
tid_x,
layout_tv=tiled_copy.layout_dst_tv_tiled,
tiler=cute.core._pack_tile(tiled_copy.tiler_mn),
)
# Handle copy where copy atom is used for both partition and copy during smem to rmem and rmem to smem copies
if type(tiled_copy.op) in [
cute.nvgpu.warp.LdMatrix8x8x16bOp,
cute.nvgpu.warp.LdMatrix16x16x8bOp,
cute.nvgpu.warp.StMatrix8x8x16bOp,
cute.nvgpu.warp.StMatrix16x8x8bOp,
]:
copy(
src_partitioned,
dst_partitioned,
copy_atom=tiled_copy,
loc=loc,
ip=ip,
)
# The rest handles copy where copy atom is used for partition
elif (
src.memspace,
dst.memspace,
) in [
(cute.AddressSpace.rmem, cute.AddressSpace.smem),
(cute.AddressSpace.smem, cute.AddressSpace.rmem),
(cute.AddressSpace.rmem, cute.AddressSpace.gmem),
(cute.AddressSpace.gmem, cute.AddressSpace.rmem),
]:
simt_auto_vec_copy(src_partitioned, dst_partitioned, loc=loc, ip=ip)
elif (
src.memspace == cute.AddressSpace.gmem
and dst.memspace == cute.AddressSpace.smem
):
simt_auto_vec_copy(
src_partitioned, dst_partitioned, async_op=True, loc=loc, ip=ip
)
# Handle copy where copy atom is used for partition and copy
else:
copy(
src_partitioned,
dst_partitioned,
copy_atom=tiled_copy,
loc=loc,
ip=ip,
)
python/CuTeDSL/cutlass/cute/experimental/core.py
+245
View File
@@ -0,0 +1,245 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
from cutlass.cutlass_dsl import dsl_user_op
from cutlass._mlir.dialects import lir as cutlass_lir_ir, nvvm as _nvvm
from cutlass._mlir import ir
from cutlass.cutlass_dsl import lru_cache_ir
from cutlass._mlir.dialects.core import OperationTypeEnum
from cutlass import cute
@dsl_user_op
def elect_sync(loc=None, ip=None):
"""
Elects one predicated thread within a warp.
"""
return _nvvm.elect_sync(loc=loc, ip=ip)
@dsl_user_op
def get_mbarrier(stage_token, loc=None, ip=None):
"""
Returns the mbarrier pointer for a given stage token.
"""
return cutlass_lir_ir.GetMbarrierOp(stage_token, loc=loc, ip=ip)
@ir.register_value_caster(cutlass_lir_ir.PipelineStateType.get_static_typeid())
class PipelineState(ir.Value):
def __init__(self, value):
if isinstance(value, ir.Value):
self.value = value
else:
raise TypeError(f"Expected ir.Value, got {type(value)}")
super().__init__(value)
@property
@lru_cache_ir()
def type(self) -> ir.Type:
return self.value.type
@classmethod
def __new_from_mlir_values__(cls, values):
assert len(values) == 1, f"Expected 1 value, but got {len(values)}"
return PipelineState(values[0])
@dsl_user_op
def create_pipeline(
stage: cute.Int32,
producer: OperationTypeEnum,
consumer: OperationTypeEnum,
producer_arv_count: cute.Int32,
consumer_arv_count: cute.Int32,
loc=None,
ip=None,
) -> tuple[PipelineState, PipelineState, PipelineState]:
"""
Creates an abstraction for a circular buffer of synchronizatoin primitives
indexed by stage count.
:param stage: Stage count
:type stage: cute.Int32
:param producer: Producer operation type
:type producer: OperationTypeEnum
:param consumer: Consumer operation type
:type consumer: OperationTypeEnum
:param producer_arv_count: Producer arrival count
:type producer_arv_count: cute.Int32
:param consumer_arv_count: Consumer arrival count
:type consumer_arv_count: cute.Int32
"""
if isinstance(producer_arv_count, int):
producer_arv_count = cute.Int32(producer_arv_count)
if isinstance(consumer_arv_count, int):
consumer_arv_count = cute.Int32(consumer_arv_count)
result = ir.Type.parse(f"!lir.pipeline<{stage}, {producer} -> {consumer}>")
op = cutlass_lir_ir.CreatePipelineOp(
result,
producer_arv_count.ir_value(),
consumer_arv_count.ir_value(),
loc=loc,
ip=ip,
)
pipeline = op.result
result = ir.Type.parse(f"!lir.pipeline_state<{stage}>")
op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip)
producer_state = op.result
result = ir.Type.parse(f"!lir.pipeline_state<{stage}>")
op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip)
consumer_state = op.result
return pipeline, producer_state, consumer_state
@dsl_user_op
def create_pipeline_with_mask(
stage: cute.Int32,
producer: OperationTypeEnum,
consumer: OperationTypeEnum,
producer_arv_count: cute.Int32,
consumer_arv_count: cute.Int32,
arrival_mask: cute.Int16,
loc=None,
ip=None,
) -> tuple[PipelineState, PipelineState, PipelineState]:
"""
Creates a pipeline with an arrival mask for cluster-scoped synchronization.
:param stage: Pipeline stage count.
:param producer: Producer operation type (e.g. SM90_TMA_LOAD_MULTICAST).
:param consumer: Consumer operation type (e.g. SM100_MMA_2SM_SS).
:param producer_arv_count: Producer arrival count for the pipeline barriers.
:param consumer_arv_count: Consumer arrival count for the pipeline barriers.
:param arrival_mask: Bitmask that selects participating peers (e.g. CTAs in a
cluster). This is attached to the pipeline value and is consulted by some
pipeline lowerings to generate cluster-scoped synchronization
"""
if isinstance(producer_arv_count, int):
producer_arv_count = cute.Int32(producer_arv_count)
if isinstance(consumer_arv_count, int):
consumer_arv_count = cute.Int32(consumer_arv_count)
if isinstance(arrival_mask, int):
arrival_mask = cute.Int16(arrival_mask)
result = ir.Type.parse(f"!lir.pipeline<{stage}, {producer} -> {consumer}>")
op = cutlass_lir_ir.CreatePipelineWithMaskOp(
result,
producer_arv_count.ir_value(),
consumer_arv_count.ir_value(),
arrival_mask.ir_value(),
loc=loc,
ip=ip,
)
pipeline = op.result
result = ir.Type.parse(f"!lir.pipeline_state<{stage}>")
op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip)
producer_state = op.result
result = ir.Type.parse(f"!lir.pipeline_state<{stage}>")
op = cutlass_lir_ir.CreatePipelineStateOp(result, pipeline, loc=loc, ip=ip)
consumer_state = op.result
return pipeline, producer_state, consumer_state
@dsl_user_op
def pipeline_advance_iterator(pipe, state, loc=None, ip=None):
"""
Advances a pipeline iterator to the next stage.
"""
op = cutlass_lir_ir.PipelineAdvanceIteratorOp(pipe, state, loc=loc, ip=ip)
return op.result
@dsl_user_op
def producer_acquire(pipe, state, loc=None, ip=None):
"""
Acquires exclusive access to a pipeline.
"""
op = cutlass_lir_ir.ProducerAcquireOp(pipe, state, loc=loc, ip=ip)
return op.result
@dsl_user_op
def producer_commit(pipe, state, loc=None, ip=None):
"""
Commits results to a pipeline.
"""
op = cutlass_lir_ir.ProducerCommitOp(pipe, state, loc=loc, ip=ip)
return op.result
@dsl_user_op
def consumer_wait(pipe, state, loc=None, ip=None):
"""
Waits for a pipeline to transition to `full`.
"""
op = cutlass_lir_ir.ConsumerWaitOp(pipe, state, loc=loc, ip=ip)
return op.result
@dsl_user_op
def consumer_release(pipe, state, loc=None, ip=None):
"""
Releases a pipeline that has been consumed.
"""
op = cutlass_lir_ir.ConsumerReleaseOp(pipe, state, loc=loc, ip=ip)
return op.result
@dsl_user_op
def consumer_tail(pipe, state, loc=None, ip=None):
"""
Called by the consumer to block until asynchronous tasks have completed.
"""
op = cutlass_lir_ir.ConsumerTailOp(pipe, state, loc=loc, ip=ip)
return op.result
@dsl_user_op
def get_pipeline_produce_stage(pipeline, state, loc=None, ip=None):
"""
Gets a pipeline produce stage.
"""
stage_token_type = ir.Type.parse(f"!lir.stage_token<{pipeline.type}>")
stage_idx = ir.IntegerType.get_signless(32)
op = cutlass_lir_ir.GetPipelineProduceStageOp(
stage_token=stage_token_type,
stage_index=stage_idx,
pipeline=pipeline,
pipelineState=state,
loc=loc,
ip=ip,
)
return op.stage_token, op.stage_index
@dsl_user_op
def get_pipeline_consume_stage(pipeline, state, loc=None, ip=None):
"""
Creates a pipeline consume stage.
"""
stage_token_type = ir.Type.parse(f"!lir.stage_token<{pipeline.type}>")
stage_idx = ir.IntegerType.get_signless(32)
op = cutlass_lir_ir.GetPipelineConsumeStageOp(
stage_token=stage_token_type,
stage_index=stage_idx,
pipeline=pipeline,
pipelineState=state,
loc=loc,
ip=ip,
)
return op.stage_token, op.stage_index
python/CuTeDSL/cutlass/cute/experimental/math.py
+84
View File
@@ -0,0 +1,84 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 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 cutlass import cute
from cutlass.cutlass_dsl import dsl_user_op
from cutlass._mlir.dialects import lir as cutlass_lir
@dsl_user_op
def dot_block_scaled(
mma_atom: cute.MmaAtom,
a: cute.Tensor,
sfa: cute.Tensor,
b: cute.Tensor,
sfb: cute.Tensor,
c: cute.Tensor,
loc=None,
ip=None,
):
"""
Computes the dot product of two tensors with block scaling and accumulates the result into a third tensor.
:param mma_atom: MMA atom
:type mma_atom: cute.MmaAtom
:param a: First tensor
:type a: cute.Tensor
:param sfa: First scale factor tensor
:type sfa: cute.Tensor
:param b: Second tensor
:type b: cute.Tensor
:param sfb: Second scale factor tensor
:type sfb: cute.Tensor
:param c: Result tensor
:type c: cute.Tensor
"""
cutlass_lir.DotBlockScaledOp(
mma_atom._unpack(),
a.value,
sfa.value,
b.value,
sfb.value,
c.value,
loc=loc,
ip=ip,
)
@dsl_user_op
def dot(
mma_atom: cute.MmaAtom,
a: cute.Tensor,
b: cute.Tensor,
c: cute.Tensor,
loc=None,
ip=None,
):
"""
Computes the dot product of two tensors and accumulates the result into a third tensor.
:param mma_atom: MMA atom
:type mma_atom: cute.MmaAtom
:param a: First tensor
:type a: cute.Tensor
:param b: Second tensor
:type b: cute.Tensor
:param c: Result tensor
:type c: cute.Tensor
"""
cutlass_lir.DotOp(
mma_atom._unpack(),
a.value,
b.value,
c.value,
loc=loc,
ip=ip,
)
python/CuTeDSL/cutlass/cute/experimental/memory.py
+256
View File
@@ -0,0 +1,256 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 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 typing import Type, Optional
from cutlass._mlir.dialects import cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
from cutlass.cutlass_dsl import dsl_user_op
from cutlass._mlir.dialects import (
lir as cutlass_lir,
cute as _cute_ir,
)
from cutlass._mlir.dialects.core import OperationTypeEnum
from cutlass import cute
def _get_tma_load_kind(tma_operation_type: OperationTypeEnum):
"""Convert OperationTypeEnum to TiledTmaLoadEnum."""
if tma_operation_type == OperationTypeEnum.SM100_TMA_LOAD_2SM_MULTICAST:
return _cute_ir.TiledTmaLoadEnum.sm_100_2sm_multicast
if tma_operation_type == OperationTypeEnum.SM90_TMA_LOAD_MULTICAST:
return _cute_ir.TiledTmaLoadEnum.sm_90_multicast
if tma_operation_type == OperationTypeEnum.SM100_TMA_LOAD_2SM:
return _cute_ir.TiledTmaLoadEnum.sm_100_2sm
if tma_operation_type == OperationTypeEnum.SM90_TMA_LOAD:
return _cute_ir.TiledTmaLoadEnum.sm_90
raise ValueError(f"Unsupported TMA operation type: {tma_operation_type}")
@dsl_user_op
def allocate(
type: Type[cute.Numeric],
address_space: cute.AddressSpace,
layout: cute.Layout | cute.ComposedLayout,
alignment: cute.Int32,
is2cta: bool = False,
loc=None,
ip=None,
) -> cute.Tensor:
"""
Allocate a buffer of the given type and layout.
:param type: The type of the buffer
:type type: cute.Tensor
:param layout: The layout of the buffer
:type layout: cute.Layout
:param address_space: The address space of the buffer
:type address_space: str
:param alignment: The alignment of the buffer
:type alignment: cute.Int32
:param is2cta: Whether TMEM allocation should span a CTA pair (2CTA TMEM)
:type is2cta: bool
"""
swizzle = None
if isinstance(layout, cute.ComposedLayout):
swizzle = layout.inner
layout = layout.outer
# Handle SparseElemType (pass through) vs regular types (get mlir_type)
if isinstance(type, _cute_ir.SparseElemType):
pass
else:
type = type.mlir_type
ptr_ty = _cute_ir.PtrType.get(
type,
address_space,
alignment,
swizzle.type.attribute if swizzle else None,
)
buffer_type = _cute_ir.MemRefType.get(ptr_ty, layout.type)
# `is2cta` is a UnitAttr flag in the IR:
# present => true, absent => false.
is2cta_attr = ir.UnitAttr.get() if is2cta else None
buffer_op = cutlass_lir.AllocateBufferOp(
buffer_type, is2cta=is2cta_attr, loc=loc, ip=ip
)
return buffer_op.result
@dsl_user_op
def tma_load(
src: cute.Tensor,
dst: cute.Tensor,
mbar,
*,
cta_v_map,
tma_operation_type: Optional[OperationTypeEnum] = None,
internal_type=None,
update_expect_tx: bool = True,
loc=None,
ip=None,
):
"""
Copies a tensor pointed by a !cute.memref into a Buffer using TMA.
update_expect_tx (bool): controls whether this operation increments the mbarrier's transaction bytes with the TMA copy size.
When used with Cute DSL pipelines, it must be set to False as the pipeline already initializes the mbarrier's transaction bytes.
tma_operation_type (optional): specifies the TMA operation type (SM90_TMA_LOAD, SM100_TMA_LOAD_2SM, etc.)
internal_type (optional): selects the TMA transfer's internal element encoding used by hardware.
Does not change src/dst memref types. For structured sparsity, use base storage types:
Float16 for 2:4 FP16 sparse element type, Uint8 for 8:1 uint8 sparse element type.
:param src: Source tensor in global memory
:type src: cute.Tensor
:param dst: Destination tensor in shared memory
:type dst: cute.Tensor
:param mbar: Memory barrier for synchronization
:type mbar: cute.core.Mbarrier
:param cta_v_map: CTA V-map for the tensor
:type cta_v_map: cute.core.CtaVMap
:param tma_operation_type: TMA operation type (e.g., SM90_TMA_LOAD, SM100_TMA_LOAD_2SM, etc.)
:type tma_operation_type: OperationTypeEnum
:param internal_type: Internal type of the TMA transfer
:type internal_type: cute.core.InternalType
:param update_expect_tx: Whether to update expected transaction bytes
:type update_expect_tx: bool
"""
if tma_operation_type is not None:
kind = _get_tma_load_kind(tma_operation_type)
else:
kind = _cute_ir.TiledTmaLoadEnum.sm_90
kwargs = {
"cta_v_map": cta_v_map.type.attribute,
"kind": kind,
"loc": loc,
"ip": ip,
}
# Map internal_type to tma_format per updated API
if internal_type is not None:
internal_mlir_ty = (
internal_type.mlir_type
if hasattr(internal_type, "mlir_type")
else internal_type
)
kwargs["tma_format"] = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_ty, False)
)
if update_expect_tx:
kwargs["update_expect_tx"] = True
cutlass_lir.TmaLoadOp(src.value, dst.value, mbar, **kwargs)
@dsl_user_op
def tma_load_multicast(
src: cute.Tensor,
dst: cute.Tensor,
mbar,
*,
vmnk_layout: cute.Layout,
cta_v_map,
tma_operation_type: OperationTypeEnum,
multicast_mode: int,
update_expect_tx: bool = True,
loc=None,
ip=None,
):
"""
Copies a tensor pointed by a !cute.memref into a Buffer using TMA with multicast.
:param src: Source tensor in global memory
:param dst: Destination tensor in shared memory
:param mbar: Memory barrier for synchronization
:param vmnk_layout: Layout describing the cluster configuration
:param cta_v_map: CTA V-map for the tensor
:param tma_operation_type: TMA operation type (e.g., SM90_TMA_LOAD_MULTICAST, SM100_TMA_LOAD_2SM_MULTICAST)
:param multicast_mode: Multicast projection mode (1=column, 2=row)
:param update_expect_tx: Whether to update expected transaction bytes
"""
kind = _get_tma_load_kind(tma_operation_type)
kwargs = {
"cta_v_map": cta_v_map.type.attribute,
"kind": kind,
"vmnk_layout": vmnk_layout,
"multicast_mode": multicast_mode,
"loc": loc,
"ip": ip,
}
if update_expect_tx:
kwargs["update_expect_tx"] = True
cutlass_lir.TmaLoadMulticastOp(
src.value,
dst.value,
mbar,
**kwargs,
)
@dsl_user_op
def tma_store(
src: cute.Tensor,
dst: cute.Tensor,
*,
cta_v_map,
internal_type=None,
loc=None,
ip=None,
):
"""
Copies a tensor from a Buffer to a tensor pointed to by a !cute.memref.
internal_type (optional): selects the TMA transfer's internal element encoding used by hardware.
Does not change src/dst memref types. For structured sparsity, use base storage types:
Float16 for 2:4 FP16 sparse element type, Uint8 for 8:1 uint8 sparse element type.
:param src: Source tensor in shared memory
:type src: cute.Tensor
:param dst: Destination tensor in global memory
:type dst: cute.Tensor
:param cta_v_map: CTA V-map for the tensor
:type cta_v_map: cute.core.CtaVMap
:param internal_type: Internal type of the TMA transfer
:type internal_type: cute.core.InternalType
"""
kwargs = {
"cta_v_map": cta_v_map.type.attribute,
"loc": loc,
"ip": ip,
}
# Map internal_type to tma_format per updated API
if internal_type is not None:
internal_mlir_ty = (
internal_type.mlir_type
if hasattr(internal_type, "mlir_type")
else internal_type
)
kwargs["tma_format"] = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_ty, False)
)
cutlass_lir.TmaStoreOp(src.value, dst.value, **kwargs)
@dsl_user_op
def copy(src: cute.Tensor, dst: cute.Tensor, *, copy_atom, loc=None, ip=None):
"""
Copy a tensor from src to dst using a given copy atom.
"""
copy_atom = ir.Attribute.parse(f"{copy_atom.type}")
cutlass_lir.CopyOp(src.value, dst.value, copy_atom=copy_atom, loc=loc, ip=ip)
python/CuTeDSL/cutlass/cute/experimental/pipeline.py
+684
View File
@@ -0,0 +1,684 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
#
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
# property and proprietary rights in and to this material, related
# documentation and any modifications thereto. Any use, reproduction,
# disclosure or distribution of this material and related documentation
# without an express license agreement from NVIDIA CORPORATION or
# its affiliates is strictly prohibited.
"""
Convenience pipeline classes that hide elect_one synchronization complexity
"""
from dataclasses import dataclass
from typing import Optional
import cutlass
import cutlass.cute as cute
from cutlass._mlir.dialects import lir as cutlass_lir_ir
from cutlass.base_dsl.typing import Int32
from cutlass._mlir.dialects.core import OperationTypeEnum
from cutlass.cute.experimental.core import (
create_pipeline,
create_pipeline_with_mask,
producer_acquire,
get_pipeline_produce_stage,
get_pipeline_consume_stage,
producer_commit,
consumer_release,
pipeline_advance_iterator,
consumer_wait,
consumer_tail,
)
from cutlass.cutlass_dsl import CuteExperimentalDSL
class GenericPipelineBase:
"""Base class for pipeline convenience wrappers"""
def __init__(
self,
raw_pipeline,
num_stages,
producer_state,
consumer_state,
):
self.raw_pipeline = raw_pipeline
self.num_stages = num_stages
# For convenience class, we always manage state internally
self.producer_state = producer_state
self.consumer_state = consumer_state
def __extract_mlir_values__(self):
"""Extract MLIR values for DynamicExpression protocol."""
# raw_pipeline is always ir.OpResult from create_pipeline (no __extract_mlir_values__)
pipeline_values = [self.raw_pipeline]
# Create DSL types and extract their underlying MLIR values
num_stages_dsl = Int32(self.num_stages)
# Pipeline states are already MLIR values (PipelineState objects)
producer_state_values = [self.producer_state]
consumer_state_values = [self.consumer_state]
return (
pipeline_values
+ [
num_stages_dsl.__extract_mlir_values__()[0],
]
+ producer_state_values
+ consumer_state_values
)
@classmethod
def __new_from_mlir_values__(cls, values):
"""Reconstruct object from MLIR values."""
# Parse the known structure: [pipeline] + [num_stages, producer_flag, consumer_flag] + [producer_state] + [consumer_state]
# All lir_* objects are single MLIR values
raw_pipeline = values[0] # Always single ir.OpResult
num_stages_val = values[1]
producer_state = values[2] # Always single PipelineState
consumer_state = values[3] # Always single PipelineState
# Create temporary DSL objects and extract Python values
temp_num_stages = Int32(0)
num_stages_dsl = temp_num_stages.__new_from_mlir_values__([num_stages_val])
return cls(
raw_pipeline,
(
num_stages_dsl.value
if hasattr(num_stages_dsl, "value")
else int(num_stages_dsl)
),
producer_state,
consumer_state,
)
def producer_acquire(self):
"""Acquire producer state."""
producer_acquire(self.raw_pipeline, self.producer_state)
return self
def get_producer_stage(self):
"""Get producer stage."""
return get_pipeline_produce_stage(self.raw_pipeline, self.producer_state)
def get_consumer_stage(self):
"""Get consumer stage."""
return get_pipeline_consume_stage(self.raw_pipeline, self.consumer_state)
# Instance methods that can now be used directly in kernel context
def producer_acquire_and_get_stage(self):
"""Combined producer acquire + get_stage with automatic elect_one using internal state."""
self.producer_acquire()
return get_pipeline_produce_stage(self.raw_pipeline, self.producer_state)
def producer_commit(self):
"""Commit producer state."""
producer_commit(self.raw_pipeline, self.producer_state)
return self
def consumer_release(self):
"""Release consumer state."""
consumer_release(self.raw_pipeline, self.consumer_state)
return self
def producer_commit_and_advance(self):
"""Combined producer commit + advance with automatic elect_one using internal state."""
self.producer_commit()
# Update internal state in-place for better performance
self.producer_state = pipeline_advance_iterator(
self.raw_pipeline, self.producer_state
)
return self
def consumer_wait_and_get_stage(self):
"""Combined consumer wait + get_stage with automatic elect_one using internal state."""
self.consumer_wait()
return get_pipeline_consume_stage(self.raw_pipeline, self.consumer_state)
def consumer_wait(self):
"""Wait for consumer to be ready."""
consumer_wait(self.raw_pipeline, self.consumer_state)
return self
def consumer_release_and_advance(self):
"""Combined consumer release + advance with automatic elect_one using internal state."""
self.consumer_release()
# Update internal state in-place for better performance
self.consumer_state = pipeline_advance_iterator(
self.raw_pipeline, self.consumer_state
)
return self
def consumer_tail(self):
"""Combined consumer tail with automatic elect_one using internal state."""
consumer_tail(self.raw_pipeline, self.consumer_state)
return self
class GenericPipeline(GenericPipelineBase):
"""
Generic pipeline for any combination of producer and consumer.
"""
@staticmethod
def create(
*,
producer: OperationTypeEnum,
consumer: OperationTypeEnum,
producer_arv_count: cute.Int32,
consumer_arv_count: cute.Int32,
num_stages: cute.Int32,
):
"""
Create a generic pipeline with parameterized producer and consumer.
Args:
producer: Producer operation type
consumer: Consumer operation type
producer_arv_count: Producer arrival count
consumer_arv_count: Consumer arrival count
num_stages: Number of pipeline stages
"""
raw_pipeline, producer_state, consumer_state = create_pipeline(
num_stages,
producer,
consumer,
producer_arv_count=producer_arv_count,
consumer_arv_count=consumer_arv_count,
)
return GenericPipeline(
raw_pipeline,
num_stages,
producer_state,
consumer_state,
)
def _validate_umma_operation_type(operation_type: OperationTypeEnum):
if operation_type not in [
OperationTypeEnum.SM100_MMA_1SM_SS,
OperationTypeEnum.SM100_MMA_1SM_TS,
OperationTypeEnum.SM100_MMA_2SM_SS,
OperationTypeEnum.SM100_MMA_2SM_TS,
OperationTypeEnum.SM100_MMA_SCALED_1SM_SS,
OperationTypeEnum.SM100_MMA_SCALED_1SM_TS,
OperationTypeEnum.SM100_MMA_SCALED_2SM_SS,
OperationTypeEnum.SM100_MMA_SCALED_2SM_TS,
]:
raise ValueError(f"Invalid UMMA operation type: {operation_type}")
def _is_2sm_umma_operation_type(operation_type: OperationTypeEnum) -> bool:
"""Check if the operation type is a 2SM UMMA operation."""
return operation_type in [
OperationTypeEnum.SM100_MMA_2SM_SS,
OperationTypeEnum.SM100_MMA_2SM_TS,
OperationTypeEnum.SM100_MMA_SCALED_2SM_SS,
OperationTypeEnum.SM100_MMA_SCALED_2SM_TS,
]
class TMAToUMMAPipeline(GenericPipelineBase):
"""
Pipeline for TMA to UMMA.
"""
@staticmethod
def create(
*,
num_stages: cute.Int32,
mma_operation_type: OperationTypeEnum,
tma_operation_type: Optional[OperationTypeEnum] = None,
cluster_layout_vmnk: Optional[cute.Layout] = None,
):
"""
Create a TMA to UMMA pipeline.
For 2SM MMA with TMA_LOAD_2SM, provide cluster_layout_vmnk for proper mask computation.
"""
_validate_umma_operation_type(
mma_operation_type,
)
# Default to SM90_TMA_LOAD if not specified
if tma_operation_type is None:
tma_operation_type = OperationTypeEnum.SM90_TMA_LOAD
if tma_operation_type == OperationTypeEnum.SM100_TMA_LOAD_2SM:
if cluster_layout_vmnk is None:
raise ValueError(
"cluster_layout_vmnk is required if using 2CTA MMA with TMA"
)
# If using 2CTA MMA, need consumer_mask == local_cta | peer_cta
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
cta_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
cta_rank_in_cluster
)
arrival_mask = cute.make_layout_image_mask(
cluster_layout_vmnk, cta_in_cluster_coord_vmnk, mode=0
)
raw_pipeline, producer_state, consumer_state = create_pipeline_with_mask(
num_stages,
tma_operation_type,
mma_operation_type,
producer_arv_count=1,
consumer_arv_count=1,
arrival_mask=arrival_mask,
)
else:
raw_pipeline, producer_state, consumer_state = create_pipeline(
num_stages,
tma_operation_type,
mma_operation_type,
producer_arv_count=1,
consumer_arv_count=1,
)
return TMAToUMMAPipeline(
raw_pipeline,
num_stages,
producer_state,
consumer_state,
)
@staticmethod
def create_with_mask(
*,
num_stages: cute.Int32,
tma_operation_type: OperationTypeEnum,
mma_operation_type: OperationTypeEnum,
cluster_layout_vmnk: cute.Layout,
):
"""
Create a TMA to UMMA pipeline with multicast mask for 2CTA operations.
"""
_validate_umma_operation_type(
mma_operation_type,
)
# Calculate TMA multicasting masks
tma_mcast_proj_A = 2 # multicast across CTAs in same row
tma_mcast_proj_B = 1 # multicast across CTAs in same column
cta_rank_in_cluster = cute.arch.make_warp_uniform(
cute.arch.block_idx_in_cluster()
)
cta_in_cluster_coord_vmnk = cluster_layout_vmnk.get_flat_coord(
cta_rank_in_cluster
)
# For 2CTA MMA (v-size==2), the peer CTA is the other v-slice (xor 1).
# For 1CTA MMA (v-size==1), the peer is the local CTA (no flip).
v_size = cute.size(cluster_layout_vmnk.shape[0])
peer_v = (
(cta_in_cluster_coord_vmnk[0] ^ 1)
if cutlass.const_expr(v_size > 1)
else cta_in_cluster_coord_vmnk[0]
)
cta_in_cluster_coord_vmnk_peer = (
peer_v,
*cta_in_cluster_coord_vmnk[1:],
)
arrival_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, cta_in_cluster_coord_vmnk, tma_mcast_proj_A
)
arrival_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
cluster_layout_vmnk, cta_in_cluster_coord_vmnk, tma_mcast_proj_B
)
arrival_mask_a_peer = cute.nvgpu.cpasync.create_tma_multicast_mask(
cluster_layout_vmnk,
cta_in_cluster_coord_vmnk_peer,
mcast_mode=tma_mcast_proj_A,
)
arrival_mask_b_peer = cute.nvgpu.cpasync.create_tma_multicast_mask(
cluster_layout_vmnk,
cta_in_cluster_coord_vmnk_peer,
mcast_mode=tma_mcast_proj_B,
)
# if 1SM MMA, arrival_mask_a_peer==arrival_mask_a && arrival_mask_b==arrival_mask_b_peer
arrival_mask_c = (
arrival_mask_a | arrival_mask_a_peer | arrival_mask_b | arrival_mask_b_peer
)
num_mcast_ctas_a = cute.size(cluster_layout_vmnk.shape[2])
num_mcast_ctas_b = cute.size(cluster_layout_vmnk.shape[1])
num_mcast_participants = num_mcast_ctas_a + num_mcast_ctas_b - 1
raw_pipeline, producer_state, consumer_state = create_pipeline_with_mask(
num_stages,
tma_operation_type,
mma_operation_type,
producer_arv_count=1,
consumer_arv_count=num_mcast_participants,
arrival_mask=arrival_mask_c,
)
return TMAToUMMAPipeline(
raw_pipeline, num_stages, producer_state, consumer_state
)
def producer_commit(self):
"""Commit producer state."""
with cute.arch.elect_one():
super().producer_commit()
return self
def consumer_release(self):
"""Release consumer state."""
with cute.arch.elect_one():
super().consumer_release()
return self
class TMAToAsyncPipeline(GenericPipelineBase):
"""
Pipeline for TMA to * (except UMMA).
"""
@staticmethod
def create(
*,
num_stages: cute.Int32,
consumer: OperationTypeEnum,
consumer_arv_count: cute.Int32,
):
"""
Create a TMA to * (except UMMA) pipeline.
"""
raw_pipeline, producer_state, consumer_state = create_pipeline(
num_stages,
OperationTypeEnum.SM90_TMA_LOAD,
consumer,
producer_arv_count=1,
consumer_arv_count=consumer_arv_count,
)
return TMAToAsyncPipeline(
raw_pipeline,
num_stages,
producer_state,
consumer_state,
)
def producer_commit(self):
"""Commit producer state."""
with cute.arch.elect_one():
super().producer_commit()
return self
class AsyncToUMMAPipeline(GenericPipelineBase):
"""
Pipeline for * (except TMA) to UMMA.
"""
@staticmethod
def create(
*,
num_stages: cute.Int32,
producer: OperationTypeEnum,
producer_arv_count: cute.Int32,
mma_operation_type: OperationTypeEnum,
):
"""
Create a * (except TMA) to UMMA pipeline.
"""
_validate_umma_operation_type(
mma_operation_type,
)
if producer == OperationTypeEnum.SM90_TMA_LOAD:
raise ValueError("TMA to UMMA is not supported.")
raw_pipeline, producer_state, consumer_state = create_pipeline(
num_stages,
producer,
mma_operation_type,
producer_arv_count=producer_arv_count,
consumer_arv_count=1,
)
return AsyncToUMMAPipeline(
raw_pipeline,
num_stages,
producer_state,
consumer_state,
)
def consumer_release(self):
"""Release consumer state."""
with cute.arch.elect_one():
super().consumer_release()
return self
class UMMAtoAsyncPipeline(GenericPipelineBase):
"""
Pipeline for UMMA to * (except TMA).
"""
@staticmethod
def create(
*,
num_stages: cute.Int32,
consumer: OperationTypeEnum,
consumer_arv_count: cute.Int32,
mma_operation_type: OperationTypeEnum,
cluster_layout_vmnk: Optional[cute.Layout] = None,
):
"""
Create a UMMA to * (except TMA) pipeline.
For 2SM MMA, provide cluster_layout_vmnk for proper mask computation.
"""
_validate_umma_operation_type(
mma_operation_type,
)
if consumer == OperationTypeEnum.SM90_TMA_LOAD:
raise ValueError("UMMA to TMA is not supported.")
if _is_2sm_umma_operation_type(mma_operation_type):
if cluster_layout_vmnk is None:
raise ValueError("cluster_layout_vmnk cannot be None if using 2SM MMA")
return UMMAtoAsyncPipeline.create_with_mask(
num_stages=num_stages,
consumer_type=consumer,
consumer_arv_count=consumer_arv_count,
mma_operation_type=mma_operation_type,
cluster_layout_vmnk=cluster_layout_vmnk,
)
else: # 1SM MMA
raw_pipeline, producer_state, consumer_state = create_pipeline(
num_stages,
mma_operation_type,
consumer,
producer_arv_count=1,
consumer_arv_count=consumer_arv_count,
)
return UMMAtoAsyncPipeline(
raw_pipeline,
num_stages,
producer_state,
consumer_state,
)
@staticmethod
def create_with_mask(
*,
num_stages: cute.Int32,
consumer_type: OperationTypeEnum,
consumer_arv_count: cute.Int32,
mma_operation_type: OperationTypeEnum,
cluster_layout_vmnk: cute.Layout,
):
"""
Create a UMMA to * pipeline with arrival mask for 2CTA operations.
"""
tmem_sync_mask = cutlass.pipeline.PipelineUmmaAsync._compute_tmem_sync_mask(
cta_layout_vmnk=cluster_layout_vmnk
)
raw_pipeline, producer_state, consumer_state = create_pipeline_with_mask(
num_stages,
mma_operation_type,
consumer_type,
producer_arv_count=1,
consumer_arv_count=consumer_arv_count,
arrival_mask=tmem_sync_mask,
)
return UMMAtoAsyncPipeline(
raw_pipeline,
num_stages,
producer_state,
consumer_state,
)
def producer_commit(self):
"""Commit producer state."""
with cute.arch.elect_one():
super().producer_commit()
return self
@dataclass
class TMAStorePipeline:
"""
TMA Store Pipeline modeling SMEM producer to TMA consumer pipeline.
A number of epilogue warps participate in the pipeline as producers, and one of them is designated as the consumer to perform TMA store.
Named barrier is used to synchronize all warps so that producers write SMEM after the pipeline stage is available, and the consumer waits for all producers before issuing TMA store.
The canonical pipeline flow is:
1. acquire_sync(): wait for pipeline stage availability + barrier
2. Each producer performs SMEM writes
3. commit_sync(): fence SMEM writes + barrier
4. Consumer performs TMA store
5. release_advance(): commit TMA store + advance stage
Args:
stages: Number of pipeline stages (type parameter)
arv_count: Number of threads participating in barriers
barrier_id: Barrier ID for synchronization
tma_warp_id: Which warp issues TMA stores (None = no TMA operations)
index: Initial stage index
"""
stages: cutlass.Constexpr[int]
arv_count: int
barrier_id: int
tma_warp_id: int
index: int = 0
def get_num_stages(self):
return self.stages
def acquire_sync(self):
"""
Acquire pipeline stage and synchronize all warps.
TMA warp waits for previous TMA operation to the same stage to complete (allowing writes to other stages to be in flight).
All warps then synchronize before producers write to SMEM.
"""
@CuteExperimentalDSL.jit
def acquire_sync_impl():
# Only TMA warp needs to wait for bulk async operations
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
# Use Python if with @Cutlass_LIR.jit preprocessor
if warp_idx == self.tma_warp_id:
# Allow N-1 TMA operations in flight for pipelining
# Now we can use the compile-time constant from type parameter
num_stages = self.get_num_stages()
wait_count = num_stages - 1 if num_stages > 1 else 0
cute.arch.cp_async_bulk_wait_group(wait_count, read=True)
# All warps must synchronize before producers write to SMEM
self._barrier()
return self
return acquire_sync_impl()
def commit_sync(self):
"""
Fence SMEM writes and synchronize all warps.
All warps fence their SMEM writes to make them visible to consumer
All warps then synchronize before TMA store operation.
"""
# All warps fence their SMEM writes for TMA visibility
cute.arch.fence_proxy("async.shared", space="cta")
# All warps synchronize before TMA store
self._barrier()
return self
def release_advance(self):
"""
Release current stage and advance to next stage.
TMA warp commits the TMA store operations to a bulk group.
All warps advance to the next pipeline stage.
"""
@CuteExperimentalDSL.jit
def release_advance_impl():
# Only TMA warp commits the TMA operations
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
# Use Python if with @Cutlass_LIR.jit preprocessor
if warp_idx == self.tma_warp_id:
cute.arch.cp_async_bulk_commit_group()
# All warps advance to next stage
self.index = (self.index + 1) % self.get_num_stages()
return self
return release_advance_impl()
def get_index(self):
"""Get current pipeline stage index."""
return self.index
def tail(self):
"""
Wait for all remaining TMA operations to complete.
Should be called at the end of the pipeline to ensure all TMA stores finish.
"""
@CuteExperimentalDSL.jit
def tail_impl():
warp_idx = cute.arch.warp_idx()
warp_idx = cute.arch.make_warp_uniform(warp_idx)
# Use Python if with @Cutlass_LIR.jit preprocessor
if warp_idx == self.tma_warp_id:
# Wait for all TMA operations to complete
cute.arch.cp_async_bulk_wait_group(0, read=True)
self._barrier()
return self
return tail_impl()
def _barrier(self):
"""Internal barrier synchronization."""
cute.arch.barrier(
barrier_id=self.barrier_id,
number_of_threads=self.arv_count,
)
python/CuTeDSL/cutlass/cute/experimental/utils.py
+162
View File
@@ -0,0 +1,162 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 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 cutlass import cute
def get_cta_v_map_ab(
gmem_tensor,
mma_tiler_mnk,
tiled_mma,
input_operand,
*,
loc=None,
ip=None,
):
"""
Build the **CTA-to-value map** (aka **CTA V-map**) layout for a TMA load of A/B
(and scale-factor variants SFA/SFB).
In practice, `cta_v_map` is a `cute.Layout` that tells TMA how this CTA’s
portion of a global tensor tile maps onto the values being transferred into
shared memory.
:param gmem_tensor: Global-memory tensor being loaded by TMA.
:type gmem_tensor: cute.Tensor
:param mma_tiler_mnk: The (M,N,K,...) tiler describing the CTA tile shape.
:type mma_tiler_mnk: tuple
:param tiled_mma: The tiled MMA object used to derive the per-operand thread/value mapping.
:type tiled_mma: cute.core.TiledMma
:param input_operand: One of {"A","B","SFA","SFB"} selecting which operand mapping to use.
:type input_operand: str
:returns: A layout suitable to pass as `cta_v_map=...` to `tma_load` / `tma_load_multicast`.
:rtype: cute.Layout
"""
ident = cute.core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
mode = 0 if (input_operand in ("A", "SFA")) else 1
mma_tiler_mk = (mma_tiler_mnk[mode], *mma_tiler_mnk[2:])
g_tile = cute.core.composition(ident, mma_tiler_mk, loc=loc, ip=ip)
if input_operand in ("A", "SFA"):
cta_v_map = tiled_mma._thrfrg_A(g_tile)
if input_operand in ("B", "SFB"):
cta_v_map = tiled_mma._thrfrg_B(g_tile)
cta_v_map = cute.core.get(cta_v_map, mode=[1])
cta_v_map = cute.core.dice(cta_v_map, (1, (1,) * cute.core.rank(g_tile)))
return cta_v_map
def get_cta_v_map_c(
gmem_tensor,
epi_tile,
*,
loc=None,
ip=None,
):
"""
Build the **CTA-to-value map** (aka **CTA V-map**) layout for a TMA store/load
of the output tensor C/D.
This returns an identity layout over the global tensor composed with the
epilogue tile, yielding a `cute.Layout` that describes which global indices
this CTA is responsible for.
:param gmem_tensor: Global-memory tensor being stored/loaded by TMA.
:type gmem_tensor: cute.Tensor
:param epi_tile: Epilogue tile layout describing the CTA's output tile shape.
:type epi_tile: cute.Layout
:returns: A layout suitable to pass as `cta_v_map=...` to `tma_store` / `tma_load`.
:rtype: cute.Layout
"""
ident = cute.core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
return cute.core.composition(ident, epi_tile, loc=loc, ip=ip)
def make_tmem_layout_acc(
tiled_mma,
mnk_tiler,
acc_stage,
*,
loc=None,
ip=None,
):
"""Return TMEM accumulator buffer layout for a tiled MMA.
This is a small helper around ``tiled_mma.make_fragment_C(...).layout`` to
keep example code fragment-free at the call site.
:param tiled_mma: The MMA tiler (``cute.TiledMma``).
:type tiled_mma: cute.TiledMma
:param mnk_tiler: Full MNK tiler; only the MN components are used for C.
:type mnk_tiler: tuple
:param acc_stage: Accumulator pipeline stages.
:param loc: Optional location for DSL ops.
:param ip: Optional insertion point for DSL ops.
:return: Layout for the accumulator TMEM buffer.
:rtype: cute.Layout
"""
acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2], loc=loc, ip=ip)
acc_shape_staged = cute.append(acc_shape, acc_stage, loc=loc, ip=ip)
return tiled_mma.make_fragment_C(acc_shape_staged, loc=loc, ip=ip).layout
def make_tmem_layout_a(
tiled_mma,
mk_tiler,
stage,
*,
loc=None,
ip=None,
):
"""Return TMEM A operand buffer layout for a tiled MMA.
:param tiled_mma: The MMA tiler (``cute.TiledMma``).
:type tiled_mma: cute.TiledMma
:param mk_tiler: MK tiler used to shape the A operand.
:type mk_tiler: tuple
:param stage: Pipeline stages for the A operand buffer.
:param loc: Optional location for DSL ops.
:param ip: Optional insertion point for DSL ops.
:return: Layout for the A operand TMEM buffer.
:rtype: cute.Layout
"""
a_shape = tiled_mma.partition_shape_A(mk_tiler, loc=loc, ip=ip)
a_shape_staged = cute.append(a_shape, stage, loc=loc, ip=ip)
return tiled_mma.make_fragment_A(a_shape_staged, loc=loc, ip=ip).layout
def make_t2r_rmem_layout(
tiled_copy_t2r,
gC_mnl_epi,
tidx,
*,
loc=None,
ip=None,
):
"""Return RMEM buffer layout for the T2R epilogue destination.
Computes the per-thread RMEM buffer layout produced by a TMEM->RMEM copy
for a single epilogue iteration.
:param tiled_copy_t2r: The TMEM->RMEM tiled copy op (``cute.TiledCopy``).
:type tiled_copy_t2r: cute.TiledCopy
:param gC_mnl_epi: Global C tensor partitioned by epilogue tile.
:type gC_mnl_epi: cute.Tensor
:param tidx: Thread index for the copy slice.
:param loc: Optional location for DSL ops.
:param ip: Optional insertion point for DSL ops.
:return: Layout for the RMEM buffer.
:rtype: cute.Layout
"""
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi, loc=loc, ip=ip)
return cute.make_fragment_like(
tTR_gC[(None, None, None, 0, 0)].layout, loc=loc, ip=ip
)
python/CuTeDSL/cutlass/cute/export/aot_config.py
-1
View File
@@ -169,4 +169,3 @@ Examples:
if __name__ == "__main__":
main()
python/CuTeDSL/cutlass/cute/math.py
-2
View File
@@ -16,8 +16,6 @@ from .tensor import TensorSSA
from cutlass._mlir.dialects import math, arith
from typing import Callable, Union
def _math_op(func: Callable, fastmath: bool, *args, **kwargs):
"""Dispatch the function to either a TensorSSA or a Numeric(Float).
python/CuTeDSL/cutlass/cute/nvgpu/common.py
+30 -3
View File
@@ -24,6 +24,7 @@ from ..typing import Float16, Float32, Float64, Numeric
__all__ = [
"OpError",
"normalize_field_to_ir_name",
"MmaUniversalOp",
"MmaUniversalTrait",
"CopyUniversalOp",
@@ -33,6 +34,33 @@ __all__ = [
"CacheEvictionPriority",
]
def normalize_field_to_ir_name(field, admissible_fields) -> str:
"""
Normalize a field specifier to its IR logical field name.
Accepted inputs:
- Enum value present in admissible_fields (must expose _to_ir_field_name()).
- Exact string IR name (e.g., "accum_c", "neg_a", "sf_a").
Any other form is rejected.
"""
# Enum path
if any(field is f for f in admissible_fields):
return field._to_ir_field_name()
# String path (must match exactly one of the IR names exposed by admissible_fields)
if isinstance(field, str):
allowed = {f._to_ir_field_name() for f in admissible_fields}
if field in allowed:
return field
# Otherwise, reject
allowed_pretty = [f._to_ir_field_name() for f in admissible_fields]
raise ValueError(
f"invalid field, must be one of {allowed_pretty} or their enum counterparts, but got {field}"
)
class OpError(DSLBaseError):
"""
An exception class for Op construction errors.
@@ -178,8 +206,8 @@ class CopyUniversalOp(atom.CopyOp):
op = cute.nvgpu.CopyUniversalOp()
atom = cute.make_copy_atom(
op,
tensor_dtype,
op,
tensor_dtype,
num_bits_per_copy=64,
l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_NORMAL
)
@@ -195,7 +223,6 @@ class CopyUniversalOp(atom.CopyOp):
- ``invariant`` is a kw argument specifying whether the load is invariant (read-only data \
that never changes). This enables compiler optimizations like instruction reordering. \
Defaults to ``False`` if not provided.
"""
def __str__(self) -> str:
python/CuTeDSL/cutlass/cute/nvgpu/cpasync/__init__.py
+1 -1
View File
@@ -24,6 +24,7 @@ __all__ = [
"CopyBulkTensorTileG2SMulticastOp",
"CopyBulkTensorTileS2GOp",
"CopyReduceBulkTensorTileS2GOp",
"CopyDsmemStoreOp",
#
# helpers.py
#
@@ -36,5 +37,4 @@ __all__ = [
"fence_tma_desc_acquire",
"cp_fence_tma_desc_release",
"fence_tma_desc_release",
"group_bulk_copy_modes",
]
python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py
+12 -8
View File
@@ -21,7 +21,7 @@ from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp
from cutlass._mlir import ir
from ...atom import CopyOp, Trait, make_atom
from ...typing import Int16, Int64, Pointer, Integer, Numeric
from ...typing import Int16, Int32, Int64, Pointer, Integer, Numeric
from ..common import OpError
from ..tcgen05.mma import CtaGroup
@@ -112,6 +112,7 @@ TMA_MBAR_PTR_FIELD_NAME = "tma_bar"
TMA_MCAST_MASK_FIELD_NAME = "mcast_mask"
TMA_DESC_PTR_FIELD_NAME = "tma_descriptor_ptr"
TMA_BYTE_MASK_FIELD_NAME = "byte_mask"
TMA_CTA_RANK_FIELD_NAME = "cta_rank"
TMA_CACHE_POLICY_FIELD_NAME = "cache_policy"
@@ -249,6 +250,7 @@ class CopyBulkTensorTileG2SNonExecTrait(Trait):
class CopyBulkTensorTileG2STrait(Trait):
pass
#
# TMA GMEM -> SMEM multicast copies
#
@@ -374,6 +376,7 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
)
return exec_value
class CopyBulkTensorTileG2SMulticastTrait(Trait):
pass
@@ -457,10 +460,6 @@ class CopyBulkTensorTileS2GTrait(Trait):
pass
class CopyBulkTensorTileS2GTrait(Trait):
pass
@dataclass
class CopyReduceBulkTensorTileS2GOp(TmaCopyOp):
"""
@@ -800,7 +799,7 @@ class CopyBulkS2GByteMaskOp(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
arch: Arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_100:
raise OpError(
self,
@@ -874,7 +873,7 @@ class CopyBulkS2SOp(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
arch: Arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
@@ -958,7 +957,7 @@ class CopyDsmemStoreOp(CopyOp):
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
arch: Arch = BaseDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
@@ -984,6 +983,11 @@ class CopyDsmemStoreOp(CopyOp):
"expects a 'num_bits_per_copy' kw argument of type int that is non-negative "
f"when creating a copy Atom for {self.__class__.__name__}"
)
if num_bits_per_copy not in [0, 32, 64, 128]:
raise ValueError(
"expects a 'num_bits_per_copy' kw argument that is one of {0, 32, 64, 128} "
f"when creating a copy Atom for {self.__class__.__name__}"
)
ty = _cute_nvgpu_ir.CopyAtomDsmemStoreType.get(
copy_internal_type.mlir_type, num_bits_per_copy
)
python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py
+33 -25
View File
@@ -10,7 +10,6 @@
# is strictly prohibited.
from typing import Optional, Tuple, Type, Union
from typing_extensions import deprecated
from cutlass.cutlass_dsl import dsl_user_op
@@ -47,11 +46,12 @@ TMAOp = Union[
CopyReduceBulkTensorTileS2GOp,
]
@dsl_user_op
def make_tiled_tma_atom(
op: TMAOp,
gmem_tensor: Tensor,
smem_layout: Union[Layout, ComposedLayout],
smem_layout_: Union[Layout, ComposedLayout],
cta_tiler: Tiler,
num_multicast: int = 1,
*,
@@ -84,7 +84,7 @@ def make_tiled_tma_atom(
:type op: TMAOp
: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
:param smem_layout: The SMEM layout to construct the Copy Atom, either w/ or w/o the stage mode
:type smem_layout: Union[Layout, ComposedLayout]
:param cta_tiler: The CTA Tiler to use
:type cta_tiler: Tiler
@@ -95,6 +95,26 @@ def make_tiled_tma_atom(
:return: A TMA Copy Atom associated with the TMA tensor
:rtype: Tuple[atom.CopyAtom, Tensor]
"""
smem_rank = core.rank(smem_layout_)
tiler_rank = core.rank(cta_tiler)
assert smem_rank == tiler_rank or smem_rank == tiler_rank + 1, (
f"smem_layout must be non-staged (rank(smem_layout) == rank(cta_tiler)) "
f"or staged (rank(smem_layout) == rank(cta_tiler) + 1)"
)
# Set the smem_layout on the operation for later retrieval
op.smem_layout = (
smem_layout_.value
if isinstance(smem_layout_, core._ComposedLayout)
else smem_layout_
)
# Slice the smem_layout if it is staged
if smem_rank == tiler_rank + 1:
smem_layout = core.select(smem_layout_, mode=list(range(tiler_rank)))
else:
smem_layout = smem_layout_
cta_v_map = core.composition(
core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip),
cta_tiler,
@@ -105,22 +125,21 @@ def make_tiled_tma_atom(
if isinstance(smem_layout, core._ComposedLayout):
smem_layout = smem_layout.value
# Set the smem_layout on the operation for later retrieval
op.smem_layout = (
smem_layout.value
if isinstance(smem_layout, core._ComposedLayout)
else smem_layout
)
tma_format = None
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}")
use_unpack = (internal_type.width == 8 and
isinstance(gmem_tensor.element_type, NumericMeta) and
gmem_tensor.element_type.width < 8)
internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type
use_unpack = (
internal_type.width == 8
and isinstance(gmem_tensor.element_type, NumericMeta)
and gmem_tensor.element_type.width < 8
)
internal_mlir_type = (
gmem_tensor.element_type.mlir_type
if use_unpack
else internal_type.mlir_type
)
tma_format = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack)
)
@@ -380,14 +399,3 @@ def fence_tma_desc_release(*, loc=None, ip=None) -> None:
loc=loc,
ip=ip,
)
@dsl_user_op
@deprecated("`group_bulk_copy_modes` is deprecated, use `group_modes` instead")
def group_bulk_copy_modes(src: Tensor, dst: Tensor, loc=None, ip=None) -> Tuple:
"""
Copy async bulk need group mode 0, acquiring whole tensor for bulk copy
"""
mSrc = core.group_modes(src, 0, core.rank(src))
mDst = core.group_modes(dst, 0, core.rank(dst))
return (mSrc, mDst)
python/CuTeDSL/cutlass/cute/nvgpu/helpers.py
+20 -23
View File
@@ -17,7 +17,6 @@ import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from .. import core, atom
from ..typing import Shape, Layout, ComposedLayout, Tensor, Numeric, NumericMeta
from ...impl_utils import check_type_in
from .cpasync.copy import (
CopyBulkTensorTileG2SOp,
CopyBulkTensorTileG2SNonExecTrait,
@@ -96,13 +95,6 @@ def make_tiled_tma_atom_A(
"""
check_type_in(
op,
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
"op",
"make_tiled_tma_atom_A",
)
# Set the smem_layout on the operation for later retrieval
op.smem_layout = (
smem_layout.value
@@ -136,10 +128,16 @@ def make_tiled_tma_atom_A(
if not isinstance(internal_type, NumericMeta):
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
use_unpack = (internal_type.width == 8 and
isinstance(gmem_tensor.element_type, NumericMeta) and
gmem_tensor.element_type.width < 8)
internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type
use_unpack = (
internal_type.width == 8
and isinstance(gmem_tensor.element_type, NumericMeta)
and gmem_tensor.element_type.width < 8
)
internal_mlir_type = (
gmem_tensor.element_type.mlir_type
if use_unpack
else internal_type.mlir_type
)
tma_format = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack)
)
@@ -224,13 +222,6 @@ def make_tiled_tma_atom_B(
"""
check_type_in(
op,
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
"op",
"make_tiled_tma_atom_B",
)
# Set the smem_layout on the operation for later retrieval
op.smem_layout = (
smem_layout.value
@@ -264,10 +255,16 @@ def make_tiled_tma_atom_B(
if not isinstance(internal_type, NumericMeta):
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
use_unpack = (internal_type.width == 8 and
isinstance(gmem_tensor.element_type, NumericMeta) and
gmem_tensor.element_type.width < 8)
internal_mlir_type = gmem_tensor.element_type.mlir_type if use_unpack else internal_type.mlir_type
use_unpack = (
internal_type.width == 8
and isinstance(gmem_tensor.element_type, NumericMeta)
and gmem_tensor.element_type.width < 8
)
internal_mlir_type = (
gmem_tensor.element_type.mlir_type
if use_unpack
else internal_type.mlir_type
)
tma_format = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack)
)
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/__init__.py
+2
View File
@@ -19,6 +19,7 @@ __all__ = [
# copy.py
#
"Repetition",
"TmemLoadRedOp",
"Pack",
"Unpack",
"Ld16x64bOp",
@@ -60,4 +61,5 @@ __all__ = [
"make_tmem_copy",
"make_s2t_copy",
"get_s2t_smem_desc_tensor",
"make_umma_smem_desc",
]
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py
+107
View File
@@ -26,6 +26,22 @@ from ...typing import Numeric
from .mma import CtaGroup
class TmemLoadRedOp(enum.Enum):
"""
An enumeration for the possible reduce operations for TMEM load operations.
"""
MAX = _cute_nvgpu_ir.TmemLoadRedOp.max
MAXABS = _cute_nvgpu_ir.TmemLoadRedOp.maxabs
MIN = _cute_nvgpu_ir.TmemLoadRedOp.min
MINABS = _cute_nvgpu_ir.TmemLoadRedOp.minabs
def __str__(self) -> str:
return f"{self.__class__.__name__}.{self.name}"
def __repr__(self) -> str:
return f"<{self.__class__.__name__}.{self.name}>"
class Repetition(enum.Enum):
"""
An enumeration for the number of repetitions of a given TMEM copy within the instruction.
@@ -390,6 +406,97 @@ class Ld32x32bTrait(Trait):
pass
@dataclass(frozen=True)
class LdRed16x32bx2Op(_LdBase):
"""
16x32bx2 TMEM load Reduce Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
This Operation corresponds to the ``.red`` and ``.16x32bx2`` qualifiers.
"""
redOp: TmemLoadRedOp = TmemLoadRedOp.MAX
nan: bool = False
half_split_off: int = 0
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "LdRed16x32bx2Trait":
"""
Create a trait object for the 16x32bx2 TMEM load Reduce operation.
:param copy_internal_type: The data type for the copy operation
:type copy_internal_type: Type[Numeric]
:param loc: MLIR location information for debugging, defaults to None
:type loc: optional
:param ip: MLIR insertion point for code generation, defaults to None
:type ip: optional
:param kwargs: Additional keyword arguments
:type kwargs: dict
:return: A trait object for this load operation
:rtype: LdRed16x32bx2Trait
"""
ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get(
copy_internal_type.mlir_type,
16,
32,
self.repeat.value,
self.redOp.value,
ir.UnitAttr.get() if self.nan else None,
ir.IntegerAttr.get(ir.IntegerType.get_signless(32), self.half_split_off),
)
return LdRed16x32bx2Trait(make_atom(ty, loc=loc, ip=ip))
class LdRed16x32bx2Trait(Trait):
pass
@dataclass(frozen=True)
class LdRed32x32bOp(_LdBase):
"""
32x32b TMEM load Reduce Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
This Operation corresponds to the ``red`` and ``.32x32`` qualifiers.
"""
redOp: TmemLoadRedOp = TmemLoadRedOp.MAX
nan: bool = False
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "LdRed32x32bTrait":
"""
Create a trait object for the 32x32b TMEM load Reduce operation.
:param copy_internal_type: The data type for the copy operation
:type copy_internal_type: Type[Numeric]
:param loc: MLIR location information for debugging, defaults to None
:type loc: optional
:param ip: MLIR insertion point for code generation, defaults to None
:type ip: optional
:param kwargs: Additional keyword arguments
:type kwargs: dict
:return: A trait object for this load operation
:rtype: LdRed32x32bTrait
"""
ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get(
copy_internal_type.mlir_type,
32,
32,
self.repeat.value,
self.redOp.value,
ir.UnitAttr.get() if self.nan else None,
None,
)
return LdRed32x32bTrait(make_atom(ty, loc=loc, ip=ip))
class LdRed32x32bTrait(Trait):
pass
@dataclass(frozen=True)
class _StBase(CopyOp):
"""
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py
+90 -12
View File
@@ -9,14 +9,16 @@
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
from typing import overload, Type, Tuple, Union
from typing import overload, Type, Tuple, Union, Optional
from cutlass.cutlass_dsl import dsl_user_op
from cutlass._mlir import ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir.dialects import nvvm
from cutlass._mlir.dialects import nvvm, builtin
from ...typing import (
Pointer,
Shape,
IntTuple,
Layout,
@@ -27,6 +29,7 @@ from ...typing import (
NumericMeta,
Int16,
Int32,
Int64,
)
from ... import core
from ...tensor import recast_tensor
@@ -102,17 +105,27 @@ def make_smem_layout_atom(
SmemLayoutAtomKind.MN_SW128_32B,
):
# M/N-major layout
outer = core.make_layout(
(num_contiguous_elems, 8), stride=(1, num_contiguous_elems), loc=loc, ip=ip
return core.make_composed_layout(
sw,
0,
core.make_layout(
(num_contiguous_elems, 8), stride=(1, num_contiguous_elems)
),
loc=loc,
ip=ip,
)
else:
# K-major layout
outer = core.make_layout(
(8, num_contiguous_elems), stride=(num_contiguous_elems, 1), loc=loc, ip=ip
return core.make_composed_layout(
sw,
0,
core.make_layout(
(8, num_contiguous_elems), stride=(num_contiguous_elems, 1)
),
loc=loc,
ip=ip,
)
return core.make_composed_layout(sw, 0, outer, loc=loc, ip=ip)
@overload
def tile_to_mma_shape(
@@ -190,14 +203,27 @@ def commit(
mbar_ptr = mbar_ptr.llvm_ptr
if mask is not None:
mask = Int16(mask).ir_value(loc=loc, ip=ip)
nvvm.tcgen05_commit_arrive(
mbar_ptr, multicast_mask=mask, group=group, loc=loc, ip=ip
)
nvvm.tcgen05_commit(mbar_ptr, multicast_mask=mask, group=group, loc=loc, ip=ip)
else:
nvvm.tcgen05_commit_arrive(mbar_ptr, group=group, loc=loc, ip=ip)
nvvm.tcgen05_commit(mbar_ptr, group=group, loc=loc, ip=ip)
return
@dsl_user_op
def int_to_smem_descriptor(i, *, loc=None, ip=None) -> ir.Value:
desc_type = _cute_nvgpu_ir.SmemDescType.get()
return builtin.unrealized_conversion_cast(
[desc_type], [Int64(i).ir_value(loc=loc, ip=ip)], loc=loc, ip=ip
)
@dsl_user_op
def smem_descriptor_to_int(desc: ir.Value, *, loc=None, ip=None) -> Int64:
return Int64(
builtin.unrealized_conversion_cast([Int64.mlir_type], [desc], loc=loc, ip=ip)
)
####################################################################################################
#
# Helper functions for Copies
@@ -324,3 +350,55 @@ def get_s2t_smem_desc_tensor(
atom._trait.value, smem_tensor.value, loc=loc, ip=ip
)
return smem_desc_tensor
def make_umma_smem_desc(
src: Pointer,
layout: Layout,
major: str,
next_src: Optional[Pointer] = None,
*,
loc=None,
ip=None,
):
"""
Construct shared memory descriptor for UMMA.
The `make_umma_smem_desc` operation accepts an input cute.ptr (optionally a nextSrc
pointer for the second buffer in a circular buffer scheme), alongside a cute.layout
and a major attr, then constructs the shared memory descriptor and returns it.
The layout must be describing the buffer pointed to by the input pointer and the
iterator must carry valid swizzle information.
There are 5 supported swizzle variants:
- S<0, 4, 3> | SWIZZLE_NONE
- S<1, 4, 3> | SWIZZLE_32B
- S<2, 4, 3> | SWIZZLE_64B
- S<3, 4, 3> | SWIZZLE_128B
- S<2, 5, 2> | SWIZZLE_128B_BASE32B
The cute.ptr must carry shared address space and must be aligned to 16B.
:param src: The source pointer to shared memory
:type src: Pointer
:param layout: The layout describing the buffer
:type layout: Layout
:param major: The major mode attribute
:type major: str
:param next_src: Optional next source pointer for circular buffer scheme
:type next_src: Optional[Pointer]
:return: The shared memory descriptor
:rtype: SmemDescType
"""
src = src.value
if next_src is not None:
next_src = next_src.value
return _cute_nvgpu_ir.make_umma_smem_desc(
src=src,
layout=layout.type.attribute,
major=major,
next_src=next_src,
loc=loc,
ip=ip,
)
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py
+247 -41
View File
@@ -20,7 +20,7 @@ import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
from ..common import OpError
from ..common import OpError, normalize_field_to_ir_name
from ... import core, atom
from ...core import _pack_shape, rank, depth
from ...typing import (
@@ -141,6 +141,7 @@ class Field(enum.Enum):
return self.value
# Base class for all tcgen05 MMA Ops with syntax `tcgen05.mma.cta_group.kind` used to factor out some internal code
@dataclass(frozen=True)
class MmaOp(Tcgen05MmaOp):
@@ -268,26 +269,30 @@ class MmaTraits(Trait):
admissible_fields = [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]
def set(self, field, value, *, loc=None, ip=None) -> None:
if field not in self.admissible_fields:
raise ValueError(
f"expects field to be one of {self.admissible_fields}, but got {field}"
field_ir = normalize_field_to_ir_name(field, self.admissible_fields)
bool_val = Boolean(value).ir_value(loc=loc, ip=ip)
try:
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, field_ir, bool_val, loc=loc, ip=ip
)
except (TypeError, AttributeError):
# Legacy fallback
attr = ir.Attribute.parse(f"#cute_nvgpu.atom_mma_field_sm100<{field_ir}>")
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, attr, bool_val, loc=loc, ip=ip
)
field_name = f"#cute_nvgpu.atom_mma_field_sm100<{field._to_ir_field_name()}>"
attr = ir.Attribute.parse(field_name)
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
)
def get(self, field, *, loc=None, ip=None) -> Any:
if field not in self.admissible_fields:
raise ValueError(
f"expects field to be one of {self.admissible_fields}, but got {field}"
field_ir = normalize_field_to_ir_name(field, self.admissible_fields)
try:
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, field_ir, loc=loc, ip=ip
)
except (TypeError, AttributeError):
attr = ir.Attribute.parse(f"#cute_nvgpu.atom_mma_field_sm100<{field_ir}>")
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
)
field_name = f"#cute_nvgpu.atom_mma_field_sm100<{field._to_ir_field_name()}>"
attr = ir.Attribute.parse(field_name)
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
)
# Base class for all tcgen05 BlockScaled MMA Ops with syntax `tcgen05.mma.cta_group.kind.block_scale` used to factor out some internal code
@@ -420,33 +425,58 @@ class BlockScaledMmaTraits(Trait):
]
def set(self, field, value, *, loc=None, ip=None) -> None:
if field not in self.admissible_fields:
raise ValueError(
f"expects field to be one of {self.admissible_fields}, but got {field}"
)
if field in [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]:
value = Boolean(value).ir_value(loc=loc, ip=ip)
elif field in [Field.SFA, Field.SFB]:
field_ir = normalize_field_to_ir_name(field, self.admissible_fields)
# Derive boolean/pointer IR names from enum values, no hard-coded strings.
bool_field_ir = {
f._to_ir_field_name()
for f in self.admissible_fields
if f in (Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B)
}
ptr_field_ir = {
f._to_ir_field_name()
for f in self.admissible_fields
if f in (Field.SFA, Field.SFB)
}
# Coerce value based on field kind
if field_ir in bool_field_ir:
val = Boolean(value).ir_value(loc=loc, ip=ip)
elif field_ir in ptr_field_ir:
if not isinstance(value, Pointer):
raise ValueError(
f"expects value to be a pointer for {field}, but got {type(value).__name__}"
f"expects value to be a pointer for {field_ir}, but got {type(value).__name__}"
)
value = value.value
field_name = f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field._to_ir_field_name()}>"
attr = ir.Attribute.parse(field_name)
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, attr, value, loc=loc, ip=ip
)
val = value.value
else:
raise ValueError(f"unsupported field: {field_ir}")
try:
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, field_ir, val, loc=loc, ip=ip
)
except (TypeError, AttributeError):
attr = ir.Attribute.parse(
f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field_ir}>"
)
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, attr, val, loc=loc, ip=ip
)
def get(self, field, *, loc=None, ip=None) -> Any:
if field not in [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]:
raise ValueError(f"the get method for {field} is not supported")
field_name = f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field._to_ir_field_name()}>"
attr = ir.Attribute.parse(field_name)
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
)
# Only boolean-returning fields supported for get. Derive from admissible_fields.
gettable_fields = [
f for f in self.admissible_fields if f not in (Field.SFA, Field.SFB)
]
field_ir = normalize_field_to_ir_name(field, gettable_fields)
try:
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, field_ir, loc=loc, ip=ip
)
except (TypeError, AttributeError):
attr = ir.Attribute.parse(
f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field_ir}>"
)
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
)
#
@@ -802,6 +832,7 @@ class MmaFP8Trait(MmaTraits):
pass
#
# MXF8F6F4 MMA
#
@@ -946,7 +977,7 @@ class MmaMXF4Op(BlockScaledMmaOp):
f"but got {self.shape_mnk[2]}",
)
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF8Trait":
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4Trait":
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get(
shape_mnk.type.attribute,
@@ -1039,7 +1070,7 @@ class MmaMXF4NVF4Op(BlockScaledMmaOp):
f"but got {self.shape_mnk[2]}",
)
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF8Trait":
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4NVF4Trait":
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get(
shape_mnk.type.attribute,
@@ -1077,6 +1108,181 @@ class MmaMXF4NVF4Trait(BlockScaledMmaTraits):
pass
#
# SM103 MXF4 MMA
#
@dataclass(frozen=True)
class SM103MmaMXF4Op(BlockScaledMmaOp):
"""
SM103 MXF4 tcgen05 BlockScaled MMA Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma>`__.
This Operation corresponds to the ``.kind::mxf4`` qualifier.
This Operation is for SM103.
"""
descriptive_name = "tcgen05 SM103 MXF4 BlockScaled MMA Operation"
def __init__(
self,
instruction_shape: Shape,
cta_group: CtaGroup,
a_src: OperandSource,
) -> None:
super().__init__(
Float4E2M1FN,
Float4E2M1FN,
Float32,
Float8E8M0FNU,
32,
instruction_shape,
cta_group,
a_src,
OperandMajorMode.K,
OperandMajorMode.K,
)
self._verify()
def _verify(self) -> None:
# Instruction shape verification
instruction_k = 96
if rank(self.shape_mnk) == 2:
object.__setattr__(self, "shape_mnk", (*self.shape_mnk, instruction_k))
if self.shape_mnk[2] != instruction_k:
raise OpError(
self,
f"expects the instruction extent in the K-mode to be {instruction_k}, "
f"but got {self.shape_mnk[2]}",
)
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4Trait":
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get(
shape_mnk.type.attribute,
self.cta_group.value,
self.a_major_mode._to_ir(),
self.b_major_mode._to_ir(),
self.a_dtype.mlir_type,
self.b_dtype.mlir_type,
self.acc_dtype.mlir_type,
self.sf_dtype.mlir_type,
self.a_src._to_ir(),
self.sf_vec_size,
1030,
)
return MmaMXF4Trait(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
Boolean(False).ir_value(loc=loc, ip=ip),
Boolean(False).ir_value(loc=loc, ip=ip),
core.make_ptr(
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
).value,
core.make_ptr(
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
).value,
),
loc=loc,
ip=ip,
)
)
#
# SM103 MXF4NVF4 MMA
#
@dataclass(frozen=True)
class SM103MmaMXF4NVF4Op(BlockScaledMmaOp):
"""
SM103 MXF4NVF4 tcgen05 BlockScaled MMA Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma>`__.
This Operation corresponds to the ``.kind::mxf4nvf4`` qualifier.
This Operation is for SM103.
"""
descriptive_name = "tcgen05 SM103 MXF4NVF4 BlockScaled MMA Operation"
def __init__(
self,
sf_dtype: Type[Numeric],
instruction_shape: Shape,
cta_group: CtaGroup,
a_src: OperandSource,
) -> None:
super().__init__(
Float4E2M1FN,
Float4E2M1FN,
Float32,
sf_dtype,
16,
instruction_shape,
cta_group,
a_src,
OperandMajorMode.K,
OperandMajorMode.K,
)
self._verify()
def _verify(self) -> None:
# Scale Factor data type verification
if self.sf_dtype not in [Float8E8M0FNU, Float8E4M3FN]:
raise OpError(
self,
"expects the 'sf_dtype' Op parameter to be one of Float8E8M0FNU",
)
# Instruction shape verification
instruction_k = 96
if rank(self.shape_mnk) == 2:
object.__setattr__(self, "shape_mnk", (*self.shape_mnk, instruction_k))
if self.shape_mnk[2] != instruction_k:
raise OpError(
self,
f"expects the instruction extent in the K-mode to be {instruction_k}, "
f"but got {self.shape_mnk[2]}",
)
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaMXF4NVF4Trait":
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
ty = _cute_nvgpu_ir.MmaAtomSM100UMMABlockScaledType.get(
shape_mnk.type.attribute,
self.cta_group.value,
self.a_major_mode._to_ir(),
self.b_major_mode._to_ir(),
self.a_dtype.mlir_type,
self.b_dtype.mlir_type,
self.acc_dtype.mlir_type,
self.sf_dtype.mlir_type,
self.a_src._to_ir(),
self.sf_vec_size,
1030,
)
return MmaMXF4NVF4Trait(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
Boolean(False).ir_value(loc=loc, ip=ip),
Boolean(False).ir_value(loc=loc, ip=ip),
core.make_ptr(
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
).value,
core.make_ptr(
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
).value,
),
loc=loc,
ip=ip,
)
)
####################################################################################################
#
# SMEM layout atoms
python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py
+15 -7
View File
@@ -82,6 +82,7 @@ class LdMatrix8x8x16bOp(BaseOp):
class LdMatrix8x8x16bTrait(Trait):
pass
@dataclass(frozen=True)
class LdMatrix8x16x8bOp(BaseOp):
"""
@@ -102,15 +103,20 @@ class LdMatrix8x16x8bOp(BaseOp):
self,
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
)
if self.unpack_bits not in [4, 6]:
raise OpError(self, "Op unpack bits must be 4 or 6")
if self.unpack_bits not in [None, 4, 6]:
raise OpError(self, "Op unpack bits must be 4 or 6 or None")
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "LdMatrix8x16x8bTrait":
mode = _pack_shape((8, 16), loc=loc, ip=ip)
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8
if self.unpack_bits == 6:
# LdMatrix8x16x8b without unpacking doesn't exist
# but is equivalent to LdMatrix8x8x16b
mode_n = 8 if self.unpack_bits is None else 16
mode = _pack_shape((8, mode_n), loc=loc, ip=ip)
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u16
if self.unpack_bits == 4:
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8
elif self.unpack_bits == 6:
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8
ty = _cute_nvgpu_ir.CopyAtomLdsmType.get(
copy_internal_type.mlir_type,
@@ -125,11 +131,12 @@ class LdMatrix8x16x8bOp(BaseOp):
class LdMatrix8x16x8bTrait(Trait):
pass
@dataclass(frozen=True)
class LdMatrix16x8x8bOp(BaseOp):
"""
16x8 8b ``ldmatrix`` Operation with transpose
There is no direct PTX correspondance to this Op.
This actually lowers to ldmatrix with the ``.m16n16`` qualifier and
additional address and value permutations to match stmatrix.m16n8.trans.
@@ -166,6 +173,7 @@ class LdMatrix16x8x8bOp(BaseOp):
)
return LdMatrix16x8x8bTrait(make_atom(ty, loc=loc, ip=ip))
class LdMatrix16x8x8bTrait(Trait):
pass
@@ -176,7 +184,7 @@ class LdMatrix16x16x8bOp(BaseOp):
16x16 ``ldmatrix`` Operation with transpose and optional unpacking to 8b container.
Packed source container is 16x4b elements with 64b padding
or 16x6b elements with 32b padding (total 128b per 16 elements)
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-load-instruction-ldmatrix>`__.
This operation corresponds to the ``.m16n16`` and the ``.b4x16_p64``,``.b6x16_p32``,``.b8`` qualifiers.
"""
python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py
+4 -2
View File
@@ -15,7 +15,7 @@ from typing import Type, Any
import enum
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import CuTeDSL
from cutlass.cutlass_dsl import BaseDSL
from ..common import OpError
@@ -134,7 +134,7 @@ class MmaSM120BlockScaledOp(MmaOp):
def __post_init__(self) -> None:
# Verify arch
arch = CuTeDSL._get_dsl().get_arch_enum()
arch = BaseDSL._get_dsl().get_arch_enum()
if not arch == Arch.sm_120a:
raise OpError(
self,
@@ -174,6 +174,7 @@ class MmaSM120BlockScaledOp(MmaOp):
self,
"expects the 'sf_vec_size' Op parameter to be 16 or 32",
)
def __str__(self) -> str:
return (
"warp-level MXF4/MXF4NVF4 MMA Operation"
@@ -190,6 +191,7 @@ class MmaSM120BlockScaledOp(MmaOp):
def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None):
pass
class Field(enum.Enum):
"""
An enumeration for the fields of the MMA Atom that can be modified at runtime.
python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py
+35 -17
View File
@@ -15,12 +15,13 @@ from typing import Type, Any
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL, T
from typing_extensions import deprecated
import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
from ..common import OpError
from ..common import OpError, normalize_field_to_ir_name
from ...core import _pack_shape, rank, depth
from ...typing import (
Shape,
@@ -208,27 +209,44 @@ class MmaOp(WarpGroupMmaOp):
class MmaTraits(Trait):
admissible_fields = [Field.ACCUMULATE]
def _normalize_field_name(self, field: Any) -> str:
"""
Normalize a field specifier (enum or string) into the IR logical field name.
Accepted inputs:
- Field.ACCUMULATE
- "accum_c"
"""
return normalize_field_to_ir_name(field, self.admissible_fields)
def set(self, field, value, *, loc=None, ip=None) -> None:
if field not in self.admissible_fields:
raise ValueError(
f"invalid field, must be {Field.ACCUMULATE}, but got {field}"
field_ir_name = self._normalize_field_name(field)
# Prefer the newer builder that accepts a logical field name, but keep
# a fallback for legacy attribute-based construction to avoid breaking changes.
bool_val = Boolean(value).ir_value(loc=loc, ip=ip)
try:
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, field_ir_name, bool_val, loc=loc, ip=ip
)
except (TypeError, AttributeError):
# Legacy path: construct the per-arch field attribute explicitly
attr_asm = f"#cute_nvgpu.atom_mma_field_sm90<{field_ir_name}>"
attr = ir.Attribute.parse(attr_asm)
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, attr, bool_val, loc=loc, ip=ip
)
field_name = f"#cute_nvgpu.atom_mma_field_sm90<{field._to_ir_field_name()}>"
attr = ir.Attribute.parse(field_name)
self.value = _cute_nvgpu_ir.atom_set_value(
self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
)
def get(self, field, *, loc=None, ip=None) -> Any:
if field not in self.admissible_fields:
raise ValueError(
f"invalid field, must be {Field.ACCUMULATE}, but got {field}"
field_ir_name = self._normalize_field_name(field)
try:
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, field_ir_name, loc=loc, ip=ip
)
except (TypeError, AttributeError):
attr_asm = f"#cute_nvgpu.atom_mma_field_sm90<{field_ir_name}>"
attr = ir.Attribute.parse(attr_asm)
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
)
field_name = f"#cute_nvgpu.atom_mma_field_sm90<{field._to_ir_field_name()}>"
attr = ir.Attribute.parse(field_name)
return _cute_nvgpu_ir.atom_get_value(
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
)
@dataclass(frozen=True)
python/CuTeDSL/cutlass/cute/runtime.py
+1 -1
View File
@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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
python/CuTeDSL/cutlass/cute/tensor.py
+50 -53
View File
@@ -20,6 +20,7 @@ from cutlass.cutlass_dsl import (
T,
cutlass_arith,
_binary_op_type_promote,
MLIR_DYNAMIC,
BaseDSL,
)
from cutlass._mlir import ir
@@ -75,35 +76,8 @@ from .core import (
recast_layout,
)
from .typing import (
IntTuple,
Coord,
Shape,
Stride,
Pointer,
Layout,
ComposedLayout,
Tensor,
AddressSpace,
is_integer,
is_int_tuple,
as_numeric,
)
from .typing import (
Numeric,
Integer,
Boolean,
Int4,
Uint8,
Int8,
Int32,
Float4E2M1FN,
Float16,
Float32,
BFloat16,
)
from .tuple import transform_leaf, product, product_like, flatten_to_tuple
from .arch import cvt_i8_bf16_intrinsic, cvt_i4_bf16_intrinsic, cvt_f4e2m1_f16_intrinsic
from .arch import cvt_i8_bf16_intrinsic, cvt_i4_bf16_intrinsic
__all__ = [
@@ -439,10 +413,9 @@ class _Tensor(Tensor):
return _cute_ir.get_layout(self.value, loc=loc, ip=ip)
@property
@dsl_user_op
@lru_cache_ir()
def shape(self, *, loc=None, ip=None) -> Shape:
return self.layout.shape_method(loc=loc, ip=ip)
def shape(self) -> Shape:
return self.layout.shape
@property
@lru_cache_ir()
@@ -480,12 +453,23 @@ class _Tensor(Tensor):
raise ValueError(f"{self} doesn't have memspace")
@dsl_user_op
def load(self, *, loc=None, ip=None) -> "TensorSSA":
def load(
self,
*,
mask: Optional["TensorSSA"] = None,
pass_thru: Optional["TensorSSA"] = None,
loc=None,
ip=None,
) -> "TensorSSA":
"""Load tensor elements as a vector.
Loads all elements of the tensor into a vector representation, assuming the tensor
has a static shape and is in a memory space that supports load operations.
:param mask: Mask vector, defaults to None
:type mask: Optional[TensorSSA]
:param pass_thru: Pass through vector, defaults to None
:type pass_thru: Optional[TensorSSA]
:param loc: Source location for MLIR operation tracking, defaults to None
:type loc: Optional[Location]
:param ip: Insertion point for MLIR operation, defaults to None
@@ -501,9 +485,15 @@ class _Tensor(Tensor):
if not is_static(self.shape):
raise ValueError("dynamic layout doesn't support load")
self._check_can_load_store()
self._check_can_load_store(vectorized=True)
res_vect = _cute_ir.memref_load_vec(self.value, loc=loc, ip=ip)
mask_val = None if mask is None else mask.ir_value(loc=loc, ip=ip)
pass_thru_val = (
None if pass_thru is None else self._cvt_to_dest(pass_thru, loc=loc, ip=ip)
)
res_vect = _cute_ir.memref_load_vec(
self.value, mask=mask_val, pass_thru=pass_thru_val, loc=loc, ip=ip
)
if self.element_type is Boolean:
assert res_vect.type.element_type == T.i8(), (
f"Boolean tensor must be stored as i8 in memory, but got {res_vect.type.element_type}"
@@ -515,7 +505,14 @@ class _Tensor(Tensor):
return TensorSSA(res_vect, self.shape, self.element_type)
@dsl_user_op
def store(self, data: "TensorSSA", *, loc=None, ip=None):
def store(
self,
data: "TensorSSA",
*,
mask: Optional["TensorSSA"] = None,
loc=None,
ip=None,
):
"""Store vector data into tensor.
Stores vector data into the tensor, assuming matching shapes and a memory space
@@ -523,6 +520,8 @@ class _Tensor(Tensor):
:param data: Vector data to store into tensor
:type data: TensorSSA
:param mask: Mask vector, defaults to None
:type mask: Optional[TensorSSA]
:param loc: Source location for MLIR operation tracking, defaults to None
:type loc: Optional[Location]
:param ip: Insertion point for MLIR operation, defaults to None
@@ -538,7 +537,7 @@ class _Tensor(Tensor):
if not is_static(self.shape):
raise ValueError("Dynamic layout doesn't support vectorized store")
self._check_can_load_store()
self._check_can_load_store(vectorized=True)
n_elems = size(self.shape, loc=loc, ip=ip)
if n_elems != size(data.shape, loc=loc, ip=ip):
@@ -556,7 +555,11 @@ class _Tensor(Tensor):
# Implicit upcast to wider type
new_data = self._cvt_to_dest(data, loc=loc, ip=ip)
return _cute_ir.memref_store_vec(new_data, self.value, loc=loc, ip=ip)
mask_val = None if mask is None else mask.ir_value(loc=loc, ip=ip)
return _cute_ir.memref_store_vec(
new_data, self.value, mask=mask_val, loc=loc, ip=ip
)
@dsl_user_op
def fill(self, value: Numeric, *, loc=None, ip=None) -> None:
@@ -585,7 +588,7 @@ class _Tensor(Tensor):
# Fill tensor with constant value
tensor.fill(0.5) # All elements become 0.5
"""
self._check_can_load_store()
self._check_can_load_store(vectorized=True)
sz = size(self, loc=loc, ip=ip)
if type(sz) is not int:
@@ -599,7 +602,7 @@ class _Tensor(Tensor):
)
self.store(vect_val, loc=loc, ip=ip)
def _check_can_load_store(self):
def _check_can_load_store(self, vectorized: bool = False):
if not isinstance(self.type, _cute_ir.MemRefType) or self.memspace not in (
AddressSpace.rmem,
AddressSpace.smem,
@@ -608,9 +611,9 @@ class _Tensor(Tensor):
):
raise ValueError(f"{self} doesn't support load and store")
if self.type.is_swizzled:
if vectorized and isinstance(self.layout, ComposedLayout):
raise NotImplementedError(
f"load & store swizzled memory is not supported yet: {self}"
"vectorized load/store on tensor with composed layout is not supported yet"
)
def _check_can_dereference(self):
@@ -1038,8 +1041,10 @@ def print_tensor(
signed = tensor.element_type.signed
else:
signed = False
else:
elif isinstance(tensor.type, _cute_ir.CoordTensorType):
signed = True
else:
raise ValueError(f"unsupported tensor type for print_tensor, got {tensor.type}")
_cute_ir.print_view(tensor.value, verbose=verbose, is_signed=signed, loc=loc, ip=ip)
@@ -1750,7 +1755,8 @@ class TensorSSA(cutlass_arith.ArithValue):
idx = crd2idx(crd, self._layout, loc=loc, ip=ip)
assert not isinstance(idx, tuple), "index must be scalar"
idx_val = as_numeric(idx).ir_value(loc=loc, ip=ip)
res_val = vector.extractelement(self, position=idx_val, loc=loc, ip=ip)
idx_val = arith.index_cast(T.index(), idx_val, loc=loc, ip=ip)
res_val = vector.extract(self, [idx_val], [MLIR_DYNAMIC], loc=loc, ip=ip)
return self.dtype(res_val)
if not is_static(crd):
@@ -1817,16 +1823,7 @@ class TensorSSA(cutlass_arith.ArithValue):
# maybe downcast can lose signedness
src = self.maybe_downcast().with_signedness(self.signed)
if src_dtype.is_float and dtype.is_float:
if src_dtype == Float4E2M1FN and dtype in (Float16, Float32):
res_vect = cvt_f4e2m1_f16_intrinsic(
src, size(self.shape), loc=loc, ip=ip
)
if dtype == Float32:
res_vect = cutlass_arith.cvtf(
res_vect, dtype.mlir_type, loc=loc, ip=ip
)
else:
res_vect = cutlass_arith.cvtf(src, dtype.mlir_type, loc=loc, ip=ip)
res_vect = cutlass_arith.cvtf(src, dtype.mlir_type, loc=loc, ip=ip)
elif src_dtype.is_float and issubclass(dtype, Integer):
res_vect = cutlass_arith.fptoi(
src, dtype.signed, dtype.mlir_type, loc=loc, ip=ip
python/CuTeDSL/cutlass/cute/testing.py
+6 -6
View File
@@ -20,15 +20,12 @@ from typing import Type, Union, Callable, Optional, Dict, List, Any
import cuda.bindings.driver as cuda_driver
import cuda.bindings.runtime as cuda_runtime
import cutlass
import cutlass.base_dsl.jit_executor
import cutlass.cutlass_dsl.cuda_jit_executor
from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op, const_expr
from .typing import Numeric, Int8, Boolean, Tensor, Layout, Shape
from . import nvgpu
from .core import recast_layout, make_layout, composition, get, rank, size, zipped_divide
from .core import recast_layout, make_layout, composition, get, rank, size
from .tuple import elem_less
from .tensor import (
make_rmem_tensor,
@@ -39,6 +36,7 @@ from .tensor import (
)
from .atom import make_copy_atom
from .algorithm import copy
from .core import zipped_divide
from .runtime import from_dlpack
from cutlass._mlir.dialects import builtin, cf, nvvm, vector
@@ -76,7 +74,7 @@ class _CompileTimeAssertion(Assertion):
def __init__(
self,
tensor: _Tensor,
tensor: Tensor,
num_assertions: int = 1,
msgs=None,
device=None,
@@ -849,7 +847,9 @@ def get_workspace_count(
:return: Number of workspaces needed
:rtype: int
"""
num_l2_cache_bytes = cutlass.utils.HardwareInfo().get_l2_cache_size_in_bytes()
from cutlass.utils import HardwareInfo
num_l2_cache_bytes = HardwareInfo().get_l2_cache_size_in_bytes()
num_workspaces = (num_l2_cache_bytes * 3) // one_workspace_bytes + 1
num_iters = warmup_iterations + iterations
return num_iters if num_iters < num_workspaces else num_workspaces
python/CuTeDSL/cutlass/cute/typing.py
+5 -3
View File
@@ -12,7 +12,6 @@
from abc import ABC, abstractmethod
import ctypes
from typing import ForwardRef, Tuple, Union, Any, Type, List, Optional, Literal
from functools import lru_cache
from cutlass.base_dsl.typing import *
@@ -28,9 +27,13 @@ class SymInt:
def __init__(self, width: Literal[32, 64] = 32, *, divisibility=1):
if width not in [32, 64]:
raise ValueError(f"Unsupported width: {width}")
self._width = width
self._divisibility = divisibility
def __hash__(self):
return hash((self._width, self._divisibility))
@property
def width(self):
return self._width
@@ -80,6 +83,7 @@ class SymInt:
else:
assert False, f"Unsupported width: {self.width}"
return self
def sym_int(width: Literal[32, 64] = 32, *, divisibility=1) -> SymInt:
return SymInt(width, divisibility=divisibility)
@@ -403,6 +407,4 @@ __all__ = [
"XTuple",
"is_integer",
"is_int_tuple",
"Pointer",
"Tensor",
]