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

v4.4 update. (#2979)

Browse Source
This commit is contained in:
Junkai-Wu
2026-01-24 11:46:17 -05:00
committed by GitHub
parent 2fafefb7b9
commit 9fba3195f9
293 changed files with 46343 additions and 2995 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/CuTeDSL/cutlass/cute/__init__.py
+17 -1
View File
@@ -57,7 +57,6 @@ from .core import (
make_composed_layout,
make_layout_tv,
make_swizzle,
make_sparse_elem,
recast_ptr,
get,
select,
@@ -99,6 +98,8 @@ from .core import (
local_partition,
local_tile,
printf,
get_nonswizzle_portion,
get_swizzle_portion,
# Wrapper classes
Swizzle,
E,
@@ -118,6 +119,8 @@ from .core import (
# FastDivmod operations
FastDivmodDivisor,
fast_divmod_create_divisor,
basis_value,
basis_get,
)
from .tuple import (
@@ -130,6 +133,9 @@ from .tuple import (
product_like,
product_each,
elem_less,
tuple_cat,
transform_apply,
filter_tuple,
)
from .tensor import (
TensorSSA,
@@ -178,6 +184,8 @@ from .atom import (
)
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
@@ -215,6 +223,7 @@ _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__,
"AddressSpace",
"CacheEvictionPriority",
"Tensor",
@@ -252,6 +261,8 @@ __all__ = [
"make_composed_layout",
"make_layout_tv",
"make_layout_image_mask",
"get_nonswizzle_portion",
"get_swizzle_portion",
# Tensor functions
"make_ptr",
"make_tensor",
@@ -271,6 +282,8 @@ __all__ = [
"find",
"find_if",
"transform_leaf",
"basis_value",
"basis_get",
"coalesce",
"group_modes",
"cosize",
@@ -287,6 +300,9 @@ __all__ = [
"prepend_ones",
"append_ones",
"elem_less",
"tuple_cat",
"transform_apply",
"filter_tuple",
# Math operations
"ceil_div",
"round_up",
python/CuTeDSL/cutlass/cute/_tvm_ffi_args_spec_converter.py
+2
View File
@@ -353,6 +353,8 @@ def _convert_single_arg(
elem_param = _convert_single_arg(elem, elem_name, None, ctx)
tuple_params.append(elem_param)
return spec.TupleParam(arg_name, tuple_params)
elif isinstance(arg, bool):
return spec.Var(arg_name, NumericToTVMFFIDtype[Boolean])
elif isinstance(arg, int):
# in cute.compile, unannotated const int is converted to int32
return spec.Var(arg_name, NumericToTVMFFIDtype[Int32])
python/CuTeDSL/cutlass/cute/arch/__init__.py
+20
View File
@@ -16,6 +16,7 @@ from .nvvm_wrappers import *
from .smem import *
from .tmem import *
from .numeric_conversion import *
from .clc import *
# __all__ is required here for documentation generation
__all__ = [
@@ -73,12 +74,24 @@ __all__ = [
"vote_any_sync",
"vote_all_sync",
"vote_uni_sync",
"atomic_add",
"atomic_and",
"atomic_or",
"atomic_xor",
"atomic_max",
"atomic_min",
"atomic_exch",
"atomic_cas",
"store",
"load",
"popc",
"fence_proxy",
"fence_view_async_tmem_load",
"fence_view_async_tmem_store",
"warpgroup_reg_alloc",
"warpgroup_reg_dealloc",
"setmaxregister_increase",
"setmaxregister_decrease",
"fma_packed_f32x2",
"mul_packed_f32x2",
"add_packed_f32x2",
@@ -100,6 +113,8 @@ __all__ = [
#
# tmem.py
#
"get_max_tmem_alloc_cols",
"get_min_tmem_alloc_cols",
"retrieve_tmem_ptr",
"alloc_tmem",
"relinquish_tmem_alloc_permit",
@@ -115,4 +130,9 @@ __all__ = [
"cvt_i8x2_to_f32x2",
"cvt_i8_bf16",
"cvt_f32x2_bf16x2",
#
# clc.py
#
"issue_clc_query",
"clc_response",
]
python/CuTeDSL/cutlass/cute/arch/clc.py
+116
View File
@@ -0,0 +1,116 @@
# 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 Tuple
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass._mlir.dialects import nvvm, vector
from ..typing import Int32, Pointer, Int128
@dsl_user_op
def issue_clc_query(
mbar_ptr: Pointer,
clc_response_ptr: Pointer,
loc=None,
ip=None,
) -> None:
"""
The clusterlaunchcontrol.try_cancel instruction requests atomically cancelling the launch
of a cluster that has not started running yet. It asynchronously writes an opaque response
to shared memory indicating whether the operation succeeded or failed. On success, the
opaque response contains the ctaid of the first CTA of the canceled cluster.
:param mbar_ptr: A pointer to the mbarrier address in SMEM
:type mbar_ptr: Pointer
:param clc_response_ptr: A pointer to the cluster launch control response address in SMEM
:type clc_response_ptr: Pointer
"""
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,
)
@dsl_user_op
def clc_response(
result_addr: Pointer, loc=None, ip=None
) -> Tuple[Int32, Int32, Int32, Int32]:
"""
After loading response from clusterlaunchcontrol.try_cancel instruction into 16-byte
register, it can be further queried using clusterlaunchcontrol.query_cancel instruction.
If the cluster is canceled successfully, predicate p is set to true; otherwise, it is
set to false. If the request succeeded, clusterlaunchcontrol.query_cancel.get_first_ctaid
extracts the CTA id of the first CTA in the canceled cluster. By default, the instruction
returns a .v4 vector whose first three elements are the x, y and z coordinate of first CTA
in canceled cluster.
:param result_addr: A pointer to the cluster launch control response address in SMEM
:type result_addr: Pointer
"""
from cutlass.cute import recast_ptr, make_tensor, make_layout
clc_ptr_i128 = recast_ptr(result_addr, dtype=Int128, loc=loc, ip=ip)
clc_tensor = make_tensor(
clc_ptr_i128, make_layout(1, loc=loc, ip=ip), loc=loc, ip=ip
)
# Load the 128-bit value from shared memory
clc_result_vec = clc_tensor.load(loc=loc, ip=ip)
# Extract the i128 scalar from the vector<1xi128>
clc_result_i128 = vector.extract(
clc_result_vec.ir_value(loc=loc, ip=ip),
[],
[0],
)
# Query if the cluster was canceled
pred = nvvm.clusterlaunchcontrol_query_cancel_is_canceled(
T.bool(),
clc_result_i128,
loc=loc,
ip=ip,
)
is_valid = Int32(pred)
# 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,
)
# 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,
)
# 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,
)
m_idx = Int32(m_idx_i32)
n_idx = Int32(n_idx_i32)
l_idx = Int32(l_idx_i32)
return m_idx, n_idx, l_idx, is_valid
python/CuTeDSL/cutlass/cute/arch/elect.py
+2
View File
@@ -69,6 +69,8 @@ def elect_one(*, loc=None, ip=None) -> IfOpRegion:
# Only one thread in the warp executes the code in this context
pass
"""
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())
if_op = scf.IfOp(is_thread_leader, loc=loc, ip=ip)
python/CuTeDSL/cutlass/cute/arch/mbar.py
+11 -13
View File
@@ -35,7 +35,10 @@ def mbarrier_init(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None:
:type cnt: Int
"""
nvvm.mbarrier_init_shared(
mbar_ptr.llvm_ptr, Int32(cnt).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
Int32(cnt).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
@@ -65,7 +68,7 @@ def mbarrier_arrive_and_expect_tx(
"""
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
if peer_cta_rank_in_cluster is not None:
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
mbar_llvm_ptr.type,
@@ -105,7 +108,7 @@ def mbarrier_expect_tx(
"""
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
if peer_cta_rank_in_cluster is not None:
mbar_llvm_ptr = nvvm.mapa(
mbar_llvm_ptr.type,
@@ -144,7 +147,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.llvm_ptr,
mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
Int32(phase).ir_value(loc=loc, ip=ip),
Int32(timeout_ns).ir_value(loc=loc, ip=ip),
loc=loc,
@@ -169,7 +172,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.llvm_ptr,
mbar_ptr.to_llvm_ptr(loc=loc, ip=ip),
Int32(phase).ir_value(loc=loc, ip=ip),
nvvm.MBarrierWaitKind.TRY,
loc=loc,
@@ -223,7 +226,7 @@ def mbarrier_arrive(
the mbarrier is converted to a remote address in the peer CTA's
SMEM.
"""
mbar_llvm_ptr = mbar_ptr.llvm_ptr
mbar_llvm_ptr = mbar_ptr.to_llvm_ptr(loc=loc, ip=ip)
if peer_cta_rank_in_cluster is not None:
BaseDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
@@ -261,10 +264,5 @@ 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.llvm_ptr
nvvm.cp_async_mbarrier_arrive_shared(
mbar_llvm_ptr,
noinc=True,
loc=loc,
ip=ip,
)
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)
python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py
+33
View File
@@ -30,6 +30,7 @@ from .nvvm_wrappers import (
cvt_f4e2m1x4_to_f16x4,
cvt_f4e2m1x2_to_f16x2,
cvt_f4e2m1_f16,
sext_unpacked_i4x4_to_i8x4,
)
from ..typing import (
@@ -272,6 +273,38 @@ def cvt_f4e2m1_f16_intrinsic(vec_f4e2m1, length, *, loc=None, ip=None):
return vec_dst
@dsl_user_op
def sext_unpacked_i4_i8_intrinsic(vec_unpacked_i4, length, *, loc=None, ip=None):
"""
Sign extend vector of int4 unpacked in 8b containers to packed int8
:param vec_unpacked_i4: The input vector of unpacked int4.
:type vec_unpacked_i4: 1D vector of unpacked int4
:param length: The length of the input vector.
:type length: int
:return: The output 1D vector of int8 with the same length as the input vector.
:rtype: 1D vector of int8
"""
assert length % 4 == 0, "unsupported length"
vec_i8x4_type = ir.VectorType.get([4], Int8.mlir_type, loc=loc)
vec_i8_type = ir.VectorType.get([length], Int8.mlir_type, loc=loc)
vec_i8 = llvm.mlir_zero(vec_i8_type, loc=loc, ip=ip)
for pos in range(0, length, 4):
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_i8 = vector.insert_strided_slice(
vec_i8x4, vec_i8, [pos], [1], loc=loc, ip=ip
)
return vec_i8
# Expose supported architectures via the intrinsic symbol
cvt_i8_bf16_intrinsic.supported_archs = (
*Arch.AmpereArchs(),
python/CuTeDSL/cutlass/cute/arch/nvvm_wrappers.py
+788 -34
View File
@@ -10,19 +10,16 @@
# is strictly prohibited.
from functools import partial
from typing import Optional, Tuple, Union, Callable, TYPE_CHECKING
from typing import Optional, Tuple, Union, Callable, Literal
from typing_extensions import deprecated
from cutlass.cutlass_dsl import T, dsl_user_op, cutlass_arith
from cutlass.cutlass_dsl import T, dsl_user_op
import cutlass.cutlass_dsl as cutlass_dsl
from cutlass._mlir import ir
from cutlass._mlir.dialects import arith, llvm, nvvm, vector
if TYPE_CHECKING:
from cutlass.tensor import TensorSSA
# Forward nvvm enums
from cutlass._mlir.dialects.nvvm import (
ProxyKind,
@@ -54,6 +51,81 @@ WARP_SIZE = 32
FULL_MASK = 0xFFFFFFFF
# ============================================================================
# Enum String Mapping Helper
# ============================================================================
# This section provides a helper to convert string literals to NVVM enum types
# by introspecting the enum's __str__() method. Each function imports and
# enhances only the enums it needs, avoiding namespace pollution.
#
# Usage within functions:
# MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind)
# sem = MemOrderKind.from_str("relaxed")
## ============================================================================
def _enhance_enum_with_str_mapping(enum_class):
"""
Enhance an IntEnum class with automatic string-to-enum conversion.
Builds a reverse mapping from __str__() output to enum members and adds
a from_str() class method for conversion. Safe to call multiple times
(idempotent - won't re-enhance if already enhanced).
:param enum_class: The enum class to enhance
:return: The enhanced enum class (for chaining)
"""
# Skip if already enhanced
if hasattr(enum_class, "from_str"):
return enum_class
# Build reverse mapping from string representation to enum member
str_to_enum_map = {}
for member in enum_class:
str_repr = str(member)
if str_repr in str_to_enum_map:
raise ValueError(
f"Duplicate string representation '{str_repr}' in {enum_class.__name__}"
)
str_to_enum_map[str_repr] = member
# Add from_str class method
@classmethod
def from_str(cls, s):
"""
Convert a string literal to the corresponding enum member.
:param s: String representation of the enum member, or an enum member itself (deprecated)
:return: The enum member (or None if s is None)
:raises ValueError: If the string is not a valid enum member
"""
import warnings
if s is None:
return None
# 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,
)
return s
if s not in str_to_enum_map:
valid_options = sorted(str_to_enum_map.keys())
raise ValueError(
f"Invalid {cls.__name__} string: '{s}'. "
f"Valid options are: {valid_options}"
)
return str_to_enum_map[s]
enum_class.from_str = from_str
return enum_class
@dsl_user_op
def lane_idx(*, loc=None, ip=None) -> Int32:
"""
@@ -374,7 +446,6 @@ 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),
@@ -389,13 +460,34 @@ 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)
nvvm.barrier(
barrier_id=barrier_id, number_of_threads=number_of_threads, 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,
)
@dsl_user_op
@@ -404,6 +496,8 @@ 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(
@@ -411,8 +505,14 @@ def barrier_arrive(
)
number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip)
nvvm.barrier_arrive(
barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip
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,
)
@@ -682,7 +782,7 @@ def popc(value: Numeric, *, loc=None, ip=None) -> Numeric:
@dsl_user_op
def fence_view_async_tmem_op(
kind: Tcgen05WaitKind,
kind: Literal["load", "store"],
*,
loc=None,
ip=None,
@@ -715,18 +815,20 @@ def fence_view_async_tmem_op(
```
:param kind: The kind of fence operation to perform including LOAD and STORE.
:type kind: Tcgen05WaitKind
:param kind: The kind of fence operation to perform ("load", "store").
:type kind: Literal["load", "store"]
"""
nvvm.tcgen05_wait(kind, loc=loc, ip=ip)
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)
nvvm.tcgen05_wait(kind=kind, loc=loc, ip=ip)
fence_view_async_tmem_load = partial(
fence_view_async_tmem_op, kind=Tcgen05WaitKind.LOAD
)
fence_view_async_tmem_store = partial(
fence_view_async_tmem_op, kind=Tcgen05WaitKind.STORE
)
fence_view_async_tmem_load = partial(fence_view_async_tmem_op, kind="load")
fence_view_async_tmem_store = partial(fence_view_async_tmem_op, kind="store")
@dsl_user_op
@@ -751,23 +853,45 @@ def fence_view_async_shared(
@dsl_user_op
def warpgroup_reg_realloc_op(
def setmaxregister_increase(
reg_count: int,
*,
loc=None,
ip=None,
):
return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip)
@dsl_user_op
def setmaxregister_decrease(
reg_count: int,
*,
loc=None,
ip=None,
):
return nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip)
@dsl_user_op
@deprecated("API is deprecated, use setmaxregister_increase instead")
def warpgroup_reg_alloc(
reg_count: int,
kind: SetMaxRegisterAction,
*,
loc=None,
ip=None,
) -> None:
nvvm.setmaxregister(reg_count, kind, loc=loc, ip=ip)
nvvm.setmaxregister(reg_count, SetMaxRegisterAction.increase, loc=loc, ip=ip)
warpgroup_reg_alloc = partial(
warpgroup_reg_realloc_op, kind=SetMaxRegisterAction.increase
)
warpgroup_reg_dealloc = partial(
warpgroup_reg_realloc_op, kind=SetMaxRegisterAction.decrease
)
@dsl_user_op
@deprecated("API is deprecated, use setmaxregister_decrease instead")
def warpgroup_reg_dealloc(
reg_count: int,
*,
loc=None,
ip=None,
) -> None:
nvvm.setmaxregister(reg_count, SetMaxRegisterAction.decrease, loc=loc, ip=ip)
@dsl_user_op
def calc_packed_f32x2_op(
@@ -776,11 +900,17 @@ def calc_packed_f32x2_op(
src_c: Optional[Tuple[Float32, Float32]],
calc_func: Callable,
*,
rnd=RoundingModeKind.RZ,
ftz=True,
rnd: Optional[Literal["rn", "rz", "rm", "rp", "none"]] = "rn",
ftz=None,
loc=None,
ip=None,
) -> Tuple[Float32, Float32]:
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)
vec_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc)
vec_src_a = vector.from_elements(
vec_type,
@@ -1259,6 +1389,15 @@ def cvt_i4x8_to_bf16x8(src_vec8, *, loc=None, ip=None):
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):
imm_u32 = arith.constant(Uint32.mlir_type, 0x78787878, loc=loc, ip=ip)
src_u32 = llvm.bitcast(Uint32.mlir_type, src_vec4, loc=loc, ip=ip)
dst_u32 = arith.addi(src_u32, imm_u32, loc=loc, ip=ip)
dst_u32 = arith.xori(dst_u32, imm_u32, loc=loc, ip=ip)
return llvm.bitcast(src_vec4.type, dst_u32, loc=loc, ip=ip)
@dsl_user_op
def log2_of_pow2_int(a: Int32, *, loc=None, ip=None) -> Int32:
@@ -1346,6 +1485,621 @@ def griddepcontrol_launch_dependents(*, loc=None, ip=None) -> None:
def _normalize_ptr(addr, *, loc=None, ip=None) -> ir.Value:
"""
Helper function to normalize pointer types to MLIR ir.Value.
Supports:
- ir.Value (LLVM pointer): returned as-is
- cute.ptr (_Pointer instance): converted via to_llvm_ptr()
:param addr: Address in various pointer formats
:return: Normalized MLIR pointer value
:rtype: ir.Value
"""
# If it's already an MLIR ir.Value, return as-is
if isinstance(addr, ir.Value):
return addr
# If it has to_llvm_ptr method (cute._Pointer instances)
if hasattr(addr, "to_llvm_ptr") and callable(addr.to_llvm_ptr):
return addr.to_llvm_ptr(loc=loc, ip=ip)
# If none of the above, return as-is and let NVVM handle it
# This allows for future pointer types without breaking existing code
return addr
def _atomic(
ptr,
val: Union[Numeric, ir.Value],
*,
op: Literal[
"add",
"fadd",
"max",
"min",
"and",
"or",
"xor",
"exch",
],
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Union[Numeric, ir.Value]:
"""
General atomic operation function.
Atomically adds `val` to the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location. Supports:
- ir.Value (LLVM pointer)
- cute.ptr (_Pointer instance)
:param val: Value to add (scalar Numeric or vector ir.Value)
:type val: Union[Numeric, ir.Value]
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:param op: Atomic operation ("add", "fadd", "max", "min", "and", "or", "xor", "exch")
:type op: Literal["add", "fadd", "max", "min", "and", "or", "xor", "exch"]
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Union[Numeric, ir.Value]
"""
from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind
from cutlass.utils.version_info import CUDA_VERSION
# Enhance enums and convert string literals to enum types
AtomicOpKind = _enhance_enum_with_str_mapping(AtomicOpKind)
MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind)
MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind)
op = AtomicOpKind.from_str(op)
sem = MemOrderKind.from_str(sem)
scope = MemScopeKind.from_str(scope)
# Normalize pointer type to MLIR ir.Value
ptr = _normalize_ptr(ptr, loc=loc, ip=ip)
# * Handle `val` Type - scalar Numeric or vector ir.Value
is_vector = isinstance(val, ir.Value) and isinstance(val.type, ir.VectorType)
if is_vector:
# Vector type atomic - val is already an ir.Value
val_ir = val
val_type = val.type
# Check if it's a floating-point vector type
elem_type = val.type.element_type
is_float_vector = (
elem_type == Float16.mlir_type
or elem_type == BFloat16.mlir_type
or elem_type == Float32.mlir_type
)
# Vector atomics for f16/bf16/f32 only support ADD (FADD)
if is_float_vector and op == AtomicOpKind.ADD:
op = AtomicOpKind.FADD
else:
# Scalar type atomic - convert to Numeric
if not isinstance(val, Numeric):
val = as_numeric(val)
val_type = type(val)
val_ir = val.ir_value(loc=loc, ip=ip)
# * Float
# For .f32, .f64, .f16, .bf16, .f16x2, .bf16x2, only .add (FADD) is supported
# For .u32 .u64, .s32, .s64, .add .and .or .xor .cas .exch .min .max are supported
if val_type.is_float:
# For floating-point types, only ADD is supported
if op == AtomicOpKind.ADD:
# Convert ADD to FADD for floating-point types
op = AtomicOpKind.FADD
# * NVVM call based on nvvm version
if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9:
# Old API: requires explicit result type as first positional argument
# For vectors: pass val_type (ir.VectorType), for scalars: pass val_type.mlir_type
result_type = val_type if is_vector else val_type.mlir_type
result = nvvm.atomicrmw(
result_type,
op=op,
ptr=ptr,
a=val_ir,
mem_order=sem,
syncscope=scope,
loc=loc,
ip=ip,
)
else:
# New API: infers result type automatically
result = nvvm.atomicrmw(
op=op,
ptr=ptr,
a=val_ir,
mem_order=sem,
syncscope=scope,
loc=loc,
ip=ip,
)
# Return raw result for vectors, wrapped for scalars
return result if is_vector else val_type(result)
def atomic_add(
ptr,
val: Union[Numeric, ir.Value],
*,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Union[Numeric, ir.Value]:
"""
Performs an atomic addition operation.
Atomically adds `val` to the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location
:param val: Value to add (scalar Numeric or vector ir.Value)
:type val: Union[Numeric, ir.Value]
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Union[Numeric, ir.Value]
"""
return _atomic(ptr, val, op="add", sem=sem, scope=scope, loc=loc, ip=ip)
def atomic_and(
ptr,
val: Numeric,
*,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Numeric:
"""
Performs an atomic bitwise AND operation.
Atomically computes bitwise AND of `val` with the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location
:param val: Value for AND operation
:type val: Numeric
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Numeric
"""
return _atomic(ptr, val, op="and", sem=sem, scope=scope, loc=loc, ip=ip)
def atomic_or(
ptr,
val: Numeric,
*,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Numeric:
"""
Performs an atomic bitwise OR operation.
Atomically computes bitwise OR of `val` with the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location
:param val: Value for OR operation
:type val: Numeric
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Numeric
"""
return _atomic(ptr, val, op="or", sem=sem, scope=scope, loc=loc, ip=ip)
def atomic_xor(
ptr,
val: Numeric,
*,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Numeric:
"""
Performs an atomic bitwise XOR operation.
Atomically computes bitwise XOR of `val` with the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location
:param val: Value for XOR operation
:type val: Numeric
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Numeric
"""
return _atomic(ptr, val, op="xor", sem=sem, scope=scope, loc=loc, ip=ip)
def atomic_max(
ptr,
val: Numeric,
*,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Numeric:
"""
Performs an atomic maximum operation.
Atomically computes maximum of `val` and the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location
:param val: Value for MAX operation
:type val: Numeric
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Numeric
"""
return _atomic(ptr, val, op="max", sem=sem, scope=scope, loc=loc, ip=ip)
def atomic_min(
ptr,
val: Numeric,
*,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Numeric:
"""
Performs an atomic minimum operation.
Atomically computes minimum of `val` and the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location
:param val: Value for MIN operation
:type val: Numeric
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Numeric
"""
return _atomic(ptr, val, op="min", sem=sem, scope=scope, loc=loc, ip=ip)
def atomic_exch(
ptr,
val: Numeric,
*,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Numeric:
"""
Performs an atomic exchange operation.
Atomically exchanges `val` with the value at memory location `ptr` and returns the old value.
:param ptr: Pointer to memory location
:param val: Value to exchange
:type val: Numeric
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Numeric
"""
return _atomic(ptr, val, op="exch", sem=sem, scope=scope, loc=loc, ip=ip)
@dsl_user_op
def atomic_cas(
ptr,
*,
cmp: Numeric,
val: Numeric,
sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> Numeric:
"""
Performs an atomic compare-and-swap (CAS) operation.
Atomically compares the value at the memory location with `cmp`. If they are equal,
stores `val` at the memory location and returns the old value.
:param ptr: Pointer to memory location. Supports:
- ir.Value (LLVM pointer)
- cute.ptr (_Pointer instance)
:param cmp: Value to compare against current memory value
:type cmp: Numeric
:param val: Value to store if comparison succeeds
:type val: Numeric
:param sem: Memory semantic ("relaxed", "release", "acquire", "acq_rel")
:type sem: Optional[Literal["relaxed", "release", "acquire", "acq_rel"]]
:param scope: Memory scope ("gpu", "cta", "cluster", "sys")
:type scope: Optional[Literal["gpu", "cta", "cluster", "sys"]]
:return: Old value at memory location
:rtype: Numeric
"""
from cutlass._mlir.dialects.nvvm import AtomicOpKind, MemOrderKind, MemScopeKind
from cutlass.utils.version_info import CUDA_VERSION
# Enhance enums and convert string literals to enum types
MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind)
MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind)
sem = MemOrderKind.from_str(sem)
scope = MemScopeKind.from_str(scope)
# Normalize pointer type to MLIR ir.Value
ptr = _normalize_ptr(ptr, loc=loc, ip=ip)
# * Hanldle `val`, `cmp` Numeric Type
if not isinstance(cmp, Numeric):
cmp = as_numeric(cmp)
if not isinstance(val, Numeric):
val = as_numeric(val)
cmp_type = type(cmp)
cmp_ir = cmp.ir_value(loc=loc, ip=ip)
val_ir = val.ir_value(loc=loc, ip=ip)
# * NVVM call based on nvvm version
if CUDA_VERSION.major == 12 and CUDA_VERSION.minor == 9:
result = nvvm.atomicrmw(
cmp_type.mlir_type,
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,
ptr=ptr,
a=cmp_ir,
b=val_ir,
mem_order=sem,
syncscope=scope,
loc=loc,
ip=ip,
)
return cmp_type(result)
@dsl_user_op
def store(
ptr,
val: Union[Numeric, ir.Value],
*,
level1_eviction_priority: Optional[
Literal[
"evict_normal",
"evict_first",
"evict_last",
"evict_no_allocate",
"evict_unchanged",
]
] = None,
cop: Optional[Literal["wb", "cg", "cs", "wt"]] = None,
ss: Optional[Literal["cta", "cluster"]] = None,
sem: Optional[Literal["relaxed", "release"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
loc=None,
ip=None,
) -> None:
"""
Store a value to a memory location.
:param ptr: Pointer to store to. Supports:
- ir.Value (LLVM pointer)
- cute.ptr (_Pointer instance)
:param val: Value to store (scalar Numeric or vector ir.Value)
:type val: Union[Numeric, ir.Value]
:param level1_eviction_priority: L1 cache eviction policy string literal:
"evict_normal" : .level1::eviction_priority = .L1::evict_normal
"evict_first" : .level1::eviction_priority = .L1::evict_first
"evict_last" : .level1::eviction_priority = .L1::evict_last
"evict_no_allocate" : .level1::eviction_priority = .L1::no_allocate
"evict_unchanged" : .level1::eviction_priority = .L1::evict_unchanged
:param cop: Store cache modifier string literal:
:param ss: Shared memory space string literal:
"cta" : .ss = .shared::cta
"cluster" : .ss = .shared::cluster
None : .ss = .global
:param sem: Memory semantic string literal:
:param scope: Memory scope string literal:
"""
from cutlass._mlir.dialects.nvvm import (
MemOrderKind,
MemScopeKind,
StoreCacheModifierKind,
EvictKind,
SharedSpace,
)
# Enhance enums and convert string literals to enum types
MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind)
MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind)
StoreCacheModifierKind = _enhance_enum_with_str_mapping(StoreCacheModifierKind)
EvictKind = _enhance_enum_with_str_mapping(EvictKind)
SharedSpace = _enhance_enum_with_str_mapping(SharedSpace)
sem = MemOrderKind.from_str(sem)
scope = MemScopeKind.from_str(scope)
cop = StoreCacheModifierKind.from_str(cop)
level1_eviction_priority = EvictKind.from_str(level1_eviction_priority)
ss = SharedSpace.from_str(ss)
# Normalize pointer type to MLIR ir.Value
ptr = _normalize_ptr(ptr, loc=loc, ip=ip)
# Handle both scalar Numeric and vector ir.Value
is_vector = isinstance(val, ir.Value) and isinstance(val.type, ir.VectorType)
if is_vector:
# Vector type store - val is already an ir.Value
val_ir = val
else:
# Scalar type store - ensure val is a Numeric and convert to MLIR Value
if not isinstance(val, Numeric):
val = as_numeric(val)
val_ir = val.ir_value(loc=loc, ip=ip)
nvvm.store_ext(
val_ir,
ptr,
order=sem,
scope=scope,
evict=level1_eviction_priority,
cache_modifier=cop,
shared_space=ss,
loc=loc,
ip=ip,
)
@dsl_user_op
def load(
ptr,
dtype: Union[type[Numeric], ir.VectorType],
*,
sem: Optional[Literal["relaxed", "acquire"]] = None,
scope: Optional[Literal["gpu", "cta", "cluster", "sys"]] = None,
level1_eviction_priority: Optional[
Literal[
"evict_normal",
"evict_first",
"evict_last",
"evict_no_allocate",
"evict_unchanged",
]
] = None,
cop: Optional[Literal["ca", "cg", "cs", "lu", "cv"]] = None,
ss: Optional[Literal["cta", "cluster"]] = None,
level_prefetch_size: Optional[Literal["size_64b", "size_128b", "size_256b"]] = None,
loc=None,
ip=None,
) -> Union[Numeric, ir.Value]:
"""
Load a value from a memory location.
:param ptr: Pointer to load from. Supports:
- ir.Value (LLVM pointer)
- cute.ptr (_Pointer instance)
:param dtype: Data type to load. Can be:
- Scalar: Numeric type class (Int8, Uint8, Int32, Float32, etc.)
- Vector: ir.VectorType for vectorized load (e.g., ir.VectorType.get([4], Int64.mlir_type))
:type dtype: Union[type[Numeric], ir.VectorType]
:param sem: Memory semantic string literal:
:param scope: Memory scope string literal:
:param level1_eviction_priority: L1 cache eviction policy string literal:
"evict_normal" : .level1::eviction_priority = .L1::evict_normal
"evict_first" : .level1::eviction_priority = .L1::evict_first
"evict_last" : .level1::eviction_priority = .L1::evict_last
"evict_no_allocate" : .level1::eviction_priority = .L1::no_allocate
"evict_unchanged" : .level1::eviction_priority = .L1::evict_unchanged
:param cop: Load cache modifier string literal:
:param ss: Shared memory space string literal:
"cta" : .ss = .shared::cta
"cluster" : .ss = .shared::cluster
None : .ss = .global
:param level_prefetch_size: L2 cache prefetch size hint string literal:
"size_64b" : .level::prefetch_size = .L2::64B
"size_128b" : .level::prefetch_size = .L2::128B
"size_256b" : .level::prefetch_size = .L2::256B
:return: Loaded value (scalar Numeric or vector ir.Value)
:rtype: Union[Numeric, ir.Value]
"""
from cutlass._mlir.dialects.nvvm import (
MemOrderKind,
MemScopeKind,
LoadCacheModifierKind,
EvictKind,
SharedSpace,
L2PrefetchSize,
)
# Enhance enums and convert string literals to enum types
MemOrderKind = _enhance_enum_with_str_mapping(MemOrderKind)
MemScopeKind = _enhance_enum_with_str_mapping(MemScopeKind)
LoadCacheModifierKind = _enhance_enum_with_str_mapping(LoadCacheModifierKind)
EvictKind = _enhance_enum_with_str_mapping(EvictKind)
SharedSpace = _enhance_enum_with_str_mapping(SharedSpace)
L2PrefetchSize = _enhance_enum_with_str_mapping(L2PrefetchSize)
sem = MemOrderKind.from_str(sem)
scope = MemScopeKind.from_str(scope)
cop = LoadCacheModifierKind.from_str(cop)
level1_eviction_priority = EvictKind.from_str(level1_eviction_priority)
ss = SharedSpace.from_str(ss)
level_prefetch_size = L2PrefetchSize.from_str(level_prefetch_size)
# Normalize pointer type to MLIR ir.Value
ptr = _normalize_ptr(ptr, loc=loc, ip=ip)
# Determine if dtype is a vector type or scalar type
is_vector = isinstance(dtype, ir.VectorType) and isinstance(dtype, ir.VectorType)
if is_vector:
# Vector load: dtype is already an ir.VectorType
mlir_type = dtype
scalar_dtype = None # We don't need to wrap the result
else:
# Scalar load: dtype is a Numeric type class
mlir_type = dtype.mlir_type
scalar_dtype = dtype
result = nvvm.load_ext(
res=mlir_type,
addr=ptr,
order=sem,
scope=scope,
evict=level1_eviction_priority,
cache_modifier=cop,
shared_space=ss,
prefetch=level_prefetch_size,
loc=loc,
ip=ip,
)
# Return raw ir.Value for vectors, wrapped Numeric for scalars
if is_vector:
return result
else:
return scalar_dtype(result)
@dsl_user_op
def cvt_f4e2m1_f16(src, *, loc=None, ip=None):
# 0 padding for upper 4 bits
python/CuTeDSL/cutlass/cute/arch/tmem.py
+69 -9
View File
@@ -16,11 +16,61 @@ from cutlass.cutlass_dsl import dsl_user_op
import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from ..typing import Pointer, Int, Int32, Numeric, NumericMeta
from ..typing import Pointer, Int, Int32, Numeric, NumericMeta, Tensor
SM100_TMEM_CAPACITY_COLUMNS = (
512 # deprecated; use get_max_tmem_alloc_cols(arch="sm_100") instead
)
SM100_TMEM_MIN_ALLOC_COLUMNS = (
32 # deprecated; use get_min_tmem_alloc_cols(arch="sm_100") instead
)
TMEM_MAX_ALLOC_COLUMNS_MAP = {
"sm_120": 512,
"sm_103": 512,
"sm_100": 512,
}
TMEM_MIN_ALLOC_COLUMNS_MAP = {
"sm_120": 32,
"sm_103": 32,
"sm_100": 32,
}
SM100_TMEM_CAPACITY_COLUMNS = 512
SM100_TMEM_MIN_ALLOC_COLUMNS = 32
def get_max_tmem_alloc_cols(compute_capability: str) -> int:
"""Get the tensor memory capacity in columns for a given compute capability.
Returns the maximum TMEM capacity in columns available for the specified
GPU compute capability.
:param compute_capability: The compute capability string (e.g. "sm_100", "sm_103")
:type compute_capability: str
:return: The TMEM capacity in columns
:rtype: int
:raises ValueError: If the compute capability is not supported
"""
if compute_capability not in TMEM_MAX_ALLOC_COLUMNS_MAP:
raise ValueError(f"Unsupported compute capability: {compute_capability}")
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.
Returns the minimum TMEM allocation columns available for the specified
GPU compute capability.
:param compute_capability: The compute capability string (e.g. "sm_100", "sm_103")
:type compute_capability: str
:return: The minimum TMEM allocation columns
:rtype: int
:raises ValueError: If the compute capability is not supported
"""
if compute_capability not in TMEM_MIN_ALLOC_COLUMNS_MAP:
raise ValueError(f"Unsupported compute capability: {compute_capability}")
return TMEM_MIN_ALLOC_COLUMNS_MAP[compute_capability]
@dsl_user_op
@@ -64,6 +114,7 @@ def alloc_tmem(
smem_ptr_to_write_address: Pointer,
is_two_cta=None,
*,
arch: str = "sm_100",
loc=None,
ip=None,
) -> None:
@@ -76,15 +127,19 @@ def alloc_tmem(
to
:type smem_ptr_to_write_address: Pointer
: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)
if isinstance(num_columns, int):
if (
num_columns < SM100_TMEM_MIN_ALLOC_COLUMNS
or num_columns > SM100_TMEM_CAPACITY_COLUMNS
num_columns < tmem_min_alloc_cols
or num_columns > tmem_max_alloc_cols
or not (num_columns & (num_columns - 1) == 0)
):
raise ValueError(
f"num_columns must be between 32 and 512, and must be pow of 2, but got {num_columns}"
f"num_columns must be between {tmem_min_alloc_cols} and {tmem_max_alloc_cols}, and must be pow of 2, but got {num_columns}"
)
_cute_nvgpu_ir.arch_sm100_alloc_tmem(
Int32(num_columns).ir_value(loc=loc, ip=ip),
@@ -112,6 +167,7 @@ def dealloc_tmem(
num_columns: Int,
is_two_cta=None,
*,
arch: str = "sm_100",
loc=None,
ip=None,
) -> None:
@@ -123,15 +179,19 @@ 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)
if isinstance(num_columns, int):
if (
num_columns < SM100_TMEM_MIN_ALLOC_COLUMNS
or num_columns > SM100_TMEM_CAPACITY_COLUMNS
num_columns < tmem_min_alloc_cols
or num_columns > tmem_max_alloc_cols
or not (num_columns & (num_columns - 1) == 0)
):
raise ValueError(
f"num_columns must be between 32 and 512, and must be pow of 2, but got {num_columns}"
f"num_columns must be between {tmem_min_alloc_cols} and {tmem_max_alloc_cols}, and must be pow of 2, but got {num_columns}"
)
_cute_nvgpu_ir.arch_sm100_dealloc_tmem(
tmem_ptr.value,
python/CuTeDSL/cutlass/cute/atom.py
+49 -12
View File
@@ -136,10 +136,15 @@ class Atom(ABC):
self._trait = trait
def __extract_mlir_values__(self):
return extract_mlir_values(self._trait)
return extract_mlir_values(self._trait) + extract_mlir_values(self._op)
def __new_from_mlir_values__(self, values):
return self.__class__(self.op, new_from_mlir_values(self._trait, values))
traits_value = values[: len(extract_mlir_values(self._trait))]
op_value = values[len(extract_mlir_values(self._trait)) :]
new_trait = new_from_mlir_values(self._trait, traits_value)
new_op = new_from_mlir_values(self._op, op_value)
return self.__class__(new_op, new_trait)
@property
def op(self) -> Op:
@@ -318,24 +323,29 @@ class TiledMma(MmaAtom):
#
@property
def tv_layout_A_tiled(self) -> Layout:
return static(self._trait.value.type.layout_a_tv_tiled)
@dsl_user_op
def tv_layout_A_tiled(self, *, loc=None, ip=None) -> Layout:
return static(self._trait.value.type.layout_a_tv_tiled, loc=loc, ip=ip)
@property
def tv_layout_B_tiled(self) -> Layout:
return static(self._trait.value.type.layout_b_tv_tiled)
@dsl_user_op
def tv_layout_B_tiled(self, *, loc=None, ip=None) -> Layout:
return static(self._trait.value.type.layout_b_tv_tiled, loc=loc, ip=ip)
@property
def tv_layout_C_tiled(self) -> Layout:
return static(self._trait.value.type.layout_c_tv_tiled)
@dsl_user_op
def tv_layout_C_tiled(self, *, loc=None, ip=None) -> Layout:
return static(self._trait.value.type.layout_c_tv_tiled, loc=loc, ip=ip)
@property
def permutation_mnk(self) -> Tile:
return _unpack_x_tuple(self._trait.value.type.permutation_mnk)
@dsl_user_op
def permutation_mnk(self, *, loc=None, ip=None) -> Tile:
return _unpack_x_tuple(self._trait.value.type.permutation_mnk, loc=loc, ip=ip)
@property
def thr_layout_vmnk(self) -> Layout:
return static(self._trait.value.type.thr_layout_vmnk)
@dsl_user_op
def thr_layout_vmnk(self, *, loc=None, ip=None) -> Layout:
return static(self._trait.value.type.thr_layout_vmnk, loc=loc, ip=ip)
@property
def size(self) -> int:
@@ -598,6 +608,33 @@ class CopyAtom(Atom):
def layout_dst_tv(self) -> Layout:
return static(self._trait.value.type.layout_dst_tv)
@property
def smem_layout(self):
"""
Convenience property to access the SMEM layout for TMA copy atoms.
This is a shortcut for ``atom.op.smem_layout`` that checks if the operation
is a TMA operation and provides a clearer error message if not.
:return: The SMEM layout
:rtype: Layout or ComposedLayout
:raises TypeError: If the operation is not a TMA operation
:raises ValueError: If the SMEM layout is not set
Example:
>>> layout = tma_atom.smem_layout # Instead of tma_atom.op.smem_layout
"""
# Import here to avoid circular dependency
from .nvgpu.cpasync.copy import TmaCopyOp
if not isinstance(self.op, TmaCopyOp):
raise TypeError(
f"smem_layout is only available for TMA copy operations, "
f"but this atom uses {type(self.op).__name__}"
)
return self.op.smem_layout
class TiledCopy(CopyAtom):
"""
python/CuTeDSL/cutlass/cute/core.py
+428 -63
View File
@@ -17,7 +17,7 @@ from typing import Any, Dict, List, Optional, Tuple, Type, Union, overload
from typing_extensions import deprecated
from cutlass._mlir import ir
from cutlass._mlir.dialects import builtin, llvm
from cutlass._mlir.dialects import builtin, llvm, vector
from cutlass._mlir.dialects import cute as _cute_ir
from cutlass._mlir.dialects.cute import (
Ratio as _Ratio,
@@ -66,6 +66,91 @@ from .typing import (
is_integer,
)
__all__ = [
# Classes
"IntValue",
"Swizzle",
"struct",
# Utility functions
"E",
"get_divisibility",
"is_valid_leaf",
"is_static",
"has_underscore",
"has_scaled_basis",
"pretty_str",
"printf",
# Layout operations
"front",
"is_major",
"assume",
"make_swizzle",
"static",
"get_leaves",
"depth",
"rank",
"is_congruent",
"is_weakly_congruent",
"get",
"select",
"group_modes",
"slice_",
"dice",
"prepend",
"append",
"prepend_ones",
"append_ones",
"repeat_as_tuple",
"repeat",
"repeat_like",
"flatten",
"filter_zeros",
"filter",
"size",
"shape_div",
"ceil_div",
"round_up",
"make_layout",
"make_identity_layout",
"make_ordered_layout",
"make_layout_like",
"make_composed_layout",
"cosize",
"size_in_bytes",
"coalesce",
"crd2idx",
"idx2crd",
"recast_layout",
"slice_and_offset",
"shape",
"recast_ptr",
"make_ptr",
"composition",
"complement",
"right_inverse",
"left_inverse",
"logical_product",
"zipped_product",
"tiled_product",
"flat_product",
"raked_product",
"blocked_product",
"logical_divide",
"zipped_divide",
"tiled_divide",
"flat_divide",
"max_common_layout",
"max_common_vector",
"tile_to_shape",
"local_partition",
"local_tile",
"make_layout_image_mask",
"leading_dim",
"make_layout_tv",
"get_nonswizzle_portion",
"get_swizzle_portion",
]
####################################################################################################
#
# Internal IntTuple helpers
@@ -131,8 +216,8 @@ def _pack_tile(tile: Tile, *, loc=None, ip=None) -> ir.Value:
leaves = []
for e in tile:
if isinstance(e, _Layout):
leaves.extend(list(flatten_to_tuple(e.shape)))
leaves.extend(list(flatten_to_tuple(e.stride)))
leaves.extend(list(flatten_to_tuple(e.shape_method(loc=loc, ip=ip))))
leaves.extend(list(flatten_to_tuple(e.stride_method(loc=loc, ip=ip))))
else:
leaves.append(e)
return leaves
@@ -143,6 +228,7 @@ def _pack_tile(tile: Tile, *, loc=None, ip=None) -> ir.Value:
_get_typed_value(x) for x in dyn_elems if isinstance(x, (Integer, ir.Value))
]
tile = transform_leaf(_get_typed_value, tile)
res_ty = _cute_ir.pack_tile(tile)
return _cute_ir.make_tile(res_ty, dyn_elems, loc=loc, ip=ip)
@@ -332,16 +418,18 @@ class IntValue(cutlass_arith.ArithValue):
# Dispatch to `__rmul__` of `other`
return NotImplemented
return IntValue(op(self, other_val, **kwargs))
return IntValue(
op(self, other_val, **kwargs),
loc=kwargs.get("loc"),
ip=kwargs.get("ip"),
)
return wrapper
@dsl_user_op
@_binary_op
def __add__(self, other, *, loc=None, ip=None):
return _cute_ir.add_offset(
self.get_typed_value(loc=loc, ip=ip), other, loc=loc, ip=ip
)
return _cute_ir.tuple_add(self.get_typed_value(), other, loc=loc, ip=ip)
@dsl_user_op
@_binary_op
@@ -373,10 +461,8 @@ class IntValue(cutlass_arith.ArithValue):
@dsl_user_op
@_binary_op
def __radd__(self, other, *, loc=None, ip=None) -> "IntValue":
return _cute_ir.add_offset(
other, self.get_typed_value(loc=loc, ip=ip), loc=loc, ip=ip
)
def __radd__(self, other, *, loc=None, ip=None):
return _cute_ir.tuple_add(other, self.get_typed_value(), loc=loc, ip=ip)
@dsl_user_op
@_binary_op
@@ -579,6 +665,9 @@ class ScaledBasis:
if isinstance(self._value, Integer):
scale = self._value.ir_value(loc=loc, ip=ip)
return _ScaledBasis(scale, self._mode, get_divisibility(scale))
elif isinstance(self._value, cutlass_arith.ArithValue):
scale = self._value
return _ScaledBasis(scale, self._mode, get_divisibility(scale))
else:
scale = self._value
return _ScaledBasis(scale, self._mode)
@@ -658,6 +747,27 @@ class ScaledBasis:
return ScaledBasis(scale * value, self.mode) # type: ignore
def __mul__(
self, scale: Union[Int, ir.Value, Ratio], *, loc=None, ip=None
) -> "ScaledBasis":
"""Multiplication by a scale factor.
This operation is used in layout algebra to scale basis elements,
which is essential for operations like composition and partitioning.
:param scale: The scale factor
:type scale: Union[Int, ir.Value, Ratio]
:param loc: The source location for the operation, defaults to None
:type loc: Location, optional
:param ip: The insertion point for the operation, defaults to None
:type ip: InsertionPoint, optional
:return: A new scaled basis element
:rtype: ScaledBasis
:raises TypeError: If scale is not of a supported type
:raises NotImplementedError: If scaling a basis element with a ratio value
"""
return self.__rmul__(scale, loc=loc, ip=ip)
def __extract_mlir_values__(self):
if isinstance(self.value, Ratio):
# Ratio is always static
@@ -716,6 +826,79 @@ def get_divisibility(x: Union[int, Integer]) -> int:
return 1
def basis_value(e: Union[ScaledBasis, Any]) -> Union[Int, ir.Value, Ratio]:
"""Extract the value from a ScaledBasis or return the input as-is.
If the input is a ScaledBasis, returns its value component.
Otherwise, returns the input unchanged.
:param e: The input element (ScaledBasis or any other type)
:type e: Any
:return: The value of the ScaledBasis or the input itself
:rtype: Any
**Examples:**
.. code-block:: python
>>> basis_value(ScaledBasis(5, 0))
5
>>> basis_value(42)
42
"""
if isinstance(e, ScaledBasis):
return e.value
else:
return e
@dsl_user_op
def basis_get(
basis: Union[ScaledBasis, Numeric, int],
t: Union[XTuple, Layout, ComposedLayout],
*,
loc=None,
ip=None,
) -> Union[XTuple, Layout, ComposedLayout]:
"""Apply the mode indices from a ScaledBasis to get an element from a tuple, layout, or composed layout.
If the basis is a ScaledBasis or Numeric with mode indices, this function uses those
indices to extract the corresponding element from the tuple using hierarchical
indexing. If the basis is not a ScaledBasis or has no modes, returns the tuple, layout, or composed layout as-is.
:param basis: The basis element (ScaledBasis)
:type basis: ScaledBasis
:param t: The tuple, layout, or composed layout to index into
:type t: Union[XTuple, Layout, ComposedLayout]
:return: The element at the position specified by the basis modes, or t itself
:rtype: Union[XTuple, Layout, ComposedLayout]
**Examples:**
.. code-block:: python
>>> basis_get(ScaledBasis(2, 1), (10, 20, 30))
20
>>> basis_get(ScaledBasis(2, [0, 1]), ((10, 20), (30, 40)))
20
>>> basis_get(5, (10, 20, 30)) # Non-basis returns tuple as-is
(10, 20, 30)
"""
if isinstance(basis, ScaledBasis):
modes = basis.mode
if len(modes) == 0:
return t
else:
# Use hierarchical indexing with the mode list
return get(t, modes, loc=loc, ip=ip)
elif isinstance(basis, (Numeric, int)):
return t
else:
raise TypeError(
f"basis must be a ScaledBasis or Numeric, but got {type(basis)}"
)
@ir.register_value_caster(_cute_ir.SwizzleType.get_static_typeid(), replace=True)
class Swizzle(ir.Value):
"""
@@ -751,6 +934,41 @@ class Swizzle(ir.Value):
# Cut off the MLIR type's string for making pretty_str more concise
return self.type.__str__()[15 : 15 + 8]
def __eq__(self, other) -> Union[bool, Boolean]:
"""Check if this Swizzle is equal to another Swizzle. Since num_bits, num_base, and num_shift are static,
this is a constant expression.
Two Swizzles are equal if they have the same num_bits, num_base, and num_shift.
:param other: The Swizzle to compare with.
:return: True if Swizzles are equal, False otherwise.
"""
if isinstance(other, Swizzle):
return self.type == other.type
else:
return False
@property
def num_bits(self) -> int:
"""
Returns the number of bits in the mask (B in Sw<B,M,S>).
"""
return self.type.num_bits
@property
def num_base(self) -> int:
"""
Returns the number of least-significant bits to keep constant (M in Sw<B,M,S>).
"""
return self.type.num_base
@property
def num_shift(self) -> int:
"""
Returns the distance to shift the mask (S in Sw<B,M,S>).
"""
return self.type.num_shift
@ir.register_value_caster(_cute_ir.LayoutType.get_static_typeid(), replace=True)
class _Layout(Layout):
@@ -794,12 +1012,26 @@ class _Layout(Layout):
"""
super().__init__(op_result)
def __str__(self) -> str:
def __repr__(self, *, loc=None, ip=None) -> str:
return self.__str__(loc=loc, ip=ip)
def __str__(self, *, loc=None, ip=None) -> str:
"""Return a string representation of the layout.
:return: A string in the format "shape:stride".
"""
return f"{pretty_str(self.shape)}:{pretty_str(self.stride)}"
type_str = self.type.__str__()
return type_str[type_str.find("<") + 2 : type_str.rfind(">") - 1]
@lru_cache_ir()
def shape_method(self, *, loc=None, ip=None) -> Shape:
return _unpack_x_tuple(_cute_ir.get_shape(self, loc=loc, ip=ip), loc=loc, ip=ip)
@lru_cache_ir()
def stride_method(self, *, loc=None, ip=None) -> Stride:
return _unpack_x_tuple(
_cute_ir.get_stride(self, loc=loc, ip=ip), loc=loc, ip=ip
)
@property
@dsl_user_op
@@ -812,7 +1044,7 @@ class _Layout(Layout):
:return: The hierarchical shape of the layout.
"""
return _unpack_x_tuple(_cute_ir.get_shape(self, loc=loc, ip=ip), loc=loc, ip=ip)
return self.shape_method(loc=loc, ip=ip)
@property
@dsl_user_op
@@ -824,9 +1056,7 @@ class _Layout(Layout):
:return: The hierarchical stride of the layout.
"""
return _unpack_x_tuple(
_cute_ir.get_stride(self, loc=loc, ip=ip), loc=loc, ip=ip
)
return self.stride_method(loc=loc, ip=ip)
@property
def max_alignment(self) -> int:
@@ -977,6 +1207,10 @@ 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
)
@@ -1088,11 +1322,15 @@ class _Pointer(Pointer):
@property
@lru_cache_ir()
def dtype(self) -> Union[Type[Numeric], _cute_ir.SparseElemType]:
if isinstance(self.value.type.value_type, _cute_ir.SparseElemType):
return self.value.type.value_type
else:
return Numeric.from_mlir_type(self.value.type.value_type)
def dtype(
self,
) -> Union[
Type[Numeric],
]:
ret_type = None
if ret_type is None:
ret_type = Numeric.from_mlir_type(self.value.type.value_type)
return ret_type
@property
def alignment(self) -> int:
@@ -1121,6 +1359,25 @@ 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
"""
@@ -1212,16 +1469,16 @@ class _Pointer(Pointer):
####################################################################################################
def _op_wrapper(op_fn, input):
def _op_wrapper(op_fn, input, *, loc=None, ip=None):
from .tensor import _Tensor
if isinstance(input, Tensor):
res = op_fn(input.value)
return _Tensor(res, dtype=input.element_type)
res = op_fn(input.value, loc=loc, ip=ip)
return _Tensor(res, dtype=input.element_type, loc=loc, ip=ip)
elif isinstance(input, _ComposedLayout):
return op_fn(input.value)
return op_fn(input.value, loc=loc, ip=ip)
else:
return op_fn(input)
return op_fn(input, loc=loc, ip=ip)
#
@@ -1486,17 +1743,14 @@ def assume(src, divby=None, *, loc=None, ip=None):
@dsl_user_op
def make_swizzle(b, m, s, *, loc=None, ip=None):
# canonicalize to <0, 4, 3> for identity swizzle (as compiler assumes <0, 4, 3>)
if not isinstance(b, int) or not isinstance(m, int) or not isinstance(s, int):
raise ValueError("b, m, and s must be int")
if b == 0:
m, s = 4, 3
ty = ir.Type.parse(f'!cute.swizzle<"S<{b},{m},{s}>">')
return Swizzle(static(ty, loc=loc, ip=ip))
@dsl_user_op
def make_sparse_elem(num_logical, num_phys, elem_type, *, loc=None, ip=None):
return _cute_ir.SparseElemType.get(num_logical, num_phys, elem_type.mlir_type)
@dsl_user_op
def static(value, *, loc=None, ip=None):
return _cute_ir.static(value, loc=loc, ip=ip)
@@ -1846,7 +2100,7 @@ def group_modes(input, begin: int, end: Optional[int] = None, *, loc=None, ip=No
return (*input[:begin], (input[begin:end]), *input[end:])
return _op_wrapper(
partial(_cute_ir.group_modes, begin=begin, end=end, loc=loc, ip=ip), input
partial(_cute_ir.group_modes, begin=begin, end=end), input, loc=loc, ip=ip
)
@@ -1939,7 +2193,7 @@ def slice_(src, coord: Coord, *, loc=None, ip=None):
return ()
coord_val = _pack_coord(coord, loc=loc, ip=ip)
return _op_wrapper(partial(_cute_ir.slice, coord=coord_val, loc=loc, ip=ip), src)
return _op_wrapper(partial(_cute_ir.slice, coord=coord_val), src, loc=loc, ip=ip)
@overload
@@ -2015,7 +2269,7 @@ def dice(src, dicer, *, loc=None, ip=None):
dicer_val = _pack_coord(dicer, loc=loc, ip=ip)
return _op_wrapper(
partial(_cute_ir.dice, coord=dicer_val.type.attribute, loc=loc, ip=ip), src
partial(_cute_ir.dice, coord=dicer_val.type.attribute), src, loc=loc, ip=ip
)
@@ -2030,7 +2284,7 @@ def _extend(func, input, elem, up_to_rank, loc, ip):
raise TypeError(f"Input type of elem ({type(elem)}) is not accepted!")
N = rank(input) + 1 if up_to_rank is None else up_to_rank
return _op_wrapper(partial(func, N, element=elem, loc=loc, ip=ip), input)
return _op_wrapper(partial(func, N, element=elem), input, loc=loc, ip=ip)
if is_valid_leaf(input) or isinstance(input, tuple):
if elem is None:
@@ -2717,9 +2971,7 @@ class _ComposedLayoutWithInnerFunc(ComposedLayout):
delta = self._outer(coord)
delta_val = _pack_int_tuple(delta, loc=loc, ip=ip)
offset_val_new = _cute_ir.add_offset(
self._offset_val, delta_val, loc=loc, ip=ip
)
offset_val_new = _cute_ir.tuple_add(self._offset_val, delta_val, loc=loc, ip=ip)
offset_new = _unpack_x_tuple(offset_val_new, loc=loc, ip=ip)
return self._inner(offset_new)
@@ -2906,7 +3158,7 @@ def coalesce(input, *, target_profile: Coord = None, loc=None, ip=None):
profile_val = None
return _op_wrapper(
partial(_cute_ir.coalesce, target_profile=profile_val, loc=loc, ip=ip), input
partial(_cute_ir.coalesce, target_profile=profile_val), input, loc=loc, ip=ip
)
@@ -2989,15 +3241,13 @@ def idx2crd(idx, shape, *, loc=None, ip=None):
import cutlass.cute as cute
@cute.jit
def foo():
coord = cute.idx2crd(11, (5,4))
coord = cute.idx2crd(11, (5, 4))
# idx2crd is always col-major
# For shape (m, n, l, ...), coord = (idx % m, idx // m % n, idx // m // n % l, ...
# Computed as: (11 % 5, 11 // 5 % 4) = (1, 2)
print(coord)
foo() # Expected output: (1, 2)
**Note:**
Python DSL is aligned with C++ DSL.
foo() # Expected output: (1, 2)
"""
if is_integer(idx) and is_integer(shape):
return idx
@@ -3008,7 +3258,56 @@ def idx2crd(idx, shape, *, loc=None, ip=None):
@dsl_user_op
def recast_layout(new_type_bits, old_type_bits, src_layout, *, loc=None, ip=None):
def recast_layout(
new_type_bits: int,
old_type_bits: int,
src_layout: Union[Layout, ComposedLayout],
*,
loc=None,
ip=None,
):
"""
Recast a layout from one data type to another.
:param new_type_bits: The new data type bits
:type new_type_bits: int
:param old_type_bits: The old data type bits
:type old_type_bits: int
:param src_layout: The layout to recast
:type src_layout: Union[Layout, ComposedLayout]
:param loc: Optional location information for IR diagnostics.
:type loc: optional
:param ip: Optional instruction pointer or context for underlying IR functions.
:type ip: optional
:return: The recast layout
:rtype: Layout or ComposedLayout
**Example:**
.. code-block:: python
import cutlass.cute as cute
@cute.jit
def foo():
# Create a layout
L = cute.make_layout((2, 3, 4))
# Recast the layout to a different data type
L_recast = cute.recast_layout(16, 8, L)
print(L_recast)
foo() # Expected output: (2, 3, 4)
"""
if not isinstance(new_type_bits, int):
raise TypeError(
f"new_type_bits must be an integer instead got {type(new_type_bits)}"
)
if not isinstance(old_type_bits, int):
raise TypeError(
f"old_type_bits must be an integer instead got {type(old_type_bits)}"
)
if not isinstance(src_layout, (Layout, ComposedLayout)):
raise TypeError(
f"src_layout must be a layout or composed layout instead got {type(src_layout)}"
)
if isinstance(src_layout, _ComposedLayout):
src_layout = src_layout.value
return _cute_ir.recast_layout(
@@ -3086,15 +3385,14 @@ def recast_ptr(
loc=None,
ip=None,
) -> Pointer:
cvt_type = None
if dtype is not None:
if isinstance(dtype, _cute_ir.SparseElemType):
# use SparseElemType as dtype
pass
else:
if cvt_type is None:
if not isclass(dtype) or not issubclass(dtype, Numeric):
raise TypeError(f"dtype must be a type of Numeric, but got {dtype}")
dtype = dtype.mlir_type
cvt_type = dtype.mlir_type
dtype = cvt_type
value_type = ptr.type.value_type if dtype is None else dtype
swizzle = swizzle_.type.attribute if swizzle_ is not None else None
res_ty = _cute_ir.PtrType.get(value_type, ptr.memspace, ptr.alignment, swizzle)
@@ -3405,7 +3703,7 @@ def logical_divide(target, tiler: Tiler, *, loc=None, ip=None):
if isinstance(tiler, tuple):
tiler = _pack_tile(tiler, loc=loc, ip=ip) # type: ignore
return _op_wrapper(
partial(_cute_ir.logical_divide, tiler=tiler, loc=loc, ip=ip), target
partial(_cute_ir.logical_divide, tiler=tiler), target, loc=loc, ip=ip
)
@@ -3420,7 +3718,7 @@ def zipped_divide(target, tiler: Tiler, *, loc=None, ip=None):
if isinstance(tiler, tuple):
tiler = _pack_tile(tiler, loc=loc, ip=ip) # type: ignore
return _op_wrapper(
partial(_cute_ir.zipped_divide, tiler=tiler, loc=loc, ip=ip), target
partial(_cute_ir.zipped_divide, tiler=tiler), target, loc=loc, ip=ip
)
@@ -3435,7 +3733,7 @@ def tiled_divide(target, tiler: Tiler, *, loc=None, ip=None):
if isinstance(tiler, tuple):
tiler = _pack_tile(tiler, loc=loc, ip=ip)
return _op_wrapper(
partial(_cute_ir.tiled_divide, tiler=tiler, loc=loc, ip=ip), target
partial(_cute_ir.tiled_divide, tiler=tiler), target, loc=loc, ip=ip
)
@@ -3450,7 +3748,7 @@ def flat_divide(target, tiler: Tile, *, loc=None, ip=None):
if isinstance(tiler, tuple):
tiler = _pack_tile(tiler, loc=loc, ip=ip)
return _op_wrapper(
partial(_cute_ir.flat_divide, tiler=tiler, loc=loc, ip=ip), target
partial(_cute_ir.flat_divide, tiler=tiler), target, loc=loc, ip=ip
)
@@ -3654,7 +3952,6 @@ 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
@@ -3732,11 +4029,76 @@ def make_layout_tv(
right_inverse(layout_mn, loc=loc, ip=ip), tmp, loc=loc, ip=ip
)
tiler_mn = product_each(layout_mn.shape, loc=loc, ip=ip)
tiler_mn = product_each(layout_mn.shape_method(loc=loc, ip=ip), loc=loc, ip=ip)
return (tiler_mn, layout_tv)
@dsl_user_op
def get_nonswizzle_portion(
layout: Union[Layout, ComposedLayout], *, loc=None, ip=None
) -> Union[Layout, ComposedLayout]:
"""
Extract the non-swizzle portion from a layout.
For a simple Layout, the entire layout is considered non-swizzled and is returned as-is.
For a ComposedLayout, the inner layout (non-swizzled portion) is extracted and returned,
effectively separating the base layout from any swizzle transformation that may be applied.
:param layout: A Layout or ComposedLayout from which to extract the non-swizzle portion.
:type layout: Union[Layout, ComposedLayout]
:param loc: Optional location information for IR diagnostics.
:type loc: optional
:param ip: Optional
:type ip: optional
:returns: The non-swizzle portion of the input layout. For Layout objects, returns the layout itself.
For ComposedLayout objects, returns the outer layout component.
:rtype: Layout
:raises TypeError: If the layout is neither a Layout nor a ComposedLayout.
"""
if isinstance(layout, Layout):
return layout
elif isinstance(layout, ComposedLayout):
return layout.outer
else:
raise TypeError(f"expects a Layout or ComposedLayout, but got {type(layout)}")
@dsl_user_op
def get_swizzle_portion(
layout: Union[Layout, ComposedLayout], *, loc=None, ip=None
) -> Swizzle:
"""
Extract or create the swizzle portion from a layout.
For a simple Layout (which has no explicit swizzle), a default identity swizzle is created.
For a ComposedLayout, the outer layout is checked and returned if it is a Swizzle object.
Otherwise, a default identity swizzle is created. The default identity swizzle has parameters
(0, 4, 3), which represents a no-op swizzle transformation.
:param layout: A Layout or ComposedLayout from which to extract the swizzle portion.
:type layout: Union[Layout, ComposedLayout]
:param loc: Optional location information for IR diagnostics.
:type loc: optional
:param ip: Optional
:type ip: optional
:returns: The swizzle portion of the layout. For Layout objects or ComposedLayout objects without
a Swizzle outer component, returns a default identity swizzle (0, 4, 3). For ComposedLayout
objects with a Swizzle outer component, returns that swizzle.
:rtype: Swizzle
:raises TypeError: If the layout is neither a Layout nor a ComposedLayout.
"""
if isinstance(layout, Layout):
return make_swizzle(0, 4, 3, loc=loc, ip=ip)
elif isinstance(layout, ComposedLayout):
if isinstance(layout.inner, Swizzle):
return layout.inner
else:
return make_swizzle(0, 4, 3, loc=loc, ip=ip)
else:
raise TypeError(f"expects a Layout or ComposedLayout, but got {type(layout)}")
##############################################################################
# User defined struct
##############################################################################
@@ -3870,7 +4232,8 @@ class struct:
self._size = size
self._base = base
def data_ptr(self):
@dsl_user_op
def data_ptr(self, *, loc=None, ip=None):
"""
Returns start pointer to the data in this memory range.
@@ -3878,9 +4241,10 @@ class struct:
:raises AssertionError: If the size of the memory range is negative.
"""
assert self._size >= 0
return recast_ptr(self._base, dtype=self._dtype)
return recast_ptr(self._base, dtype=self._dtype, loc=loc, ip=ip)
def get_tensor(self, layout, swizzle=None, dtype=None):
@dsl_user_op
def get_tensor(self, layout, swizzle=None, dtype=None, *, loc=None, ip=None):
"""
Creates a tensor from the memory range.
@@ -3898,8 +4262,8 @@ class struct:
if isinstance(layout, ComposedLayout) and (swizzle is not None):
raise TypeError("incompatible layout with swizzle")
elem_type = self._dtype if dtype is None else dtype
ptr = recast_ptr(self._base, swizzle, dtype=elem_type)
res = make_tensor(ptr, layout)
ptr = recast_ptr(self._base, swizzle, dtype=elem_type, loc=loc, ip=ip)
res = make_tensor(ptr, layout, loc=loc, ip=ip)
return res
def __getitem__(self, index: int) -> Any:
@@ -4046,7 +4410,8 @@ class struct:
self._size_of = self.align_offset(offset, alignment)
# create the __init__ method for decorated struct
def __call__(self, base: Any) -> None:
@dsl_user_op
def __call__(self, base: Any, *, loc=None, ip=None) -> None:
"""
Creates a new instance of the decorated struct.
@@ -4065,7 +4430,7 @@ class struct:
if isinstance(obj, struct._AlignMeta):
obj = obj.dtype
if struct._is_scalar_type(obj):
new_obj = recast_ptr(base + off, dtype=obj)
new_obj = recast_ptr(base + off, dtype=obj, loc=loc, ip=ip)
setattr(cls, name, new_obj)
elif isinstance(obj, struct._MemRangeMeta):
new_obj = struct._MemRangeData(obj._dtype, obj._size, base + off)
python/CuTeDSL/cutlass/cute/experimental/__init__.py
Executable
+14
View File
@@ -0,0 +1,14 @@
# 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.
raise NotImplementedError(
"CuTe Experimental module is only supported on Cuda toolkit 13.1 and above!"
)
python/CuTeDSL/cutlass/cute/export/__init__.py
+29 -18
View File
@@ -11,28 +11,39 @@
from .c_header_generator import CuteCHeaderGenerator
from ...base_dsl.export import (
get_export_module,
dump_to_object as _dump_to_object,
export_to_c as _export_to_c,
)
from ...cutlass_dsl import CuTeDSL
from functools import partial as _partial
from ...cutlass_dsl.cuda_jit_executor import CudaDialectJitCompiledFunction
from .export import object_file_version as _object_file_version
from .export import CuteArgsSpecProcessor as _CuteArgsSpecProcessor
dump_to_object = _partial(
_dump_to_object,
dsl=CuTeDSL._get_dsl(),
from ...base_dsl.jit_executor import ExportProvider as _ExportProvider
from ...cutlass_dsl import CuTeDSL as _CuTeDSL
from ...cutlass_dsl.cuda_jit_executor import (
CudaDialectJitCompiledFunction as _CudaDialectJitCompiledFunction,
)
export_to_c = _partial(
_export_to_c,
dsl=CuTeDSL._get_dsl(),
from ..._mlir._mlir_libs._cutlass_ir import _mlirExecutionEngine
_CudaDialectJitCompiledFunction.export_provider = _ExportProvider(
dsl=_CuTeDSL,
arg_spec_processor=_CuteArgsSpecProcessor(),
c_header_generator=CuteCHeaderGenerator(),
use_gpu_dialect=False,
object_file_version=_object_file_version,
mlirExecutionEngine=_mlirExecutionEngine,
)
from ...base_dsl.export import ExternalBinaryModule as _ExternalBinaryModule
from ...base_dsl.export import LoadProvider as _LoadProvider
from .load import version_checker as _version_checker
from ..._mlir._mlir_libs._cutlass_ir._execution_engine import (
BinaryExecutionEngine as _BinaryExecutionEngine,
)
_ExternalBinaryModule.load_provider = _LoadProvider(
dsl=_CuTeDSL,
args_spec_processor=_CuteArgsSpecProcessor(),
version_checker=_version_checker,
execution_engine_constructor=_BinaryExecutionEngine,
jit_function_constructor=_CudaDialectJitCompiledFunction,
)
__all__ = [
"CuteCHeaderGenerator",
"get_export_module",
"dump_to_object",
"export_to_c",
]
python/CuTeDSL/cutlass/cute/export/aot_config.py
+172
View File
@@ -0,0 +1,172 @@
# 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.
"""
CLI tool to help with AOT compilation configuration.
Similar to tvm-ffi-config or llvm-config, this tool provides compiler flags
for linking against CuTe DSL runtime libraries.
Usage:
python -m cutlass.cute.export.aot_config --libdir # Returns the library directory path
python -m cutlass.cute.export.aot_config --ldflags # Returns -L flags for linking
python -m cutlass.cute.export.aot_config --libs # Returns -l flags for linking
Examples:
# Compile and link a shared library using shell substitution
g++ -shared -o kernel.so kernel.o \\
$(python -m cutlass.cute.export.aot_config --ldflags) \\
$(python -m cutlass.cute.export.aot_config --libs)
# Or using backticks
g++ -shared -o kernel.so kernel.o `python -m cutlass.cute.export.aot_config --ldflags` `python -m cutlass.cute.export.aot_config --libs`
"""
import argparse
import sys
from pathlib import Path
def get_libdir() -> str:
"""
Get the library directory path containing libcuda_dialect_runtime.so.
:return: Path to the library directory
:rtype: str
"""
from ..runtime import find_runtime_libraries
libs = find_runtime_libraries(enable_tvm_ffi=False)
if libs:
# Return the directory containing the first library found
return str(Path(libs[0]).parent)
return ""
def get_libs(enable_tvm_ffi: bool = False) -> str:
"""
Get the -l flags needed for AOT compilation linking.
Similar to `tvm-ffi-config --libs` which returns `-ltvm_ffi`,
this returns `-lcuda_dialect_runtime` (and `-ltvm_ffi` if TVM-FFI is enabled).
:param enable_tvm_ffi: Whether to include TVM-FFI library
:return: Space-separated -l flags (e.g., "-lcuda_dialect_runtime -ltvm_ffi")
:rtype: str
"""
from ..runtime import find_runtime_libraries
libs = find_runtime_libraries(enable_tvm_ffi=enable_tvm_ffi)
# Convert full paths to -l flags
# e.g., /path/to/libcuda_dialect_runtime.so -> -lcuda_dialect_runtime
flags = []
for lib in libs:
lib_path = Path(lib)
lib_name = lib_path.stem # e.g., "libcuda_dialect_runtime"
if lib_name.startswith("lib"):
lib_name = lib_name[3:]
flags.append(f"-l{lib_name}")
return " ".join(flags)
def get_lib_paths(enable_tvm_ffi: bool = False) -> list[str]:
"""
Get the full paths to runtime libraries.
:param enable_tvm_ffi: Whether to include TVM-FFI library
:return: List of full library paths
:rtype: list[str]
"""
from ..runtime import find_runtime_libraries
return find_runtime_libraries(enable_tvm_ffi=enable_tvm_ffi)
def get_ldflags() -> str:
"""
Get the -L flags for the linker.
Similar to `tvm-ffi-config --ldflags` which returns `-L<libdir>`.
:return: -L flag with library directory path
:rtype: str
"""
libdir = get_libdir()
if libdir:
return f"-L{libdir}"
return ""
def main():
parser = argparse.ArgumentParser(
description="AOT configuration helper for CuTe DSL (similar to tvm-ffi-config)",
formatter_class=argparse.RawDescriptionHelpFormatter,
epilog="""
Examples:
# Get library directory path
python -m cutlass.cute.export.aot_config --libdir
# Get -L flags for linking
python -m cutlass.cute.export.aot_config --ldflags
# Get -l flags for linking
python -m cutlass.cute.export.aot_config --libs
# Compile a shared library
g++ -shared -o kernel.so kernel.o \\
$(python -m cutlass.cute.export.aot_config --ldflags) \\
$(python -m cutlass.cute.export.aot_config --libs)
""",
)
parser.add_argument(
"--libdir",
action="store_true",
help="Print the library directory path containing runtime libraries",
)
parser.add_argument(
"--ldflags",
action="store_true",
help="Print -L flags for linking (e.g., -L/path/to/lib)",
)
parser.add_argument(
"--libs",
action="store_true",
help="Print -l flags for linking (e.g., -lcuda_dialect_runtime)",
)
parser.add_argument(
"--with-tvm-ffi",
action="store_true",
help="Include TVM-FFI library in --libs output (disabled by default)",
)
args = parser.parse_args()
if not args.libdir and not args.ldflags and not args.libs:
parser.print_help()
sys.exit(1)
enable_tvm_ffi = args.with_tvm_ffi
if args.libdir:
print(get_libdir())
if args.ldflags:
print(get_ldflags())
if args.libs:
print(get_libs(enable_tvm_ffi=enable_tvm_ffi))
if __name__ == "__main__":
main()
python/CuTeDSL/cutlass/cute/export/c_header_generator.py
+68 -43
View File
@@ -10,7 +10,7 @@
# is strictly prohibited.
from cutlass.cute.typing import NumericMeta, Integer
from cutlass.base_dsl.export import CHeaderGenerator
from cutlass.base_dsl.export import CHeaderGenerator, CHeaderArguments
from cutlass.base_dsl.dsl import is_dynamic_expression
from cutlass.base_dsl.common import DSLRuntimeError
from cutlass.base_dsl.jit_executor import ExecutionArgs
@@ -29,6 +29,22 @@ import cuda.bindings.driver as cuda
class CuteCHeaderGenerator(CHeaderGenerator):
"""This class provides a Export C Header Generator for cute c/cpp AOT support."""
includes = """
#pragma once
#include <cuda_runtime.h>
#include <cuda_fp16.h>
#include <stdio.h>
#include <stdint.h>
"""
cuda_error_check = r"""_CUDA_ERROR_CHECK(err) { \
if ((err) != cudaSuccess) { \
printf("Got Cuda Error %s: %s\n", cudaGetErrorName(err), cudaGetErrorString(err)); \
} \
}
"""
def _get_cute_algebra_type(self, arg_type: Any, arg: Any) -> str:
"""Judge if the dynamic elements of the cute algebra type are same(Int32 or Int64).
If so, generate the corresponding C type. Otherwise, refuse to generate the argument
@@ -65,55 +81,50 @@ class CuteCHeaderGenerator(CHeaderGenerator):
"""
return ""
def _generate_kernel_metadata(
def _generate_kernel_module(
self, symbol_prefix: str, kernel_info: Dict[str, List], dsl_name: str
):
"""
Generate the kernel metadata for the compiled function.
Generate the kernel module for the compiled function.
"""
kernel_metadata_struct = f"""
kernel_module_struct = f"""
typedef struct {{
CUlibrary module;
}} {symbol_prefix}_Kernel_Metadata_t;
cudaLibrary_t module;
}} {symbol_prefix}_Kernel_Module_t;
"""
kernel_metadata_load = f"""
kernel_module_load = f"""
#ifdef __cplusplus
extern "C" {{
#endif
void _mlir_{symbol_prefix}_cuda_init(void **);
void _mlir_{symbol_prefix}_cuda_load(void **);
static inline void {symbol_prefix}_Kernel_Metadata_Load({symbol_prefix}_Kernel_Metadata_t *metadata) {{
CUlibrary *libraryPtr = &(metadata->module);
int32_t ret;
void _mlir_{symbol_prefix}_cuda_load_to_device(void **);
static inline void {symbol_prefix}_Kernel_Module_Load({symbol_prefix}_Kernel_Module_t *module) {{
cudaLibrary_t *libraryPtr = &(module->module);
cudaError_t ret;
struct {{
CUlibrary **libraryPtr;
int32_t *ret;
cudaLibrary_t **libraryPtr;
cudaError_t *ret;
}} initArgs = {{&libraryPtr, &ret}};
_mlir_{symbol_prefix}_cuda_init((void **)(&initArgs));
{dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret));
{dsl_name}_CUDA_ERROR_CHECK(ret);
int32_t device_id = 0;
struct {{
CUlibrary *library;
int32_t *ret;
}} loadArgs = {{libraryPtr, &ret}};
_mlir_{symbol_prefix}_cuda_load((void **)(&loadArgs));
{dsl_name}_CUDA_ERROR_CHECK((CUresult)(ret));
CUdevice device;
{dsl_name}_CUDA_ERROR_CHECK(cuCtxGetDevice(&device));
int max_shared_memory_per_block_optin;
{dsl_name}_CUDA_ERROR_CHECK(cuDeviceGetAttribute(&max_shared_memory_per_block_optin, CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN, device));
unsigned int num_kernels;
{dsl_name}_CUDA_ERROR_CHECK(cuLibraryGetKernelCount(&num_kernels, metadata->module));
CUkernel *kernels = (CUkernel *)malloc(num_kernels * sizeof(CUkernel));
{dsl_name}_CUDA_ERROR_CHECK(cuLibraryEnumerateKernels(kernels, num_kernels, metadata->module));
for (unsigned int i = 0; i < num_kernels; i++) {{
{dsl_name}_CUDA_ERROR_CHECK(cuKernelSetAttribute(CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES, max_shared_memory_per_block_optin, kernels[i], device));
cudaLibrary_t **library;
int32_t *device_id;
cudaError_t *ret;
}} loadArgs = {{&libraryPtr, &device_id, &ret}};
int32_t device_count;
{dsl_name}_CUDA_ERROR_CHECK(cudaGetDeviceCount(&device_count));
for (int32_t i = 0; i < device_count; i++) {{
device_id = i;
_mlir_{symbol_prefix}_cuda_load_to_device((void **)(&loadArgs));
{dsl_name}_CUDA_ERROR_CHECK(ret);
}}
free(kernels);
}}
"""
kernel_metadata_unload = f"""
static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel_Metadata_t *metadata) {{
{dsl_name}_CUDA_ERROR_CHECK(cuLibraryUnload(metadata->module));
kernel_module_unload = f"""
static inline void {symbol_prefix}_Kernel_Module_Unload({symbol_prefix}_Kernel_Module_t *module) {{
{dsl_name}_CUDA_ERROR_CHECK(cudaLibraryUnload(module->module));
}}
#ifdef __cplusplus
@@ -121,7 +132,7 @@ static inline void {symbol_prefix}_Kernel_Metadata_Unload({symbol_prefix}_Kernel
#endif
"""
return kernel_metadata_struct + kernel_metadata_load + kernel_metadata_unload
return kernel_module_struct + kernel_module_load + kernel_module_unload
def _generate_arguments(
self,
@@ -174,13 +185,13 @@ typedef struct {{
elif isinstance(arg_type, NumericMeta):
arguments.append(self._generate_numeric_argument(arg_name, arg_type))
packed_args.append("&" + arg_name)
elif is_cute_algebra_type(arg_type):
elif is_cute_algebra_type(arg_type) or isinstance(arg, (tuple, list)):
c_type = self._get_cute_algebra_type(arg_type, arg)
arguments.append(f"{c_type}*{arg_name}")
for i in range(self._count_dynamic_expression(arg)):
packed_args.append("&" + arg_name + "[" + str(i) + "]")
elif isclass(arg_type) and issubclass(arg_type, cuda.CUstream):
arguments.append("CUstream " + arg_name)
arguments.append("cudaStream_t " + arg_name)
packed_args.append("&" + arg_name)
else:
raise DSLRuntimeError(
@@ -191,24 +202,38 @@ typedef struct {{
def _generate_wrapper_function(
self,
dsl_name: str,
symbol_prefix: str,
args_spec: ExecutionArgs,
function_name: str,
kernel_info: Dict[str, List],
dynamic_args: list,
dynamic_kwargs: dict,
c_header_arguments: CHeaderArguments,
):
"""
Generate the wrapper function for the compiled function which is provided to users as the entry point.
It uses the `symbol_prefix` as the function name for identification. The host/device symbols are hidden under the bytecode.
"""
# 1. Get the name of the function wrapper
wrapper_function_name = f"{symbol_prefix}_wrapper"
wrapper_function_name = f"{dsl_name.lower()}_{symbol_prefix}_wrapper"
capi_function_name = f"_mlir_{symbol_prefix}__mlir_ciface_{function_name}"
# 2. Generate the signature of the wrapper function
arguments, packed_args, declarations = self._generate_arguments(
symbol_prefix, args_spec, dynamic_args, dynamic_kwargs
)
if c_header_arguments.error_msg is not None:
raise DSLRuntimeError(
f"Error generating c header arguments: {c_header_arguments.error_msg}"
)
arguments = [
arg.replace(c_header_arguments.dummy_prefix_name, symbol_prefix)
for arg in c_header_arguments.arguments
]
packed_args = [
arg.replace(c_header_arguments.dummy_prefix_name, symbol_prefix)
for arg in c_header_arguments.packed_args
]
declarations = [
declaration.replace(c_header_arguments.dummy_prefix_name, symbol_prefix)
for declaration in c_header_arguments.declarations
]
# 3. Get the return type of the wrapper function.
# Note that this requires the return type to be properly annotated in python.
return_type = args_spec.args_spec.annotations.get("return", None)
@@ -227,7 +252,7 @@ extern "C"
#endif
void {capi_function_name}(void **args, int32_t num_args);
static inline {return_type} {wrapper_function_name}({symbol_prefix}_Kernel_Metadata_t *metadata, {", ".join(arguments)}) {{
static inline {return_type} {wrapper_function_name}({symbol_prefix}_Kernel_Module_t *module, {", ".join(arguments)}) {{
{return_type} ret;
void *args[{len(packed_args) + 1}] = {{
{", ".join(packed_args)},
python/CuTeDSL/cutlass/cute/export/export.py
+55
View File
@@ -0,0 +1,55 @@
# 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.
import os
import pickle
import copy
from ..typing import IntTuple, Shape, Stride, Coord, Tile
from inspect import FullArgSpec
from cutlass.base_dsl.export import (
ArgsSpecProcessor,
)
cute_algebra_types_dump = {
IntTuple: "IntTuple",
Shape: "Shape",
Stride: "Stride",
Coord: "Coord",
Tile: "Tile",
}
cute_algebra_types_load = {
"IntTuple": IntTuple,
"Shape": Shape,
"Stride": Stride,
"Coord": Coord,
"Tile": Tile,
}
class CuteArgsSpecProcessor(ArgsSpecProcessor):
def dumps(self, args_spec: FullArgSpec) -> bytes:
new_args_spec = copy.deepcopy(args_spec)
for arg, arg_type in new_args_spec.annotations.items():
if arg_type in cute_algebra_types_dump.keys():
new_args_spec.annotations[arg] = cute_algebra_types_dump[arg_type]
return pickle.dumps(new_args_spec)
def loads(self, args_spec_bytes: bytes) -> FullArgSpec:
args_spec = pickle.loads(args_spec_bytes)
for arg, arg_type in args_spec.annotations.items():
if arg_type in cute_algebra_types_load.keys():
args_spec.annotations[arg] = cute_algebra_types_load[arg_type]
return args_spec
# This is the version of the object file. It is used to check the version of the object file is compatible with the current dsl version or not.
object_file_version = "1.1"
python/CuTeDSL/cutlass/cute/export/load.py
+19
View File
@@ -0,0 +1,19 @@
# 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.base_dsl.common import DSLRuntimeError
def version_checker(version: str) -> bool:
"""Check the version of the object file is compatible with the current dsl version or not."""
if version not in ["1.0", "1.1"]:
raise DSLRuntimeError("Incompatible version: " + version)
return True
python/CuTeDSL/cutlass/cute/nvgpu/__init__.py
+10 -3
View File
@@ -17,10 +17,17 @@ from . import tcgen05
from .common import *
from .helpers import *
from . import common
from . import helpers
# __all__ is required here for documentation generation
__all__ = [
"OpError",
"MmaUniversalOp",
"CopyUniversalOp",
*common.__all__,
*helpers.__all__,
# submodules With namespace
"warp",
"cpasync",
"warpgroup",
"tcgen05",
]
python/CuTeDSL/cutlass/cute/nvgpu/common.py
+52 -12
View File
@@ -20,9 +20,19 @@ from cutlass._mlir import ir
from .. import atom
from ..typing import Float16, Float32, Float64, Numeric
from cutlass import cute
__all__ = [
"OpError",
"MmaUniversalOp",
"MmaUniversalTrait",
"CopyUniversalOp",
"CopyUniversalTrait",
"MemoryOrder",
"MemoryScope",
"CacheEvictionPriority",
]
class OpError(DSLBaseError):
"""
An exception class for Op construction errors.
@@ -83,7 +93,7 @@ class MmaUniversalOp(atom.MmaOp):
self.abacc_dtype.mlir_type,
self.abacc_dtype.mlir_type,
)
return MmaUniversalTrait(cute.make_atom(atom_ty, loc=loc, ip=ip))
return MmaUniversalTrait(atom.make_atom(atom_ty, loc=loc, ip=ip))
def _verify_fragment_A(self, input, *, loc=None, ip=None):
pass
@@ -140,6 +150,23 @@ class MemoryScope(enum.Enum):
return self.value
class CacheEvictionPriority(enum.Enum):
EVICT_NORMAL = _cute_ir.CacheEvictionPriority.EVICT_NORMAL
EVICT_FIRST = _cute_ir.CacheEvictionPriority.EVICT_FIRST
EVICT_LAST = _cute_ir.CacheEvictionPriority.EVICT_LAST
EVICT_UNCHANGED = _cute_ir.CacheEvictionPriority.EVICT_UNCHANGED
NO_ALLOCATE = _cute_ir.CacheEvictionPriority.NO_ALLOCATE
def __str__(self) -> str:
return f"{self.__class__.__name__}.{self.name}"
def __repr__(self) -> str:
return f"<{self.__class__.__name__}.{self.name}>"
def _to_ir(self) -> _cute_ir.CacheEvictionPriority:
return self.value
@dataclass(frozen=True)
class CopyUniversalOp(atom.CopyOp):
"""
@@ -150,7 +177,12 @@ class CopyUniversalOp(atom.CopyOp):
.. code-block:: python
op = cute.nvgpu.CopyUniversalOp()
atom = cute.make_copy_atom(op, tensor_dtype, num_bits_per_copy=64)
atom = cute.make_copy_atom(
op,
tensor_dtype,
num_bits_per_copy=64,
l1c_evict_priority=cute.nvgpu.CacheEvictionPriority.EVICT_NORMAL
)
- ``tensor_dtype`` is the data type used to build the reference TV Layout (either the source \
or the destination TV Layout) in unit of tensor elements and is used for partitioning by \
@@ -158,6 +190,12 @@ class CopyUniversalOp(atom.CopyOp):
- ``num_bits_per_copy`` is a kw argument specifying the number of bits to copy per Atom \
execution. This can be larger than the width of the above data type. When not provided, \
the compiler will do a best effort at auto-vectorizing.
- ``l1c_evict_priority`` is a kw argument specifying the L1 cache eviction priority hint for \
the copy operation. Defaults to ``EVICT_NORMAL`` if not provided.
- ``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:
@@ -167,25 +205,27 @@ class CopyUniversalOp(atom.CopyOp):
self,
copy_internal_type: Type[Numeric],
*,
num_bits_per_copy: int = 0,
memory_order: MemoryOrder = MemoryOrder.WEAK,
memory_scope: MemoryScope = MemoryScope.CTA,
l1c_evict_priority: CacheEvictionPriority = CacheEvictionPriority.EVICT_NORMAL,
invariant: bool = False,
loc=None,
ip=None,
**kwargs,
) -> "CopyUniversalTrait":
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
memory_order = kwargs.get("memory_order", MemoryOrder.WEAK)
memory_scope = kwargs.get("memory_scope", MemoryScope.CTA)
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
if not isinstance(num_bits_per_copy, int) or num_bits_per_copy < 0:
raise ValueError(
"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__}"
f"'num_bits_per_copy' must be a non-negative int when creating a copy Atom for {self.__class__.__name__!r}"
)
ty = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
atom_type = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
copy_internal_type.mlir_type,
num_bits_per_copy,
memory_order._to_ir(),
memory_scope._to_ir(),
l1c_evict_priority._to_ir(),
invariant,
)
return CopyUniversalTrait(cute.make_atom(ty, loc=loc, ip=ip))
return CopyUniversalTrait(atom.make_atom(atom_type, loc=loc, ip=ip))
class CopyUniversalTrait(atom.Trait):
python/CuTeDSL/cutlass/cute/nvgpu/cpasync/copy.py
+265 -13
View File
@@ -13,15 +13,14 @@ import enum
from dataclasses import dataclass
from typing import Optional, Type
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp
from cutlass._mlir import ir
from ...atom import CopyOp, Trait
from ...tensor import ReductionOp
from ...atom import CopyOp, Trait, make_atom
from ...typing import Int16, Int64, Pointer, Integer, Numeric
from ..common import OpError
from ..tcgen05.mma import CtaGroup
@@ -29,7 +28,7 @@ from ..tcgen05.mma import CtaGroup
####################################################################################################
#
# Aynchronous copies
# Asynchronous copies
#
####################################################################################################
@@ -96,7 +95,7 @@ class CopyG2SOp(CopyOp):
ty = _cute_nvgpu_ir.CopyAtomSIMTAsyncCopyType.get(
copy_internal_type.mlir_type, self.cache_mode._to_ir(), num_bits_per_copy
)
return CopyG2STrait(cute.make_atom(ty, loc=loc, ip=ip))
return CopyG2STrait(make_atom(ty, loc=loc, ip=ip))
class CopyG2STrait(Trait):
@@ -121,7 +120,16 @@ class TmaCopyOp(CopyOp):
Base class for all TMA copy operations.
"""
pass
def __init__(self, smem_layout: Optional[ir.Value] = None) -> None:
self.smem_layout = smem_layout
def __extract_mlir_values__(self):
return [self.smem_layout]
def __new_from_mlir_values__(self, values):
res = self.__class__()
res.smem_layout = values[0]
return res
#
@@ -129,7 +137,7 @@ class TmaCopyOp(CopyOp):
#
@dataclass(frozen=True)
@dataclass
class CopyBulkTensorTileG2SOp(TmaCopyOp):
"""
Bulk tensor asynchrnous GMEM to SMEM Copy Operation using the TMA unit.
@@ -246,7 +254,7 @@ class CopyBulkTensorTileG2STrait(Trait):
#
@dataclass(frozen=True)
@dataclass
class CopyBulkTensorTileG2SMulticastOp(TmaCopyOp):
"""
Bulk tensor asynchrnous multicast GMEM to SMEM Copy Operation using the TMA unit.
@@ -375,7 +383,7 @@ class CopyBulkTensorTileG2SMulticastTrait(Trait):
#
@dataclass(frozen=True)
@dataclass
class CopyBulkTensorTileS2GOp(TmaCopyOp):
"""
Bulk tensor asynchronous SMEM to GMEM Copy Operation using the TMA unit.
@@ -449,7 +457,11 @@ class CopyBulkTensorTileS2GTrait(Trait):
pass
@dataclass(frozen=True)
class CopyBulkTensorTileS2GTrait(Trait):
pass
@dataclass
class CopyReduceBulkTensorTileS2GOp(TmaCopyOp):
"""
Bulk tensor asynchronous SMEM to GMEM Reduction Operation using the TMA unit.
@@ -585,7 +597,7 @@ class CopyBulkG2SOp(CopyOp):
ty = _cute_nvgpu_ir.CopyAtomBulkCopyG2SType.get(
copy_internal_type.mlir_type, num_bits_per_copy, False
)
return CopyBulkG2STrait(cute.make_atom(ty, loc=loc, ip=ip))
return CopyBulkG2STrait(make_atom(ty, loc=loc, ip=ip))
class CopyBulkG2STrait(Trait):
@@ -670,7 +682,7 @@ class CopyBulkG2SMulticastOp(CopyOp):
ty = _cute_nvgpu_ir.CopyAtomBulkCopyG2SType.get(
copy_internal_type.mlir_type, num_bits_per_copy, True
)
return CopyBulkG2SMulticastTrait(cute.make_atom(ty, loc=loc, ip=ip))
return CopyBulkG2SMulticastTrait(make_atom(ty, loc=loc, ip=ip))
class CopyBulkG2SMulticastTrait(Trait):
@@ -764,8 +776,248 @@ class CopyBulkS2GOp(CopyOp):
ty = _cute_nvgpu_ir.CopyAtomBulkCopyS2GType.get(
copy_internal_type.mlir_type, num_bits_per_copy, False
)
return CopyBulkS2GTrait(cute.make_atom(ty, loc=loc, ip=ip))
return CopyBulkS2GTrait(make_atom(ty, loc=loc, ip=ip))
class CopyBulkS2GTrait(Trait):
pass
#
# Bulk SMEM -> GMEM mask copies
#
@dataclass(frozen=True)
class CopyBulkS2GByteMaskOp(CopyOp):
"""
Bulk copy asynchrnous SMEM to GMEM Copy Operation with mask.
The i-th bit in the 16-bit wide byteMask operand specifies whether
the i-th byte of each 16-byte wide chunk of source data is copied to the destination.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk>`__.
"""
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_100:
raise OpError(
self,
f"expects arch to be at least {Arch.sm_100.name}, but got {arch.name}",
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
)
def __str__(self) -> str:
res = "cp.async SMEM -> GMEM bulk copy Operation"
return res
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "CopyBulkS2GByteMaskTrait":
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
raise ValueError(
"expects a 'num_bits_per_copy' kw argument of type int that is positive "
f"when creating a copy Atom for {self.__class__.__name__}"
)
ty = _cute_nvgpu_ir.CopyAtomBulkCopyS2GType.get(
copy_internal_type.mlir_type, num_bits_per_copy, True
)
return CopyBulkS2GByteMaskTrait(make_atom(ty, loc=loc, ip=ip))
class CopyBulkS2GByteMaskTrait(Trait):
def unpack(
self,
*,
loc=None,
ip=None,
byte_mask=None,
**kwargs,
):
"""
Custom implementation of unpack for bulk copy store with mask.
The bulk store with mask requires `byte_mask` keyword argument to be provided when
using `copy`. Any other kw arguments will be ignored instead of triggering an error.
"""
if not isinstance(byte_mask, Integer):
raise ValueError(
"expects a byte mask to be provided via the byte_mask kw argument"
)
# Support for .cp_mask qualifier requires sm_100 or higher.
attr_str = f"#cute_nvgpu.atom_copy_field_bulks2g<{TMA_BYTE_MASK_FIELD_NAME}>"
attr = ir.Attribute.parse(attr_str)
val = _cute_nvgpu_ir.atom_set_value(
self.value,
attr,
Int16(byte_mask).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return val
#
# Bulk SMEM CTA to Cluster copies
#
@dataclass(frozen=True)
class CopyBulkS2SOp(CopyOp):
"""
Bulk copy asynchrnous SMEM CTA to Cluster Copy Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk>`__.
"""
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
)
def __str__(self) -> str:
res = "cp.async CTA -> Cluster bulk copy Operation"
return res
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "CopyBulkS2STrait":
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
raise ValueError(
"expects a 'num_bits_per_copy' kw argument of type int that is positive "
f"when creating a copy Atom for {self.__class__.__name__}"
)
ty = _cute_nvgpu_ir.CopyAtomBulkCopyS2SType.get(
copy_internal_type.mlir_type, num_bits_per_copy
)
return CopyBulkS2STrait(make_atom(ty, loc=loc, ip=ip))
class CopyBulkS2STrait(Trait):
def unpack(
self,
*,
loc=None,
ip=None,
mbar_ptr: Optional[Pointer] = None,
cta_rank: Optional[Integer] = None,
**kwargs,
):
"""
Custom implementation of unpack for bulk copy cta to cluster.
The bulk cta to cluster copy requires a `mbar_ptr` and `cta_rank` keyword argument to be provided
when using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error.
"""
if not isinstance(mbar_ptr, Pointer):
raise ValueError(
"expects a pointer to an mbarrier to be provided via the mbar_ptr kw argument"
)
if not isinstance(cta_rank, Integer):
raise ValueError(
"expects a cta rank of int32 to be provided via the cta_rank kw argument"
)
attr_str = f"#cute_nvgpu.atom_copy_field_bulks2s<{TMA_MBAR_PTR_FIELD_NAME}>"
attr = ir.Attribute.parse(attr_str)
val = _cute_nvgpu_ir.atom_set_value(
self.value, attr, mbar_ptr.value, loc=loc, ip=ip
)
attr_str = f"#cute_nvgpu.atom_copy_field_bulks2s<{TMA_CTA_RANK_FIELD_NAME}>"
attr = ir.Attribute.parse(attr_str)
val = _cute_nvgpu_ir.atom_set_value(
val, attr, Int32(cta_rank).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
)
return val
####################################################################################################
#
# Aynchronous distributed shared memory stores
#
####################################################################################################
MBAR_PTR_FIELD_NAME = "mbar_ptr"
@dataclass(frozen=True)
class CopyDsmemStoreOp(CopyOp):
"""
Asynchronous Store operation to DSMEM with explicit synchronization.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/index.html#data-movement-and-conversion-instructions-st-async>`__.
"""
def __post_init__(self) -> None:
# Arch verification
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
if not arch >= Arch.sm_90:
raise OpError(
self,
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
)
def __str__(self) -> str:
res = "st.async RMEM -> DSMEM copy Operation"
return res
def _make_trait(
self,
copy_internal_type: Type[Numeric],
*,
loc=None,
ip=None,
**kwargs,
) -> "CopyDsmemStoreTrait":
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
raise ValueError(
"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__}"
)
ty = _cute_nvgpu_ir.CopyAtomDsmemStoreType.get(
copy_internal_type.mlir_type, num_bits_per_copy
)
return CopyDsmemStoreTrait(make_atom(ty, loc=loc, ip=ip))
class CopyDsmemStoreTrait(Trait):
def unpack(
self,
*,
loc=None,
ip=None,
mbar_ptr: Optional[Pointer] = None,
**kwargs,
):
"""
Custom implementation of unpack for dsmem async copy.
The dsmem async copy requires `mbar_ptr` keyword argument to be provided when using `cute.copy`.
Any other kw arguments will be ignored instead of triggering an error.
"""
if not isinstance(mbar_ptr, Pointer):
raise ValueError(
"expects a pointer to an mbarrier to be provided via the mbar_ptr kw argument",
)
attr_str = f"#cute_nvgpu.atom_copy_field_dsmem_store<{MBAR_PTR_FIELD_NAME}>"
attr = ir.Attribute.parse(attr_str)
val = _cute_nvgpu_ir.atom_set_value(
self.value,
attr,
mbar_ptr.value,
loc=loc,
ip=ip,
)
return val
python/CuTeDSL/cutlass/cute/nvgpu/cpasync/helpers.py
+48 -26
View File
@@ -10,6 +10,7 @@
# is strictly prohibited.
from typing import Optional, Tuple, Type, Union
from typing_extensions import deprecated
from cutlass.cutlass_dsl import dsl_user_op
@@ -39,15 +40,16 @@ from .copy import (
CopyReduceBulkTensorTileS2GNonExecTrait,
)
TMAOp = Union[
CopyBulkTensorTileG2SOp,
CopyBulkTensorTileG2SMulticastOp,
CopyBulkTensorTileS2GOp,
CopyReduceBulkTensorTileS2GOp,
]
@dsl_user_op
def make_tiled_tma_atom(
op: Union[
CopyBulkTensorTileG2SOp,
CopyBulkTensorTileG2SMulticastOp,
CopyBulkTensorTileS2GOp,
CopyReduceBulkTensorTileS2GOp,
],
op: TMAOp,
gmem_tensor: Tensor,
smem_layout: Union[Layout, ComposedLayout],
cta_tiler: Tiler,
@@ -69,37 +71,30 @@ def make_tiled_tma_atom(
this function figures out the bulk tensor asynchronous copy instruction to use with the maximum
"TMA vector length" to copy tiles of the GMEM tensor to/from an SMEM buffer with the provided
layout and consistent with the provided Tiler.
layout while maintaining consistency with the provided Tiler.
This function returns two results:
1. the Copy Atom
2. the so-called TMA tensor used to map logical coordinates of the GMEM tensor to coordinates \
that the TMA unit can consume. TMA tensors have so-called basis stride elements so that the \
associated layout can output coordinates. Otherwise, TMA tensors can be partitioned \
similarly to any other CuTe tensors using the algebra.
2. a TMA tensor that maps logical coordinates of the GMEM tensor to coordinates consumed by the \
TMA unit. TMA tensors contain basis stride elements that enable their associated layout to \
compute coordinates. Like other CuTe tensors, TMA tensors can be partitioned.
:param op: The Copy Operation to construct an Atom for
:type op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileS2GOp, CopyReduceBulkTensorTileS2GOp]
:param op: The TMA Copy Operation to construct an 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 for
:param smem_layout: The SMEM layout to construct the Copy Atom
:type smem_layout: Union[Layout, ComposedLayout]
:param cta_tiler: The CTA Tiler to use
:type cta_tiler: Tiler
:param num_multicast: The multicast factor
:type num_multicast: int
:param internal_type: An optional parameter for the internal data type to use when the actual data type is not supported by the TMA unit
:param internal_type: Optional internal data type to use when the tensor data type is not supported by the TMA unit
:type internal_type: Type[Numeric]
:return: A Copy Atom for this Operation and the associated TMA tensor
:return: A TMA Copy Atom associated with the TMA tensor
:rtype: Tuple[atom.CopyAtom, Tensor]
"""
if internal_type is not None:
if not isinstance(internal_type, NumericMeta):
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
internal_type = internal_type.mlir_type
cta_v_map = core.composition(
core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip),
cta_tiler,
@@ -110,6 +105,26 @@ 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
tma_format = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack)
)
if isinstance(op, CopyBulkTensorTileG2SOp):
if num_multicast != 1:
raise ValueError(
@@ -122,7 +137,7 @@ def make_tiled_tma_atom(
cta_v_map,
op._to_ir(),
num_multicast=num_multicast,
internal_type=internal_type,
tma_format=tma_format,
loc=loc,
ip=ip,
)
@@ -139,7 +154,7 @@ def make_tiled_tma_atom(
cta_v_map,
op._to_ir(),
num_multicast=num_multicast,
internal_type=internal_type,
tma_format=tma_format,
loc=loc,
ip=ip,
)
@@ -152,7 +167,7 @@ def make_tiled_tma_atom(
gmem_tensor.value,
smem_layout,
cta_v_map,
internal_type=internal_type,
tma_format=tma_format,
loc=loc,
ip=ip,
)
@@ -163,7 +178,7 @@ def make_tiled_tma_atom(
smem_layout,
cta_v_map,
op._to_ir(),
internal_type=internal_type,
tma_format=tma_format,
loc=loc,
ip=ip,
)
@@ -314,6 +329,8 @@ def fence_tma_desc_acquire(
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@@ -342,6 +359,8 @@ def cp_fence_tma_desc_release(
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
loc=loc,
ip=ip,
)
@@ -358,10 +377,13 @@ def fence_tma_desc_release(*, loc=None, ip=None) -> None:
has_side_effects=True,
is_align_stack=False,
asm_dialect=llvm.AsmDialect.AD_ATT,
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
python/CuTeDSL/cutlass/cute/nvgpu/helpers.py
+47 -16
View File
@@ -26,6 +26,11 @@ from .cpasync.copy import (
)
__all__ = [
"make_tiled_tma_atom_A",
"make_tiled_tma_atom_B",
]
####################################################################################################
#
# TMA creation helpers for tcgen05 MMAs
@@ -91,10 +96,6 @@ def make_tiled_tma_atom_A(
"""
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}")
internal_type = internal_type.mlir_type
check_type_in(
op,
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
@@ -102,6 +103,13 @@ def make_tiled_tma_atom_A(
"make_tiled_tma_atom_A",
)
# 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
)
ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
mma_tiler_mk = (mma_tiler_mnk[0], *mma_tiler_mnk[2:])
g_tile = core.composition(ident, mma_tiler_mk, loc=loc, ip=ip)
@@ -123,6 +131,19 @@ def make_tiled_tma_atom_A(
if isinstance(smem_layout, core._ComposedLayout):
smem_layout = smem_layout.value
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
tma_format = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack)
)
# res[0] = the IR Value for the non-executable atom instance
# res[1] = the IR Value for the associated TMA tensor
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load(
@@ -131,7 +152,7 @@ def make_tiled_tma_atom_A(
cta_v_map,
op._to_ir(),
num_multicast=num_multicast,
internal_type=internal_type,
tma_format=tma_format,
loc=loc,
ip=ip,
)
@@ -203,10 +224,6 @@ def make_tiled_tma_atom_B(
"""
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}")
internal_type = internal_type.mlir_type
check_type_in(
op,
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
@@ -214,6 +231,13 @@ def make_tiled_tma_atom_B(
"make_tiled_tma_atom_B",
)
# 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
)
ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
mma_tiler_nk = (mma_tiler_mnk[1], *mma_tiler_mnk[2:])
g_tile = core.composition(ident, mma_tiler_nk, loc=loc, ip=ip)
@@ -235,6 +259,19 @@ def make_tiled_tma_atom_B(
if isinstance(smem_layout, core._ComposedLayout):
smem_layout = smem_layout.value
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
tma_format = _cute_nvgpu_ir.TmaDataFormat(
_cute_nvgpu_ir.get_default_tma_format(internal_mlir_type, use_unpack)
)
# res[0] = the IR Value for the non-executable atom instance
# res[1] = the IR Value for the associated TMA tensor
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load(
@@ -243,7 +280,7 @@ def make_tiled_tma_atom_B(
cta_v_map,
op._to_ir(),
num_multicast=num_multicast,
internal_type=internal_type,
tma_format=tma_format,
loc=loc,
ip=ip,
)
@@ -255,9 +292,3 @@ def make_tiled_tma_atom_B(
atom.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
res[1],
)
__all__ = [
"make_tiled_tma_atom_A",
"make_tiled_tma_atom_B",
]
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py
+17 -18
View File
@@ -13,7 +13,6 @@ import enum
from dataclasses import dataclass
from typing import Type
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL
@@ -21,7 +20,7 @@ import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
from ..common import OpError
from ...atom import CopyOp, Trait
from ...atom import CopyOp, Trait, make_atom
from ...typing import Numeric
from .mma import CtaGroup
@@ -188,7 +187,7 @@ class Ld16x64bOp(_LdBase):
self.repeat.value,
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
)
return Ld16x64bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Ld16x64bTrait(make_atom(ty, loc=loc, ip=ip))
class Ld16x64bTrait(Trait):
@@ -246,7 +245,7 @@ class Ld16x128bOp(_LdBase):
self.repeat.value,
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
)
return Ld16x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Ld16x128bTrait(make_atom(ty, loc=loc, ip=ip))
class Ld16x128bTrait(Trait):
@@ -304,7 +303,7 @@ class Ld16x256bOp(_LdBase):
self.repeat.value,
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
)
return Ld16x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Ld16x256bTrait(make_atom(ty, loc=loc, ip=ip))
class Ld16x256bTrait(Trait):
@@ -344,7 +343,7 @@ class Ld16x32bx2Op(_LdBase):
self.repeat.value,
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
)
return Ld16x32bx2Trait(cute.make_atom(ty, loc=loc, ip=ip))
return Ld16x32bx2Trait(make_atom(ty, loc=loc, ip=ip))
class Ld16x32bx2Trait(Trait):
@@ -384,7 +383,7 @@ class Ld32x32bOp(_LdBase):
self.repeat.value,
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
)
return Ld32x32bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Ld32x32bTrait(make_atom(ty, loc=loc, ip=ip))
class Ld32x32bTrait(Trait):
@@ -478,7 +477,7 @@ class St16x64bOp(_StBase):
self.repeat.value,
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
)
return St16x64bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return St16x64bTrait(make_atom(ty, loc=loc, ip=ip))
class St16x64bTrait(Trait):
@@ -513,7 +512,7 @@ class St16x128bOp(_StBase):
self.repeat.value,
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
)
return St16x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return St16x128bTrait(make_atom(ty, loc=loc, ip=ip))
class St16x128bTrait(Trait):
@@ -548,7 +547,7 @@ class St16x256bOp(_StBase):
self.repeat.value,
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
)
return St16x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return St16x256bTrait(make_atom(ty, loc=loc, ip=ip))
class St16x256bTrait(Trait):
@@ -574,7 +573,7 @@ class St16x32bx2Op(_StBase):
self.repeat.value,
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
)
return St16x32bx2Trait(cute.make_atom(ty, loc=loc, ip=ip))
return St16x32bx2Trait(make_atom(ty, loc=loc, ip=ip))
class St16x32bx2Trait(Trait):
@@ -600,7 +599,7 @@ class St32x32bOp(_StBase):
self.repeat.value,
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
)
return St32x32bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return St32x32bTrait(make_atom(ty, loc=loc, ip=ip))
class St32x32bTrait(Trait):
@@ -681,7 +680,7 @@ class Cp128x256bOp(_S2TCopyBase):
self.cta_group.value,
_cute_nvgpu_ir.CopyS2TBroadcast.none,
)
return Cp128x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Cp128x256bTrait(make_atom(ty, loc=loc, ip=ip))
class Cp128x256bTrait(Trait):
@@ -707,7 +706,7 @@ class Cp128x128bOp(_S2TCopyBase):
self.cta_group.value,
_cute_nvgpu_ir.CopyS2TBroadcast.none,
)
return Cp128x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Cp128x128bTrait(make_atom(ty, loc=loc, ip=ip))
class Cp128x128bTrait(Trait):
@@ -733,7 +732,7 @@ class Cp4x256bOp(_S2TCopyBase):
self.cta_group.value,
_cute_nvgpu_ir.CopyS2TBroadcast.none,
)
return Cp4x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Cp4x256bTrait(make_atom(ty, loc=loc, ip=ip))
class Cp4x256bTrait(Trait):
@@ -759,7 +758,7 @@ class Cp4x32x128bOp(_S2TCopyBase):
self.cta_group.value,
_cute_nvgpu_ir.CopyS2TBroadcast.x4,
)
return Cp4x32x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return Cp4x32x128bTrait(make_atom(ty, loc=loc, ip=ip))
class Cp4x32x128bTrait(Trait):
@@ -785,7 +784,7 @@ class Cp2x64x128b0213Op(_S2TCopyBase):
self.cta_group.value,
_cute_nvgpu_ir.CopyS2TBroadcast.lw_0213,
)
return Cp2x64x128b0213Trait(cute.make_atom(ty, loc=loc, ip=ip))
return Cp2x64x128b0213Trait(make_atom(ty, loc=loc, ip=ip))
class Cp2x64x128b0213Trait(Trait):
@@ -811,7 +810,7 @@ class Cp2x64x128b0123Op(_S2TCopyBase):
self.cta_group.value,
_cute_nvgpu_ir.CopyS2TBroadcast.lw_0123,
)
return Cp2x64x128b0123Trait(cute.make_atom(ty, loc=loc, ip=ip))
return Cp2x64x128b0123Trait(make_atom(ty, loc=loc, ip=ip))
class Cp2x64x128b0123Trait(Trait):
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/helpers.py
+6 -16
View File
@@ -102,27 +102,17 @@ def make_smem_layout_atom(
SmemLayoutAtomKind.MN_SW128_32B,
):
# M/N-major layout
return core.make_composed_layout(
sw,
0,
core.make_layout(
(num_contiguous_elems, 8), stride=(1, num_contiguous_elems)
),
loc=loc,
ip=ip,
outer = core.make_layout(
(num_contiguous_elems, 8), stride=(1, num_contiguous_elems), loc=loc, ip=ip
)
else:
# K-major layout
return core.make_composed_layout(
sw,
0,
core.make_layout(
(8, num_contiguous_elems), stride=(num_contiguous_elems, 1)
),
loc=loc,
ip=ip,
outer = 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(
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/mma.py
+13 -274
View File
@@ -13,7 +13,6 @@ import enum
from dataclasses import dataclass
from typing import Type, Any
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL, T
@@ -24,9 +23,9 @@ from cutlass._mlir import ir
from ..common import OpError
from ... import core, atom
from ...core import _pack_shape, rank, depth
from ...tensor import _Tensor
from ...typing import (
Shape,
Tensor,
Float4E2M1FN,
Float8E8M0FNU,
Float8E5M2,
@@ -43,9 +42,7 @@ from ...typing import (
AddressSpace,
Pointer,
)
from ...atom import Trait
from ..warp.mma import SparseMetadataFormat
from ...atom import Trait, make_atom
####################################################################################################
@@ -242,7 +239,7 @@ class MmaOp(Tcgen05MmaOp):
+ f"\n Instruction shape MNK = {self.shape_mnk}"
)
def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
@@ -254,7 +251,7 @@ class MmaOp(Tcgen05MmaOp):
)
return True
def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
@@ -388,7 +385,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp):
+ f"\n Instruction shape MNK = {self.shape_mnk}"
)
def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
@@ -400,7 +397,7 @@ class BlockScaledMmaOp(Tcgen05MmaOp):
)
return True
def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
@@ -452,161 +449,6 @@ class BlockScaledMmaTraits(Trait):
)
# Base class for all tcgen05 Sparse MMA Ops with syntax `tcgen05.mma.cta_group.kind.sparse` used to factor out some internal code
@dataclass(frozen=True)
class SparseMmaOp(Tcgen05MmaOp):
a_dtype: Type[Numeric]
b_dtype: Type[Numeric]
acc_dtype: Type[Numeric]
shape_mnk: Shape
cta_group: CtaGroup
a_src: OperandSource
a_major_mode: OperandMajorMode
b_major_mode: OperandMajorMode
sparse_metadata_format: SparseMetadataFormat
admissible_archs = Arch.filter(
lambda arch: arch.is_family_of(Arch.sm_100f) or arch.is_family_of(Arch.sm_110f)
)
def __post_init__(self) -> None:
# Verify arch
arch = BaseDSL._get_dsl().get_arch_enum()
if arch not in self.admissible_archs:
raise OpError(
self,
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
)
# Verify that the user provided enum values
if not isinstance(self.cta_group, CtaGroup):
raise OpError(
self,
"expects the 'cta_group' Op parameter to be a tcgen05.CtaGroup instance",
)
if not isinstance(self.a_src, OperandSource):
raise OpError(
self,
"expects the 'a_src' Op parameter to be a tcgen05.OperandSource instance",
)
if not isinstance(self.a_major_mode, OperandMajorMode):
raise OpError(
self,
"expects the 'a_major_mode' Op parameter to be a tcgen05.OperandMajorMode instance",
)
if not isinstance(self.b_major_mode, OperandMajorMode):
raise OpError(
self,
"expects the 'b_major_mode' Op parameter to be a tcgen05.OperandMajorMode instance",
)
if not isinstance(self.sparse_metadata_format, SparseMetadataFormat):
raise OpError(
self,
"expects the 'sparse_metadata_format' Op parameter to be a tcgen05.SparseMetadataFormat instance",
)
# Verify the instruction shape
if (rank(self.shape_mnk) not in [2, 3]) or (depth(self.shape_mnk) != 1):
raise OpError(
self,
f"expected a flat rank 2 or 3 tuple for the 'shape_mnk' Op parameter, "
f"but got {self.shape_mnk}",
)
m, n = self.shape_mnk[0], self.shape_mnk[1]
# For sparse MMA, the shape validation follows the same rules as dense MMA
# but the K dimension is typically doubled in the derived classes
if self.cta_group == CtaGroup.ONE:
if m not in [64, 128]:
raise OpError(self, f"expects the M-mode to be 64 or 128, but got {m}")
if m == 64:
if (n < 8) or (n > 256) or (n % 8 != 0):
raise OpError(
self,
f"expects the N-mode to satisfy 8 <= N <= 256 and N % 8 == 0, but got {n}",
)
elif m == 128:
if (n < 16) or (n > 256) or (n % 16 != 0):
raise OpError(
self,
f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}",
)
else:
if m not in [128, 256]:
raise OpError(self, f"expects the M-mode to be 128 or 256, but got {m}")
if (n < 32) or (n > 256) or (n % 32 != 0):
raise OpError(
self,
f"expects the N-mode to satisfy 32 <= N <= 256 and N % 32 == 0, but got {n}",
)
def __str__(self) -> str:
return (
self.__class__.descriptive_name # type: ignore
+ f"\n A data type = {self.a_dtype}"
+ f"\n B data type = {self.b_dtype}"
+ f"\n Accumulator data type = {self.acc_dtype}"
+ f"\n CTA group = {self.cta_group}"
+ f"\n A source location = {self.a_src}"
+ f"\n A major mode = {self.a_major_mode}"
+ f"\n B major mode = {self.b_major_mode}"
+ f"\n Instruction shape MNK = {self.shape_mnk}"
+ f"\n Sparse metadata format = {self.sparse_metadata_format}"
)
def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
raise OpError(
self,
f"Expected affine layout for {self._make_trait()}'s operand A, "
f"but got composed layout instead: {input.layout}"
f"\nPlease use recast_ptr(ptr, {input.layout.inner}, element_type) operation to move swizzle to the ptr",
)
return True
def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
raise OpError(
self,
f"Expected affine layout for {self._make_trait()}'s operand B, "
f"but got composed layout instead: {input.layout}"
f"\nPlease use recast_ptr(ptr, {input.layout.inner}, element_type) operation to move swizzle to the ptr",
)
return True
class SparseMmaTraits(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_name = (
f"#cute_nvgpu.atom_mma_field_sm100_sparse<{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_name = (
f"#cute_nvgpu.atom_mma_field_sm100_sparse<{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
)
#
# TF32 MMA
#
@@ -669,7 +511,7 @@ class MmaTF32Op(MmaOp):
0,
)
return MmaTF32Trait(
cute.make_atom(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
@@ -763,7 +605,7 @@ class MmaF16BF16Op(MmaOp):
0,
)
return MmaF16BF16Trait(
cute.make_atom(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
@@ -780,109 +622,6 @@ class MmaF16BF16Trait(MmaTraits):
pass
@dataclass(frozen=True)
class MmaF16BF16SparseOp(SparseMmaOp):
"""
F16/BF16 tcgen05 Sparse MMA Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma-sp>`__.
This Operation corresponds to the ``.kind::f16`` qualifier with sparse support.
"""
descriptive_name = "tcgen05 F16/BF16 Sparse MMA Operation"
def __init__(
self,
ab_dtype: Type[Numeric],
acc_dtype: Type[Numeric],
instruction_shape: Shape,
cta_group: CtaGroup,
a_src: OperandSource,
a_major_mode: OperandMajorMode,
b_major_mode: OperandMajorMode,
sparse_metadata_format: SparseMetadataFormat,
) -> None:
super().__init__(
ab_dtype,
ab_dtype,
acc_dtype,
instruction_shape,
cta_group,
a_src,
a_major_mode,
b_major_mode,
sparse_metadata_format,
)
self._verify()
def _verify(self) -> None:
# Input data type verification
if self.a_dtype not in [Float16, BFloat16]:
raise OpError(
self,
"expects the 'ab_dtype' Op parameter to be one of Float16 or BFloat16",
)
assert self.b_dtype == self.a_dtype, "a_dtype and b_dtype must be the same"
# Accumulator data type verification
if self.acc_dtype not in [Float16, Float32]:
raise OpError(
self,
"expects the 'acc_dtype' Op parameter to be one of Float16 or Float32",
)
# Instruction shape verification
instruction_k = 32 # For sparse, K is doubled compared to dense F16/BF16
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) -> "MmaF16BF16SparseTrait":
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
ty = _cute_nvgpu_ir.MmaAtomSM100UMMASparseType.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,
T.ui8(),
self.sparse_metadata_format._to_ir(),
self.a_src._to_ir(),
0, # cScaleExp
)
def get_e_ptr():
ptr_type = _cute_ir.PtrType.get(T.ui8(), _cute_ir.AddressSpace.tmem, 8)
address_value = Int32(0).ir_value(loc=loc, ip=ip)
aligned_ty = _cute_ir.ConstrainedIntType.get(8, 32)
aligned_intptr = _cute_ir.assume(aligned_ty, address_value, loc=loc, ip=ip)
ui8_tmem_ptr = _cute_ir.inttoptr(ptr_type, aligned_intptr, loc=loc, ip=ip)
return ui8_tmem_ptr
return MmaF16BF16SparseTrait(
cute.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),
get_e_ptr().value,
),
loc=loc,
ip=ip,
)
)
class MmaF16BF16SparseTrait(SparseMmaTraits):
pass
#
# I8 MMA
#
@@ -953,7 +692,7 @@ class MmaI8Op(MmaOp):
0,
)
return MmaI8Trait(
cute.make_atom(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
@@ -1046,7 +785,7 @@ class MmaFP8Op(MmaOp):
0,
)
return MmaFP8Trait(
cute.make_atom(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
@@ -1136,7 +875,7 @@ class MmaMXF8Op(BlockScaledMmaOp):
self.sf_vec_size,
)
return MmaMXF8Trait(
cute.make_atom(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
@@ -1222,7 +961,7 @@ class MmaMXF4Op(BlockScaledMmaOp):
self.sf_vec_size,
)
return MmaMXF4Trait(
cute.make_atom(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
@@ -1315,7 +1054,7 @@ class MmaMXF4NVF4Op(BlockScaledMmaOp):
self.sf_vec_size,
)
return MmaMXF4NVF4Trait(
cute.make_atom(
make_atom(
ty,
(
Boolean(False).ir_value(loc=loc, ip=ip),
python/CuTeDSL/cutlass/cute/nvgpu/warp/__init__.py
+4
View File
@@ -16,9 +16,13 @@ from .mma import *
# __all__ is required here for documentation generation
__all__ = [
# mma.py
"Field",
"MmaF16BF16Op",
"MmaMXF4Op",
"MmaMXF4NVF4Op",
# copy.py
"LdMatrix8x8x16bOp",
"LdMatrix16x8x8bOp",
"LdMatrix16x16x8bOp",
"StMatrix8x8x16bOp",
"StMatrix16x8x8bOp",
python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py
+123 -23
View File
@@ -12,20 +12,20 @@
from dataclasses import dataclass
from typing import Type
from cutlass import cute
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
from ..common import OpError
from ...core import _pack_shape
from ...typing import Numeric
from ...atom import CopyOp, Trait
from ...typing import Numeric, Optional
from ...atom import CopyOp, Trait, make_atom
@dataclass(frozen=True)
class BaseOp(CopyOp):
transpose: bool = False
num_matrices: int = 1
unpack_bits: Optional[int] = None
def __post_init__(self) -> None:
if not isinstance(self.transpose, bool):
@@ -41,6 +41,8 @@ class BaseOp(CopyOp):
)
if self.transpose:
res += "\n transposed"
if self.unpack_bits is not None:
res += f"\n unpack {self.unpack_bits}b to 8b"
return res
@@ -60,6 +62,8 @@ class LdMatrix8x8x16bOp(BaseOp):
self,
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
)
if self.unpack_bits is not None:
raise OpError(self, "Op doesn't support unpacking")
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
@@ -72,46 +76,140 @@ class LdMatrix8x8x16bOp(BaseOp):
self.num_matrices,
ir.UnitAttr.get() if self.transpose else None,
)
return LdMatrix8x8x16bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return LdMatrix8x8x16bTrait(make_atom(ty, loc=loc, ip=ip))
class LdMatrix8x8x16bTrait(Trait):
pass
@dataclass(frozen=True)
class LdMatrix8x16x8bOp(BaseOp):
"""
8x16 ``ldmatrix`` Operation with 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 ``.m8n16`` and the ``.b4x16_p64``, ``.b6x16_p32`` qualifiers.
"""
def __post_init__(self) -> None:
super().__post_init__()
if self.transpose:
raise OpError(self, "Op doesn't support transpose")
if self.num_matrices not in [1, 2, 4]:
raise OpError(
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")
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:
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8
ty = _cute_nvgpu_ir.CopyAtomLdsmType.get(
copy_internal_type.mlir_type,
mode.type.attribute,
sz_pattern,
self.num_matrices,
None,
)
return LdMatrix8x16x8bTrait(make_atom(ty, loc=loc, ip=ip))
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.
Useful for vectorizing with Ampere-style 8x8 matrix thread-value layouts
"""
def __post_init__(self) -> None:
super().__post_init__()
if not self.transpose:
raise OpError(self, "Op only supports transpose")
if self.num_matrices not in [2, 4]:
raise OpError(
self,
"expects the 'num_matrices' Op parameter to be one of [2,4]",
)
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
) -> "LdMatrix16x8x8bTrait":
mode = _pack_shape((16, 8), loc=loc, ip=ip)
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u8
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,
mode.type.attribute,
sz_pattern,
self.num_matrices,
ir.UnitAttr.get(),
)
return LdMatrix16x8x8bTrait(make_atom(ty, loc=loc, ip=ip))
class LdMatrix16x8x8bTrait(Trait):
pass
@dataclass(frozen=True)
class LdMatrix16x16x8bOp(BaseOp):
"""
16x16 8-bit ``ldmatrix`` Operation.
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 ``.b16`` qualifiers.
This operation corresponds to the ``.m16n16`` and the ``.b4x16_p64``,``.b6x16_p32``,``.b8`` qualifiers.
"""
def __init__(self, num_matrices: int) -> None:
super().__init__(transpose=True, num_matrices=num_matrices)
self._verify()
def _verify(self):
assert self.transpose, "transpose must be True"
def __post_init__(self) -> None:
super().__post_init__()
if not self.transpose:
raise OpError(self, "Op only supports transpose")
if self.num_matrices not in [1, 2]:
raise OpError(
self,
"expects the 'num_matrices' Op parameter to be one of [1,2]",
)
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
) -> "LdMatrix16x16x8bTrait":
mode = _pack_shape((16, 16), loc=loc, ip=ip)
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u8
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,
mode.type.attribute,
_cute_nvgpu_ir.LdsmSzPattern.u8,
sz_pattern,
self.num_matrices,
ir.UnitAttr.get(),
)
return LdMatrix16x16x8bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return LdMatrix16x16x8bTrait(make_atom(ty, loc=loc, ip=ip))
class LdMatrix16x16x8bTrait(Trait):
@@ -134,6 +232,8 @@ class StMatrix8x8x16bOp(BaseOp):
self,
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
)
if self.unpack_bits is not None:
raise OpError(self, "Op doesn't support unpacking")
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
@@ -145,7 +245,7 @@ class StMatrix8x8x16bOp(BaseOp):
self.num_matrices,
ir.UnitAttr.get() if self.transpose else None,
)
return StMatrix8x8x16bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return StMatrix8x8x16bTrait(make_atom(ty, loc=loc, ip=ip))
class StMatrix8x8x16bTrait(Trait):
@@ -161,17 +261,17 @@ class StMatrix16x8x8bOp(BaseOp):
This operation corresponds to the ``m16n8`` qualifier.
"""
def __init__(self, num_matrices: int) -> None:
super().__init__(transpose=True, num_matrices=num_matrices)
self._verify()
def _verify(self):
def __post_init__(self) -> None:
super().__post_init__()
if not self.transpose:
raise OpError(self, "Op only supports transpose")
if self.num_matrices not in [1, 2, 4]:
assert self.transpose, "transpose must be True"
raise OpError(
self,
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
)
if self.unpack_bits is not None:
raise OpError(self, "Op doesn't support unpacking")
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
@@ -183,7 +283,7 @@ class StMatrix16x8x8bOp(BaseOp):
self.num_matrices,
ir.UnitAttr.get(),
)
return StMatrix16x8x8bTrait(cute.make_atom(ty, loc=loc, ip=ip))
return StMatrix16x8x8bTrait(make_atom(ty, loc=loc, ip=ip))
class StMatrix16x8x8bTrait(Trait):
python/CuTeDSL/cutlass/cute/nvgpu/warp/mma.py
+226 -60
View File
@@ -10,19 +10,33 @@
# is strictly prohibited.
from dataclasses import dataclass
from typing import Type
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 ..common import OpError
from ...typing import Shape, Float16, BFloat16, Float32, Numeric
from ...typing import (
Shape,
Float4E2M1FN,
Float8E8M0FNU,
Float8E4M3FN,
Float16,
BFloat16,
Float32,
Boolean,
Numeric,
Pointer,
)
from ...core import _pack_shape
from ...tensor import _Tensor
from ...atom import MmaOp, Trait
from ...typing import Tensor
from ...atom import MmaOp, Trait, make_atom
from cutlass._mlir import ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir.dialects.cute_nvgpu import SparseMetadataFormat
####################################################################################################
@@ -43,7 +57,7 @@ class WarpMmaOp(MmaOp):
@dataclass(frozen=True)
class MmaF16BF16Op(WarpMmaOp):
"""
F16/BF16 tcgen05 MMA Operation.
F16/BF16 warp-level MMA Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma>`__.
This Operation covers the instructions using the ``.f16`` or ``.bf16`` qualifiers for the input operands.
@@ -83,7 +97,7 @@ class MmaF16BF16Op(WarpMmaOp):
self.ab_dtype.mlir_type,
self.acc_dtype.mlir_type,
)
return MmaF16BF16Trait(cute.make_atom(ty, loc=loc, ip=ip))
return MmaF16BF16Trait(make_atom(ty, loc=loc, ip=ip))
def __str__(self) -> str:
return (
@@ -93,10 +107,10 @@ class MmaF16BF16Op(WarpMmaOp):
+ f"\n Instruction shape MNK = {self.shape_mnk}"
)
def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None):
pass
def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None):
pass
@@ -104,12 +118,86 @@ class MmaF16BF16Trait(Trait):
pass
class SparseMetadataFormat(enum.Enum):
# Base class for SM120 Blockscaled MMA Ops
@dataclass(frozen=True)
class MmaSM120BlockScaledOp(MmaOp):
ab_dtype: Type[Numeric]
acc_dtype: Type[Numeric]
shape_mnk: Shape
sf_type: Type[Numeric]
sf_vec_size: int
use_sf_layout_TV: bool = False
admissible_archs = [
"sm_120a",
]
def __post_init__(self) -> None:
# Verify arch
arch = CuTeDSL._get_dsl().get_arch_enum()
if not arch == Arch.sm_120a:
raise OpError(
self,
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
)
if self.ab_dtype != Float4E2M1FN:
raise OpError(
self,
"expects the 'ab_dtype' Op parameter to be Float4E2M1FN",
)
if self.acc_dtype != Float32:
raise OpError(
self,
"expects the 'acc_dtype' Op parameter to be Float32",
)
if self.shape_mnk != (16, 8, 64):
raise OpError(
self,
"expects the 'shape_mnk' Op parameter to be (16,8,64)",
)
if self.sf_vec_size == 16:
if self.sf_type != Float8E4M3FN:
raise OpError(
self,
"expects the 'sf_type' Op parameter to be Float8E4M3FN",
)
elif self.sf_vec_size == 32:
if self.sf_type != Float8E8M0FNU:
raise OpError(
self,
"expects the 'sf_type' Op parameter to be Float8E8M0FNU",
)
else:
raise OpError(
self,
"expects the 'sf_vec_size' Op parameter to be 16 or 32",
)
def __str__(self) -> str:
return (
"warp-level MXF4/MXF4NVF4 MMA Operation"
+ f"\n A/B data type = {self.ab_dtype}"
+ f"\n Accumulator data type = {self.acc_dtype}"
+ f"\n Instruction shape MNK = {self.shape_mnk}"
+ f"\n Vector size = {self.sf_vec_size}"
+ f"\n SF data type = {self.sf_type}"
)
def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None):
pass
def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None):
pass
class Field(enum.Enum):
"""
An enumeration for the sparse metadata format of the MMA.
An enumeration for the fields of the MMA Atom that can be modified at runtime.
"""
TID = SparseMetadataFormat.tid
ACCUMULATE = "accum_c"
SFA = "sf_a"
SFB = "sf_b"
def __str__(self) -> str:
return f"{self.__class__.__name__}.{self.name}"
@@ -117,66 +205,144 @@ class SparseMetadataFormat(enum.Enum):
def __repr__(self) -> str:
return f"<{self.__class__.__name__}.{self.name}>"
def _to_ir(self) -> _cute_nvgpu_ir.SparseMetadataFormat:
def _to_ir_field_name(self) -> str:
return self.value
class MmaBlockScaledTrait(Trait):
admissible_fields = [
Field.ACCUMULATE,
Field.SFA,
Field.SFB,
]
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 == Field.ACCUMULATE:
value = Boolean(value).ir_value(loc=loc, ip=ip)
elif field in [Field.SFA, Field.SFB]:
if not isinstance(value, Pointer):
raise ValueError(
f"expects value to be a pointer for {field}, but got {type(value).__name__}"
)
value = value.value
field_name = f"#cute_nvgpu.atom_mma_field_sm120_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
)
def get(self, field, *, loc=None, ip=None) -> Any:
if field not in [Field.ACCUMULATE]:
raise ValueError(f"the get method for {field} is not supported")
field_name = f"#cute_nvgpu.atom_mma_field_sm120_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
)
#
# MXF4 MMA
#
@dataclass(frozen=True)
class MmaF16BF16SparseOp(WarpMmaOp):
ab_dtype: Type[Numeric]
acc_dtype: Type[Numeric]
shape_mnk: Shape
sparse_metadata_format: SparseMetadataFormat
class MmaMXF4Op(MmaSM120BlockScaledOp):
"""
MXF4 warp-level MMA Operation.
def __post_init__(self) -> None:
# verify field after initialization
if not isinstance(self.sparse_metadata_format, SparseMetadataFormat):
raise OpError(
self,
"expects the 'sparse_metadata_format' Op parameter to be a SparseMetadataFormat instance",
)
# verify the instruction shape
if self.ab_dtype not in [Float16, BFloat16]:
raise OpError(
self,
"expects the 'ab_dtype' Op parameter to be one of Float16 or BFloat16",
)
if self.acc_dtype not in [Float16, Float32]:
raise OpError(
self,
"expects the 'acc_dtype' Op parameter to be one of Float16 or Float32",
)
if (self.ab_dtype == BFloat16) and (self.acc_dtype != Float32):
raise OpError(
self,
"expects the 'acc_dtype' Op parameter to be Float32 when 'ab_dtype' is BFloat16",
)
if self.shape_mnk not in [(16, 8, 16), (16, 8, 32)]:
raise OpError(
self,
"expects the 'shape_mnk' Op parameter to be one of (16,8,16) or (16,8,32)",
)
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma>`__.
This Operation covers the instructions using the ``.e2m1`` qualifiers for the input operands.
.kind = {.kind::mxf4};
.scale_vec_size = {.scale_vec::2X};
.stype = {.ue8m0};
"""
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaF16BF16SparseTrait":
descriptive_name = "warp-level MXF4 MMA Operation"
def __init__(
self,
ab_dtype: Type[Numeric],
acc_dtype: Type[Numeric],
sf_type: Type[Numeric],
) -> None:
super().__init__(
ab_dtype,
acc_dtype,
(16, 8, 64),
sf_type,
32,
)
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.MmaAtomSM80SparseType.get(
ty = _cute_nvgpu_ir.MmaAtomSM120BlockScaledType.get(
shape_mnk.type.attribute,
32,
False,
self.ab_dtype.mlir_type,
self.ab_dtype.mlir_type,
self.acc_dtype.mlir_type,
self.sparse_metadata_format._to_ir(),
)
return MmaF16BF16SparseTrait(cute.make_atom(ty, loc=loc, ip=ip))
def __str__(self) -> str:
return (
"warp-level F16/BF16 Sparse MMA Operation"
+ f"\n A/B data type = {self.ab_dtype}"
+ f"\n Accumulator data type = {self.acc_dtype}"
+ f"\n Instruction shape MNK = {self.shape_mnk}"
+ f"\n Sparse metadata format = {self.sparse_metadata_format}"
self.sf_type.mlir_type,
)
return MmaMXF4Trait(make_atom(ty, loc=loc, ip=ip))
class MmaF16BF16SparseTrait(Trait):
class MmaMXF4Trait(MmaBlockScaledTrait):
pass
#
# MXF4NVF4 MMA
#
@dataclass(frozen=True)
class MmaMXF4NVF4Op(MmaSM120BlockScaledOp):
"""
MXF4NVF4 warp-level MMA Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma>`__.
This Operation covers the instructions using the ``.e2m1`` qualifiers for the input operands.
.kind = {.kind::mxf4nvf4};
.scale_vec_size = {.scale_vec::2X, .scale_vec::4X};
.stype = {.ue8m0, .ue4m3};
"""
descriptive_name = "warp-level MXF4NVF4 MMA Operation"
def __init__(
self,
ab_dtype: Type[Numeric],
acc_dtype: Type[Numeric],
sf_type: Type[Numeric],
) -> None:
super().__init__(
ab_dtype,
acc_dtype,
(16, 8, 64),
sf_type,
16,
)
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.MmaAtomSM120BlockScaledType.get(
shape_mnk.type.attribute,
16,
False,
self.ab_dtype.mlir_type,
self.ab_dtype.mlir_type,
self.acc_dtype.mlir_type,
self.sf_type.mlir_type,
)
return MmaMXF4NVF4Trait(make_atom(ty, loc=loc, ip=ip))
class MmaMXF4NVF4Trait(MmaBlockScaledTrait):
pass
python/CuTeDSL/cutlass/cute/nvgpu/warpgroup/mma.py
+7 -23
View File
@@ -13,7 +13,6 @@ import enum
from dataclasses import dataclass
from typing import Type, Any
from cutlass import cute
from cutlass.base_dsl.arch import Arch
from cutlass.cutlass_dsl import BaseDSL, T
@@ -23,9 +22,9 @@ from cutlass._mlir import ir
from ..common import OpError
from ...core import _pack_shape, rank, depth
from ...tensor import _Tensor
from ...typing import (
Shape,
Tensor,
Float16,
BFloat16,
Float32,
@@ -38,7 +37,7 @@ from ...typing import (
Numeric,
AddressSpace,
)
from ...atom import MmaOp, Trait
from ...atom import MmaOp, Trait, make_atom
####################################################################################################
@@ -181,7 +180,7 @@ class MmaOp(WarpGroupMmaOp):
+ f"\n Instruction shape MNK = {self.shape_mnk}"
)
def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_A(self, input: Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
@@ -193,7 +192,7 @@ class MmaOp(WarpGroupMmaOp):
)
return True
def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None):
def _verify_fragment_B(self, input: Tensor, *, loc=None, ip=None):
if input.memspace == AddressSpace.smem and isinstance(
input.layout.type, _cute_ir.ComposedLayoutType
):
@@ -305,12 +304,7 @@ class MmaF16BF16Op(MmaOp):
self.a_src._to_ir(),
)
return MmaF16BF16Trait(
cute.make_atom(
ty,
(Boolean(False).ir_value(loc=loc, ip=ip),),
loc=loc,
ip=ip,
)
make_atom(ty, (Boolean(False).ir_value(loc=loc, ip=ip),), loc=loc, ip=ip)
)
@@ -391,12 +385,7 @@ class MmaF8Op(MmaOp):
self.a_src._to_ir(),
)
return MmaF8Trait(
cute.make_atom(
ty,
(Boolean(False).ir_value(loc=loc, ip=ip),),
loc=loc,
ip=ip,
)
make_atom(ty, (Boolean(False).ir_value(loc=loc, ip=ip),), loc=loc, ip=ip)
)
@@ -486,12 +475,7 @@ class MmaI8Op(MmaOp):
self.a_src._to_ir(),
)
return MmaI8Trait(
cute.make_atom(
ty,
(Boolean(False).ir_value(loc=loc, ip=ip),),
loc=loc,
ip=ip,
)
make_atom(ty, (Boolean(False).ir_value(loc=loc, ip=ip),), loc=loc, ip=ip)
)
python/CuTeDSL/cutlass/cute/runtime.py
+82 -53
View File
@@ -25,9 +25,19 @@ import cutlass._mlir.dialects.cuda as _cuda_dialect
from cutlass.cutlass_dsl import JitArgAdapterRegistry, CuTeDSL as _CuTeDSL
from cutlass.base_dsl.common import DSLRuntimeError
from cutlass.base_dsl.export import ExternalBinaryModule
# Local modules imports
from .typing import AddressSpace, Layout, Tensor, Pointer, Numeric, SymInt
from .typing import (
AddressSpace,
Layout,
Tensor,
Pointer,
Numeric,
SymInt,
Float32,
TFloat32,
)
from . import core
from .tensor import _Tensor as CoreTensor
@@ -69,7 +79,10 @@ class _Pointer(Pointer):
else:
self._assumed_align = assumed_align
self._c_pointer = None
self._desc = ctypes.c_void_p(int(self._pointer))
self._c_pointer = ctypes.addressof(self._desc)
self._c_pointers_cache = [self._c_pointer]
assert int(self._pointer) % self._assumed_align == 0, (
f"pointer must be {self._assumed_align} bytes aligned"
)
@@ -85,10 +98,7 @@ class _Pointer(Pointer):
return self._pointer
def __c_pointers__(self):
if self._c_pointer is None:
self._desc = ctypes.c_void_p(int(self._pointer))
self._c_pointer = ctypes.addressof(self._desc)
return [self._c_pointer]
return self._c_pointers_cache
def __new_from_mlir_values__(self, values):
assert len(values) == 1
@@ -151,29 +161,24 @@ class _Tensor(Tensor):
self._memref_desc = None
self._dtype = None
self._use_32bit_stride = use_32bit_stride
self._c_pointers_cache = None
@property
def __class__(self) -> Type[Tensor]:
# Cheat to let `type(_Tensor())` to return cute.Tensor
return Tensor
def lazily_load_dltensor(func):
"""Decorator to lazily load the DLTensorWrapper.
def load_dltensor(self):
"""Lazily load the DLTensorWrapper.
This decorator loads the DLTensorWrapper when needed,
This function loads the DLTensorWrapper when needed,
avoiding overhead in the critical path of calling JIT functions.
"""
if self._dltensor_wrapper is None:
self._dltensor_wrapper = _cute_ir.DLTensorWrapper(
self._dlpack_data, self._use_32bit_stride
)
def wrapper(self, *args, **kwargs):
if self._dltensor_wrapper is None:
self._dltensor_wrapper = _cute_ir.DLTensorWrapper(
self._dlpack_data, self._use_32bit_stride
)
return func(self, *args, **kwargs)
return wrapper
@lazily_load_dltensor
def mark_layout_dynamic(self, leading_dim: Optional[int] = None):
"""Marks the tensor layout as dynamic based on the leading dimension.
@@ -193,10 +198,10 @@ class _Tensor(Tensor):
:return: The tensor with dynamic layout
:rtype: _Tensor
"""
self.load_dltensor()
self._dltensor_wrapper.mark_layout_dynamic(leading_dim)
return self
@lazily_load_dltensor
def mark_compact_shape_dynamic(
self,
mode: int,
@@ -208,6 +213,7 @@ class _Tensor(Tensor):
:param mode: The mode of the compact shape, defaults to 0
:type mode: int
:param stride_order: Consistent with `torch.Tensor.dim_order`. Defaults to None.
Indicates the order of the modes (dimensions) if the current layout were converted to row-major order.
It starts from the outermost to the innermost dimension.
:type stride_order: tuple[int, ...], optional
@@ -228,17 +234,18 @@ class _Tensor(Tensor):
Using `torch.Tensor.dim_order()` to get the stride order of the torch tensor.
.. code-block:: python
a = torch.empty(3, 4)
t = cute.runtime.from_dlpack(a)
t = t.mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order())
a = torch.empty(3, 4)
t = cute.runtime.from_dlpack(a)
t = t.mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order())
"""
self.load_dltensor()
self._dltensor_wrapper.mark_compact_shape_dynamic(
mode, stride_order, divisibility
)
return self
@property
@lazily_load_dltensor
def element_type(self) -> Type[Numeric]:
self.load_dltensor()
if self._dtype is None:
self._dtype = self._dltensor_wrapper.dtype
return self._dtype
@@ -274,24 +281,24 @@ class _Tensor(Tensor):
self._dtype = new_type
@property
@lazily_load_dltensor
def memspace(self):
self.load_dltensor()
return self._dltensor_wrapper.address_space
@property
@lazily_load_dltensor
def size_in_bytes(self) -> int:
self.load_dltensor()
return self._dltensor_wrapper.size_in_bytes()
@property
@lazily_load_dltensor
def mlir_type(self) -> ir.Type:
self.load_dltensor()
return self._dltensor_wrapper.get_type(
self.element_type.mlir_type, self._assumed_align
)
@lazily_load_dltensor
def __str__(self) -> str:
self.load_dltensor()
return f"Tensor<0x{self._dltensor_wrapper.str}>"
def __repr__(self):
@@ -304,8 +311,8 @@ class _Tensor(Tensor):
raise TypeError("runtime._Tensor is not indexable")
@property
@lazily_load_dltensor
def iterator(self):
self.load_dltensor()
return _Pointer(
self._dltensor_wrapper.data_ptr,
self.element_type,
@@ -320,13 +327,13 @@ class _Tensor(Tensor):
)
@property
@lazily_load_dltensor
def shape(self):
self.load_dltensor()
return self._dltensor_wrapper.shape
@property
@lazily_load_dltensor
def stride(self):
self.load_dltensor()
strides = self._dltensor_wrapper.stride
if strides is None:
strides = itertools.accumulate(
@@ -356,28 +363,30 @@ class _Tensor(Tensor):
raise TypeError("fill function is not supported in runtime")
@property
@lazily_load_dltensor
def data_ptr(self):
self.load_dltensor()
return self._dltensor_wrapper.data_ptr
@property
@lazily_load_dltensor
def dynamic_shapes_mask(self):
"""Get the mask of dynamic shapes in the tensor."""
self.load_dltensor()
return self._dltensor_wrapper.get_dynamic_shapes_mask()
@property
@lazily_load_dltensor
def dynamic_strides_mask(self):
"""Get the mask of dynamic strides in the tensor."""
self.load_dltensor()
return self._dltensor_wrapper.get_dynamic_strides_mask()
@lazily_load_dltensor
def __c_pointers__(self):
self._memref_desc = self._dltensor_wrapper.build_memref_desc(
self._assumed_align
)
return [_cute_ir.pycapsule_get_pointer(self._memref_desc)]
if self._c_pointers_cache is None:
self.load_dltensor()
self._memref_desc = self._dltensor_wrapper.build_memref_desc(
self._assumed_align
)
self._c_pointers_cache = [_cute_ir.pycapsule_get_pointer(self._memref_desc)]
return self._c_pointers_cache
def __get_mlir_types__(self):
return [self.mlir_type]
@@ -480,6 +489,12 @@ class _FakeCompactTensor(Tensor):
def stride(self):
return self._stride
@property
def leading_dim(self):
for dim, order in enumerate(self._stride_order):
if order == 0:
return dim
@property
def dynamic_shapes_mask(self):
return tuple(1 if isinstance(e, SymInt) else 0 for e in self._shape)
@@ -618,7 +633,8 @@ def make_fake_compact_tensor(
:param shape: Shape of the tensor.
:type shape: tuple[int, ...]
:param stride_order: Order in which strides (memory layout) are assigned to the tensor dimensions.
If None, the default layout is col-major. Otherwise, it should be a permutation of the dimension indices.
If None, the default layout is left-to-right order (known as column-major order for flatten layout).
Otherwise, it should be a permutation order of the dimension indices.
:type stride_order: tuple[int, ...], optional
:param memspace: Memory space where the fake tensor resides. Optional.
:type memspace: str, optional
@@ -646,6 +662,9 @@ def make_fake_compact_tensor(
# Compiled function will take a tensor with the type:
# tensor<ptr<f32, generic> o (100,?{div=8}):(?{i32 div=8},1)>
compiled_foo = cute.compile(foo, x)
# Default stride order is left-to-right order: (1, 8)
y = make_fake_compact_tensor(cutlass.Float32, (8, 3))
"""
return _FakeCompactTensor(
@@ -730,6 +749,7 @@ def from_dlpack(
use_32bit_stride=False,
*,
enable_tvm_ffi=False,
force_tf32=False,
) -> Tensor:
"""Convert from tensor object supporting __dlpack__() to a CuTe Tensor.
@@ -745,6 +765,8 @@ def from_dlpack(
:param enable_tvm_ffi: Whether to enable TVM-FFI, defaults to False. When True, the tensor will be converted to
a TVM-FFI function compatible tensor.
:type enable_tvm_ffi: bool, optional
:param force_tf32: Whether to force the element type to TFloat32 if the element type is Float32.
:type force_tf32: bool, optional
:return: A CuTe Tensor object
:rtype: Tensor
@@ -764,12 +786,15 @@ def from_dlpack(
# If the environment variable `CUTE_DSL_ENABLE_TVM_FFI` is set to True, the tensor will be converted to
# a TVM-FFI function compatible tensor.
enable_tvm_ffi = enable_tvm_ffi or _CuTeDSL._get_dsl().envar.enable_tvm_ffi
return _Tensor(
res = _Tensor(
tensor_dlpack,
assumed_align=assumed_align,
use_32bit_stride=use_32bit_stride,
enable_tvm_ffi=enable_tvm_ffi,
)
if force_tf32 and res.element_type == Float32:
res.element_type = TFloat32
return res
def make_ptr(
@@ -851,15 +876,21 @@ class TensorAdapter:
def __init__(self, arg):
self._arg = from_dlpack(arg).mark_layout_dynamic()
self._c_pointers_cache = None
self._mlir_types_cache = None
def __new_from_mlir_values__(self, values):
return self._arg.__new_from_mlir_values__(values)
def __c_pointers__(self):
return self._arg.__c_pointers__()
if self._c_pointers_cache is None:
self._c_pointers_cache = self._arg.__c_pointers__()
return self._c_pointers_cache
def __get_mlir_types__(self):
return self._arg.__get_mlir_types__()
if self._mlir_types_cache is None:
self._mlir_types_cache = self._arg.__get_mlir_types__()
return self._mlir_types_cache
def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]:
@@ -872,7 +903,7 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]:
:rtype: list
"""
def _get_cuda_dialect_runtime_path():
def _get_cute_dsl_runtime_path():
libs = get_prefix_dsl_libs("CUTE_DSL")
if libs is None:
return None
@@ -884,15 +915,15 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]:
libs = libs.split(":")
for path in libs:
if path.endswith("libcuda_dialect_runtime.so"):
if path.endswith("libcute_dsl_runtime.so"):
return path
return None
libs = []
cuda_dialect_runtime_path = _get_cuda_dialect_runtime_path()
if cuda_dialect_runtime_path:
libs.append(cuda_dialect_runtime_path)
cute_dsl_runtime_path = _get_cute_dsl_runtime_path()
if cute_dsl_runtime_path:
libs.append(cute_dsl_runtime_path)
if enable_tvm_ffi:
import tvm_ffi
@@ -905,8 +936,8 @@ def find_runtime_libraries(*, enable_tvm_ffi: bool = True) -> List[str]:
_LOAD_MODULE_LIBS_CACHE = []
def load_module(file_path: str, *, enable_tvm_ffi: bool = True):
"""Load a module from a file path. Today only support TVM-FFI module.
def load_module(file_path: str, *, enable_tvm_ffi: bool = False):
"""Load a module from a file path.
:param file_path: The path to the module file
:type file_path: str
@@ -937,9 +968,7 @@ def load_module(file_path: str, *, enable_tvm_ffi: bool = True):
# compatible with tvm-ffi < 0.1.6
return tvm_ffi.load_module(file_path)
else:
raise DSLRuntimeError(
"Unimplemented, please load the module with enable_tvm_ffi=True."
)
return ExternalBinaryModule(file_path)
# -------------------------------------------------------------------------
# Try to register_jit_arg_adapter for TensorAdapter
python/CuTeDSL/cutlass/cute/tensor.py
+73 -11
View File
@@ -28,6 +28,30 @@ from cutlass._mlir.dialects.cute import ReductionOp as ReductionOp
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir.dialects import vector, arith
from .typing import (
Numeric,
Integer,
Boolean,
Int4,
Uint8,
Int8,
Int32,
Int64,
BFloat16,
IntTuple,
Coord,
Shape,
Stride,
Pointer,
Layout,
ComposedLayout,
Tensor,
AddressSpace,
is_integer,
is_int_tuple,
as_numeric,
)
from .core import (
_unpack_x_tuple,
_pack_int_tuple,
@@ -82,6 +106,29 @@ 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
__all__ = [
"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_",
]
@ir.register_value_caster(_cute_ir.MemRefType.get_static_typeid(), replace=True)
@ir.register_value_caster(_cute_ir.CoordTensorType.get_static_typeid(), replace=True)
@ir.register_value_caster(
@@ -132,8 +179,9 @@ class _Tensor(Tensor):
iter_val = _cute_ir.get_iter(self.value, loc=loc, ip=ip)
if isinstance(iter_val, Pointer):
self._iterator = iter_val
elif isinstance(iter_val.type, _cute_ir.IntTupleType):
self._iterator = _unpack_x_tuple(iter_val)
elif isinstance(iter_val.type, _cute_ir.ArithTupleIteratorType):
itup_val = _cute_ir.deref_arith_tuple_iter(iter_val)
self._iterator = _unpack_x_tuple(itup_val)
elif isinstance(iter_val, ir.Value):
# Example: SMEM descriptor iterator, not well supported today
self._iterator = iter_val
@@ -152,6 +200,9 @@ class _Tensor(Tensor):
else:
raise TypeError(f"unsupported iterator type, got {type(self.iterator)}")
def __repr__(self):
return self.__str__()
def __str__(self):
from .core import pretty_str
@@ -257,7 +308,8 @@ class _Tensor(Tensor):
res = _cute_ir.get_iter(
slice_(self, crd, loc=loc, ip=ip).value, loc=loc, ip=ip
)
return _unpack_x_tuple(res, loc=loc, ip=ip)
itup_val = _cute_ir.deref_arith_tuple_iter(res)
return _unpack_x_tuple(itup_val)
else:
self._check_can_load_store()
self._check_can_dereference()
@@ -387,9 +439,10 @@ class _Tensor(Tensor):
return _cute_ir.get_layout(self.value, loc=loc, ip=ip)
@property
@dsl_user_op
@lru_cache_ir()
def shape(self) -> Shape:
return self.layout.shape
def shape(self, *, loc=None, ip=None) -> Shape:
return self.layout.shape_method(loc=loc, ip=ip)
@property
@lru_cache_ir()
@@ -633,7 +686,7 @@ def make_tensor(
"""
if isinstance(layout, _ComposedLayoutWithInnerFunc):
raise ValueError(
"CuTe DSL tensor does not support composed layouts with inner functions: {layout}"
f"CuTe DSL tensor does not support composed layouts with inner functions: {layout}"
)
if not isinstance(layout, (Layout, ComposedLayout)):
@@ -644,8 +697,12 @@ def make_tensor(
res_ty = None
if is_integer(iterator) or isinstance(iterator, tuple):
iterator = _pack_int_tuple(iterator, loc=loc, ip=ip)
res_ty = _cute_ir.CoordTensorType.get(iterator.type, layout.type)
itup_val = _pack_int_tuple(iterator, loc=loc, ip=ip)
iter_ty = _cute_ir.ArithTupleIteratorType.get(itup_val.type)
iterator = _cute_ir.make_arith_tuple_iter(
iter=iter_ty, value=itup_val, loc=loc, ip=ip
)
res_ty = _cute_ir.CoordTensorType.get(itup_val.type, layout.type)
elif isinstance(iterator, Pointer):
iterator = iterator.value
res_ty = _cute_ir.MemRefType.get(iterator.type, layout.type)
@@ -773,7 +830,7 @@ def make_fragment(
@dsl_user_op
def make_rmem_tensor_like(
src: Union[Layout, ComposedLayout, Tensor],
src: Union[Layout, ComposedLayout, Tensor, "TensorSSA"],
dtype: Optional[Type[Numeric]] = None,
*,
loc=None,
@@ -826,7 +883,7 @@ def make_rmem_tensor_like(
create register storage for intermediate results.
"""
if not isinstance(src, (Layout, ComposedLayout, Tensor)):
if not isinstance(src, (Layout, ComposedLayout, Tensor, TensorSSA)):
raise TypeError(
f"src must be a Layout or ComposedLayout or Tensor, got {type(src)}"
)
@@ -844,6 +901,9 @@ def make_rmem_tensor_like(
else:
res_dtype = dtype or src.element_type
src_layout = src.layout
elif isinstance(src, TensorSSA):
res_dtype = dtype or src.element_type
src_layout = make_layout(src.shape, loc=loc, ip=ip)
else:
if dtype is None:
raise ValueError("dtype must be provided when src is a layout")
@@ -918,7 +978,7 @@ def domain_offset(coord: Coord, tensor: Tensor, *, loc=None, ip=None) -> Tensor:
ip=ip,
)
elif is_integer(tensor.iterator) or isinstance(tensor.iterator, tuple):
new_iter = _cute_ir.add_offset(
new_iter = _cute_ir.tuple_add(
_pack_int_tuple(tensor.iterator, loc=loc, ip=ip),
_pack_int_tuple(offset, loc=loc, ip=ip),
loc=loc,
@@ -1136,10 +1196,12 @@ class TensorSSA(cutlass_arith.ArithValue):
def _apply_op(
self, op, other: "TensorSSA", flip=False, *, loc, ip
) -> "TensorSSA": ...
@overload
def _apply_op(
self, op, other: cutlass_arith.ArithValue, flip=False, *, loc, ip
) -> "TensorSSA": ...
@overload
def _apply_op(
self, op, other: Union[int, float, bool], flip=False, *, loc, ip
python/CuTeDSL/cutlass/cute/testing.py
+309 -53
View File
@@ -22,12 +22,24 @@ import cuda.bindings.runtime as cuda_runtime
import cutlass
import cutlass.base_dsl.jit_executor
from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op
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
from .typing import Numeric, Int8, Boolean, Tensor, Layout, Shape
import cutlass.cute as cute
from cutlass.cute import nvgpu
from . import nvgpu
from .core import recast_layout, make_layout, composition, get, rank, size, zipped_divide
from .tuple import elem_less
from .tensor import (
make_rmem_tensor,
recast_tensor,
make_identity_tensor,
TensorSSA,
_Tensor,
)
from .atom import make_copy_atom
from .algorithm import copy
from .runtime import from_dlpack
from cutlass._mlir.dialects import builtin, cf, nvvm, vector
@@ -37,14 +49,252 @@ def assert_(cond, msg=None, *, loc=None, ip=None):
cf.assert_(Boolean(cond).ir_value(), msg if msg else "", loc=loc, ip=ip)
def _maybe_recast_tensor_from_f4(src: cute.Tensor, tv_layout: cute.Layout):
################################################
# Runtime Assertion Helper Utilities For Testing
################################################
class AssertionError(RuntimeError):
"""Custom assertion error for runtime assertions."""
pass
class Assertion:
"""Base class for runtime assertion."""
pass
class _CompileTimeAssertion(Assertion):
"""Compile-time assertion helper that tracks assertion results during execution.
This assertion is used internally when RuntimeAssertion is passed through
JIT compilation. It stores assertion results in a tensor and provides compile-time
tracking of assertion results.
"""
def __init__(
self,
tensor: _Tensor,
num_assertions: int = 1,
msgs=None,
device=None,
disable: bool = False,
init_value: bool = False,
used_indices: set = None,
):
"""Initialize _CompileTimeAssertion.
:param tensor: Tensor to store assertion results
:param num_assertions: Number of assertions to support
:param msgs: List of assertion messages
:param device: Device to run assertions on
:param disable: If True, assertions are disabled
:param init_value: Initial value for assertion tensor
:param used_indices: Set of used assertion indices
"""
if msgs is None:
msgs = []
self._tensor = tensor
self._num_assertions = num_assertions
self._device = device
self._disable = disable
self._msgs = msgs
self._init_value = init_value
self._used_indices = used_indices
def __new_from_mlir_values__(self, values):
if self._disable:
return _CompileTimeAssertion(
None,
self._num_assertions,
self._msgs,
self._device,
self._disable,
self._init_value,
self._used_indices,
)
return _CompileTimeAssertion(
_Tensor(values[0], dtype=Boolean),
self._num_assertions,
self._msgs,
self._device,
self._disable,
self._init_value,
self._used_indices,
)
def __extract_mlir_values__(self):
if self._disable:
return []
return self._tensor.__extract_mlir_values__()
@dsl_user_op
@CuTeDSL.jit
def store(self, idx: Constexpr, pred: Boolean, msg: str = "", *, loc=None, ip=None):
"""Assert a predicate condition.
:param idx: Assertion index
:type idx: int
:param pred: Predicate condition to assert
:type pred: Boolean
:param msg: Assertion message
:type msg: str, optional
: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
"""
if const_expr(self._disable):
return
if const_expr(not isinstance(idx, int)):
raise ValueError(f"expects idx to be 'int', but got {type(idx)}")
if const_expr(idx >= self._num_assertions):
raise ValueError(f"please increase the number of assertions!!!")
if const_expr(self._init_value is True):
self._tensor[idx] = pred and self._tensor[idx]
else:
self._tensor[idx] = pred
self._msgs[idx] = f"{msg}\nAt {loc}"
self._used_indices.add(idx)
def __enter__(self):
"""Enter context manager."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Exit context manager and verify assertions if no exception occurred."""
# Only verify if there was no exception in the with block
if exc_type is None and not self._disable:
# _CompileTimeAssertion doesn't have verify method as it's checked at compile time
pass
return False # Don't suppress exceptions
class RuntimeAssertion(Assertion):
"""Runtime assertion helper that verifies conditions at runtime.
```python
There are two modes to use RuntimeAssertion:
1. Manual mode - explicitly call verify():
```python
@cute.jit
def jit_func(assertions: Assertion):
assertions.store(0, pred, "assertion failed")
assertions = cute.testing.RuntimeAssertion(num_assertions=1)
jit_func(assertions)
assertions.verify()
```
2. Context manager mode - automatically verifies on exit:
```python
with cute.testing.RuntimeAssertion(num_assertions=1) as assertions:
jit_func(assertions)
# verify() is called automatically after the with block
```
"""
def __init__(
self,
num_assertions: int = 1,
device=None,
disable: bool = False,
init_value: bool = False,
):
"""Initialize _RuntimeAssertion.
:param num_assertions: Number of assertions to support
:param device: Device to run assertions on (None for CPU, "cuda" for GPU)
:param disable: If True, assertions are disabled
:param init_value: Initial value for assertion tensor
"""
self._num_assertions = num_assertions
self._device = device
self._disable = disable
self._msgs = [""] * num_assertions
self._init_value = init_value
self._used_indices = set()
if self._disable:
return
import torch
self._torch_tensor = torch.full(
(self._num_assertions,),
device=self._device,
dtype=torch.bool,
fill_value=init_value,
)
self._tensor = from_dlpack(self._torch_tensor)
def __c_pointers__(self):
"""Get C pointers for passing to JIT functions."""
if self._disable:
return []
return self._tensor.__c_pointers__()
def __get_mlir_types__(self):
"""Get MLIR types for code generation."""
if self._disable:
return []
return self._tensor.__get_mlir_types__()
def __new_from_mlir_values__(self, values):
"""Create new instance from MLIR values (for JIT compilation)."""
if self._disable:
return _CompileTimeAssertion(
None,
self._num_assertions,
self._msgs,
self._device,
self._disable,
self._init_value,
self._used_indices,
)
return _CompileTimeAssertion(
_Tensor(values[0], dtype=Boolean),
self._num_assertions,
self._msgs,
self._device,
self._disable,
self._init_value,
self._used_indices,
)
def verify(self):
"""Verify all assertions have passed."""
if self._disable:
return
import torch
if self._device is not None:
torch.cuda.synchronize()
false_indices = torch.where(self._torch_tensor == False)[0].tolist()
valid_indices = [idx for idx in false_indices if idx in self._used_indices]
if len(valid_indices) > 0:
# emit the first assertion error.
raise AssertionError(self._msgs[valid_indices[0]])
def __enter__(self):
"""Enter the context manager, returns self for use in 'with' statement."""
return self
def __exit__(self, exc_type, exc_val, exc_tb):
"""Exit the context manager, automatically calls verify()."""
if exc_type is None:
# Only verify if no exception occurred in the with block
self.verify()
# Return False to propagate any exception that occurred
return False
def _maybe_recast_tensor_from_f4(src: Tensor, tv_layout: Layout):
if src.element_type.width == 4:
tv_layout = cute.recast_layout(8, 4, tv_layout)
src = cute.recast_tensor(src, dtype=Int8)
tv_layout = recast_layout(8, 4, tv_layout)
src = recast_tensor(src, dtype=Int8)
return src, tv_layout
def _maybe_recast_to_f4(input: cute.TensorSSA, dtype: Type[Numeric]):
def _maybe_recast_to_f4(input: TensorSSA, dtype: Type[Numeric]):
"""Conditionally recasts the tensor to 4-bit type if the destination type is 4-bit.
:param input: The input tensor to recast.
@@ -56,18 +306,18 @@ def _maybe_recast_to_f4(input: cute.TensorSSA, dtype: Type[Numeric]):
raise TypeError(f"dst_ty must be a type of Numeric, but got {dtype}")
if dtype.width == 4:
recast_shape = cute.recast_layout(4, 8, cute.make_layout(input.shape)).shape
recast_shape = recast_layout(4, 8, make_layout(input.shape)).shape
i4_vec = vector.bitcast(
T.vector(input.type.shape[0] * 2, T.i(4)), input.maybe_downcast()
)
res_vect = builtin.unrealized_conversion_cast(
[T.vector(i4_vec.type.shape[0], dtype.mlir_type)], [i4_vec]
)
return cute.TensorSSA(res_vect, recast_shape, dtype)
return TensorSSA(res_vect, recast_shape, dtype)
return input
def _maybe_recast_from_f4(input: cute.TensorSSA, src_dtype: Type[Numeric]):
def _maybe_recast_from_f4(input: TensorSSA, src_dtype: Type[Numeric]):
"""Conditionally recasts the tensor from 4-bit type if the source type is 4-bit.
:param input: The input tensor to recast.
@@ -79,23 +329,23 @@ def _maybe_recast_from_f4(input: cute.TensorSSA, src_dtype: Type[Numeric]):
raise TypeError(f"src_ty must be a type of Numeric, but got {src_dtype}")
if src_dtype.width == 4:
recast_shape = cute.recast_layout(8, 4, cute.make_layout(input.shape)).shape
recast_shape = recast_layout(8, 4, make_layout(input.shape)).shape
i4_vec = builtin.unrealized_conversion_cast(
[T.vector(input.type.shape[0], T.i(4))], [input.maybe_downcast()]
)
res_vect = vector.bitcast(T.vector(i4_vec.type.shape[0] // 2, T.i8()), i4_vec)
return cute.TensorSSA(res_vect, recast_shape, Int8)
return TensorSSA(res_vect, recast_shape, Int8)
return input
@CuTeDSL.kernel
def _convert_kernel(
gSrc: cute.Tensor,
gDst: cute.Tensor,
cSrc: cute.Tensor,
src_tv_layout: cute.Layout,
dst_tv_layout: cute.Layout,
src_shape: cute.Shape,
gSrc: Tensor,
gDst: Tensor,
cSrc: Tensor,
src_tv_layout: Layout,
dst_tv_layout: Layout,
src_shape: Shape,
src_ty,
dst_ty,
):
@@ -111,9 +361,9 @@ def _convert_kernel(
# compose with CTA TV layout
# tid, vid -> address
tidfrgSrc = cute.composition(ctaSrc, src_tv_layout) # (T,V)
tidfrgDst = cute.composition(ctaDst, dst_tv_layout) # (T,V)
tidfrgCSrc = cute.composition(ctaCSrc, src_tv_layout) # (T,V)
tidfrgSrc = composition(ctaSrc, src_tv_layout) # (T,V)
tidfrgDst = composition(ctaDst, dst_tv_layout) # (T,V)
tidfrgCSrc = composition(ctaCSrc, src_tv_layout) # (T,V)
# print(f"tidfrgSrc = {tidfrgSrc.type}")
# slice for threads
@@ -124,19 +374,19 @@ def _convert_kernel(
# print(f"thrSrc = {thrSrc.type}")
# predicate
if cute.elem_less(thrCSrc[0], src_shape):
if elem_less(thrCSrc[0], src_shape):
# allocate fragments for gmem->rmem
frgSrc = cute.make_rmem_tensor(
cute.get(src_tv_layout, mode=[1]), gSrc.element_type
frgSrc = make_rmem_tensor(
get(src_tv_layout, mode=[1]), gSrc.element_type
) # (V)
frgDst = cute.make_rmem_tensor(
cute.get(dst_tv_layout, mode=[1]), gDst.element_type
frgDst = make_rmem_tensor(
get(dst_tv_layout, mode=[1]), gDst.element_type
) # (V)
# print(f"frgSrc = {frgSrc.type}")
# Move data to reg address space
copy_atom_load = cute.make_copy_atom(nvgpu.CopyUniversalOp(), gSrc.element_type)
cute.copy(copy_atom_load, thrSrc, frgSrc)
copy_atom_load = make_copy_atom(nvgpu.CopyUniversalOp(), gSrc.element_type)
copy(copy_atom_load, thrSrc, frgSrc)
vec_src = frgSrc.load()
vec_src = _maybe_recast_to_f4(vec_src, src_ty)
@@ -145,14 +395,14 @@ def _convert_kernel(
frgDst.store(vec_dst)
# Copy the results back to c
copy_atom_stg = cute.make_copy_atom(nvgpu.CopyUniversalOp(), gDst.element_type)
cute.copy(copy_atom_stg, frgDst, thrDst)
copy_atom_stg = make_copy_atom(nvgpu.CopyUniversalOp(), gDst.element_type)
copy(copy_atom_stg, frgDst, thrDst)
@CuTeDSL.jit(preprocess=False)
def _convert(
src: cute.Tensor,
dst: cute.Tensor,
src: Tensor,
dst: Tensor,
leading_mode: Constexpr,
elem_per_copy: Constexpr,
):
@@ -160,35 +410,29 @@ def _convert(
src_ty = src.element_type
dst_ty = dst.element_type
tv_layout = cute.make_layout((128, elem_per_copy), stride=(elem_per_copy, 1))
tv_layout = make_layout((128, elem_per_copy), stride=(elem_per_copy, 1))
# Step 2. maybe recast from f4 tensor
src, src_tv_layout = _maybe_recast_tensor_from_f4(src, tv_layout)
dst, dst_tv_layout = _maybe_recast_tensor_from_f4(dst, tv_layout)
src_shape = src.shape
# predicate tensor
idA = cute.make_identity_tensor(src.shape)
idA = make_identity_tensor(src.shape)
# Step 3. select a proper tiling pattern as (...,TileV, ...)
src_cta_tiler = [
1,
] * cute.rank(src.layout)
src_cta_tiler[leading_mode] = cute.size(src_tv_layout) # (...,TileV,...)
] * rank(src.layout)
src_cta_tiler[leading_mode] = size(src_tv_layout) # (...,TileV,...)
dst_cta_tiler = [
1,
] * cute.rank(dst.layout)
dst_cta_tiler[leading_mode] = cute.size(dst_tv_layout) # (...,TileV,...)
] * rank(dst.layout)
dst_cta_tiler[leading_mode] = size(dst_tv_layout) # (...,TileV,...)
# Step 4. partition input and output tensor by cta tiler.
gS = cute.zipped_divide(
src, tuple(src_cta_tiler)
) # ((...,TileV,...),(...,RestV,...))
cS = cute.zipped_divide(
idA, tuple(src_cta_tiler)
) # ((...,TileV,...),(...,RestV,...))
gD = cute.zipped_divide(
dst, tuple(dst_cta_tiler)
) # ((...,TileV,...),(...,RestV,...))
gS = zipped_divide(src, tuple(src_cta_tiler)) # ((...,TileV,...),(...,RestV,...))
cS = zipped_divide(idA, tuple(src_cta_tiler)) # ((...,TileV,...),(...,RestV,...))
gD = zipped_divide(dst, tuple(dst_cta_tiler)) # ((...,TileV,...),(...,RestV,...))
# print(f"{gS.type=}")
_convert_kernel(
@@ -201,8 +445,8 @@ def _convert(
src_ty,
dst_ty,
).launch(
grid=[cute.size(gS, mode=[1]), 1, 1],
block=[cute.size(src_tv_layout, mode=[0]), 1, 1],
grid=[size(gS, mode=[1]), 1, 1],
block=[size(src_tv_layout, mode=[0]), 1, 1],
)
@@ -210,7 +454,7 @@ def _convert(
# And when src or dst dtype is narrow precision(Float4E2M1FN/Float8E8M0FNU/Float8E4M3FN), the shape of
# their leading dimension should be 4(fp8)/8(fp4) element align. (nvgpu.cvt_fptrunc/cvt_fpext
# needs 32-bits aligned input/output)
def convert(src: cute.Tensor, dst: cute.Tensor):
def convert(src: Tensor, dst: Tensor):
assert len(src.shape) == len(dst.shape), (
"Shape of src and dst tensors should be the same rank."
)
@@ -292,6 +536,14 @@ class JitArguments:
def __init__(self, *args, **kwargs):
self.args = args
self.kwargs = kwargs
self.references = list()
def add_to_scope(self, references: Any) -> None:
"""
Keeps references to external variables (e.g., Torch tensors when taking a view)
in the scope of the lifetime of the JitArguments object.
"""
self.references.extend(references)
def _cuda_success(
@@ -428,6 +680,9 @@ def benchmark(
:rtype: float
"""
import cutlass.base_dsl.jit_executor as jit_executor
import cutlass.cutlass_dsl.cuda_jit_executor as cuda_jit_executor
if stream is None:
stream = cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT)
@@ -697,7 +952,7 @@ def _benchmark_for_autotune(
_cuda_success(err, "Error on querying event")
execution_time_ms.append(elapsed_time)
# unit: us
time_us = sum(execution_time_ms) / len(execution_time_ms)
time_us = sum(execution_time_ms) * 1e3 / len(execution_time_ms)
except Exception as e:
print(f"This config execution error: {e}")
time_us = float("inf")
@@ -775,6 +1030,7 @@ class autotune_jit:
Returns:
Decorated wrapper function
"""
from cutlass.cute import compile
# Initialize autotune parameters
if not hasattr(func, "_autotune_params"):
@@ -825,7 +1081,7 @@ class autotune_jit:
# For example, if current_config contains "cluster_shape_mn": (2, 1)
# It will override func's default parameter value
merged_kwargs = {**kwargs, **current_config}
compiled_func = cute.compile(
compiled_func = compile(
func._original_func, *args, **merged_kwargs
)
python/CuTeDSL/cutlass/cute/tuple.py
+165 -3
View File
@@ -17,7 +17,17 @@ from cutlass.cutlass_dsl import is_dynamic_expression, dsl_user_op
from cutlass._mlir import ir
import cutlass._mlir.dialects.cute as _cute_ir
from .typing import XTuple, IntTuple, Shape, Coord, Boolean, is_integer
from .typing import (
ComposedLayout,
Layout,
Stride,
XTuple,
IntTuple,
Shape,
Coord,
Boolean,
is_integer,
)
def wrap(x) -> Tuple[Any, ...]:
@@ -184,7 +194,7 @@ def product_each(a: IntTuple, *, loc=None, ip=None) -> IntTuple:
def find_if(
t: Union[tuple, ir.Value, int],
pred_fn: Callable[[int, Tuple[int, ...]], bool],
pred_fn: Callable[[Union[tuple, ir.Value, int], int], bool],
*,
loc=None,
ip=None,
@@ -197,7 +207,8 @@ def find_if(
:type t: Union[tuple, ir.Value, int]
:param pred_fn: A callable object (lambda, function, etc.) that predicates the value and position in t.
It takes the current leaf value and position, returns True if the value or position is satisfied.
:type pred_fn: Callable[[int, Tuple[int, ...]], bool]
The type must be compatible with rank(t).
:type pred_fn: Callable[[Union[tuple, ir.Value, int], int], bool]
:return: Index if found at top level, tuple of indices showing nested position, or None if not found
:rtype: Union[int, Tuple[int, ...], None]
@@ -329,3 +340,154 @@ def elem_less(
lhs_val = _pack_coord(lhs, loc=loc, ip=ip)
rhs_val = _pack_coord(rhs, loc=loc, ip=ip)
return Boolean(_cute_ir.elem_less(lhs_val, rhs_val, loc=loc, ip=ip))
def tuple_cat(*tuples):
"""Concatenate multiple tuples into a single tuple.
This function takes any number of tuples and concatenates them into a single tuple.
Non-tuple arguments are treated as single-element tuples.
:param tuples: Variable number of tuples to concatenate
:type tuples: tuple or any
:return: A single concatenated tuple
:rtype: tuple
**Examples:**
.. code-block:: python
>>> tuple_cat((1, 2), (3, 4))
(1, 2, 3, 4)
>>> tuple_cat((1,), (2, 3), (4,))
(1, 2, 3, 4)
>>> tuple_cat(1, (2, 3))
(1, 2, 3)
"""
result = ()
for t in tuples:
if isinstance(t, tuple):
result += t
else:
result += (t,)
return result
def transform_apply(*args, f: Callable, g: Callable):
"""Transform elements of tuple(s) with f, then apply g to all results.
This function applies f to corresponding elements across input tuple(s),
then applies g to all transformed results. It mimics the C++ CuTe implementation.
Supports multiple signatures:
- transform_apply(t, f, g): For single tuple, computes g(f(t[0]), f(t[1]), ...)
- transform_apply(t0, t1, f, g): For two tuples, computes g(f(t0[0], t1[0]), f(t0[1], t1[1]), ...)
- transform_apply(t0, t1, t2, ..., f, g): For multiple tuples of same length
For non-tuple inputs, f is applied to the input(s) and g is applied to that single result.
:param args: One or more tuples (or non-tuples) to transform
:param f: The function to apply to each element (or corresponding elements across tuples)
:type f: Callable
:param g: The function to apply to all transformed elements
:type g: Callable
:param loc: Source location for MLIR, defaults to None
:type loc: optional
:param ip: Insertion point, defaults to None
:type ip: optional
:return: The result of applying g to all transformed elements
:rtype: any
**Examples:**
.. code-block:: python
>>> transform_apply((1, 2, 3), f=lambda x: x * 2, g=lambda *args: sum(args))
12 # (1*2 + 2*2 + 3*2) = 12
>>> transform_apply((1, 2), f=lambda x: (x, x+1), g=tuple_cat)
(1, 2, 2, 3)
>>> transform_apply((1, 2), (3, 4), f=lambda x, y: x + y, g=lambda *args: args)
(4, 6)
"""
if not isinstance(f, Callable):
raise TypeError(f"f must be callable, but got {type(f)}")
if not isinstance(g, Callable):
raise TypeError(f"g must be callable, but got {type(g)}")
if not args:
raise ValueError("transform_apply requires at least one argument")
# Check if first argument is a tuple to determine behavior
if isinstance(args[0], tuple):
# Verify all args are tuples of the same length
if not all(isinstance(arg, tuple) for arg in args):
raise TypeError("All arguments must be tuples or all must be non-tuples")
tuple_length = len(args[0])
for i, arg in enumerate(args[1:], 1):
if len(arg) != tuple_length:
raise ValueError(
f"All tuple arguments must have the same length. "
f"arg[0] has length {tuple_length}, but arg[{i}] has length {len(arg)}"
)
# Apply f to corresponding elements across all tuples: g(f(args[0][i], args[1][i], ...), ...)
transformed_results = tuple(
f(*(arg[i] for arg in args)) for i in range(tuple_length)
)
return g(*transformed_results)
else:
# Non-tuple case: apply f to all args, then g to that single result
result = f(*args)
return g(result)
def filter_tuple(*args, f: Callable):
"""Filter and flatten tuple elements by applying a function.
The function f should return tuples, which are then concatenated together
to produce the final result. This is useful for filtering and transforming
tuple structures in a single pass.
:param t: The tuple to filter
:type t: Union[tuple, ir.Value, int]
:param f: The function to apply to each element of t
:type f: Callable
:param loc: Source location for MLIR, defaults to None
:type loc: optional
:param ip: Insertion point, defaults to None
:type ip: optional
:return: A concatenated tuple of all results
:rtype: tuple
**Examples:**
.. code-block:: python
>>> # Keep only even numbers, wrapped in tuples
>>> filter_tuple((1, 2, 3, 4), lambda x: (x,) if x % 2 == 0 else ())
(2, 4)
>>> # Duplicate each element
>>> filter_tuple((1, 2, 3), lambda x: (x, x))
(1, 1, 2, 2, 3, 3)
"""
if not isinstance(f, Callable):
raise TypeError(f"f must be callable, but got {type(f)}")
return transform_apply(*args, f=f, g=lambda *args: tuple_cat(*args))
__all__ = [
"transform_leaf",
"find_if",
"find",
"flatten_to_tuple",
"unflatten",
"product",
"product_like",
"product_each",
"elem_less",
"tuple_cat",
"transform_apply",
"filter_tuple",
]
python/CuTeDSL/cutlass/cute/typing.py
+14 -1
View File
@@ -52,6 +52,15 @@ class SymInt:
[self._width == other._width, self._divisibility == other._divisibility]
)
def __mod__(self, other: int) -> Union["SymInt", int]:
if self._divisibility % other != 0:
from math import gcd
div = gcd(self._divisibility, other)
return SymInt(self._width, divisibility=div)
else:
return 0
def __c_pointers__(self):
return [ctypes.c_void_p(0).value]
@@ -381,10 +390,14 @@ __all__ = [
"Float6E2M3FN",
"Float6E3M2FN",
"IntTuple",
"Layout",
"ScaledBasis",
"Coord",
"Shape",
"Stride",
"Layout",
"ComposedLayout",
"Pointer",
"Tensor",
"Tile",
"Tiler",
"XTuple",