v4.2 tag release. (#2638)
This commit is contained in:
Junkai-Wu
@@ -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",
|
||||
|
||||
@@ -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,
|
||||
)
|
||||
|
||||
@@ -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]:
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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",
|
||||
|
||||
@@ -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))
|
||||
|
||||
|
||||
@@ -23,6 +23,7 @@ __all__ = [
|
||||
"CopyBulkTensorTileG2SOp",
|
||||
"CopyBulkTensorTileG2SMulticastOp",
|
||||
"CopyBulkTensorTileS2GOp",
|
||||
"CopyReduceBulkTensorTileS2GOp",
|
||||
#
|
||||
# helpers.py
|
||||
#
|
||||
|
||||
@@ -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",
|
||||
]
|
||||
|
||||
@@ -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}")
|
||||
|
||||
|
||||
@@ -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__()
|
||||
|
||||
@@ -31,9 +31,12 @@ from .helpers import (
|
||||
|
||||
from .sm90 import (
|
||||
PipelineAsync,
|
||||
PipelineCpAsync,
|
||||
PipelineTmaAsync,
|
||||
PipelineTmaMultiConsumersAsync,
|
||||
PipelineTmaStore,
|
||||
PipelineProducer,
|
||||
PipelineConsumer,
|
||||
)
|
||||
|
||||
from .sm100 import (
|
||||
@@ -53,10 +56,13 @@ __all__ = [
|
||||
"PipelineUserType",
|
||||
"PipelineState",
|
||||
"PipelineAsync",
|
||||
"PipelineCpAsync",
|
||||
"PipelineTmaAsync",
|
||||
"PipelineTmaUmma",
|
||||
"PipelineTmaMultiConsumersAsync",
|
||||
"PipelineAsyncUmma",
|
||||
"PipelineUmmaAsync",
|
||||
"PipelineTmaStore",
|
||||
"PipelineProducer",
|
||||
"PipelineConsumer",
|
||||
]
|
||||
|
||||
@@ -89,6 +89,8 @@ class PipelineOp(enum.Enum):
|
||||
TmaStore = enum.auto()
|
||||
# Composite of multiple PipelineOps
|
||||
Composite = enum.auto()
|
||||
# Async load without TMA
|
||||
AsyncLoad = enum.auto()
|
||||
|
||||
|
||||
def _get_pipeline_op(type_str):
|
||||
@@ -226,6 +228,8 @@ class MbarrierArray(SyncObject):
|
||||
self.arrive_tcgen05mma(index, dst, cta_group)
|
||||
elif self.op_type in [PipelineOp.TmaLoad]:
|
||||
self.arrive_and_expect_tx(index, self.tx_count)
|
||||
elif self.op_type is PipelineOp.AsyncLoad:
|
||||
self.arrive_cp_async_mbarrier(index)
|
||||
else:
|
||||
assert (
|
||||
False
|
||||
@@ -237,6 +241,9 @@ class MbarrierArray(SyncObject):
|
||||
else:
|
||||
cute.arch.mbarrier_arrive(self.get_barrier(index), dst_rank)
|
||||
|
||||
def arrive_cp_async_mbarrier(self, index: int):
|
||||
cute.arch.cp_async_mbarrier_arrive_noinc(self.get_barrier(index))
|
||||
|
||||
def arrive_tcgen05mma(
|
||||
self, index: int, mask: Optional[int], cta_group: cute.nvgpu.tcgen05.CtaGroup
|
||||
) -> None:
|
||||
|
||||
@@ -19,6 +19,7 @@ import cutlass.cute as cute
|
||||
from cutlass.cutlass_dsl import Boolean, if_generate
|
||||
|
||||
from cutlass.pipeline import (
|
||||
Agent,
|
||||
CooperativeGroup,
|
||||
PipelineOp,
|
||||
PipelineState,
|
||||
@@ -106,9 +107,9 @@ class PipelineTmaUmma(PipelineAsync):
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:param producer_group: `CooperativeGroup` for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:param consumer_group: `CooperativeGroup` for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
|
||||
:type tx_count: int
|
||||
@@ -258,9 +259,9 @@ class PipelineAsyncUmma(PipelineAsync):
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:param producer_group: `CooperativeGroup` for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:param consumer_group: `CooperativeGroup` for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param cta_layout_vmnk: Layout of the cluster shape
|
||||
:type cta_layout_vmnk: cute.Layout | None
|
||||
@@ -368,9 +369,9 @@ class PipelineUmmaAsync(PipelineAsync):
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:param producer_group: `CooperativeGroup` for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:param consumer_group: `CooperativeGroup` for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param cta_layout_vmnk: Layout of the cluster shape
|
||||
:type cta_layout_vmnk: cute.Layout | None
|
||||
|
||||
@@ -10,15 +10,18 @@
|
||||
# is strictly prohibited.
|
||||
|
||||
import enum
|
||||
from typing import Type, Tuple
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Union
|
||||
import warnings
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cutlass_dsl import Boolean, Int32, if_generate
|
||||
|
||||
from cutlass.pipeline import (
|
||||
Agent,
|
||||
CooperativeGroup,
|
||||
PipelineOp,
|
||||
SyncObject,
|
||||
@@ -91,6 +94,30 @@ class PipelineAsync:
|
||||
- D: Data ready (producer has written data to buffer)
|
||||
- R: Consumer reading (consumer is consuming data from buffer)
|
||||
|
||||
**Example:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Create pipeline with 5 stages
|
||||
pipeline = PipelineAsync.create(
|
||||
num_stages=5, # number of pipeline stages
|
||||
producer_group=producer_warp,
|
||||
consumer_group=consumer_warp
|
||||
barrier_storage=smem_ptr, # smem pointer for array of mbarriers in shared memory
|
||||
)
|
||||
|
||||
producer, consumer = pipeline.make_participants()
|
||||
# Producer side
|
||||
for i in range(num_iterations):
|
||||
handle = producer.acquire_and_advance() # Wait for buffer to be empty & Move index to next stage
|
||||
# Write data to pipeline buffer
|
||||
handle.commit() # Signal buffer is full
|
||||
|
||||
# Consumer side
|
||||
for i in range(num_iterations):
|
||||
handle = consumer.wait_and_advance() # Wait for buffer to be full & Move index to next stage
|
||||
# Read data from pipeline buffer
|
||||
handle.release() # Signal buffer is empty
|
||||
"""
|
||||
|
||||
sync_object_full: SyncObject
|
||||
@@ -114,6 +141,7 @@ class PipelineAsync:
|
||||
PipelineOp.TmaLoad,
|
||||
PipelineOp.TCGen05Mma,
|
||||
PipelineOp.Composite,
|
||||
PipelineOp.AsyncLoad,
|
||||
]:
|
||||
return MbarrierArray(
|
||||
barrier_storage=barrier_storage,
|
||||
@@ -232,6 +260,74 @@ class PipelineAsync:
|
||||
state.advance()
|
||||
self.producer_acquire(state)
|
||||
|
||||
# Util methods to manage produer and consumer
|
||||
def make_producer(self):
|
||||
state = make_pipeline_state(PipelineUserType.Producer, self.num_stages)
|
||||
return PipelineProducer(self, state, self.sync_object_full.cg)
|
||||
|
||||
def make_consumer(self):
|
||||
state = make_pipeline_state(PipelineUserType.Consumer, self.num_stages)
|
||||
return PipelineConsumer(self, state, self.sync_object_empty.cg)
|
||||
|
||||
def make_participants(self):
|
||||
return self.make_producer(), self.make_consumer()
|
||||
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineCpAsync(PipelineAsync):
|
||||
"""
|
||||
PipelineCpAsync is used for CpAsync producers and AsyncThread consumers (e.g. Hopper non-TMA mainloops).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
barrier_storage: cute.Pointer,
|
||||
num_stages: Int32,
|
||||
producer_group: CooperativeGroup,
|
||||
consumer_group: CooperativeGroup,
|
||||
producer_mask: Int32 = None,
|
||||
consumer_mask: Int32 = None,
|
||||
):
|
||||
"""
|
||||
This helper function computes any necessary attributes and returns an instance of PipelineAsync.
|
||||
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param producer_mask: Mask for signaling arrives for the producer agent
|
||||
:type producer_mask: Int32 | None
|
||||
:param consumer_mask: Mask for signaling arrives for the consumer agent
|
||||
:type consumer_mask: Int32 | None
|
||||
"""
|
||||
producer_type = PipelineOp.AsyncLoad
|
||||
consumer_type = PipelineOp.AsyncThread
|
||||
|
||||
producer = (producer_type, producer_group)
|
||||
consumer = (consumer_type, consumer_group)
|
||||
|
||||
sync_object_array_full = PipelineCpAsync._make_sync_object(
|
||||
barrier_storage.align(min_align=8), num_stages, producer
|
||||
)
|
||||
sync_object_array_empty = PipelineCpAsync._make_sync_object(
|
||||
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
|
||||
)
|
||||
|
||||
pipeline_init_wait()
|
||||
|
||||
return PipelineCpAsync(
|
||||
sync_object_array_full,
|
||||
sync_object_array_empty,
|
||||
num_stages,
|
||||
producer_mask,
|
||||
consumer_mask,
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineTmaAsync(PipelineAsync):
|
||||
"""
|
||||
@@ -294,9 +390,9 @@ class PipelineTmaAsync(PipelineAsync):
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:param producer_group: `CooperativeGroup` for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:param consumer_group: `CooperativeGroup` for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
|
||||
:type tx_count: int
|
||||
@@ -404,11 +500,11 @@ class PipelineTmaMultiConsumersAsync(PipelineAsync):
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:param producer_group: `CooperativeGroup` for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group_umma: CooperativeGroup for the UMMA consumer agent
|
||||
:param consumer_group_umma: `CooperativeGroup` for the UMMA consumer agent
|
||||
:type consumer_group_umma: CooperativeGroup
|
||||
:param consumer_group_async: CooperativeGroup for the AsyncThread consumer agent
|
||||
:param consumer_group_async: `CooperativeGroup` for the AsyncThread consumer agent
|
||||
:type consumer_group_async: CooperativeGroup
|
||||
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
|
||||
:type tx_count: int
|
||||
@@ -529,9 +625,10 @@ class PipelineTmaStore(PipelineAsync):
|
||||
This helper function computes any necessary attributes and returns an instance of PipelineTmaStore.
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:param producer_group: `CooperativeGroup` for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
"""
|
||||
|
||||
producer_type = PipelineOp.TmaStore
|
||||
|
||||
producer = (producer_type, producer_group)
|
||||
@@ -556,3 +653,333 @@ class PipelineTmaStore(PipelineAsync):
|
||||
self.sync_object_full.tail()
|
||||
|
||||
|
||||
#################################################################
|
||||
# Utilities to help user of pipeline to simplify the workflow
|
||||
#################################################################
|
||||
|
||||
|
||||
class ImmutableResourceHandle:
|
||||
__origin: PipelineAsync
|
||||
__immutable_state: PipelineState
|
||||
|
||||
def __init__(self, origin: PipelineAsync, immutable_state: PipelineState):
|
||||
self.__origin = origin
|
||||
self.__immutable_state = immutable_state
|
||||
|
||||
@property
|
||||
def index(self):
|
||||
"""Get the index of the current pipeline stage."""
|
||||
return self.__immutable_state.index
|
||||
|
||||
@property
|
||||
def count(self):
|
||||
"""Get the count of how many handles this producer has committed.
|
||||
This is useful for tracking the number of blocks that have been loaded from gmem.
|
||||
"""
|
||||
return self.__immutable_state.count
|
||||
|
||||
def get_origin(self):
|
||||
"""Get the original pipeline this resource handle belongs to."""
|
||||
return self.__origin
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
"""Extract MLIR values from the current state.
|
||||
|
||||
:return: List of MLIR values representing the current state
|
||||
:rtype: list
|
||||
"""
|
||||
# TODO: need to handle pipeline as well
|
||||
return self.__immutable_state.__extract_mlir_values__()
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
"""Create a new Producer instance from MLIR values.
|
||||
|
||||
:param values: MLIR values to initialize the state
|
||||
:type values: Any
|
||||
:return: New Producer instance with state initialized from values
|
||||
:rtype: Producer
|
||||
"""
|
||||
return self.__class__(
|
||||
self.__origin, self.__immutable_state.__new_from_mlir_values__(values)
|
||||
)
|
||||
|
||||
class PipelineProducer:
|
||||
"""A class representing a producer in an asynchronous pipeline.
|
||||
|
||||
The Producer class manages the producer side of an asynchronous pipeline, handling
|
||||
synchronization and state management for producing data. It provides methods for
|
||||
acquiring, committing, and advancing through pipeline stages.
|
||||
|
||||
:ivar __pipeline: The asynchronous pipeline this producer belongs to
|
||||
:type __pipeline: PipelineAsync
|
||||
:ivar __state: The current state of the producer in the pipeline
|
||||
:type __state: PipelineState
|
||||
:ivar __group: The cooperative group this producer operates in
|
||||
:type __group: CooperativeGroup
|
||||
|
||||
**Examples:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
pipeline = PipelineAsync.create(...)
|
||||
producer = pipeline.create_producer(producer_group, stages)
|
||||
for i in range(iterations):
|
||||
handle = producer.acquire_and_advance() # Wait for buffer to be empty
|
||||
# Produce data
|
||||
producer.commit(handle) # Signal data is ready
|
||||
# An alternative way to do this is:
|
||||
# handle.commit() # Signal data is ready
|
||||
"""
|
||||
|
||||
__pipeline: PipelineAsync
|
||||
__state: PipelineState
|
||||
__group: CooperativeGroup
|
||||
|
||||
class ImmutableResourceHandle(ImmutableResourceHandle):
|
||||
@property
|
||||
def barrier(self):
|
||||
"""Get the barrier pointer for the current pipeline stage.
|
||||
|
||||
:return: Pointer to the barrier for the current stage
|
||||
:rtype: cute.Pointer
|
||||
"""
|
||||
return self.get_origin().producer_get_barrier(
|
||||
self._ImmutableResourceHandle__immutable_state
|
||||
)
|
||||
|
||||
def commit(self):
|
||||
"""Signal that data production is complete for the current stage.
|
||||
This allows consumers to start processing the data.
|
||||
"""
|
||||
self.get_origin().producer_commit(
|
||||
self._ImmutableResourceHandle__immutable_state
|
||||
)
|
||||
|
||||
def __init__(self, pipeline, state, group: CooperativeGroup):
|
||||
"""Initialize a new Producer instance.
|
||||
|
||||
:param pipeline: The pipeline this producer belongs to
|
||||
:type pipeline: PipelineAsync
|
||||
:param state: Initial pipeline state
|
||||
:type state: PipelineState
|
||||
:param group: The cooperative group for synchronization
|
||||
:type group: CooperativeGroup
|
||||
"""
|
||||
self.__pipeline = pipeline
|
||||
self.__state = state
|
||||
self.__group = group
|
||||
|
||||
def acquire(
|
||||
self,
|
||||
try_acquire_token: Optional[Boolean] = None,
|
||||
) -> ImmutableResourceHandle:
|
||||
"""Wait for the current buffer to be empty before producing data.
|
||||
This is a blocking operation.
|
||||
|
||||
:param try_acquire_token: Optional token to try to acquire the buffer
|
||||
:type try_acquire_token: Optional[Boolean]
|
||||
:return: A handle to the producer for committing the data
|
||||
:rtype: ImmutableResourceHandle
|
||||
"""
|
||||
self.__pipeline.producer_acquire(self.__state, try_acquire_token)
|
||||
handle = PipelineProducer.ImmutableResourceHandle(
|
||||
self.__pipeline, self.__state.clone()
|
||||
)
|
||||
return handle
|
||||
|
||||
def advance(self):
|
||||
"""Move to the next pipeline stage."""
|
||||
self.__state.advance()
|
||||
|
||||
def acquire_and_advance(
|
||||
self, try_acquire_token: Optional[Boolean] = None
|
||||
) -> ImmutableResourceHandle:
|
||||
"""Wait for the current buffer to be empty before producing data.
|
||||
Then advance to the next stage.
|
||||
This is a blocking operation.
|
||||
|
||||
:param try_acquire_token: Optional token to try to acquire the buffer
|
||||
:type try_acquire_token: Optional[Boolean]
|
||||
:return: A handle to the producer for committing the data
|
||||
:rtype: ImmutableResourceHandle
|
||||
"""
|
||||
handle = self.acquire(try_acquire_token)
|
||||
self.advance()
|
||||
return handle
|
||||
|
||||
def try_acquire(self) -> Boolean:
|
||||
"""Try to acquire the current buffer without blocking.
|
||||
|
||||
:return: True if acquisition was successful, False otherwise
|
||||
:rtype: Boolean
|
||||
"""
|
||||
return self.__pipeline.producer_try_acquire(self.__state)
|
||||
|
||||
def commit(self, handle: Optional[ImmutableResourceHandle] = None):
|
||||
"""Signal that data production is complete for the current stage.
|
||||
This allows consumers to start processing the data.
|
||||
"""
|
||||
if handle is not None:
|
||||
assert (
|
||||
handle.get_origin() is self
|
||||
), "ResourceHandle does not belong to this PipelineProducer instance"
|
||||
handle.commit()
|
||||
else:
|
||||
self.__pipeline.producer_commit(self.__state)
|
||||
|
||||
def tail(self):
|
||||
"""Ensure all used buffers are properly synchronized before producer exit.
|
||||
This should be called before the producer finishes to avoid dangling signals.
|
||||
"""
|
||||
self.__pipeline.producer_tail(self.__state)
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
"""Extract MLIR values from the current state.
|
||||
|
||||
:return: List of MLIR values representing the current state
|
||||
:rtype: list
|
||||
"""
|
||||
# TODO: need to handle pipeline as well
|
||||
return self.__state.__extract_mlir_values__()
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
"""Create a new Producer instance from MLIR values.
|
||||
|
||||
:param values: MLIR values to initialize the state
|
||||
:type values: Any
|
||||
:return: New Producer instance with state initialized from values
|
||||
:rtype: Producer
|
||||
"""
|
||||
return PipelineProducer(
|
||||
self.__pipeline, self.__state.__new_from_mlir_values__(values), self.__group
|
||||
)
|
||||
|
||||
class PipelineConsumer:
|
||||
"""A class representing a consumer in an asynchronous pipeline.
|
||||
|
||||
The Consumer class manages the consumer side of an asynchronous pipeline, handling
|
||||
synchronization and state management for consuming data. It provides methods for
|
||||
waiting, releasing, and advancing through pipeline stages.
|
||||
|
||||
:ivar __pipeline: The asynchronous pipeline this consumer belongs to
|
||||
:type __pipeline: PipelineAsync
|
||||
:ivar __state: The current state of the consumer in the pipeline
|
||||
:type __state: PipelineState
|
||||
:ivar __group: The cooperative group this consumer operates in
|
||||
:type __group: CooperativeGroup
|
||||
|
||||
**Examples:**
|
||||
.. code-block:: python
|
||||
|
||||
pipeline = PipelineAsync.create(...)
|
||||
consumer = pipeline.create_consumer(consumer_group, stages)
|
||||
for i in range(iterations):
|
||||
handle = consumer.wait_and_advance() # Wait for data to be ready
|
||||
# Consume data
|
||||
consumer.release(handle) # Signal buffer is empty
|
||||
# An alternative way to do this is:
|
||||
# handle.release() # Signal buffer is empty
|
||||
"""
|
||||
|
||||
__pipeline: PipelineAsync
|
||||
__state: PipelineState
|
||||
__group: CooperativeGroup
|
||||
|
||||
class ImmutableResourceHandle(ImmutableResourceHandle):
|
||||
def release(self):
|
||||
"""Signal that data production is complete for the current stage.
|
||||
This allows consumers to start processing the data.
|
||||
"""
|
||||
self.get_origin().consumer_release(
|
||||
self._ImmutableResourceHandle__immutable_state
|
||||
)
|
||||
|
||||
def __init__(self, pipeline, state: PipelineState, group: CooperativeGroup):
|
||||
"""Initialize a new Consumer instance.
|
||||
|
||||
:param pipeline: The pipeline this consumer belongs to
|
||||
:type pipeline: PipelineAsync
|
||||
:param state: Initial pipeline state
|
||||
:type state: PipelineState
|
||||
:param group: The cooperative group for synchronization
|
||||
:type group: CooperativeGroup
|
||||
"""
|
||||
self.__pipeline = pipeline
|
||||
self.__group = group
|
||||
self.__state = state
|
||||
|
||||
def wait(self, try_wait_token: Optional[Boolean] = None) -> ImmutableResourceHandle:
|
||||
"""Wait for data to be ready in the current buffer.
|
||||
This is a blocking operation.
|
||||
|
||||
:param try_wait_token: Optional token to try to wait for the buffer
|
||||
:type try_wait_token: Optional[Boolean]
|
||||
:return: A handle to the consumer for releasing the data
|
||||
:rtype: PipelineConsumerHandle
|
||||
"""
|
||||
self.__pipeline.consumer_wait(self.__state, try_wait_token)
|
||||
handle = PipelineConsumer.ImmutableResourceHandle(
|
||||
self.__pipeline, self.__state.clone()
|
||||
)
|
||||
return handle
|
||||
|
||||
def advance(self):
|
||||
"""Move to the next pipeline stage."""
|
||||
self.__state.advance()
|
||||
|
||||
def wait_and_advance(
|
||||
self, try_wait_token: Optional[Boolean] = None
|
||||
) -> ImmutableResourceHandle:
|
||||
"""Wait for data to be ready in the current buffer.
|
||||
Then advance to the next stage.
|
||||
This is a blocking operation.
|
||||
|
||||
:param try_wait_token: Optional token to try to wait for the buffer
|
||||
:type try_wait_token: Optional[Boolean]
|
||||
:return: A handle to the consumer for releasing the data
|
||||
:rtype: PipelineConsumerHandle
|
||||
"""
|
||||
handle = self.wait(try_wait_token)
|
||||
self.advance()
|
||||
return handle
|
||||
|
||||
def try_wait(self) -> Boolean:
|
||||
"""Try to check if data is ready without blocking.
|
||||
|
||||
:return: True if data is ready, False otherwise
|
||||
:rtype: Boolean
|
||||
"""
|
||||
return self.__pipeline.consumer_try_wait(self.__state)
|
||||
|
||||
def release(self, handle: Optional[ImmutableResourceHandle] = None):
|
||||
"""Signal that data consumption is complete for the current stage.
|
||||
This allows producers to start producing new data.
|
||||
"""
|
||||
if handle is not None:
|
||||
assert (
|
||||
handle.get_origin() is self
|
||||
), "ResourceHandle does not belong to this PipelineConsumer instance"
|
||||
handle.release()
|
||||
else:
|
||||
self.__pipeline.consumer_release(self.__state)
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
"""Extract MLIR values from the current state.
|
||||
|
||||
:return: List of MLIR values representing the current state
|
||||
:rtype: list
|
||||
"""
|
||||
return self.__state.__extract_mlir_values__()
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
"""Create a new Consumer instance from MLIR values.
|
||||
|
||||
:param values: MLIR values to initialize the state
|
||||
:type values: Any
|
||||
:return: New Consumer instance with state initialized from values
|
||||
:rtype: Consumer
|
||||
"""
|
||||
# TODO: need to call pipeline.__new_from_mlir_values__ recursively
|
||||
return PipelineConsumer(
|
||||
self.__pipeline, self.__state.__new_from_mlir_values__(values), self.__group
|
||||
)
|
||||
|
||||
@@ -9,6 +9,8 @@
|
||||
# and related documentation outside the scope permitted by the EULA
|
||||
# is strictly prohibited.
|
||||
|
||||
import ctypes
|
||||
from math import prod
|
||||
from dataclasses import dataclass
|
||||
from enum import Enum
|
||||
from typing import Optional, Type, Union
|
||||
@@ -54,6 +56,25 @@ def dtype(ty: Type[Numeric]):
|
||||
return torch_dtype
|
||||
|
||||
|
||||
def as_tensor(pointer, shape, torch_type):
|
||||
"""Convert a pointer to a torch tensor"""
|
||||
if torch_type.itemsize == 1:
|
||||
cytype = ctypes.c_uint8
|
||||
elif torch_type.itemsize == 2:
|
||||
cytype = ctypes.c_uint16
|
||||
elif torch_type.itemsize == 4:
|
||||
cytype = ctypes.c_uint32
|
||||
elif torch_type.itemsize == 8:
|
||||
cytype = ctypes.c_uint64
|
||||
else:
|
||||
raise ValueError(f"Unsupported torch dtype: {torch_type}")
|
||||
cpointer = ctypes.cast(pointer, ctypes.POINTER(cytype))
|
||||
arr = (cpointer._type_ * prod(shape)).from_address(
|
||||
ctypes.addressof(cpointer.contents)
|
||||
)
|
||||
return torch.frombuffer(arr, dtype=torch_type).view(*shape)
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScalarInitConfig:
|
||||
"""Configuration for scalar initialization"""
|
||||
@@ -128,7 +149,7 @@ def create_and_permute_torch_tensor(
|
||||
if not isinstance(init_config, GaussianInitConfig):
|
||||
raise ValueError("init_config must be GaussianInitConfig()")
|
||||
f32_torch_tensor = init_torch_tensor.normal_(init_config.mean, init_config.std)
|
||||
f32_torch_tensor = f32_torch_tensor * (1 << init_config.scale)
|
||||
f32_torch_tensor = f32_torch_tensor * init_config.scale
|
||||
else:
|
||||
raise ValueError(f"Invalid init type: {init_type}")
|
||||
|
||||
|
||||
@@ -64,6 +64,18 @@ from .smem_capacity import (
|
||||
get_smem_capacity_in_bytes,
|
||||
)
|
||||
|
||||
from .distributed_helpers import (
|
||||
spin_lock_wait,
|
||||
spin_lock_multimem_arrive,
|
||||
multimem_ld_reduce_8xf16,
|
||||
multimem_ld_reduce_4xf32,
|
||||
multimem_ld_reduce_8xbf16,
|
||||
multimem_ld_reduce_16xe4m3,
|
||||
multimem_ld_reduce_16xe5m2,
|
||||
multimem_st_4xb32,
|
||||
sm_wise_inter_gpu_multimem_barrier,
|
||||
)
|
||||
|
||||
__all__ = [
|
||||
"get_smem_capacity_in_bytes",
|
||||
"SmemAllocator",
|
||||
|
||||
179
python/CuTeDSL/cutlass/utils/distributed_helpers.py
Normal file
179
python/CuTeDSL/cutlass/utils/distributed_helpers.py
Normal file
@@ -0,0 +1,179 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# Use of this software is governed by the terms and conditions of the
|
||||
# NVIDIA End User License Agreement (EULA), available at:
|
||||
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
|
||||
#
|
||||
# Any use, reproduction, disclosure, or distribution of this software
|
||||
# and related documentation outside the scope permitted by the EULA
|
||||
# is strictly prohibited.
|
||||
|
||||
from functools import partial
|
||||
from typing import Tuple
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cutlass_dsl import T, dsl_user_op, while_generate
|
||||
|
||||
from cutlass._mlir import ir
|
||||
from cutlass._mlir.dialects import arith, llvm, nvvm, scf
|
||||
from cutlass._mlir.dialects.nvvm import (
|
||||
MemOrderKind,
|
||||
MemScopeKind,
|
||||
AtomicOpKind,
|
||||
)
|
||||
from cutlass.cute.typing import Pointer, Int32, Boolean
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def atomicAdd(dst_ptr: Pointer, val: Int32, loc=None, ip=None) -> Int32:
|
||||
return nvvm.atomicrmw(
|
||||
T.i32(),
|
||||
AtomicOpKind.ADD,
|
||||
dst_ptr.llvm_ptr,
|
||||
val.ir_value(loc=loc, ip=ip),
|
||||
mem_order=MemOrderKind.RELAXED,
|
||||
syncscope=MemScopeKind.SYS,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@cute.jit
|
||||
def ld_bypass(input_tensor: cute.Tensor):
|
||||
fragment = cute.make_fragment(input_tensor.layout, input_tensor.element_type)
|
||||
copy_atom_load = cute.make_copy_atom(
|
||||
cute.nvgpu.CopyUniversalOp(),
|
||||
input_tensor.element_type,
|
||||
memory_order=cute.nvgpu.common.MemoryOrder.VOLATILE,
|
||||
memory_scope=cute.nvgpu.common.MemoryScope.SYS,
|
||||
)
|
||||
cute.copy(copy_atom_load, input_tensor, fragment)
|
||||
vals = fragment.load()
|
||||
return vals
|
||||
|
||||
@cute.jit
|
||||
def spin_lock_wait(lock_ptr: Pointer, expect_count: Int32, mem_order : str = "relaxed", mem_scope : str = "gpu", loc=None, ip=None) -> None:
|
||||
"""
|
||||
wait on a spin lock until the expected count is reached.
|
||||
"""
|
||||
res = 0
|
||||
while res != expect_count:
|
||||
res = nvvm.atomicrmw(
|
||||
T.i32(),
|
||||
AtomicOpKind.CAS,
|
||||
lock_ptr.llvm_ptr,
|
||||
Int32(0).ir_value(loc=loc, ip=ip),
|
||||
b=Int32(expect_count).ir_value(loc=loc, ip=ip),
|
||||
mem_order=MemOrderKind.ACQUIRE if mem_order == "acquire" else MemOrderKind.RELAXED,
|
||||
syncscope=MemScopeKind.GPU if mem_scope == "gpu" else MemScopeKind.SYS
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def multimem_red_add_sys_release(mc_ptr: Pointer, loc=None, ip=None) -> None:
|
||||
"""
|
||||
add 1 to the multimem address
|
||||
"""
|
||||
llvm.inline_asm(
|
||||
None,
|
||||
[mc_ptr.toint().ir_value()],
|
||||
"multimem.red.release.sys.global.add.u32 [$0], 1;",
|
||||
"l",
|
||||
has_side_effects=True,
|
||||
asm_dialect=0,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
@dsl_user_op
|
||||
def multimem_red_add_gpu_relaxed(mc_ptr: Pointer, loc=None, ip=None) -> None:
|
||||
"""
|
||||
add 1 to the multimem address
|
||||
"""
|
||||
llvm.inline_asm(
|
||||
None,
|
||||
[mc_ptr.toint().ir_value()],
|
||||
"multimem.red.relaxed.gpu.global.add.u32 [$0], 1;",
|
||||
"l",
|
||||
has_side_effects=True,
|
||||
asm_dialect=0,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
def spin_lock_multimem_arrive(lock_ptr: Pointer, loc=None, ip=None) -> None:
|
||||
"""
|
||||
arrive a spin lock when the lock_ptr is a multimem address.
|
||||
"""
|
||||
multimem_red_add_gpu_relaxed(lock_ptr, loc=loc, ip=ip)
|
||||
|
||||
|
||||
def sm_wise_inter_gpu_multimem_barrier(barrier : Pointer, barrier_mc : Pointer, num_ranks, loc=None, ip=None) -> None :
|
||||
"""
|
||||
barrier for inter-gpu sm-wise
|
||||
"""
|
||||
bidx, bidy, bidz = cute.arch.block_idx()
|
||||
bdimx, bdimy, _ = cute.arch.grid_dim()
|
||||
pid = bidx + bidy * bdimx + bidz * bdimx * bdimy
|
||||
multimem_red_add_sys_release(barrier_mc + pid, loc=loc, ip=ip)
|
||||
cute.arch.fence_proxy(cute.arch.ProxyKind.alias)
|
||||
spin_lock_wait(barrier + pid, num_ranks, mem_order="acquire", mem_scope="sys", loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def multimem_ld_reduce_base(
|
||||
mc_ptr: Pointer,
|
||||
*,
|
||||
ptx_string: str = "",
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[Int32, Int32, Int32, Int32]:
|
||||
# ld reduce 8xf16 elts
|
||||
mc_ptr_int = mc_ptr.toint(loc=loc, ip=ip).ir_value()
|
||||
return_struct = llvm.inline_asm(
|
||||
ir.Type.parse("!llvm.struct<(i32,i32,i32,i32)>"),
|
||||
[mc_ptr_int],
|
||||
ptx_string,
|
||||
"=r,=r,=r,=r,l",
|
||||
has_side_effects=True,
|
||||
asm_dialect=0,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return_regs = [llvm.extractvalue(T.i32(), return_struct, [i]) for i in range(4)]
|
||||
return return_regs[0], return_regs[1], return_regs[2], return_regs[3]
|
||||
|
||||
|
||||
multimem_ld_reduce_8xf16 = partial(multimem_ld_reduce_base, ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.f16x2 {$0, $1, $2, $3}, [$4];")
|
||||
multimem_ld_reduce_4xf32 = partial(multimem_ld_reduce_base, ptx_string="multimem.ld_reduce.sys.relaxed.global.add.v4.f32 {$0, $1, $2, $3}, [$4];")
|
||||
multimem_ld_reduce_8xbf16 = partial(multimem_ld_reduce_base, ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f32.v4.bf16x2 {$0, $1, $2, $3}, [$4];")
|
||||
multimem_ld_reduce_16xe4m3 = partial(multimem_ld_reduce_base, ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f16.v4.e4m3x4 {$0, $1, $2, $3}, [$4];")
|
||||
multimem_ld_reduce_16xe5m2 = partial(multimem_ld_reduce_base, ptx_string="multimem.ld_reduce.sys.relaxed.global.add.acc::f16.v4.e5m2x4 {$0, $1, $2, $3}, [$4];")
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def multimem_st_4xb32(
|
||||
mc_ptr: Pointer,
|
||||
x: Int32,
|
||||
y: Int32,
|
||||
z: Int32,
|
||||
w: Int32,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
# st 4x32 bits of data
|
||||
mc_ptr_int = mc_ptr.toint(loc=loc, ip=ip).ir_value()
|
||||
llvm.inline_asm(
|
||||
T.i32(),
|
||||
[mc_ptr_int, x, y, z, w],
|
||||
"multimem.st.sys.relaxed.global.v4.f32 [$1], {$2, $3, $4, $5};",
|
||||
"=r,l,r,r,r,r",
|
||||
has_side_effects=True,
|
||||
asm_dialect=0,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
@@ -34,18 +34,6 @@ class LayoutEnum(Enum):
|
||||
else warpgroup.OperandMajorMode.MN
|
||||
)
|
||||
|
||||
def is_k_major_a(self):
|
||||
return self == LayoutEnum.ROW_MAJOR
|
||||
|
||||
def is_m_major_a(self):
|
||||
return self == LayoutEnum.COL_MAJOR
|
||||
|
||||
def is_k_major_b(self):
|
||||
return self == LayoutEnum.COL_MAJOR
|
||||
|
||||
def is_n_major_b(self):
|
||||
return self == LayoutEnum.ROW_MAJOR
|
||||
|
||||
def is_n_major_c(self):
|
||||
return self == LayoutEnum.ROW_MAJOR
|
||||
|
||||
|
||||
@@ -11,7 +11,7 @@
|
||||
|
||||
from typing import Type, Union, overload
|
||||
|
||||
from cutlass.cutlass_dsl import Int8, Numeric, NumericMeta
|
||||
from cutlass.cutlass_dsl import Int8, Numeric, NumericMeta, CutlassBaseDSL
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.arch import get_dyn_smem, get_dyn_smem_size
|
||||
@@ -40,14 +40,17 @@ class SmemAllocator:
|
||||
"""
|
||||
self._base = get_dyn_smem(Int8, alignment=1024)
|
||||
self._allocated_bytes = 0
|
||||
CutlassBaseDSL.track_smem_allocator(self, lambda cls: cls._allocated_bytes)
|
||||
|
||||
@overload
|
||||
def allocate(self, size_or_type: int, byte_alignment: int): ...
|
||||
def allocate(self, size_or_type: int, byte_alignment: int) -> cute.Pointer: ...
|
||||
|
||||
@overload
|
||||
def allocate(self, size_or_type: cute.struct, byte_alignment: int): ...
|
||||
def allocate(
|
||||
self, size_or_type: cute.struct, byte_alignment: int
|
||||
) -> cute.Pointer: ...
|
||||
|
||||
def allocate(self, size_or_type, byte_alignment: int = 1) -> int:
|
||||
def allocate(self, size_or_type, byte_alignment: int = 1) -> cute.Pointer:
|
||||
"""Allocate a block of memory with specified size and alignment.
|
||||
|
||||
This method adjusts the base pointer to ensure proper alignment and updates
|
||||
|
||||
@@ -382,3 +382,5 @@ class StaticPersistentTileScheduler:
|
||||
@property
|
||||
def num_tiles_executed(self) -> Int32:
|
||||
return self._num_tiles_executed
|
||||
|
||||
|
||||
|
||||
Reference in New Issue
Block a user