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

v4.2 tag release. (#2638)

Browse Source
This commit is contained in:
Junkai-Wu
2025-09-16 00:21:53 +08:00
committed by GitHub
parent 56f0718a97
commit 6a35b4d22f
161 changed files with 14056 additions and 3793 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -126,6 +126,7 @@ from .core import (
basic_copy_if,
autovec_copy,
copy,
copy_atom_call,
gemm,
# Wrapper classes
ComposedLayout,
@@ -290,6 +291,7 @@ __all__ = [
"basic_copy_if",
"autovec_copy",
"copy",
"copy_atom_call",
"gemm",
# Tensor creation
"full",

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

@@ -315,3 +315,35 @@ def mbarrier_arrive(
loc=loc,
ip=ip,
)
@dsl_user_op
def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> None:
"""
Arrives on an mbarrier for async load **without incrementing** the arrival count
(`cp.async.mbarrier.arrive.shared ..., noinc=1`).
Used in the warp-specialized kernel when the non-TMA load warp(producer) is not the same
as the math/epilogue warp(consumer).
:param mbar_ptr: A pointer to the mbarrier in SMEM
:type mbar_ptr: Pointer
"""
arch = CuTeDSL._get_dsl().envar.arch
check_value_in(
arch,
[
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
],
"arch",
)
mbar_llvm_ptr = mbar_ptr.llvm_ptr
nvvm.cp_async_mbarrier_arrive_shared(
mbar_llvm_ptr,
noinc=True,
loc=loc,
ip=ip,
)

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

@@ -11,6 +11,7 @@
from functools import partial
from typing import Optional, Tuple, Union, Callable
from typing_extensions import deprecated
from cutlass.cutlass_dsl import T, dsl_user_op
@@ -642,6 +643,9 @@ def rcp_approx(a: Union[float, Float32], *, loc=None, ip=None):
@dsl_user_op
@deprecated(
"cute.arch.exp2 is deprecated, use cute.math.exp2 with `fastmath=True` instead"
)
def exp2(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
return Float32(
llvm.inline_asm(
@@ -656,15 +660,19 @@ def exp2(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
)
# TODO: add `fastmath` flag for this op
@dsl_user_op
@deprecated(
"cute.arch.exp is deprecated, use cute.math.exp with `fastmath=True` instead"
)
def exp(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
LOG2_E = 1.4426950408889634
return exp2(a * LOG2_E, loc=loc, ip=ip)
# TODO: add `fastmath` flag for this op
@dsl_user_op
@deprecated(
"cute.arch.exp_packed_f32x2 is deprecated, use cute.arch.mul_packed_f32x2 and cute.math.exp2 with `fastmath=True` instead"
)
def exp_packed_f32x2(
a: Tuple[Float32, Float32], *, loc=None, ip=None
) -> Tuple[Float32, Float32]:

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

@@ -31,7 +31,6 @@ from typing import (
Optional,
)
from enum import Enum, auto
from typing_extensions import deprecated
from cutlass.cutlass_dsl import (
const,
@@ -1662,7 +1661,9 @@ class _Tensor(Tensor):
@dsl_user_op
def print_tensor(tensor: Tensor, *, verbose: bool = False, loc=None, ip=None):
def print_tensor(
tensor: Union[Tensor, "TensorSSA"], *, verbose: bool = False, loc=None, ip=None
):
"""Print content of the tensor in human readable format.
Outputs the tensor data in a structured format showing both metadata
@@ -1693,6 +1694,11 @@ def print_tensor(tensor: Tensor, *, verbose: bool = False, loc=None, ip=None):
[ 0.9159, 0.7577, 0.6918, 0.0754, 0.0591],
[ 0.6551, 0.1626, 0.1189, 0.0292, 0.8655]])
"""
if isinstance(tensor, TensorSSA):
tmp = make_fragment(tensor.shape, tensor.dtype)
tmp.store(tensor)
tensor = tmp
if not isinstance(tensor.type, _cute_ir.MemRefType):
raise NotImplementedError(
f"printing {tensor} is not supported because it doesn't support trivial dereferencing. "
@@ -1769,7 +1775,7 @@ def is_static(x: Union[ir.Type, ir.Value, XTuple]) -> bool:
return False
elif is_dynamic_expression(x):
return _cute_ir.is_static(x.type)
elif isinstance(x, int) or x is None:
elif isinstance(x, (bool, int, float)) or x is None:
return True
elif isinstance(x, ScaledBasis):
return x.is_static()
@@ -2241,7 +2247,7 @@ def is_weakly_congruent(
* X is a non-tuple value, OR
* X and Y are both tuples of the same rank AND all corresponding elements are weakly congruent.
Weak congruence allows scalar values to match with tuples, making it useful
Weak congruence allows scalar values to match with tuples, making it useful
for determining whether an object has a hierarchical structure "up to" another.
:param a: First object to compare
@@ -2921,33 +2927,46 @@ def flatten_to_tuple(a: Union[IntTuple, Coord, Shape, Stride]) -> tuple:
return tuple(chain.from_iterable(tuple(flatten_to_tuple(x) for x in a)))
def flatten(a: Union[IntTuple, Coord, Shape, Stride, Layout, Tensor]) -> tuple:
@overload
def flatten(a: Union[IntTuple, Coord, Shape, Stride]) -> IntTuple: ...
@overload
def flatten(a: Tensor) -> Tensor: ...
@overload
def flatten(a: Layout) -> Layout: ...
def flatten(a):
"""Flattens a CuTe data structure into a simpler form.
For tuples, this function flattens the structure into a single-level tuple.
For non-tuple types, it returns the input unchanged.
For layouts, it returns a new layout with flattened shape and stride.
For tensors, it returns a new tensor with flattened layout.
For other types, it returns the input unchanged.
:param a: The structure to flatten
:type a: Union[IntTuple, Coord, Shape, Stride, Layout, Tensor]
:return: The flattened structure
:rtype: Union[tuple, Any]
:raises NotImplementedError: If input is a Layout or Tensor
**Examples:**
.. code-block:: python
flatten((1, 2, 3)) # Returns (1, 2, 3)
flatten(((1, 2), (3, 4))) # Returns (1, 2, 3, 4)
flatten(5) # Returns 5
"""
if isinstance(a, (Layout, Tensor)):
raise NotImplementedError("flatten layout and tensor is not supported")
flatten((1, 2, 3)) # Returns (1, 2, 3)
flatten(((1, 2), (3, 4))) # Returns (1, 2, 3, 4)
flatten(5) # Returns 5
flatten(Layout(shape, stride)) # Returns Layout(flatten(shape), flatten(stride))
flatten(Tensor(layout)) # Returns Tensor(flatten(layout))
if not isinstance(a, tuple):
return a
else:
"""
if isinstance(a, Tensor):
return make_tensor(a.iterator, flatten(a.layout))
elif isinstance(a, Layout):
return make_layout(flatten(a.shape), stride=flatten(a.stride))
elif isinstance(a, tuple):
return flatten_to_tuple(a)
else:
return a
def unflatten(
@@ -4120,14 +4139,14 @@ def complement(
@dsl_user_op
def right_inverse(input: Layout, *, loc=None, ip=None) -> Layout:
if not isinstance(input, Layout):
raise TypeError(f"expects input of type Layout, but got {type(Layout)}")
raise TypeError(f"expects input of type Layout, but got {type(input)}")
return _cute_ir.right_inverse(input=input, loc=loc, ip=ip)
@dsl_user_op
def left_inverse(input: Layout, *, loc=None, ip=None) -> Layout:
if not isinstance(input, Layout):
raise TypeError(f"expects input of type Layout, but got {type(Layout)}")
raise TypeError(f"expects input of type Layout, but got {type(input)}")
return _cute_ir.left_inverse(input=input, loc=loc, ip=ip)
@@ -5156,7 +5175,6 @@ def _make_tiled_copy(atom, layout_tv, tiler_mn, *, loc=None, ip=None):
return TiledCopy(atom.op, trait)
@deprecated("Use make_tiled_copy_tv instead")
def make_tiled_copy(atom, layout_tv, tiler_mn, *, loc=None, ip=None):
"""Create a tiled type given a TV partitioner and tiler.
@@ -5434,6 +5452,14 @@ def gemm(
For MMA Atoms that require single-threaded execution, the gemm op automatically handles thread
election internally. Manual thread selection is not required in such cases.
Following dispatch rules are supported:
- Dispatch [1]: (V) x (V) => (V) => (V,1,1) x (V,1,1) => (V,1,1)
- Dispatch [2]: (M) x (N) => (M,N) => (1,M,1) x (1,N,1) => (1,M,N)
- Dispatch [3]: (M,K) x (N,K) => (M,N) => (1,M,K) x (1,N,K) => (1,M,N)
- Dispatch [4]: (V,M) x (V,N) => (V,M,N) => (V,M,1) x (V,N,1) => (V,M,N)
- Dispatch [5]: (V,M,K) x (V,N,K) => (V,M,N)
:param atom: MMA atom
:type atom: MmaAtom
:param d: Destination tensor
@@ -5454,6 +5480,27 @@ def gemm(
:rtype: None
"""
a_rank = rank(a.shape)
b_rank = rank(b.shape)
c_rank = rank(c.shape)
d_rank = rank(d.shape)
if a_rank != b_rank:
raise ValueError("`a` and `b` must have the same rank")
if c_rank != d_rank:
raise ValueError("`c` and `d` must have the same rank")
if a_rank == 1:
if c_rank > 2:
raise ValueError("`c` must have rank <= 2 when `a` has rank 1")
elif a_rank == 2:
if c_rank not in (2, 3):
raise ValueError("`c` must have rank 2 or 3 when `a` has rank 2")
elif a_rank == 3:
if c_rank != 3:
raise ValueError("`c` must have rank 3 when `a` has rank 3")
value = atom._unpack(loc=loc, ip=ip, **kwargs)
return _cute_ir.gemm(value, d.value, a.value, b.value, c.value, loc=loc, ip=ip)
@@ -5645,6 +5692,76 @@ def copy(
return _cute_ir.copy(value, src.value, dst.value, pred=pred, loc=loc, ip=ip)
@dsl_user_op
def copy_atom_call(
atom: CopyAtom,
src: Tensor,
dst: Tensor,
*,
pred: Optional[Tensor] = None,
loc=None,
ip=None,
**kwargs,
) -> None:
"""
Execute a single copy atom operation.
The copy_atom_call operation executes a copy atom with the given operands.
Following src/dst layout of atom are valid:
* ((atom_v))
* (atom_v)
Note: The format ((atom_v, rest_v)) is NOT valid for copy_atom_call since it would
require multiple atom operations, which contradicts the definition of a single copy atom call.
Examples:
.. code-block:: python
# Call a copy atom operation
cute.copy_atom_call(copy_atom, src_tensor, dst_tensor)
An additional predication tensor can be provided. If the partitioned tensors have the following
logical profile ``((ATOM_V,ATOM_REST),REST_M,...)``, the predication tensor must have a profile
consistent with ``(ATOM_REST,REST_M,...)``.
"""
if isinstance(src.type, _cute_ir.MemRefType) and isinstance(
dst.type, _cute_ir.MemRefType
):
if src.element_type.width != dst.element_type.width:
raise TypeError(
"`copy_atom_call` currently only supports equal source and destination "
"element type bit width"
)
value = atom._unpack(loc=loc, ip=ip, **kwargs)
if isinstance(pred, Tensor):
pred = pred.value
return _cute_ir.copy_atom_call(
value, src.value, dst.value, pred=pred, loc=loc, ip=ip
)
def prefetch(atom: CopyAtom, src: Tensor, *, loc=None, ip=None) -> None:
"""
The Prefetch algorithm.
The "prefetch" expects source tensors to be partitioned according to the provided Copy Atom.
Prefetch is used for loading tensors from global memory to L2.
Prefetch accepts Copy Atom but not all are allowed. Currently, only support for tma load tensor prefetch.
.. code-block:: python
cute.prefetch(tma_atom, src)
For Copy Atoms that require single-threaded execution, the copy op automatically handles thread
election internally. Manual thread selection is not required in such cases.
"""
dummy_tma_bar_ptr = make_ptr(Int64, 0, AddressSpace.smem, loc=loc, ip=ip)
value = atom._unpack(loc=loc, ip=ip, tma_bar_ptr=dummy_tma_bar_ptr)
return _cute_ir.prefetch(value, src.value, loc=loc, ip=ip)
####################################################################################################
#
# TensorSSA class (experimental)
@@ -5657,6 +5774,11 @@ class ReductionOp(Enum):
MUL = auto()
MAX = auto()
MIN = auto()
INC = auto()
DEC = auto()
AND = auto()
OR = auto()
XOR = auto()
def __str__(self):
return self.name.lower()
@@ -5697,6 +5819,7 @@ class TensorSSA(cutlass_arith.ArithValue):
self._shape = shape
self._dtype = dtype
self._layout = None
@property
def dtype(self) -> Type[Numeric]:
@@ -5776,13 +5899,26 @@ class TensorSSA(cutlass_arith.ArithValue):
):
res_type = Boolean
if lhs.shape != rhs.shape:
raise ValueError(
f"lhs and rhs must have the same shape type, but got {lhs.shape} and {rhs.shape}"
)
assert isinstance(rhs, TensorSSA), f"rhs must be TensorSSA but got {rhs}"
if not isinstance(rhs, TensorSSA):
raise TypeError(f"rhs must be TensorSSA but got {rhs}")
def _broadcast(s, t):
if s == 1:
return t
elif t == 1:
return s
elif s == t:
return s
else:
raise ValueError(f"cannot broadcast {s} and {t}")
max_rank = max(rank(lhs.shape), rank(rhs.shape))
lhs_shape = append(lhs.shape, 1, up_to_rank=max_rank)
rhs_shape = append(rhs.shape, 1, up_to_rank=max_rank)
res_shape = transform_leaf(_broadcast, lhs_shape, rhs_shape)
# broadcast to the same shape
lhs = lhs.broadcast_to(res_shape)
rhs = rhs.broadcast_to(res_shape)
if (
op in (operator.add, operator.sub)
@@ -5807,6 +5943,38 @@ class TensorSSA(cutlass_arith.ArithValue):
return res
def broadcast_to(self, target_shape: Shape, *, loc=None, ip=None) -> "TensorSSA":
"""
Broadcast the tensor to the target shape.
"""
# pad source shape to the same rank
shape = append(self.shape, 1, up_to_rank=rank(target_shape))
if shape == target_shape:
return self
def _check_broadcast(s, t):
if s != t and s != 1:
raise ValueError(
f"src_shape and target_shape must be the same when src_shape is not 1, but got {s} and {t}"
)
transform_leaf(_check_broadcast, shape, target_shape)
# reshape to flatten N-D vector
flat_shp = flatten_to_tuple(shape)
temp_ty = ir.VectorType.get(list(flat_shp), self.dtype.mlir_type)
temp_vect = vector.shape_cast(temp_ty, self, loc=loc, ip=ip)
# broadcast to result N-D vector
flat_tgt_shp = flatten_to_tuple(target_shape)
temp_tgt_ty = ir.VectorType.get(list(flat_tgt_shp), self.dtype.mlir_type)
temp_tgt_vect = vector.broadcast(temp_tgt_ty, temp_vect, loc=loc, ip=ip)
res_1d_ty = ir.VectorType.get([size(target_shape)], self.dtype.mlir_type) # type: ignore
res_1d_vect = vector.shape_cast(res_1d_ty, temp_tgt_vect, loc=loc, ip=ip)
return TensorSSA(res_1d_vect, target_shape, self.dtype)
def __pow__(self, other, *, loc=None, ip=None) -> "TensorSSA":
"""
Returns the results of tensor^other.
@@ -6093,6 +6261,16 @@ class TensorSSA(cutlass_arith.ArithValue):
"""
return self._apply_op(operator.and_, other, flip=True, loc=loc, ip=ip)
def __neg__(self, *, loc=None, ip=None) -> "TensorSSA":
"""
Returns the negation of the tensor.
:return: The element-wise negation of the tensor
:rtype: TensorSSA
"""
return self._apply_op(operator.sub, 0, flip=True, loc=loc, ip=ip)
def _flatten_shape_and_coord(self, crd, *, loc=None, ip=None):
# Coalesce and flatten source layout at terminal of coordinate
# (N_0,(N_1,...), ...) -> (N_0,N_1,N_2,...)
@@ -6158,17 +6336,13 @@ class TensorSSA(cutlass_arith.ArithValue):
if crd is None:
return self
if not has_underscore(crd) or depth(crd) == 0:
idx = crd2idx(crd, make_layout(self._shape))
if is_static(idx):
res = vector.extract(
self, dynamic_position=[], static_position=[idx], loc=loc, ip=ip
)
else:
res = vector.extract(
self, dynamic_position=[crd], static_position=[], loc=loc, ip=ip
)
return self.dtype(res)
if not has_underscore(crd):
if self._layout is None:
self._layout = make_layout(self._shape, loc=loc, ip=ip)
idx = crd2idx(crd, self._layout, loc=loc, ip=ip)
idx_val = as_numeric(idx).ir_value(loc=loc, ip=ip)
res_val = vector.extractelement(self, position=idx_val, loc=loc, ip=ip)
return self.dtype(res_val)
if not is_static(crd):
raise ValueError("dynamic coordinate is not supported")
@@ -6274,7 +6448,7 @@ class TensorSSA(cutlass_arith.ArithValue):
:type op: operator
:param init_val: The initial value for the reduction
:type init_val: numeric
:param reduction_profile: Specifies which dimensions to reduce. Dimensions marked with '_' are kept.
:param reduction_profile: Specifies which dimensions to reduce. Dimensions marked with `None` are kept.
:type reduction_profile: Coord
:return: The reduced tensor
@@ -6289,9 +6463,9 @@ class TensorSSA(cutlass_arith.ArithValue):
reduce(f32 o (4, 5))
=> f32
reduce(f32 o (4, (5, 4)), reduction_profile=(_, 1))
reduce(f32 o (4, (5, 4)), reduction_profile=(None, 1))
=> f32 o (4,)
reduce(f32 o (4, (5, 4)), reduction_profile=(_, (_, 1)))
reduce(f32 o (4, (5, 4)), reduction_profile=(None, (None, 1)))
=> f32 o (4, (5,))
"""
# short-cut to no-op
@@ -6354,21 +6528,6 @@ class TensorSSA(cutlass_arith.ArithValue):
return self._build_result(res_vect, res_shp, loc=loc, ip=ip)
def _get_attr_for_type(ty, value):
if isinstance(ty, ir.IntegerType):
return ir.IntegerAttr.get(ty, value.to(int))
elif isinstance(ty, ir.FloatType):
return ir.FloatAttr.get(ty, value.to(float))
else:
raise TypeError(f"unsupported type: {ty}")
def _splat(res_ty, fill_value):
elem_attr = _get_attr_for_type(res_ty.element_type, fill_value)
vect_attr = ir.DenseElementsAttr.get_splat(res_ty, elem_attr)
return arith.constant(res_ty, vect_attr)
@dsl_user_op
def full(shape, fill_value, dtype: Type[Numeric], *, loc=None, ip=None) -> TensorSSA:
"""
@@ -6389,9 +6548,14 @@ def full(shape, fill_value, dtype: Type[Numeric], *, loc=None, ip=None) -> Tenso
if isinstance(fill_value, (ir.Value, int, float, bool)):
fill_value = dtype(fill_value)
elif isinstance(fill_value, Numeric):
fill_value = fill_value.to(dtype, loc=loc, ip=ip)
else:
raise ValueError(f"Expected fill_value be numeric type, but got {fill_value}")
res_mlir_type = T.vector(size, dtype.mlir_type)
return TensorSSA(_splat(res_mlir_type, fill_value), shape, dtype)
res_ty = T.vector(size, dtype.mlir_type)
res_val = vector.splat(res_ty, fill_value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
return TensorSSA(res_val, shape, dtype)
def full_like(
@@ -6547,7 +6711,7 @@ class struct:
**Usage:**
.. code-block::
.. code-block:: python
# Supports base_dsl scalar int/float elements, array and nested struct:
@cute.struct
@@ -6661,7 +6825,8 @@ class struct:
Initializes a new memory range.
:param dtype: The data type.
:param size: The size of the memory range in bytes.
:param size: Size of the memory range in bytes. A size of **0** is accepted, but in that
case the range can only be used for its address (e.g. as a partition marker).
:param base: The base address of the memory range.
"""
self._dtype = dtype
@@ -6673,9 +6838,9 @@ class struct:
Returns start pointer to the data in this memory range.
:return: A pointer to the start of the memory range.
:raises AssertionError: If the size of the memory range is not greater than zero.
:raises AssertionError: If the size of the memory range is negative.
"""
assert self._size > 0
assert self._size >= 0
return recast_ptr(self._base, dtype=self._dtype)
def get_tensor(self, layout, swizzle=None, dtype=None):
@@ -6716,31 +6881,48 @@ class struct:
:param v: The object to align. Must be a struct, MemRange, or a scalar type.
:param align: The alignment value to set.
:return: A copy of the object with the specified alignment.
:raises TypeError: If the object is not a struct, MemRange, or a scalar type.
:ivar _dtype: The data type to be aligned.
:ivar _align: The alignment of the data type.
"""
_dtype = None
_align = None
def __new__(cls, name, bases, dct):
return super().__new__(cls, name, bases, dct)
def __getitem__(cls, params) -> Any:
if len(params) == 2:
obj, align = params
dtype, align = params
assert align > 0
else:
raise TypeError("Invalid struct.Align Arguments")
# make a copy of type and mark alignment
if struct._is_scalar_type(obj) or isinstance(
obj, (struct, struct._MemRangeMeta)
if not struct._is_scalar_type(dtype) and not isinstance(
dtype, (struct, struct._MemRangeMeta)
):
new_obj = py_copy.copy(obj)
setattr(new_obj, "_struct_alignment_", align)
return new_obj
else:
raise TypeError(
"align only can be applied to sturct/MemRange/base_dsl scalar"
"align only can be applied to struct/MemRange/base_dsl scalar"
)
# Create new class with alignment
new_cls = type(
f"struct.Align[{dtype.__name__}, {align}]",
(struct.Align,),
{"_dtype": dtype, "_align": align},
)
return new_cls
@property
def dtype(cls):
return cls._dtype
@property
def align(cls):
return cls._align
class Align(metaclass=_AlignMeta):
"""
Aligns the given type by `Align[T, alignment]`.
@@ -6768,6 +6950,7 @@ class struct:
:raises TypeError: If the struct is empty.
"""
self._cls = cls
self.__name__ = f"struct::{cls.__name__}"
# Get the class annotations
self._annotations = cls.__annotations__
# Create a dictionary to store the offsets
@@ -6780,12 +6963,10 @@ class struct:
raise TypeError("Empty struct is not supported!")
for name, object in self._annotations.items():
# get alignment of object
def alignof(object, default: int = 1):
return getattr(object, "_struct_alignment_", default)
# alignment for the next offset
def align_offset(offset, align):
return (offset + (align - 1)) & ~(align - 1)
sub_align = 1
if isinstance(object, struct._AlignMeta):
sub_align = object.align
object = object.dtype
# switch addition order to support dynamic size
def add_offset(val):
@@ -6793,35 +6974,37 @@ class struct:
# size of scalar
if struct._is_scalar_type(object):
dtype_size = object.width // 8
sub_align = alignof(object, dtype_size)
offset = align_offset(offset, sub_align)
dtype_size = max(1, object.width // 8)
sub_align = max(dtype_size, sub_align)
offset = self.align_offset(offset, sub_align)
self._offsets[name] = offset
offset = add_offset(dtype_size)
# size of array is size_in_bytes, alignment is elem_size
elif isinstance(object, struct._MemRangeMeta):
if object.size == 0:
continue # skip empty array
sub_align = alignof(object, max(1, object.elem_width // 8))
offset = align_offset(offset, sub_align)
# Allow empty array as a free marker-only struct member.
# Use max(sub_align, ) because we might have in the future some
# object.elem_width less than 8, such as fp4, bit and others,
# and align_offset() does not support an alignment of 0.
sub_align = max(object.elem_width // 8, sub_align)
offset = self.align_offset(offset, sub_align)
self._offsets[name] = offset
offset = add_offset(object.size_in_bytes)
# size of struct
elif isinstance(object, struct):
sub_align = max(object.__alignof__(), alignof(object))
offset = align_offset(offset, sub_align)
sub_align = max(object.__alignof__(), sub_align)
offset = self.align_offset(offset, sub_align)
self._offsets[name] = offset
offset = add_offset(object.__sizeof__())
else:
raise TypeError(
f"Struct element only support sturct/array/base_dsl scalar, "
f"Struct element only support struct/array/base_dsl scalar, "
f"but got {object}"
)
# Total aligment determined by the strictest requirement
alignment = max(alignment, sub_align)
# Total size determined by alignment
self._align_of = alignment
self._size_of = align_offset(offset, alignment)
self._size_of = self.align_offset(offset, alignment)
# create the __init__ method for decorated struct
def __call__(self, base: Any) -> None:
@@ -6840,6 +7023,8 @@ class struct:
setattr(cls, "_base", base)
for name, off in self._offsets.items():
obj = self._annotations[name]
if isinstance(obj, struct._AlignMeta):
obj = obj.dtype
if struct._is_scalar_type(obj):
new_obj = recast_ptr(base + off, dtype=obj)
setattr(cls, name, new_obj)
@@ -6851,7 +7036,7 @@ class struct:
setattr(cls, name, new_obj)
else:
raise TypeError(
f"Struct element only support sturct/array/base_dsl scalar, "
f"Struct element only support struct/array/base_dsl scalar, "
f"but got {obj}"
)
return cls
@@ -6872,3 +7057,14 @@ class struct:
# get alignment
def __alignof__(self) -> int:
return self._align_of
# util func for aligning offset
@staticmethod
def align_offset(offset, align):
"""
Return the round-up offset up to the next multiple of align.
"""
assert align > 0 and not (
align & (align - 1)
), "align should be a strictly positive power of 2."
return (offset + (align - 1)) & ~(align - 1)

217
python/CuTeDSL/cutlass/cute/math.py
View File

@@ -10,16 +10,53 @@
# is strictly prohibited.
from .core import TensorSSA
from .typing import Numeric
from cutlass._mlir.dialects import math, arith
from typing import Callable, Union
def acos(a: TensorSSA) -> TensorSSA:
def _math_op(func: Callable, fastmath: bool, *args, **kwargs):
"""Dispatch the function to either a TensorSSA or a Numeric(Float).
:param func: The function to dispatch
:param args: The input tensor or scalar
:param kwargs: The input tensor or scalar
"""
arg_type = type(args[0])
for arg in args:
if not isinstance(arg, TensorSSA) and (
not isinstance(arg, Numeric) or not type(arg).is_float
):
raise TypeError(
f"Expected a TensorSSA or Numeric(Float), but got {type(arg)}"
)
if not isinstance(arg, arg_type):
raise TypeError(
f"Expected all inputs to be of type {arg_type}, but got {type(arg)}"
)
fastmath_flag = arith.FastMathFlags.fast if fastmath else arith.FastMathFlags.none
if isinstance(args[0], TensorSSA):
return TensorSSA(
func(*args, fastmath=fastmath_flag), args[0].shape, args[0].dtype
)
else:
args = [a.ir_value() for a in args]
return func(*args, fastmath=fastmath_flag)
def acos(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc cosine of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc cosine of each element in input tensor
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -29,16 +66,20 @@ def acos(a: TensorSSA) -> TensorSSA:
y = x.load() # Load values
z = acos(y) # Compute arc cosine
"""
return TensorSSA(math.acos(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.acos, fastmath, a)
def asin(a: TensorSSA) -> TensorSSA:
def asin(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc sine of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc sine of each element in input tensor
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -48,18 +89,20 @@ def asin(a: TensorSSA) -> TensorSSA:
y = x.load() # Load values
z = asin(y) # Compute arc sine
"""
return TensorSSA(math.asin(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.asin, fastmath, a)
def atan(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def atan(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc tangent of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc tangent of each element in input tensor
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -70,23 +113,25 @@ def atan(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
z = atan(y) # Compute arc tangent
"""
raise NotImplementedError("atan is not implemented")
return TensorSSA(math.atan(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.atan, fastmath, a)
def atan2(a: TensorSSA, b: TensorSSA, fastmath: bool = False) -> TensorSSA:
def atan2(
a: Union[TensorSSA, Numeric], b: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise arc tangent of two tensors.
Computes atan2(a, b) element-wise. The function atan2(a, b) is the angle in radians
between the positive x-axis and the point given by the coordinates (b, a).
:param a: First input tensor (y-coordinates)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param b: Second input tensor (x-coordinates)
:type b: TensorSSA
:type b: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the arc tangent of a/b element-wise
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -96,20 +141,20 @@ def atan2(a: TensorSSA, b: TensorSSA, fastmath: bool = False) -> TensorSSA:
x = cute.make_fragment(ptr2, layout).load() # x coordinates
theta = atan2(y, x) # Compute angles
"""
return TensorSSA(
math.atan2(a, b, fastmath=arith.FastMathFlags.none), a.shape, a.dtype
)
return _math_op(math.atan2, fastmath, a, b)
def cos(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def cos(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise cosine of the input tensor.
:param a: Input tensor (in radians)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the cosine of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -119,21 +164,23 @@ def cos(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = cos(y) # Compute cosine
"""
return TensorSSA(math.cos(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.cos, fastmath, a)
def erf(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def erf(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise error function of the input tensor.
The error function is defined as:
erf(x) = 2/√π ∫[0 to x] exp(-t²) dt
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the error function value for each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -143,18 +190,43 @@ def erf(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = erf(y) # Compute error function
"""
return TensorSSA(math.erf(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.erf, fastmath, a)
def exp2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def exp(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise exponential of the input tensor.
:param a: Input tensor
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the exponential of each element
:rtype: Union[TensorSSA, Numeric]
Example:
.. code-block::
x = cute.make_fragment(layout) # Create tensor
y = x.load() # Load values
z = exp(y) # Compute exponential
"""
return _math_op(math.exp, fastmath, a)
def exp2(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise base-2 exponential of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing 2 raised to the power of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -164,18 +236,20 @@ def exp2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = exp2(y) # Compute 2^x
"""
return TensorSSA(math.exp2(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.exp2, fastmath, a)
def log(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def log(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise natural logarithm of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the natural logarithm of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -185,18 +259,20 @@ def log(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = log(y) # Compute natural logarithm
"""
return TensorSSA(math.log(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.log, fastmath, a)
def log2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def log2(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise base-2 logarithm of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the base-2 logarithm of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -206,18 +282,20 @@ def log2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = log2(y) # Compute log base 2
"""
return TensorSSA(math.log2(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.log2, fastmath, a)
def log10(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def log10(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise base-10 logarithm of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the base-10 logarithm of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -227,20 +305,22 @@ def log10(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = log10(y) # Compute log base 10
"""
return TensorSSA(math.log10(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.log10, fastmath, a)
def rsqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def rsqrt(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise reciprocal square root of the input tensor.
Computes 1/√x element-wise.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the reciprocal square root of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -250,18 +330,20 @@ def rsqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = rsqrt(y) # Compute 1/√x
"""
return TensorSSA(math.rsqrt(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.rsqrt, fastmath, a)
def sin(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def sin(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise sine of the input tensor.
:param a: Input tensor (in radians)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the sine of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -271,18 +353,20 @@ def sin(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = sin(y) # Compute sine
"""
return TensorSSA(math.sin(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.sin, fastmath, a)
def sqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def sqrt(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise square root of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the square root of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -292,16 +376,20 @@ def sqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = sqrt(y) # Compute square root
"""
return TensorSSA(math.sqrt(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.sqrt, fastmath, a)
def tan(a: TensorSSA) -> TensorSSA:
def tan(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise tangent of the input tensor.
:param a: Input tensor (in radians)
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the tangent of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -311,18 +399,20 @@ def tan(a: TensorSSA) -> TensorSSA:
y = x.load() # Load values
z = tan(y) # Compute tangent
"""
return TensorSSA(math.tan(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.tan, fastmath, a)
def tanh(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
def tanh(
a: Union[TensorSSA, Numeric], fastmath: bool = False
) -> Union[TensorSSA, Numeric]:
"""Compute element-wise hyperbolic tangent of the input tensor.
:param a: Input tensor
:type a: TensorSSA
:type a: Union[TensorSSA, Numeric]
:param fastmath: Enable fast math optimizations, defaults to False
:type fastmath: bool, optional
:return: Tensor containing the hyperbolic tangent of each element
:rtype: TensorSSA
:rtype: Union[TensorSSA, Numeric]
Example:
@@ -332,7 +422,7 @@ def tanh(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
y = x.load() # Load values
z = tanh(y) # Compute hyperbolic tangent
"""
return TensorSSA(math.tanh(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
return _math_op(math.tanh, fastmath, a)
__all__ = [
@@ -342,6 +432,7 @@ __all__ = [
"atan2",
"cos",
"erf",
"exp",
"exp2",
"log",
"log10",

45
python/CuTeDSL/cutlass/cute/nvgpu/common.py
View File

@@ -8,7 +8,7 @@
# Any use, reproduction, disclosure, or distribution of this software
# and related documentation outside the scope permitted by the EULA
# is strictly prohibited.
import enum
from dataclasses import dataclass
from typing import Type, Optional
@@ -101,6 +101,42 @@ class MmaUniversalTrait(core.Trait):
####################################################################################################
class MemoryOrder(enum.Enum):
WEAK = _cute_ir.MemOrderKind.WEAK
RELAXED = _cute_ir.MemOrderKind.RELAXED
ACQUIRE = _cute_ir.MemOrderKind.ACQUIRE
RELEASE = _cute_ir.MemOrderKind.RELEASE
ACQ_REL = _cute_ir.MemOrderKind.ACQ_REL
SC = _cute_ir.MemOrderKind.SC
MMIO = _cute_ir.MemOrderKind.MMIO
CONSTANT = _cute_ir.MemOrderKind.CONSTANT
VOLATILE = _cute_ir.MemOrderKind.VOLATILE
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.MemOrderKind:
return self.value
class MemoryScope(enum.Enum):
CTA = _cute_ir.MemScopeKind.CTA
CLUSTER = _cute_ir.MemScopeKind.CLUSTER
GPU = _cute_ir.MemScopeKind.GPU
SYS = _cute_ir.MemScopeKind.SYS
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.MemScopeKind:
return self.value
@dataclass(frozen=True)
class CopyUniversalOp(core.CopyOp):
"""
@@ -133,13 +169,18 @@ class CopyUniversalOp(core.CopyOp):
**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):
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.CopyAtomSIMTSyncCopyType.get(
copy_internal_type.mlir_type, num_bits_per_copy
copy_internal_type.mlir_type,
num_bits_per_copy,
memory_order._to_ir(),
memory_scope._to_ir(),
)
return CopyUniversalTrait(_cute_ir.atom(ty, loc=loc, ip=ip))

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

@@ -23,6 +23,7 @@ __all__ = [
"CopyBulkTensorTileG2SOp",
"CopyBulkTensorTileG2SMulticastOp",
"CopyBulkTensorTileS2GOp",
"CopyReduceBulkTensorTileS2GOp",
#
# helpers.py
#

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

@@ -19,7 +19,7 @@ import cutlass._mlir.dialects.cute as _cute_ir
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
from cutlass._mlir import ir
from ...core import CopyOp, Trait
from ...core import CopyOp, Trait, ReductionOp
from ...typing import Int16, Pointer, Integer, Numeric
from ..common import OpError
from ..tcgen05.mma import CtaGroup
@@ -80,6 +80,12 @@ class CopyG2SOp(CopyOp):
**kwargs,
) -> "CopyG2STrait":
num_bits_per_copy = kwargs.get("num_bits_per_copy", None)
# Verify that the user provided enum values
if not isinstance(self.cache_mode, LoadCacheMode):
raise OpError(
self,
"expects the 'cache_mode' Op parameter to be a LoadCacheMode instance",
)
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 "
@@ -330,7 +336,7 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
@dataclass(frozen=True)
class CopyBulkTensorTileS2GOp(CopyOp):
"""
Bulk tensor asynchrnous SMEM to GMEM Copy Operation using the TMA unit.
Bulk tensor asynchronous SMEM to GMEM Copy Operation using the TMA unit.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor>`__.
This Operation uses TMA in the ``.tile`` mode.
@@ -379,3 +385,87 @@ class CopyBulkTensorTileS2GTrait(Trait):
exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
)
return exec_value
@dataclass(frozen=True)
class CopyReduceBulkTensorTileS2GOp(CopyOp):
"""
Bulk tensor asynchronous SMEM to GMEM Reduction Operation using the TMA unit.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-reduce-async-bulk>`__.
This Operation uses TMA in the ``.tile`` mode.
"""
reduction_kind: ReductionOp = ReductionOp.ADD
admissible_archs = [
"sm_90",
"sm_90a",
"sm_100a",
"sm_100f",
]
def __post__init__(self):
# Arch verification
arch = CuTeDSL.__get_dsl().envar.arch
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",
)
def __str__(self) -> str:
return "cp.async SMEM -> GMEM bulk tensor reduction Operation"
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "CopyReduceBulkTensorTileS2GTrait":
raise NotImplementedError(
"Use cpasync.make_tiled_tma_atom to obtain a copy Atom for TMA"
)
def _to_ir(self) -> _cute_nvgpu_ir.ReductionKind:
if self.reduction_kind == ReductionOp.ADD:
return _cute_nvgpu_ir.ReductionKind.ADD
elif self.reduction_kind == ReductionOp.MIN:
return _cute_nvgpu_ir.ReductionKind.MIN
elif self.reduction_kind == ReductionOp.MAX:
return _cute_nvgpu_ir.ReductionKind.MAX
elif self.reduction_kind == ReductionOp.INC:
return _cute_nvgpu_ir.ReductionKind.INC
elif self.reduction_kind == ReductionOp.DEC:
return _cute_nvgpu_ir.ReductionKind.DEC
elif self.reduction_kind == ReductionOp.AND:
return _cute_nvgpu_ir.ReductionKind.AND
elif self.reduction_kind == ReductionOp.OR:
return _cute_nvgpu_ir.ReductionKind.OR
elif self.reduction_kind == ReductionOp.XOR:
return _cute_nvgpu_ir.ReductionKind.XOR
else:
assert False, "unrecognized self.reduction_kind"
class CopyReduceBulkTensorTileS2GTrait(Trait):
def unpack(self, *, loc=None, ip=None, tma_desc_ptr: Optional[Pointer] = None):
"""
Custom implementation of unpack for non-executable TMAs.
"""
exec_value = _cute_nvgpu_ir.atom_make_exec_tma(self.value, loc=loc, ip=ip)
if isinstance(tma_desc_ptr, Pointer):
attr_str = (
f"#cute_nvgpu.atom_copy_field_tmareduce<{TMA_DESC_PTR_FIELD_NAME}>"
)
attr = ir.Attribute.parse(attr_str)
exec_value = _cute_nvgpu_ir.atom_set_value(
exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
)
return exec_value
__all__ = [
"LoadCacheMode",
"CopyG2SOp",
"CopyBulkTensorTileG2SOp",
"CopyBulkTensorTileG2SMulticastOp",
"CopyBulkTensorTileS2GOp",
"CopyReduceBulkTensorTileS2GOp",
]

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

@@ -22,9 +22,11 @@ from .copy import (
CopyBulkTensorTileG2SOp,
CopyBulkTensorTileG2SMulticastOp,
CopyBulkTensorTileS2GOp,
CopyReduceBulkTensorTileS2GOp,
CopyBulkTensorTileG2SNonExecTrait,
CopyBulkTensorTileG2SMulticastNonExecTrait,
CopyBulkTensorTileS2GTrait,
CopyReduceBulkTensorTileS2GTrait,
)
@@ -34,6 +36,7 @@ def make_tiled_tma_atom(
CopyBulkTensorTileG2SOp,
CopyBulkTensorTileG2SMulticastOp,
CopyBulkTensorTileS2GOp,
CopyReduceBulkTensorTileS2GOp,
],
gmem_tensor: Tensor,
smem_layout: Union[Layout, core.ComposedLayout],
@@ -67,7 +70,7 @@ def make_tiled_tma_atom(
similarly to any other CuTe tensors using the algebra.
:param op: The Copy Operation to construct an Atom for
:type op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileS2GOp]
:type op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileS2GOp, CopyReduceBulkTensorTileS2GOp]
: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
@@ -141,6 +144,17 @@ def make_tiled_tma_atom(
ip=ip,
)
return core.CopyAtom(op, CopyBulkTensorTileS2GTrait(res[0])), res[1]
elif isinstance(op, CopyReduceBulkTensorTileS2GOp):
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_reduce(
gmem_tensor.value,
smem_layout,
cta_v_map,
op._to_ir(),
internal_type=internal_type,
loc=loc,
ip=ip,
)
return core.CopyAtom(op, CopyReduceBulkTensorTileS2GTrait(res[0])), res[1]
else:
raise ValueError(f"expects a bulk tensor (TMA) Copy Op, but got {op}")

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

@@ -21,7 +21,7 @@ from cutlass._mlir import ir
import cutlass._mlir.dialects.cute as _cute_ir
from cutlass.base_dsl.dsl import is_dynamic_expression
from cutlass.cutlass_dsl import TensorFormat, JitArgAdapterRegistry
from cutlass.cutlass_dsl import JitArgAdapterRegistry
# Local modules imports
from .typing import (
@@ -82,42 +82,36 @@ class _Pointer(Pointer):
self._dtype = dtype
self._addr_space = mem_space
is_in_device = mem_space == _cute_ir.AddressSpace.gmem
if assumed_align is None:
if is_in_device:
self._assumed_align = 32
else:
self._assumed_align = dtype.width // 8
self._assumed_align = dtype.width // 8
else:
self._assumed_align = assumed_align
class PtrDescriptor(ctypes.Structure):
"""A ctype descriptor for CuTe memref ptr"""
_fields_ = [("ptr", ctypes.c_void_p)]
def __str__(self):
return f"0x{self.ptr:016x}"
self._desc = PtrDescriptor(int(self._pointer))
self._c_pointer = ctypes.cast(ctypes.pointer(self._desc), ctypes.c_void_p)
self._c_pointer = None
assert (
self._desc.ptr % self._assumed_align == 0
int(self._pointer) % self._assumed_align == 0
), f"pointer must be {self._assumed_align} bytes aligned"
def size_in_bytes(self) -> int:
self._desc = ctypes.c_void_p(int(self._pointer))
return ctypes.sizeof(self._desc)
def __get_mlir_types__(self):
return [self.mlir_type]
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]
def __new_from_mlir_values__(self, values):
assert len(values) == 1
return values[0]
def __extract_mlir_values__(self):
return [self._c_pointer]
# Move mlir Type out of __init__ to decouple with mlir Context
@property
def mlir_type(self) -> ir.Type:
@@ -145,7 +139,7 @@ class _Pointer(Pointer):
return False
def __str__(self) -> str:
return f"Ptr<0x{self._desc.ptr:016x}@{self._addr_space}>"
return f"Ptr<0x{int(self._pointer):016x}@{self._addr_space}>"
def __repr__(self):
return self.__str__()