v4.3 update. (#2709)
* v4.3 update. * Update the cute_dsl_api changelog's doc link * Update version to 4.3.0 * Update the example link * Update doc to encourage user to install DSL from requirements.txt --------- Co-authored-by: Larry Wu <larwu@nvidia.com>
This commit is contained in:
Junkai-Wu
@@ -26,6 +26,7 @@ from .typing import (
|
||||
XTuple,
|
||||
Tiler,
|
||||
Layout,
|
||||
ComposedLayout,
|
||||
Pointer,
|
||||
Tensor,
|
||||
)
|
||||
@@ -35,47 +36,37 @@ from .typing import *
|
||||
|
||||
from .core import (
|
||||
assume,
|
||||
is_integer,
|
||||
is_int_tuple,
|
||||
is_static,
|
||||
size,
|
||||
static,
|
||||
get_leaves,
|
||||
has_underscore,
|
||||
slice_,
|
||||
make_ptr,
|
||||
make_layout,
|
||||
recast_layout,
|
||||
make_fragment_like,
|
||||
depth,
|
||||
rank,
|
||||
flatten_to_tuple,
|
||||
flatten,
|
||||
unflatten,
|
||||
product,
|
||||
product_like,
|
||||
shape,
|
||||
size_in_bytes,
|
||||
make_identity_layout,
|
||||
make_ordered_layout,
|
||||
make_layout_like,
|
||||
make_composed_layout,
|
||||
make_layout_tv,
|
||||
make_swizzle,
|
||||
make_sparse_elem,
|
||||
recast_ptr,
|
||||
make_tensor,
|
||||
make_identity_tensor,
|
||||
make_fragment,
|
||||
recast_tensor,
|
||||
get,
|
||||
select,
|
||||
front,
|
||||
is_major,
|
||||
leading_dim,
|
||||
find,
|
||||
find_if,
|
||||
coalesce,
|
||||
group_modes,
|
||||
cosize,
|
||||
dice,
|
||||
product_each,
|
||||
prepend,
|
||||
append,
|
||||
prepend_ones,
|
||||
@@ -83,9 +74,7 @@ from .core import (
|
||||
ceil_div,
|
||||
slice_and_offset,
|
||||
crd2idx,
|
||||
domain_offset,
|
||||
elem_less,
|
||||
transform_leaf,
|
||||
idx2crd,
|
||||
filter_zeros,
|
||||
filter,
|
||||
tile_to_shape,
|
||||
@@ -109,7 +98,65 @@ from .core import (
|
||||
local_partition,
|
||||
local_tile,
|
||||
printf,
|
||||
# Wrapper classes
|
||||
Swizzle,
|
||||
E,
|
||||
# User defined struct
|
||||
struct,
|
||||
pretty_str,
|
||||
make_layout_image_mask,
|
||||
repeat,
|
||||
repeat_as_tuple,
|
||||
repeat_like,
|
||||
round_up,
|
||||
is_congruent,
|
||||
is_weakly_congruent,
|
||||
ScaledBasis,
|
||||
get_divisibility,
|
||||
Ratio,
|
||||
)
|
||||
|
||||
from .tuple import (
|
||||
transform_leaf,
|
||||
find_if,
|
||||
find,
|
||||
flatten_to_tuple,
|
||||
unflatten,
|
||||
product,
|
||||
product_like,
|
||||
product_each,
|
||||
elem_less,
|
||||
)
|
||||
from .tensor import (
|
||||
TensorSSA,
|
||||
ReductionOp,
|
||||
make_tensor,
|
||||
make_identity_tensor,
|
||||
make_fragment,
|
||||
make_fragment_like,
|
||||
make_rmem_tensor_like,
|
||||
make_rmem_tensor,
|
||||
recast_tensor,
|
||||
domain_offset,
|
||||
print_tensor,
|
||||
full,
|
||||
full_like,
|
||||
empty_like,
|
||||
ones_like,
|
||||
zeros_like,
|
||||
where,
|
||||
any_,
|
||||
all_,
|
||||
)
|
||||
from .atom import (
|
||||
Atom,
|
||||
MmaAtom,
|
||||
CopyAtom,
|
||||
TiledCopy,
|
||||
TiledMma,
|
||||
ThrMma,
|
||||
ThrCopy,
|
||||
make_atom,
|
||||
# tiled mma/tiled copy
|
||||
make_mma_atom,
|
||||
make_tiled_mma,
|
||||
@@ -122,48 +169,16 @@ from .core import (
|
||||
make_tiled_copy_B,
|
||||
make_tiled_copy_C,
|
||||
make_tiled_copy_C_atom,
|
||||
basic_copy,
|
||||
basic_copy_if,
|
||||
autovec_copy,
|
||||
copy,
|
||||
make_cotiled_copy,
|
||||
copy_atom_call,
|
||||
gemm,
|
||||
# Wrapper classes
|
||||
ComposedLayout,
|
||||
Swizzle,
|
||||
E,
|
||||
Atom,
|
||||
MmaAtom,
|
||||
CopyAtom,
|
||||
TiledCopy,
|
||||
TiledMma,
|
||||
TensorSSA,
|
||||
ReductionOp,
|
||||
full,
|
||||
full_like,
|
||||
empty_like,
|
||||
ones_like,
|
||||
zeros_like,
|
||||
where,
|
||||
any_,
|
||||
all_,
|
||||
# User defined struct
|
||||
struct,
|
||||
pretty_str,
|
||||
make_layout_image_mask,
|
||||
repeat_like,
|
||||
round_up,
|
||||
is_congruent,
|
||||
is_weakly_congruent,
|
||||
ScaledBasis,
|
||||
get_divisibility,
|
||||
Ratio,
|
||||
)
|
||||
from .algorithm import gemm, copy, basic_copy, basic_copy_if, autovec_copy, prefetch
|
||||
|
||||
from . import arch
|
||||
from . import nvgpu
|
||||
from . import testing
|
||||
from . import runtime
|
||||
from . import math
|
||||
|
||||
# Export all math ops without "math."
|
||||
from .math import *
|
||||
@@ -175,7 +190,15 @@ from .. import cutlass_dsl as _dsl
|
||||
jit = _dsl.CuTeDSL.jit
|
||||
kernel = _dsl.CuTeDSL.kernel
|
||||
register_jit_arg_adapter = _dsl.JitArgAdapterRegistry.register_jit_arg_adapter
|
||||
compile = _dsl.compile
|
||||
compile = _dsl.CompileCallable()
|
||||
OptLevel = _dsl.OptLevel
|
||||
PtxasOptions = _dsl.PtxasOptions
|
||||
EnableAssertions = _dsl.EnableAssertions
|
||||
GenerateLineInfo = _dsl.GenerateLineInfo
|
||||
KeepCUBIN = _dsl.KeepCUBIN
|
||||
KeepPTX = _dsl.KeepPTX
|
||||
GPUArch = _dsl.GPUArch
|
||||
LinkLibraries = _dsl.LinkLibraries
|
||||
|
||||
# Explicitly export all symbols for documentation generation
|
||||
__all__ = [
|
||||
@@ -186,22 +209,22 @@ __all__ = [
|
||||
"ComposedLayout",
|
||||
"Swizzle",
|
||||
"E",
|
||||
"ScaledBasis",
|
||||
"Atom",
|
||||
"MmaAtom",
|
||||
"CopyAtom",
|
||||
"TiledCopy",
|
||||
"TiledMma",
|
||||
"ThrMma",
|
||||
"ThrCopy",
|
||||
"TensorSSA",
|
||||
"ReductionOp",
|
||||
# Basic utility functions
|
||||
"assume",
|
||||
"is_integer",
|
||||
"is_int_tuple",
|
||||
"is_static",
|
||||
"size",
|
||||
"has_underscore",
|
||||
"slice_",
|
||||
"depth",
|
||||
"rank",
|
||||
"shape",
|
||||
"printf",
|
||||
"print_tensor",
|
||||
@@ -211,6 +234,7 @@ __all__ = [
|
||||
"recast_layout",
|
||||
"make_identity_layout",
|
||||
"make_ordered_layout",
|
||||
"make_layout_like",
|
||||
"make_composed_layout",
|
||||
"make_layout_tv",
|
||||
"make_layout_image_mask",
|
||||
@@ -220,6 +244,8 @@ __all__ = [
|
||||
"make_identity_tensor",
|
||||
"make_fragment",
|
||||
"make_fragment_like",
|
||||
"make_rmem_tensor",
|
||||
"make_rmem_tensor_like",
|
||||
"recast_ptr",
|
||||
"recast_tensor",
|
||||
# Tensor manipulation
|
||||
@@ -230,6 +256,7 @@ __all__ = [
|
||||
"leading_dim",
|
||||
"find",
|
||||
"find_if",
|
||||
"transform_leaf",
|
||||
"coalesce",
|
||||
"group_modes",
|
||||
"cosize",
|
||||
@@ -237,6 +264,7 @@ __all__ = [
|
||||
# Tuple operations
|
||||
"flatten_to_tuple",
|
||||
"flatten",
|
||||
"unflatten",
|
||||
"product",
|
||||
"product_like",
|
||||
"product_each",
|
||||
@@ -244,6 +272,7 @@ __all__ = [
|
||||
"append",
|
||||
"prepend_ones",
|
||||
"append_ones",
|
||||
"elem_less",
|
||||
# Math operations
|
||||
"ceil_div",
|
||||
"round_up",
|
||||
@@ -251,7 +280,6 @@ __all__ = [
|
||||
"slice_and_offset",
|
||||
"crd2idx",
|
||||
"domain_offset",
|
||||
"elem_less",
|
||||
"filter_zeros",
|
||||
"filter",
|
||||
"tile_to_shape",
|
||||
@@ -280,18 +308,27 @@ __all__ = [
|
||||
"tiled_divide",
|
||||
"local_partition",
|
||||
"local_tile",
|
||||
# MMA and Copy operations
|
||||
# MMA and Copy atom operations
|
||||
"make_atom",
|
||||
"make_mma_atom",
|
||||
"make_tiled_mma",
|
||||
"make_copy_atom",
|
||||
"make_tiled_copy_tv",
|
||||
"make_tiled_copy",
|
||||
"make_tiled_copy_S",
|
||||
"make_tiled_copy_D",
|
||||
"make_tiled_copy_A",
|
||||
"make_tiled_copy_B",
|
||||
"make_tiled_copy_C",
|
||||
"make_tiled_copy_C_atom",
|
||||
"make_cotiled_copy",
|
||||
"copy_atom_call",
|
||||
# Algorithm operations
|
||||
"basic_copy",
|
||||
"basic_copy_if",
|
||||
"autovec_copy",
|
||||
"copy",
|
||||
"copy_atom_call",
|
||||
"prefetch",
|
||||
"gemm",
|
||||
# Tensor creation
|
||||
"full",
|
||||
@@ -302,8 +339,9 @@ __all__ = [
|
||||
"where",
|
||||
"any_",
|
||||
"all_",
|
||||
"repeat_as_tuple",
|
||||
"repeat",
|
||||
"repeat_like",
|
||||
"ScaledBasis",
|
||||
# User defined struct
|
||||
"struct",
|
||||
# Modules
|
||||
@@ -311,6 +349,8 @@ __all__ = [
|
||||
"nvgpu",
|
||||
"testing",
|
||||
"runtime",
|
||||
# Math utils
|
||||
*math.__all__,
|
||||
# Decorators and code generation
|
||||
"jit",
|
||||
"kernel",
|
||||
|
||||
438
python/CuTeDSL/cutlass/cute/algorithm.py
Normal file
438
python/CuTeDSL/cutlass/cute/algorithm.py
Normal file
@@ -0,0 +1,438 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# Use of this software is governed by the terms and conditions of the
|
||||
# NVIDIA End User License Agreement (EULA), available at:
|
||||
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
|
||||
#
|
||||
# Any use, reproduction, disclosure, or distribution of this software
|
||||
# and related documentation outside the scope permitted by the EULA
|
||||
# is strictly prohibited.
|
||||
|
||||
import math
|
||||
from typing import Optional, Dict, Any, List, Tuple
|
||||
|
||||
from cutlass._mlir import ir
|
||||
from cutlass.cutlass_dsl import for_generate, yield_out, if_generate, dsl_user_op
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
|
||||
from .typing import Tensor, Int64, Int16, AddressSpace
|
||||
from .core import (
|
||||
rank,
|
||||
is_static,
|
||||
size,
|
||||
make_layout,
|
||||
make_ptr,
|
||||
max_common_layout,
|
||||
logical_divide,
|
||||
append_ones,
|
||||
group_modes,
|
||||
)
|
||||
from .atom import MmaAtom, CopyAtom, make_atom
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def gemm(
|
||||
atom: MmaAtom,
|
||||
d: Tensor,
|
||||
a: Tensor,
|
||||
b: Tensor,
|
||||
c: Tensor,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
"""The GEMM algorithm.
|
||||
|
||||
Computes ``D <- A * B + C`` where ``C`` and ``D`` can alias. Note that some MMA Atoms (e.g.
|
||||
warpgroup-wide or tcgen05 MMAs) require manually setting an "accumulate" boolean field.
|
||||
|
||||
All tensors must be partitioned according to the provided MMA Atom.
|
||||
|
||||
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
|
||||
:type d: Tensor
|
||||
:param a: First source tensor
|
||||
:type a: Tensor
|
||||
:param b: Second source tensor
|
||||
:type b: Tensor
|
||||
:param c: Third source tensor
|
||||
:type c: Tensor
|
||||
:param loc: Source location for MLIR, defaults to None
|
||||
:type loc: Optional[Location], optional
|
||||
:param ip: Insertion point for MLIR, defaults to None
|
||||
:type ip: Optional[InsertionPoint], optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
:return: None
|
||||
: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)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def basic_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
|
||||
"""Performs a basic element-wise copy.
|
||||
|
||||
This functions **assumes** the following pre-conditions:
|
||||
1. `size(src) == size(dst)`
|
||||
|
||||
When the `src` and `dst` shapes are static, the pre-conditions are actually verified and the
|
||||
element-wise loop is fully unrolled.
|
||||
|
||||
:param src: Source tensor
|
||||
:type src: Tensor
|
||||
:param dst: Destination tensor
|
||||
:type dst: Tensor
|
||||
:param loc: Source location for MLIR, defaults to None
|
||||
:type loc: Optional[Location], optional
|
||||
:param ip: Insertion point, defaults to None
|
||||
:type ip: Optional[InsertionPoint], optional
|
||||
"""
|
||||
|
||||
if is_static(src.shape) and is_static(dst.shape):
|
||||
simt_copy_ty = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
|
||||
src.element_type.mlir_type, src.element_type.width
|
||||
)
|
||||
simt_copy = make_atom(simt_copy_ty, loc=loc, ip=ip)
|
||||
return _cute_ir.copy(simt_copy, src.value, dst.value, loc=loc, ip=ip)
|
||||
|
||||
s = size(dst, loc=loc, ip=ip)
|
||||
# Always generate an scf.for Op when one of the tensors is dynamic
|
||||
for i in for_generate(0, s, loc=loc, ip=ip):
|
||||
dst[i] = src[i]
|
||||
yield_out()
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def basic_copy_if(pred: Tensor, src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
|
||||
"""Performs a basic predicated element-wise copy.
|
||||
|
||||
This functions **assumes** the following pre-conditions:
|
||||
1. `size(src) == size(dst)`
|
||||
2. `size(src) == size(pred)`
|
||||
|
||||
When all shapes are static, the pre-conditions are actually verified and the element-wise loop
|
||||
is fully unrolled.
|
||||
|
||||
"""
|
||||
if src.element_type.width != dst.element_type.width:
|
||||
raise NotImplementedError(
|
||||
"basic_copy_if currently only supports equal source and destination "
|
||||
"element type bit width"
|
||||
)
|
||||
|
||||
if is_static(src.shape) and is_static(dst.shape) and is_static(pred.shape):
|
||||
return _basic_copy_if_static(pred, src, dst, loc=loc, ip=ip)
|
||||
|
||||
s = size(dst, loc=loc, ip=ip)
|
||||
# Always generate an scf.for Op when one of the tensors is dynamic
|
||||
for i in for_generate(0, s, loc=loc, ip=ip):
|
||||
if_generate(pred[i], lambda: dst.__setitem__(i, src[i]), loc=loc, ip=ip) # type: ignore
|
||||
yield_out()
|
||||
|
||||
|
||||
# Version of basic_copy_if when src and dst have static shapes
|
||||
# - verify size(src) == size(dst) == size(prd)
|
||||
# - fully unroll the loop for now
|
||||
def _basic_copy_if_static(
|
||||
pred: Tensor, src: Tensor, dst: Tensor, *, loc=None, ip=None
|
||||
) -> None:
|
||||
assert is_static(src.shape) and is_static(dst.shape) and is_static(pred.shape)
|
||||
if size(src, loc=loc, ip=ip) != size(dst, loc=loc, ip=ip):
|
||||
raise ValueError(
|
||||
"basic_copy expects the size of source, destination, and predicate tensors to match"
|
||||
)
|
||||
# Fully unrolled loop in the static case for now
|
||||
for i in range(size(dst, loc=loc, ip=ip)):
|
||||
if_generate(pred[i], lambda: dst.__setitem__(i, src[i]), loc=loc, ip=ip) # type: ignore
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Auto-vectorization SIMT copy policy.
|
||||
|
||||
Given a source and destination tensors that are statically shaped, this policy figures out the
|
||||
largest safe vector width that the copy instruction can take and performs the copy.
|
||||
"""
|
||||
if src.element_type.width != dst.element_type.width:
|
||||
raise NotImplementedError(
|
||||
"autovec_copy currently only supports equal source and destination "
|
||||
"element type bit width"
|
||||
)
|
||||
|
||||
# We are going to dispatch to copy-with-atom which requires shapes to be static
|
||||
if not is_static(src.shape) or not is_static(dst.shape):
|
||||
raise ValueError(
|
||||
"autovec_copy expects source and destination tensors to be statically shaped"
|
||||
)
|
||||
|
||||
vec_layout = max_common_layout(src, dst, loc=loc, ip=ip)
|
||||
num_common_elements = size(vec_layout, loc=loc, ip=ip)
|
||||
|
||||
# Next we construct an upper-bound on the number bits that can be vectorized by considering
|
||||
# - the maximum alignment of the layouts
|
||||
# - the maximum alignment of the pointers
|
||||
|
||||
upper_bound = math.gcd(src.layout.max_alignment, dst.layout.max_alignment)
|
||||
upper_bound = math.gcd(upper_bound, num_common_elements)
|
||||
upper_bound *= src.element_type.width
|
||||
|
||||
# For our instructions, the alignment of the pointer is an upper bound to the vector width
|
||||
# max_alignment, as opposed to alignment, takes into account possible address swizzling
|
||||
upper_bound = math.gcd(upper_bound, src.iterator.max_alignment * 8)
|
||||
upper_bound = math.gcd(upper_bound, dst.iterator.max_alignment * 8)
|
||||
|
||||
# Finally, we put a cap at 128b
|
||||
num_bits_per_copy = math.gcd(upper_bound, 128)
|
||||
|
||||
if (num_common_elements > 1) and (num_bits_per_copy % 8 == 0):
|
||||
num_common_elements = num_bits_per_copy // src.element_type.width
|
||||
|
||||
# 2 step logical divides ensuring that the divides are valid at every step
|
||||
vec_src = logical_divide(src, vec_layout, loc=loc, ip=ip)
|
||||
vec_dst = logical_divide(dst, vec_layout, loc=loc, ip=ip)
|
||||
tiled_src = logical_divide(
|
||||
vec_src, make_layout(num_common_elements, loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
tiled_dst = logical_divide(
|
||||
vec_dst, make_layout(num_common_elements, loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
|
||||
# Dispatch to copy with atom
|
||||
simt_type = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
|
||||
src.element_type.mlir_type, num_bits_per_copy
|
||||
)
|
||||
simt_copy = make_atom(simt_type, loc=loc, ip=ip)
|
||||
return _cute_ir.copy(
|
||||
simt_copy, tiled_src.value, tiled_dst.value, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
# Failed to vectorize, use a basic copy
|
||||
basic_copy(src, dst, loc=loc, ip=ip)
|
||||
|
||||
|
||||
def _parse_auto_multicast_args(
|
||||
kwargs: Dict[str, Any],
|
||||
) -> List[Tuple[str, ir.Attribute]]:
|
||||
"""
|
||||
Parse multicast-related kwargs and return a list of (attr_name, attr) pairs.
|
||||
|
||||
This function consumes the following key from kwargs if present:
|
||||
- 'auto_multicast': dict
|
||||
dict: { 'multicast_layout': str, 'use_2cta': bool }
|
||||
|
||||
Returns:
|
||||
List of (attr_name, ir.Attribute) pairs to be attached to the op.
|
||||
Recognized attributes:
|
||||
- ('multicast_layout', #cute.layout<...>) when a layout string is provided
|
||||
- ('use_2cta', unit) when use_2cta is True
|
||||
"""
|
||||
attr_pairs: List[Tuple[str, ir.Attribute]] = []
|
||||
|
||||
# Pop known keys to avoid leaking to trait unpack
|
||||
auto_multicast = kwargs.pop("auto_multicast", None)
|
||||
|
||||
use_2cta: bool = False
|
||||
layout_str: Optional[str] = None
|
||||
|
||||
if auto_multicast is not None:
|
||||
if not isinstance(auto_multicast, dict):
|
||||
raise TypeError(
|
||||
"auto_multicast must be a dict with keys 'multicast_layout' and optional 'use_2cta'"
|
||||
)
|
||||
layout_str = auto_multicast.get("multicast_layout", None)
|
||||
use_2cta = bool(auto_multicast.get("use_2cta", False))
|
||||
|
||||
if layout_str is not None:
|
||||
if not isinstance(layout_str, str):
|
||||
raise TypeError(
|
||||
"multicast_layout must be a string representing a CuTe layout, e.g. '(4,2):(1,0)'"
|
||||
)
|
||||
attr_pairs.append(
|
||||
(
|
||||
"multicast_layout",
|
||||
ir.Attribute.parse(f'#cute.layout<"{layout_str}">'),
|
||||
)
|
||||
)
|
||||
|
||||
if use_2cta:
|
||||
attr_pairs.append(("use_2cta", ir.UnitAttr.get()))
|
||||
|
||||
return attr_pairs
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def copy(
|
||||
atom: CopyAtom,
|
||||
src: Tensor,
|
||||
dst: Tensor,
|
||||
*,
|
||||
pred: Optional[Tensor] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
**kwargs,
|
||||
) -> None:
|
||||
"""Facilitates data transfer between two tensors conforming to layout profile ``(V, Rest...)``.
|
||||
|
||||
:param atom: Copy atom specifying the transfer operation
|
||||
:type atom: CopyAtom
|
||||
:param src: Source tensor with layout profile ``(V, Rest...)``
|
||||
:type src: Tensor
|
||||
:param dst: Destination tensor with layout profile ``(V, Rest...)``
|
||||
:type dst: Tensor
|
||||
:param pred: Optional predication tensor for conditional transfers, defaults to None
|
||||
:type pred: Optional[Tensor], optional
|
||||
:param loc: Source location information, defaults to None
|
||||
:type loc: Any, optional
|
||||
:param ip: Insertion point, defaults to None
|
||||
:type ip: Any, optional
|
||||
:param kwargs: Additional copy atom specific arguments
|
||||
:type kwargs: Dict[str, Any]
|
||||
:raises TypeError: If source and destination element type bit widths differ
|
||||
:raises ValueError: If source and destination ranks differ
|
||||
:raises ValueError: If source and destination mode-1 sizes differ
|
||||
:raises NotImplementedError: If ``V-mode`` rank exceeds 2
|
||||
:return: None
|
||||
:rtype: None
|
||||
|
||||
The ``V-mode`` represents either:
|
||||
|
||||
- A singular mode directly consumable by the provided Copy Atom
|
||||
- A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``,
|
||||
and src/dst layout like ``((V, Rest...), Rest...)``
|
||||
|
||||
The algorithm recursively processes the ``V-mode``, decomposing it until reaching the minimum granularity
|
||||
compatible with the provided Copy Atom's requirements.
|
||||
|
||||
Source and destination tensors must be partitioned in accordance with the Copy Atom specifications.
|
||||
Post-partitioning, both tensors will exhibit a ``(V, Rest...)`` layout profile.
|
||||
|
||||
**Precondition:** The size of mode 1 must be equal for both source and destination tensors:
|
||||
``size(src, mode=[1]) == size(dst, mode=[1])``
|
||||
|
||||
**Examples**:
|
||||
|
||||
TMA copy operation with multicast functionality:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
cute.copy(tma_atom, src, dst, tma_bar_ptr=mbar_ptr, mcast_mask=mask)
|
||||
|
||||
Optional predication is supported through an additional tensor parameter. For partitioned tensors with
|
||||
logical profile ``((ATOM_V,ATOM_REST),REST,...)``, the predication tensor must maintain profile
|
||||
compatibility with ``(ATOM_REST,REST,...)``.
|
||||
|
||||
For Copy Atoms requiring single-threaded execution, thread election is managed automatically by the
|
||||
copy operation. External thread selection mechanisms are not necessary.
|
||||
|
||||
.. note::
|
||||
|
||||
- Certain Atoms may require additional operation-specific keyword arguments.
|
||||
- Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned
|
||||
for future releases.
|
||||
|
||||
"""
|
||||
if isinstance(src.type, _cute_ir.MemRefType) and isinstance(
|
||||
dst.type, _cute_ir.MemRefType
|
||||
):
|
||||
if src.element_type.width != dst.element_type.width:
|
||||
raise TypeError(
|
||||
"`copy` currently only supports equal source and destination "
|
||||
"element type bit width"
|
||||
)
|
||||
|
||||
if rank(src) != rank(dst):
|
||||
raise ValueError(
|
||||
"Expected source and destination tensors to have the same rank, "
|
||||
f"but got {rank(src)} and {rank(dst)}"
|
||||
)
|
||||
|
||||
# Canonicalize to at least rank-2 tensors
|
||||
src = group_modes(append_ones(src, up_to_rank=2), 1)
|
||||
dst = group_modes(append_ones(dst, up_to_rank=2), 1)
|
||||
if pred is not None:
|
||||
pred = group_modes(append_ones(pred, up_to_rank=2), 1)
|
||||
|
||||
if is_static(src.shape[1]) and is_static(dst.shape[1]):
|
||||
if size(src, mode=[1]) != size(dst, mode=[1]):
|
||||
raise ValueError(
|
||||
"Expected source and destination tensors to have the same size in mode-1, "
|
||||
f"but got {size(src, mode=[1])} and {size(dst, mode=[1])}"
|
||||
)
|
||||
|
||||
multicast_attr_pairs = _parse_auto_multicast_args(kwargs)
|
||||
|
||||
value = atom._unpack(loc=loc, ip=ip, **kwargs)
|
||||
if isinstance(pred, Tensor):
|
||||
pred = pred.value
|
||||
|
||||
op = _cute_ir.copy(value, src.value, dst.value, pred=pred, loc=loc, ip=ip)
|
||||
|
||||
for name, attr in multicast_attr_pairs:
|
||||
op.attributes[name] = attr
|
||||
|
||||
return op
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
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 supports TMA prefetch.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
cute.prefetch(tma_prefetch, 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)
|
||||
dummy_mcast_mask = Int16(0)
|
||||
value = atom._unpack(
|
||||
loc=loc, ip=ip, tma_bar_ptr=dummy_tma_bar_ptr, mcast_mask=dummy_mcast_mask
|
||||
)
|
||||
return _cute_ir.prefetch(value, src.value, loc=loc, ip=ip)
|
||||
@@ -11,9 +11,11 @@
|
||||
|
||||
from .elect import *
|
||||
from .mbar import *
|
||||
from .numeric_conversion import *
|
||||
from .nvvm_wrappers import *
|
||||
from .smem import *
|
||||
from .tmem import *
|
||||
from .numeric_conversion import *
|
||||
|
||||
# __all__ is required here for documentation generation
|
||||
__all__ = [
|
||||
@@ -44,6 +46,7 @@ __all__ = [
|
||||
"grid_dim",
|
||||
"cluster_idx",
|
||||
"cluster_dim",
|
||||
"cluster_size",
|
||||
"block_in_cluster_idx",
|
||||
"block_in_cluster_dim",
|
||||
"block_idx_in_cluster",
|
||||
@@ -66,9 +69,12 @@ __all__ = [
|
||||
"cluster_wait",
|
||||
"cluster_arrive",
|
||||
"cluster_arrive_relaxed",
|
||||
"fence_proxy",
|
||||
"vote_ballot_sync",
|
||||
"vote_any_sync",
|
||||
"vote_all_sync",
|
||||
"vote_uni_sync",
|
||||
"popc",
|
||||
"fence_proxy",
|
||||
"fence_view_async_tmem_load",
|
||||
"fence_view_async_tmem_store",
|
||||
"warpgroup_reg_alloc",
|
||||
@@ -98,4 +104,15 @@ __all__ = [
|
||||
"alloc_tmem",
|
||||
"relinquish_tmem_alloc_permit",
|
||||
"dealloc_tmem",
|
||||
#
|
||||
# numeric_conversion.py
|
||||
#
|
||||
"prmt",
|
||||
"cvt_i8_bf16_intrinsic",
|
||||
"cvt_i4_bf16_intrinsic",
|
||||
"cvt_f4e2m1_f16_intrinsic",
|
||||
"cvt_i8x4_to_f32x4",
|
||||
"cvt_i8x2_to_f32x2",
|
||||
"cvt_i8_bf16",
|
||||
"cvt_f32x2_bf16x2",
|
||||
]
|
||||
|
||||
@@ -9,6 +9,7 @@
|
||||
# and related documentation outside the scope permitted by the EULA
|
||||
# is strictly prohibited.
|
||||
|
||||
from cutlass.base_dsl.arch import Arch
|
||||
from cutlass.cutlass_dsl import CuTeDSL, T, dsl_user_op
|
||||
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
@@ -16,17 +17,17 @@ from cutlass._mlir.dialects import nvvm, scf
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..typing import Int, Int32
|
||||
from ...impl_utils import check_value_in
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_warp_uniform(value: Int, *, loc=None, ip=None) -> Int32:
|
||||
"""
|
||||
Creates a warp-uniform value from the given integer input.
|
||||
Provides a compiler hint indicating that the specified value is invariant across all threads in the warp,
|
||||
which may enable performance optimizations.
|
||||
|
||||
:param value: The integer to make warp uniform.
|
||||
:param value: The integer value to be marked as warp-uniform.
|
||||
:type value: Int
|
||||
:return: The warp-uniform value equal to the input.
|
||||
:return: The input value, marked as warp-uniform.
|
||||
:rtype: Int32
|
||||
"""
|
||||
return Int32(
|
||||
@@ -68,17 +69,7 @@ def elect_one(*, loc=None, ip=None) -> IfOpRegion:
|
||||
# Only one thread in the warp executes the code in this context
|
||||
pass
|
||||
"""
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
is_thread_leader = nvvm.elect_sync(T.bool())
|
||||
if_op = scf.IfOp(is_thread_leader, loc=loc, ip=ip)
|
||||
return IfOpRegion(if_op.then_block, loc=loc, ip=ip)
|
||||
|
||||
@@ -10,14 +10,12 @@
|
||||
# is strictly prohibited.
|
||||
from typing import Optional
|
||||
|
||||
from cutlass.base_dsl.arch import Arch
|
||||
from cutlass.cutlass_dsl import CuTeDSL, T, if_generate, dsl_user_op
|
||||
|
||||
from cutlass._mlir.dialects import nvvm
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..typing import Pointer, Int, Boolean, Int32
|
||||
from ...impl_utils import check_value_in
|
||||
|
||||
from ..typing import Pointer, Int, Boolean, Int32, AddressSpace
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
@@ -46,17 +44,7 @@ def mbarrier_init_fence(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
A fence operation that applies to the mbarrier initializations.
|
||||
"""
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
nvvm.fence_mbarrier_init(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@@ -75,17 +63,7 @@ def mbarrier_arrive_and_expect_tx(
|
||||
the mbarrier is converted to a remote address in the peer CTA's
|
||||
SMEM.
|
||||
"""
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
|
||||
mbar_llvm_ptr = mbar_ptr.llvm_ptr
|
||||
if peer_cta_rank_in_cluster is not None:
|
||||
@@ -125,17 +103,7 @@ def mbarrier_expect_tx(
|
||||
the mbarrier is converted to a remote address in the peer CTA's
|
||||
SMEM.
|
||||
"""
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
|
||||
mbar_llvm_ptr = mbar_ptr.llvm_ptr
|
||||
if peer_cta_rank_in_cluster is not None:
|
||||
@@ -170,17 +138,7 @@ def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None:
|
||||
:param phase: The phase to wait for (either 0 or 1)
|
||||
:type phase: Int
|
||||
"""
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
|
||||
timeout_ns = 10000000
|
||||
# This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX
|
||||
@@ -206,17 +164,7 @@ def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Bo
|
||||
:return: A boolean value indicating whether the wait operation was successful
|
||||
:rtype: Boolean
|
||||
"""
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
|
||||
return Boolean(
|
||||
nvvm.mbarrier_wait_parity(
|
||||
@@ -245,23 +193,15 @@ def mbarrier_conditional_try_wait(
|
||||
:return: A boolean value indicating whether the wait operation was successful
|
||||
:rtype: Boolean
|
||||
"""
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
return if_generate(
|
||||
cond,
|
||||
lambda: mbarrier_try_wait(mbar_ptr, phase, loc=loc, ip=ip),
|
||||
lambda: Boolean(True).ir_value(loc=loc, ip=ip),
|
||||
None,
|
||||
[Boolean],
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@@ -284,17 +224,7 @@ def mbarrier_arrive(
|
||||
"""
|
||||
mbar_llvm_ptr = mbar_ptr.llvm_ptr
|
||||
if peer_cta_rank_in_cluster is not None:
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
check_value_in(
|
||||
arch,
|
||||
[
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
],
|
||||
"arch",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
|
||||
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
|
||||
mbar_llvm_ptr.type,
|
||||
@@ -328,17 +258,7 @@ def cp_async_mbarrier_arrive_noinc(mbar_ptr: Pointer, *, loc=None, ip=None) -> N
|
||||
: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",
|
||||
)
|
||||
CuTeDSL._get_dsl().check_arch(lambda arch: arch >= Arch.sm_90)
|
||||
|
||||
mbar_llvm_ptr = mbar_ptr.llvm_ptr
|
||||
nvvm.cp_async_mbarrier_arrive_shared(
|
||||
|
||||
263
python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py
Normal file
263
python/CuTeDSL/cutlass/cute/arch/numeric_conversion.py
Normal file
@@ -0,0 +1,263 @@
|
||||
# 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 cutlass.cutlass_dsl import dsl_user_op
|
||||
from cutlass._mlir import ir
|
||||
from cutlass._mlir.dialects import builtin, arith, llvm, vector
|
||||
|
||||
from .nvvm_wrappers import (
|
||||
cvt_i8_bf16,
|
||||
cvt_f32x2_bf16x2,
|
||||
cvt_i8x4_to_f32x4,
|
||||
cvt_i8x2_to_f32x2,
|
||||
cvt_i4x8_to_bf16x8,
|
||||
cvt_i4x4_to_bf16x4,
|
||||
cvt_i4x2_to_bf16x2,
|
||||
cvt_i4_bf16,
|
||||
cvt_f4e2m1x8_to_f16x8,
|
||||
cvt_f4e2m1x4_to_f16x4,
|
||||
cvt_f4e2m1x2_to_f16x2,
|
||||
cvt_f4e2m1_f16,
|
||||
)
|
||||
|
||||
from ..typing import (
|
||||
Int4,
|
||||
Int8,
|
||||
Int32,
|
||||
Float16,
|
||||
BFloat16,
|
||||
Float32,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cvt_i8_bf16_intrinsic(vec_i8, length, *, loc=None, ip=None):
|
||||
"""
|
||||
Convert a vector of int8 to a vector of bfloat16.
|
||||
|
||||
:param vec_i8: The input vector of int8.
|
||||
:type vec_i8: 1D vector of int8
|
||||
:param length: The length of the input vector.
|
||||
:type length: int
|
||||
:return: The output 1D vector of bfloat16 with the same length as the input vector.
|
||||
:rtype: 1D vector of bfloat16
|
||||
"""
|
||||
src_pos = 0
|
||||
vec_i8x4_type = ir.VectorType.get([4], Int8.mlir_type, loc=loc)
|
||||
vec_i8x2_type = ir.VectorType.get([2], Int8.mlir_type, loc=loc)
|
||||
vec_f32x2_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc)
|
||||
vec_dst_type = ir.VectorType.get([length], BFloat16.mlir_type, loc=loc)
|
||||
vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)
|
||||
# try to use vectorized version
|
||||
if length >= 4:
|
||||
num_vec4 = length // 4
|
||||
for _ in range(num_vec4):
|
||||
vec_i8x4 = vector.extract_strided_slice(
|
||||
vec_i8x4_type, vec_i8, [src_pos], [4], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_f32x4 = cvt_i8x4_to_f32x4(vec_i8x4, loc=loc, ip=ip)
|
||||
vec_f32x2_lo = vector.extract_strided_slice(
|
||||
vec_f32x2_type, vec_f32x4, [0], [2], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_f32x2_hi = vector.extract_strided_slice(
|
||||
vec_f32x2_type, vec_f32x4, [2], [2], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_bf16x2_lo = cvt_f32x2_bf16x2(vec_f32x2_lo, loc=loc, ip=ip)
|
||||
vec_bf16x2_hi = cvt_f32x2_bf16x2(vec_f32x2_hi, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_bf16x2_lo, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_bf16x2_hi, vec_dst, [src_pos + 2], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 4
|
||||
length -= 4
|
||||
if length >= 2:
|
||||
vec_i8x2 = vector.extract_strided_slice(
|
||||
vec_i8x2_type, vec_i8, [src_pos], [2], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_f32x2 = cvt_i8x2_to_f32x2(vec_i8x2, loc=loc, ip=ip)
|
||||
vec_bf16x2 = cvt_f32x2_bf16x2(vec_f32x2, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_bf16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 2
|
||||
length -= 2
|
||||
if length >= 1:
|
||||
val_bf16 = cvt_i8_bf16(
|
||||
vector.extractelement(
|
||||
vec_i8,
|
||||
position=arith.constant(Int32.mlir_type, src_pos),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
vec_dst = vector.insertelement(
|
||||
val_bf16,
|
||||
vec_dst,
|
||||
position=arith.constant(Int32.mlir_type, src_pos),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return vec_dst
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cvt_i4_bf16_intrinsic(vec_i4, length, *, loc=None, ip=None):
|
||||
"""
|
||||
Convert a vector of int4 to a vector of bfloat16.
|
||||
|
||||
:param vec_i4: The input vector of int4.
|
||||
:type vec_i4: 1D vector of int4
|
||||
:param length: The length of the input vector.
|
||||
:type length: int
|
||||
:return: The output 1D vector of bfloat16 with the same length as the input vector.
|
||||
:rtype: 1D vector of bfloat16
|
||||
"""
|
||||
src_pos = 0
|
||||
vec_i4x8_type = ir.VectorType.get([8], Int4.mlir_type, loc=loc)
|
||||
vec_i4x4_type = ir.VectorType.get([4], Int4.mlir_type, loc=loc)
|
||||
vec_i4x2_type = ir.VectorType.get([2], Int4.mlir_type, loc=loc)
|
||||
vec_dst_type = ir.VectorType.get([length], BFloat16.mlir_type, loc=loc)
|
||||
vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)
|
||||
# try to use vectorized version
|
||||
if length >= 8:
|
||||
num_vec8 = length // 8
|
||||
for _ in range(num_vec8):
|
||||
vec_i4x8 = vector.extract_strided_slice(
|
||||
vec_i4x8_type, vec_i4, [src_pos], [8], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_bf16x8 = cvt_i4x8_to_bf16x8(vec_i4x8, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_bf16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 8
|
||||
length -= 8
|
||||
if length >= 4:
|
||||
vec_i4x4 = vector.extract_strided_slice(
|
||||
vec_i4x4_type, vec_i4, [src_pos], [4], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_bf16x4 = cvt_i4x4_to_bf16x4(vec_i4x4, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_bf16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 4
|
||||
length -= 4
|
||||
if length >= 2:
|
||||
vec_i4x2 = vector.extract_strided_slice(
|
||||
vec_i4x2_type, vec_i4, [src_pos], [2], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_bf16x2 = cvt_i4x2_to_bf16x2(vec_i4x2, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_bf16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 2
|
||||
length -= 2
|
||||
if length >= 1:
|
||||
val_bf16 = cvt_i4_bf16(
|
||||
vector.extractelement(
|
||||
vec_i4,
|
||||
position=arith.constant(Int32.mlir_type, src_pos),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
vec_dst = vector.insertelement(
|
||||
val_bf16,
|
||||
vec_dst,
|
||||
position=arith.constant(Int32.mlir_type, src_pos),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return vec_dst
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cvt_f4e2m1_f16_intrinsic(vec_f4e2m1, length, *, loc=None, ip=None):
|
||||
"""
|
||||
Convert a vector of float4e2m1 to a vector of float16.
|
||||
|
||||
:param vec_f4e2m1: The input vector of float4e2m1.
|
||||
:type vec_f4e2m1: 1D vector of float4e2m1
|
||||
:param length: The length of the input vector.
|
||||
:type length: int
|
||||
:return: The output 1D vector of float16 with the same length as the input vector.
|
||||
:rtype: 1D vector of float16
|
||||
"""
|
||||
src_pos = 0
|
||||
vec_src_i4 = builtin.unrealized_conversion_cast(
|
||||
[ir.VectorType.get([length], Int4.mlir_type, loc=loc)],
|
||||
[vec_f4e2m1],
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
vec_i4x8_type = ir.VectorType.get([8], Int4.mlir_type, loc=loc)
|
||||
vec_i4x4_type = ir.VectorType.get([4], Int4.mlir_type, loc=loc)
|
||||
vec_i4x2_type = ir.VectorType.get([2], Int4.mlir_type, loc=loc)
|
||||
vec_dst_type = ir.VectorType.get([length], Float16.mlir_type, loc=loc)
|
||||
vec_dst = llvm.mlir_zero(vec_dst_type, loc=loc, ip=ip)
|
||||
# try to use vectorized version
|
||||
if length >= 8:
|
||||
num_vec8 = length // 8
|
||||
for _ in range(num_vec8):
|
||||
vec_f4e2m1x8 = vector.extract_strided_slice(
|
||||
vec_i4x8_type, vec_src_i4, [src_pos], [8], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_f16x8 = cvt_f4e2m1x8_to_f16x8(vec_f4e2m1x8, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_f16x8, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 8
|
||||
length -= 8
|
||||
if length >= 4:
|
||||
vec_f4e2m1x4 = vector.extract_strided_slice(
|
||||
vec_i4x4_type, vec_src_i4, [src_pos], [4], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_f16x4 = cvt_f4e2m1x4_to_f16x4(vec_f4e2m1x4, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_f16x4, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 4
|
||||
length -= 4
|
||||
if length >= 2:
|
||||
vec_f4e2m1x2 = vector.extract_strided_slice(
|
||||
vec_i4x2_type, vec_src_i4, [src_pos], [2], [1], loc=loc, ip=ip
|
||||
)
|
||||
vec_f16x2 = cvt_f4e2m1x2_to_f16x2(vec_f4e2m1x2, loc=loc, ip=ip)
|
||||
vec_dst = vector.insert_strided_slice(
|
||||
vec_f16x2, vec_dst, [src_pos], [1], loc=loc, ip=ip
|
||||
)
|
||||
src_pos += 2
|
||||
length -= 2
|
||||
if length >= 1:
|
||||
val_f16 = cvt_f4e2m1_f16(
|
||||
vector.extractelement(
|
||||
vec_src_i4,
|
||||
position=arith.constant(Int32.mlir_type, src_pos),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
vec_dst = vector.insertelement(
|
||||
val_f16,
|
||||
vec_dst,
|
||||
position=arith.constant(Int32.mlir_type, src_pos),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return vec_dst
|
||||
@@ -16,7 +16,7 @@ from typing_extensions import deprecated
|
||||
from cutlass.cutlass_dsl import T, dsl_user_op
|
||||
|
||||
from cutlass._mlir import ir
|
||||
from cutlass._mlir.dialects import llvm, nvvm, vector
|
||||
from cutlass._mlir.dialects import arith, llvm, nvvm, vector
|
||||
|
||||
# Forward nvvm enums
|
||||
from cutlass._mlir.dialects.nvvm import (
|
||||
@@ -30,11 +30,13 @@ from cutlass._mlir.dialects.nvvm import (
|
||||
from ..typing import (
|
||||
Int,
|
||||
Boolean,
|
||||
Int8,
|
||||
Int16,
|
||||
Uint16,
|
||||
Int32,
|
||||
Uint32,
|
||||
Int64,
|
||||
Float16,
|
||||
Float32,
|
||||
BFloat16,
|
||||
Numeric,
|
||||
@@ -164,6 +166,14 @@ def block_in_cluster_dim(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cluster_size(*, loc=None, ip=None) -> Int32:
|
||||
"""
|
||||
Returns the number of CTA within the cluster.
|
||||
"""
|
||||
return Int32(nvvm.read_ptx_sreg_cluster_nctarank(T.i32(), loc=loc, ip=ip))
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def block_idx_in_cluster(*, loc=None, ip=None) -> Int32:
|
||||
"""
|
||||
@@ -295,12 +305,49 @@ def shuffle_sync_op(
|
||||
shlf_res = llvm.bitcast(orig_type.mlir_type, shlf_res, loc=loc, ip=ip)
|
||||
return orig_type(shlf_res)
|
||||
|
||||
|
||||
shuffle_sync = partial(shuffle_sync_op, kind=nvvm.ShflKind.idx)
|
||||
shuffle_sync_up = partial(shuffle_sync_op, kind=nvvm.ShflKind.up)
|
||||
shuffle_sync_down = partial(shuffle_sync_op, kind=nvvm.ShflKind.down)
|
||||
shuffle_sync_bfly = partial(shuffle_sync_op, kind=nvvm.ShflKind.bfly)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def warp_reduction(
|
||||
val: Numeric, op: Callable, *, threads_in_group: int = 32, loc=None, ip=None
|
||||
) -> Numeric:
|
||||
"""warp reduction of a Numeric value(e.g.Float32) by shuffle_sync_bfly, accepts custom binary operator.
|
||||
The threads_in_group is the number of threads reduction group in a warp.
|
||||
E.g. 32 means the whole warp reduced in one group. 8 means the warp is divided into 4 thread groups, each group has 8 threads in reduction.
|
||||
|
||||
|
||||
:param val: register value
|
||||
:type val: cutlass.Numeric
|
||||
:param op: binary operator
|
||||
:type op: Callable
|
||||
:param threads_in_group: the number of threads reduction group in a warp
|
||||
:type threads_in_group: int
|
||||
:return: reduced value
|
||||
:rtype: cutlass.Numeric
|
||||
"""
|
||||
offset = threads_in_group // 2
|
||||
while offset > 0:
|
||||
val = op(
|
||||
val,
|
||||
shuffle_sync_bfly(
|
||||
val, offset=offset, mask=-1, mask_and_clamp=31, loc=loc, ip=ip
|
||||
),
|
||||
)
|
||||
offset = offset // 2
|
||||
return val
|
||||
|
||||
|
||||
warp_reduction_max = partial(
|
||||
warp_reduction, op=lambda x, y: fmax(x, y) if isinstance(x, Float32) else max(x, y)
|
||||
)
|
||||
warp_reduction_sum = partial(warp_reduction, op=lambda x, y: x + y)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def barrier(*, barrier_id=None, number_of_threads=None, loc=None, ip=None) -> None:
|
||||
"""
|
||||
@@ -473,8 +520,19 @@ def fence_proxy(
|
||||
def vote_ballot_sync(
|
||||
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
|
||||
) -> Int32:
|
||||
"""
|
||||
Performs a ballot operation across the warp.
|
||||
"""Performs a ballot operation across the warp.
|
||||
|
||||
It copies the predicate from each thread in mask into the corresponding bit position of
|
||||
destination register d, where the bit position corresponds to the thread's lane id.
|
||||
|
||||
:param pred: The predicate value for the current thread
|
||||
:type pred: Boolean
|
||||
:param mask: A 32-bit integer mask specifying which threads participate, defaults to all threads (0xFFFFFFFF)
|
||||
:type mask: Int, optional
|
||||
:return: A 32-bit integer where each bit represents a thread's predicate value
|
||||
:rtype: Int32
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`__.
|
||||
"""
|
||||
return Int32(
|
||||
nvvm.vote_ballot_sync(
|
||||
@@ -487,6 +545,97 @@ def vote_ballot_sync(
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def vote_sync_op(
|
||||
pred: Boolean, kind: str, mask: Int = FULL_MASK, *, loc=None, ip=None
|
||||
) -> Union[Int32, Boolean]:
|
||||
return_type = Boolean
|
||||
return_type_str = "pred"
|
||||
return return_type(
|
||||
llvm.inline_asm(
|
||||
T.bool(),
|
||||
[
|
||||
Boolean(pred).ir_value(loc=loc, ip=ip),
|
||||
Int32(mask).ir_value(loc=loc, ip=ip),
|
||||
],
|
||||
f"""{{\n\t
|
||||
.reg .pred ps;\n\t
|
||||
.reg .pred pd;\n\t
|
||||
setp.ne.b32 ps, $1, 0;\n\t
|
||||
vote.sync.{kind}.{return_type_str} pd, ps, $2;\n\t
|
||||
selp.b32 $0, 1, 0, pd;\n\t
|
||||
}}""",
|
||||
"=r,r,i",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def vote_any_sync(
|
||||
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
|
||||
) -> Boolean:
|
||||
"""True if source predicate is True for any non-exited threads in mask. Negate the source
|
||||
predicate to compute .not_all.
|
||||
|
||||
:param pred: The predicate value for the current thread
|
||||
:type pred: Boolean
|
||||
:param mask: A 32-bit integer mask specifying which threads participate, defaults to all
|
||||
threads (0xFFFFFFFF)
|
||||
:type mask: Int, optional
|
||||
:return: A boolean value indicating if the source predicate is True for all non-exited
|
||||
threads in mask
|
||||
:rtype: Boolean
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`__.
|
||||
"""
|
||||
return vote_sync_op(pred, "any", mask, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def vote_all_sync(
|
||||
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
|
||||
) -> Boolean:
|
||||
"""True if source predicate is True for all non-exited threads in mask. Negate the source
|
||||
predicate to compute .none.
|
||||
|
||||
:param pred: The predicate value for the current thread
|
||||
:type pred: Boolean
|
||||
:param mask: A 32-bit integer mask specifying which threads participate, defaults to all
|
||||
threads (0xFFFFFFFF)
|
||||
:type mask: Int, optional
|
||||
:return: A boolean value indicating if the source predicate is True for all non-exited
|
||||
threads in mask
|
||||
:rtype: Boolean
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-vote-sync>`__.
|
||||
"""
|
||||
return vote_sync_op(pred, "all", mask, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def vote_uni_sync(
|
||||
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
|
||||
) -> Boolean:
|
||||
"""True f source predicate has the same value in all non-exited threads in mask. Negating
|
||||
the source predicate also computes .uni
|
||||
|
||||
:param pred: The predicate value for the current thread
|
||||
:type pred: Boolean
|
||||
:param mask: A 32-bit integer mask specifying which threads participate, defaults to all
|
||||
threads (0xFFFFFFFF)
|
||||
:type mask: Int, optional
|
||||
:return: A boolean value indicating if the source predicate is True for all non-exited
|
||||
threads in mask
|
||||
:rtype: Boolean
|
||||
"""
|
||||
return vote_sync_op(pred, "uni", mask, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def popc(value: Numeric, *, loc=None, ip=None) -> Numeric:
|
||||
"""
|
||||
@@ -494,7 +643,7 @@ def popc(value: Numeric, *, loc=None, ip=None) -> Numeric:
|
||||
"""
|
||||
if not isinstance(value, Numeric):
|
||||
value = as_numeric(value)
|
||||
return type(value)(llvm.intr_ctpop(value.ir_value(), loc=loc, ip=ip))
|
||||
return type(value)(llvm.intr_ctpop(value.ir_value(loc=loc, ip=ip), loc=loc, ip=ip))
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
@@ -546,6 +695,27 @@ fence_view_async_tmem_store = partial(
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_view_async_shared(
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
Perform a fence operation on the async shared memory load or store.
|
||||
|
||||
.. note::
|
||||
This function is only available on sm_90 or higher.
|
||||
The fence is required to synchronize the shared memory load/store
|
||||
and let the pipeline release or commit the buffer.
|
||||
|
||||
This function is usually used for async execution unit (like TMA, UMMA) after the load/store operations.
|
||||
"""
|
||||
nvvm.fence_proxy(
|
||||
nvvm.ProxyKind.async_shared, space=nvvm.SharedSpace.shared_cta, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def warpgroup_reg_realloc_op(
|
||||
reg_count: int,
|
||||
@@ -569,7 +739,7 @@ warpgroup_reg_dealloc = partial(
|
||||
def calc_packed_f32x2_op(
|
||||
src_a: Tuple[Float32, Float32],
|
||||
src_b: Tuple[Float32, Float32],
|
||||
src_c: Tuple[Float32, Float32] | None,
|
||||
src_c: Optional[Tuple[Float32, Float32]],
|
||||
calc_func: Callable,
|
||||
*,
|
||||
rnd=RoundingModeKind.RZ,
|
||||
@@ -579,14 +749,23 @@ def calc_packed_f32x2_op(
|
||||
) -> Tuple[Float32, Float32]:
|
||||
vec_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc)
|
||||
vec_src_a = vector.from_elements(
|
||||
vec_type, tuple(as_numeric(a).ir_value() for a in src_a), loc=loc, ip=ip
|
||||
vec_type,
|
||||
tuple(as_numeric(a).ir_value(loc=loc, ip=ip) for a in src_a),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
vec_src_b = vector.from_elements(
|
||||
vec_type, tuple(as_numeric(b).ir_value() for b in src_b), loc=loc, ip=ip
|
||||
vec_type,
|
||||
tuple(as_numeric(b).ir_value(loc=loc, ip=ip) for b in src_b),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
if src_c is not None:
|
||||
vec_src_c = vector.from_elements(
|
||||
vec_type, tuple(as_numeric(c).ir_value() for c in src_c), loc=loc, ip=ip
|
||||
vec_type,
|
||||
tuple(as_numeric(c).ir_value(loc=loc, ip=ip) for c in src_c),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
vec_res = calc_func(
|
||||
vec_type, vec_src_a, vec_src_b, vec_src_c, rnd=rnd, ftz=ftz, loc=loc, ip=ip
|
||||
@@ -660,6 +839,418 @@ def exp2(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
|
||||
)
|
||||
|
||||
|
||||
# Convert 1 int8 value to 1 bfloat16 value
|
||||
@dsl_user_op
|
||||
def cvt_i8_bf16(src_i8, *, loc=None, ip=None):
|
||||
src_i16 = llvm.zext(Int16.mlir_type, src_i8, loc=loc, ip=ip)
|
||||
val_i16 = llvm.inline_asm(
|
||||
Uint16.mlir_type,
|
||||
[
|
||||
src_i16,
|
||||
],
|
||||
"""{\n\t
|
||||
.reg .b16 r;\n\t
|
||||
.reg .b8 s;\n\t
|
||||
mov.b16 {s,_}, $1;\n\t
|
||||
cvt.rn.bf16.s8 r, s;\n\t
|
||||
mov.b16 $0, r;\n\t
|
||||
}""",
|
||||
"=h,h",
|
||||
)
|
||||
val_bf16 = llvm.bitcast(BFloat16.mlir_type, val_i16, loc=loc, ip=ip)
|
||||
return val_bf16
|
||||
|
||||
|
||||
# Convert vector of 2 float values to vector of 2 bfloat16 values with satfinite rounding
|
||||
@dsl_user_op
|
||||
def cvt_f32x2_bf16x2(src_vec2, *, loc=None, ip=None):
|
||||
src0 = vector.extractelement(
|
||||
src_vec2, position=arith.constant(Int32.mlir_type, 0, loc=loc, ip=ip)
|
||||
)
|
||||
src1 = vector.extractelement(
|
||||
src_vec2, position=arith.constant(Int32.mlir_type, 1, loc=loc, ip=ip)
|
||||
)
|
||||
rst = llvm.inline_asm(
|
||||
T.i32(),
|
||||
[
|
||||
Float32(src1).ir_value(loc=loc, ip=ip),
|
||||
Float32(src0).ir_value(loc=loc, ip=ip),
|
||||
],
|
||||
"cvt.rn.satfinite.bf16x2.f32 $0, $1, $2;",
|
||||
"=r,f,f",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
vec_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc)
|
||||
vec_bf16x2 = llvm.bitcast(vec_type, rst, loc=loc, ip=ip)
|
||||
return vec_bf16x2
|
||||
|
||||
|
||||
# Convert 1 float32 value to 1 bfloat16 value
|
||||
@dsl_user_op
|
||||
def cvt_f32_bf16(src_f32, *, loc=None, ip=None):
|
||||
bf16_val = llvm.inline_asm(
|
||||
BFloat16.mlir_type,
|
||||
[
|
||||
src_f32,
|
||||
],
|
||||
"cvt.rn.bf16.f32 $0, $1;",
|
||||
"=h,f",
|
||||
)
|
||||
return bf16_val
|
||||
|
||||
|
||||
# Convert vector of 4 int8 values to vector of 4 float32 values
|
||||
@dsl_user_op
|
||||
def cvt_i8x4_to_f32x4(src_vec4, *, loc=None, ip=None):
|
||||
zero = arith.constant(Int32.mlir_type, 0, loc=loc, ip=ip)
|
||||
mask4 = (
|
||||
arith.constant(Int32.mlir_type, 0x00000001, loc=loc, ip=ip),
|
||||
arith.constant(Int32.mlir_type, 0x00000100, loc=loc, ip=ip),
|
||||
arith.constant(Int32.mlir_type, 0x00010000, loc=loc, ip=ip),
|
||||
arith.constant(Int32.mlir_type, 0x01000000, loc=loc, ip=ip),
|
||||
)
|
||||
src_i32 = llvm.bitcast(Int32.mlir_type, src_vec4, loc=loc, ip=ip)
|
||||
rst0 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32,
|
||||
mask4[0],
|
||||
zero,
|
||||
],
|
||||
"dp4a.s32.s32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
rst1 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32,
|
||||
mask4[1],
|
||||
zero,
|
||||
],
|
||||
"dp4a.s32.s32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
rst2 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32,
|
||||
mask4[2],
|
||||
zero,
|
||||
],
|
||||
"dp4a.s32.s32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
rst3 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32,
|
||||
mask4[3],
|
||||
zero,
|
||||
],
|
||||
"dp4a.s32.s32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
res0 = llvm.inline_asm(
|
||||
Float32.mlir_type,
|
||||
[
|
||||
rst0,
|
||||
],
|
||||
"cvt.rn.f32.s32 $0, $1;",
|
||||
"=f,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
res1 = llvm.inline_asm(
|
||||
Float32.mlir_type,
|
||||
[
|
||||
rst1,
|
||||
],
|
||||
"cvt.rn.f32.s32 $0, $1;",
|
||||
"=f,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
res2 = llvm.inline_asm(
|
||||
Float32.mlir_type,
|
||||
[
|
||||
rst2,
|
||||
],
|
||||
"cvt.rn.f32.s32 $0, $1;",
|
||||
"=f,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
res3 = llvm.inline_asm(
|
||||
Float32.mlir_type,
|
||||
[
|
||||
rst3,
|
||||
],
|
||||
"cvt.rn.f32.s32 $0, $1;",
|
||||
"=f,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
vec_f32x4_type = ir.VectorType.get([4], Float32.mlir_type, loc=loc)
|
||||
vec_f32x4 = vector.from_elements(
|
||||
vec_f32x4_type, [res0, res1, res2, res3], loc=loc, ip=ip
|
||||
)
|
||||
return vec_f32x4
|
||||
|
||||
|
||||
# Convert vector of 2 int8 values to vector of 2 float32 values
|
||||
@dsl_user_op
|
||||
def cvt_i8x2_to_f32x2(src_vec2, *, loc=None, ip=None):
|
||||
zero = arith.constant(Int32.mlir_type, 0, loc=loc, ip=ip)
|
||||
mask2 = (
|
||||
arith.constant(Int32.mlir_type, 0x00000001, loc=loc, ip=ip),
|
||||
arith.constant(Int32.mlir_type, 0x00000100, loc=loc, ip=ip),
|
||||
)
|
||||
src_i16 = llvm.bitcast(Int16.mlir_type, src_vec2, loc=loc, ip=ip)
|
||||
src_i32_pad16b = llvm.zext(Int32.mlir_type, src_i16, loc=loc, ip=ip)
|
||||
rst0 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32_pad16b,
|
||||
mask2[0],
|
||||
zero,
|
||||
],
|
||||
"dp4a.s32.s32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
rst1 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32_pad16b,
|
||||
mask2[1],
|
||||
zero,
|
||||
],
|
||||
"dp4a.s32.s32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
res0 = llvm.inline_asm(
|
||||
Float32.mlir_type,
|
||||
[
|
||||
rst0,
|
||||
],
|
||||
"cvt.rn.f32.s32 $0, $1;",
|
||||
"=f,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
res1 = llvm.inline_asm(
|
||||
Float32.mlir_type,
|
||||
[
|
||||
rst1,
|
||||
],
|
||||
"cvt.rn.f32.s32 $0, $1;",
|
||||
"=f,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
vec_f32x2_type = ir.VectorType.get([2], Float32.mlir_type, loc=loc)
|
||||
vec_f32x2 = vector.from_elements(vec_f32x2_type, [res0, res1], loc=loc, ip=ip)
|
||||
return vec_f32x2
|
||||
|
||||
|
||||
# Permute bytes from register pair.
|
||||
@dsl_user_op
|
||||
def prmt(src, src_reg_shifted, prmt_indices, *, loc=None, ip=None):
|
||||
return llvm.inline_asm(
|
||||
T.i32(),
|
||||
[
|
||||
Int32(src).ir_value(loc=loc, ip=ip),
|
||||
Int32(src_reg_shifted).ir_value(loc=loc, ip=ip),
|
||||
Int32(prmt_indices).ir_value(loc=loc, ip=ip),
|
||||
],
|
||||
"prmt.b32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
|
||||
|
||||
# Convert 1 int4 value to 1 bfloat16 value
|
||||
@dsl_user_op
|
||||
def cvt_i4_bf16(src_i4, *, loc=None, ip=None):
|
||||
# i4 -> i32 -> f32 -> bf
|
||||
src_i32 = llvm.zext(Int32.mlir_type, src_i4, loc=loc, ip=ip)
|
||||
src_f32 = llvm.sitofp(Float32.mlir_type, src_i32, loc=loc, ip=ip)
|
||||
bf16_val = cvt_f32_bf16(src_f32, loc=loc, ip=ip)
|
||||
return bf16_val
|
||||
|
||||
|
||||
# Convert multiple int4 values to bfloat16 values.
|
||||
# The number of elements to be converted must be be even as specified by num_elts.
|
||||
# Int4 values are packed into int32 values with upper bits filled with 0 if there are less than 4 int4 values.
|
||||
# Results bfloat16 values are also packed into int32 values.
|
||||
@dsl_user_op
|
||||
def cvt_i4_to_bf16_impl(src_i32, num_elts, *, loc=None, ip=None):
|
||||
c4 = arith.constant(Int32.mlir_type, 4, loc=loc, ip=ip)
|
||||
src_shr4 = llvm.lshr(src_i32, c4, loc=loc, ip=ip)
|
||||
xor_mask0 = arith.constant(Int32.mlir_type, 0x08080808, loc=loc, ip=ip)
|
||||
and_mask = arith.constant(Int32.mlir_type, 0x0F0F0F0F, loc=loc, ip=ip)
|
||||
imm_lut = arith.constant(Int32.mlir_type, 0x0000006A, loc=loc, ip=ip)
|
||||
src_i32 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32,
|
||||
and_mask,
|
||||
xor_mask0,
|
||||
imm_lut,
|
||||
],
|
||||
"lop3.b32 $0, $1, $2, $3, $4;",
|
||||
"=r,r,n,n,n",
|
||||
)
|
||||
xor_mask1 = arith.constant(Int32.mlir_type, 0x88080808, loc=loc, ip=ip)
|
||||
src_shr4 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_shr4,
|
||||
and_mask,
|
||||
xor_mask1,
|
||||
imm_lut,
|
||||
],
|
||||
"lop3.b32 $0, $1, $2, $3, $4;",
|
||||
"=r,r,n,n,n",
|
||||
)
|
||||
prmt_indices = [
|
||||
arith.constant(Int32.mlir_type, imme, loc=loc, ip=ip)
|
||||
for imme in [
|
||||
0x0000F4F0,
|
||||
0x0000F5F1,
|
||||
0x0000F6F2,
|
||||
0x0000F7F3,
|
||||
]
|
||||
]
|
||||
num_i32_elts = num_elts // 2
|
||||
rsts = []
|
||||
for i in range(num_i32_elts):
|
||||
rst = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
src_i32,
|
||||
src_shr4,
|
||||
prmt_indices[i],
|
||||
],
|
||||
"prmt.b32 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
)
|
||||
rsts.append(rst)
|
||||
mask_clear_top_bit = arith.constant(Int32.mlir_type, 0xFF7FFFFF, loc=loc, ip=ip)
|
||||
rsts[-1] = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
rsts[-1],
|
||||
mask_clear_top_bit,
|
||||
],
|
||||
"and.b32 $0, $1, $2;",
|
||||
"=r,r,r",
|
||||
)
|
||||
mul = arith.constant(Int32.mlir_type, 0x83808380, loc=loc, ip=ip)
|
||||
bias = arith.constant(Int32.mlir_type, 0xC308C308, loc=loc, ip=ip)
|
||||
for i in range(num_i32_elts):
|
||||
rsts[i] = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[
|
||||
rsts[i],
|
||||
mul,
|
||||
bias,
|
||||
],
|
||||
"fma.rn.bf16x2 $0, $1, $2, $3;",
|
||||
"=r,r,r,r",
|
||||
)
|
||||
# pack rsts into a vector
|
||||
vec_type = ir.VectorType.get([num_i32_elts], Int32.mlir_type, loc=loc)
|
||||
vec_rsts = vector.from_elements(vec_type, rsts, loc=loc, ip=ip)
|
||||
return vec_rsts
|
||||
|
||||
|
||||
# Convert 2 int4 values to 2 bfloat16 values
|
||||
@dsl_user_op
|
||||
def cvt_i4x2_to_bf16x2(src_vec2, *, loc=None, ip=None):
|
||||
# pack 2 int4 into 1 int32 value and fill upper bits with 0
|
||||
src_i8 = llvm.bitcast(Int8.mlir_type, src_vec2, loc=loc, ip=ip)
|
||||
src_i32 = llvm.zext(Int32.mlir_type, src_i8, loc=loc, ip=ip)
|
||||
rst_i32 = cvt_i4_to_bf16_impl(src_i32, 2, loc=loc, ip=ip)
|
||||
vec_bf16x2_type = ir.VectorType.get([2], BFloat16.mlir_type, loc=loc)
|
||||
vec_bf16x2 = llvm.bitcast(vec_bf16x2_type, rst_i32, loc=loc, ip=ip)
|
||||
return vec_bf16x2
|
||||
|
||||
|
||||
# Convert 4 int4 values to 4 bfloat16 values
|
||||
@dsl_user_op
|
||||
def cvt_i4x4_to_bf16x4(src_vec4, *, loc=None, ip=None):
|
||||
# pack 4 int4 into 1 int32 value and fill upper bits with 0
|
||||
src_i16 = llvm.bitcast(Int16.mlir_type, src_vec4, loc=loc, ip=ip)
|
||||
src_i32 = llvm.zext(Int32.mlir_type, src_i16, loc=loc, ip=ip)
|
||||
rst_i32 = cvt_i4_to_bf16_impl(src_i32, 4, loc=loc, ip=ip)
|
||||
vec_bf16x4_type = ir.VectorType.get([4], BFloat16.mlir_type, loc=loc)
|
||||
vec_bf16x4 = llvm.bitcast(vec_bf16x4_type, rst_i32, loc=loc, ip=ip)
|
||||
return vec_bf16x4
|
||||
|
||||
|
||||
# Convert 8 int4 values to 8 bfloat16 values
|
||||
@dsl_user_op
|
||||
def cvt_i4x8_to_bf16x8(src_vec8, *, loc=None, ip=None):
|
||||
# pack 8 int4 into 1 int32 value and fill upper bits with 0
|
||||
src_i32 = llvm.bitcast(Int32.mlir_type, src_vec8, loc=loc, ip=ip)
|
||||
rst_i32 = cvt_i4_to_bf16_impl(src_i32, 8, loc=loc, ip=ip)
|
||||
vec_bf16x8_type = ir.VectorType.get([8], BFloat16.mlir_type, loc=loc)
|
||||
vec_bf16x8 = llvm.bitcast(vec_bf16x8_type, rst_i32, loc=loc, ip=ip)
|
||||
return vec_bf16x8
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def log2_of_pow2_int(a: Int32, *, loc=None, ip=None) -> Int32:
|
||||
tmp = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[a.ir_value(loc=loc, ip=ip)],
|
||||
"brev.b32 $0, $1;",
|
||||
"=r,r",
|
||||
has_side_effects=False,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
return Int32(
|
||||
llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[tmp],
|
||||
"bfind.shiftamt.u32 $0, $1;",
|
||||
"=r,r",
|
||||
has_side_effects=False,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
@deprecated(
|
||||
"cute.arch.exp is deprecated, use cute.math.exp with `fastmath=True` instead"
|
||||
@@ -679,3 +1270,96 @@ def exp_packed_f32x2(
|
||||
LOG2_E = Float32(1.4426950408889634)
|
||||
b = mul_packed_f32x2(a, (LOG2_E, LOG2_E), loc=loc, ip=ip)
|
||||
return exp2(b[0], loc=loc, ip=ip), exp2(b[1], loc=loc, ip=ip)
|
||||
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cvt_f4e2m1_f16(src, *, loc=None, ip=None):
|
||||
# 0 padding for upper 4 bits
|
||||
zero = arith.constant(src.type, 0, loc=loc, ip=ip)
|
||||
vec2 = vector.from_elements(
|
||||
ir.VectorType.get([2], src.type, loc=loc), [src, zero], loc=loc, ip=ip
|
||||
)
|
||||
rst_vec2 = cvt_f4e2m1x2_to_f16x2(vec2, loc=loc, ip=ip)
|
||||
# only the 1st element is valid
|
||||
rst = vector.extract(
|
||||
rst_vec2, dynamic_position=[], static_position=[0], loc=loc, ip=ip
|
||||
)
|
||||
return rst
|
||||
|
||||
|
||||
# Convert 2 float4e2m1 values to 2 float16 values
|
||||
@dsl_user_op
|
||||
def cvt_f4e2m1x2_to_f16x2(src_vec2, *, loc=None, ip=None):
|
||||
# pack 2 float4e2m1 into 1 int8 value and fill upper bits with 0
|
||||
src_i8 = llvm.bitcast(Int8.mlir_type, src_vec2, loc=loc, ip=ip)
|
||||
src_i16 = llvm.zext(Int16.mlir_type, src_i8, loc=loc, ip=ip)
|
||||
rst_i32 = llvm.inline_asm(
|
||||
Int32.mlir_type,
|
||||
[src_i16],
|
||||
"""{\n\t
|
||||
.reg .b8 b;\n\t
|
||||
mov.b16 {b,_}, $1;\n\t
|
||||
cvt.rn.f16x2.e2m1x2 $0, b;\n\t
|
||||
}""",
|
||||
"=r,h",
|
||||
)
|
||||
vec_f16x2_type = ir.VectorType.get([2], Float16.mlir_type, loc=loc)
|
||||
vec_f16x2 = llvm.bitcast(vec_f16x2_type, rst_i32, loc=loc, ip=ip)
|
||||
return vec_f16x2
|
||||
|
||||
|
||||
# Convert 4 float4e2m1 values to 4 float16 values
|
||||
@dsl_user_op
|
||||
def cvt_f4e2m1x4_to_f16x4(src_vec4, *, loc=None, ip=None):
|
||||
# pack 4 float4e2m1 into 1 int16 value
|
||||
src_i16 = llvm.bitcast(Int16.mlir_type, src_vec4, loc=loc, ip=ip)
|
||||
rst_i32x2 = llvm.inline_asm(
|
||||
llvm.StructType.get_literal([T.i32(), T.i32()]),
|
||||
[src_i16],
|
||||
"""{\n\t
|
||||
.reg .b8 b0, b1;\n\t
|
||||
mov.b16 {b0, b1}, $2;\n\t
|
||||
cvt.rn.f16x2.e2m1x2 $0, b0;\n\t
|
||||
cvt.rn.f16x2.e2m1x2 $1, b1;\n\t
|
||||
}""",
|
||||
"=r,=r,h",
|
||||
)
|
||||
res0 = llvm.extractvalue(T.i32(), rst_i32x2, [0])
|
||||
res1 = llvm.extractvalue(T.i32(), rst_i32x2, [1])
|
||||
vec_f32x2_type = ir.VectorType.get([2], Int32.mlir_type, loc=loc)
|
||||
vec_f32x2 = vector.from_elements(vec_f32x2_type, [res0, res1], loc=loc, ip=ip)
|
||||
vec_f16x4_type = ir.VectorType.get([4], Float16.mlir_type, loc=loc)
|
||||
vec_f16x4 = llvm.bitcast(vec_f16x4_type, vec_f32x2, loc=loc, ip=ip)
|
||||
return vec_f16x4
|
||||
|
||||
|
||||
# Convert 8 float4e2m1 values to 8 float16 values
|
||||
@dsl_user_op
|
||||
def cvt_f4e2m1x8_to_f16x8(src_vec8, *, loc=None, ip=None):
|
||||
# pack 8 float4e2m1 into 1 int32 value and fill upper bits with 0
|
||||
src_i32 = llvm.bitcast(Int32.mlir_type, src_vec8, loc=loc, ip=ip)
|
||||
rst_i32x4 = llvm.inline_asm(
|
||||
llvm.StructType.get_literal([T.i32(), T.i32(), T.i32(), T.i32()]),
|
||||
[src_i32],
|
||||
"""{\n\t
|
||||
.reg .b8 b0, b1, b2, b3;\n\t
|
||||
mov.b32 {b0, b1, b2, b3}, $4;\n\t
|
||||
cvt.rn.f16x2.e2m1x2 $0, b0;\n\t
|
||||
cvt.rn.f16x2.e2m1x2 $1, b1;\n\t
|
||||
cvt.rn.f16x2.e2m1x2 $2, b2;\n\t
|
||||
cvt.rn.f16x2.e2m1x2 $3, b3;\n\t
|
||||
}""",
|
||||
"=r,=r,=r,=r,r",
|
||||
)
|
||||
res0 = llvm.extractvalue(T.i32(), rst_i32x4, [0])
|
||||
res1 = llvm.extractvalue(T.i32(), rst_i32x4, [1])
|
||||
res2 = llvm.extractvalue(T.i32(), rst_i32x4, [2])
|
||||
res3 = llvm.extractvalue(T.i32(), rst_i32x4, [3])
|
||||
vec_f32x4_type = ir.VectorType.get([4], Int32.mlir_type, loc=loc)
|
||||
vec_f32x4 = vector.from_elements(
|
||||
vec_f32x4_type, [res0, res1, res2, res3], loc=loc, ip=ip
|
||||
)
|
||||
vec_f16x8_type = ir.VectorType.get([8], Float16.mlir_type, loc=loc)
|
||||
vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_f32x4, loc=loc, ip=ip)
|
||||
return vec_f16x8
|
||||
|
||||
1132
python/CuTeDSL/cutlass/cute/atom.py
Normal file
1132
python/CuTeDSL/cutlass/cute/atom.py
Normal file
File diff suppressed because it is too large
Load Diff
File diff suppressed because it is too large
Load Diff
@@ -9,8 +9,11 @@
|
||||
# and related documentation outside the scope permitted by the EULA
|
||||
# is strictly prohibited.
|
||||
|
||||
from .core import TensorSSA
|
||||
from typing import Callable, Union
|
||||
|
||||
from .typing import Numeric
|
||||
from .tensor import TensorSSA
|
||||
|
||||
from cutlass._mlir.dialects import math, arith
|
||||
|
||||
from typing import Callable, Union
|
||||
@@ -62,7 +65,7 @@ def acos(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = acos(y) # Compute arc cosine
|
||||
"""
|
||||
@@ -85,7 +88,7 @@ def asin(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = asin(y) # Compute arc sine
|
||||
"""
|
||||
@@ -108,11 +111,10 @@ def atan(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = atan(y) # Compute arc tangent
|
||||
"""
|
||||
raise NotImplementedError("atan is not implemented")
|
||||
return _math_op(math.atan, fastmath, a)
|
||||
|
||||
|
||||
@@ -137,8 +139,8 @@ def atan2(
|
||||
|
||||
.. code-block::
|
||||
|
||||
y = cute.make_fragment(ptr1, layout).load() # y coordinates
|
||||
x = cute.make_fragment(ptr2, layout).load() # x coordinates
|
||||
y = cute.make_rmem_tensor(ptr1, layout).load() # y coordinates
|
||||
x = cute.make_rmem_tensor(ptr2, layout).load() # x coordinates
|
||||
theta = atan2(y, x) # Compute angles
|
||||
"""
|
||||
return _math_op(math.atan2, fastmath, a, b)
|
||||
@@ -160,7 +162,7 @@ def cos(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = cos(y) # Compute cosine
|
||||
"""
|
||||
@@ -186,7 +188,7 @@ def erf(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = erf(y) # Compute error function
|
||||
"""
|
||||
@@ -209,7 +211,7 @@ def exp(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = exp(y) # Compute exponential
|
||||
"""
|
||||
@@ -232,7 +234,7 @@ def exp2(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = exp2(y) # Compute 2^x
|
||||
"""
|
||||
@@ -255,7 +257,7 @@ def log(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = log(y) # Compute natural logarithm
|
||||
"""
|
||||
@@ -278,7 +280,7 @@ def log2(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = log2(y) # Compute log base 2
|
||||
"""
|
||||
@@ -301,7 +303,7 @@ def log10(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = log10(y) # Compute log base 10
|
||||
"""
|
||||
@@ -326,7 +328,7 @@ def rsqrt(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = rsqrt(y) # Compute 1/√x
|
||||
"""
|
||||
@@ -349,7 +351,7 @@ def sin(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = sin(y) # Compute sine
|
||||
"""
|
||||
@@ -372,7 +374,7 @@ def sqrt(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = sqrt(y) # Compute square root
|
||||
"""
|
||||
@@ -395,7 +397,7 @@ def tan(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = tan(y) # Compute tangent
|
||||
"""
|
||||
@@ -418,7 +420,7 @@ def tanh(
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
x = cute.make_rmem_tensor(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = tanh(y) # Compute hyperbolic tangent
|
||||
"""
|
||||
|
||||
@@ -18,8 +18,9 @@ import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from .. import core
|
||||
from .. import atom
|
||||
from ..typing import Float16, Float32, Float64, Numeric
|
||||
from cutlass import cute
|
||||
|
||||
|
||||
class OpError(DSLBaseError):
|
||||
@@ -28,7 +29,7 @@ class OpError(DSLBaseError):
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, op: core.Op, message: str, suggestion: Optional[str] = None
|
||||
self, op: atom.Op, message: str, suggestion: Optional[str] = None
|
||||
) -> None:
|
||||
if suggestion is None:
|
||||
# Default suggestion
|
||||
@@ -48,7 +49,7 @@ class OpError(DSLBaseError):
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaUniversalOp(core.MmaOp):
|
||||
class MmaUniversalOp(atom.MmaOp):
|
||||
"""
|
||||
The universal MMA Operation.
|
||||
|
||||
@@ -65,7 +66,7 @@ class MmaUniversalOp(core.MmaOp):
|
||||
if self.abacc_dtype not in [Float16, Float32, Float64]:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the 'abacc_dtype' Op parameter to be one of Float16, Float32, or Float64",
|
||||
"expects the 'abacc_dtype' Op parameter to be one of Float16, Float32, or Float64",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
@@ -75,14 +76,14 @@ class MmaUniversalOp(core.MmaOp):
|
||||
)
|
||||
|
||||
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaUniversalTrait":
|
||||
shape_mnk_attr = ir.Attribute.parse(f'#cute.shape<"(1,1,1)">')
|
||||
shape_mnk_attr = ir.Attribute.parse('#cute.shape<"(1,1,1)">')
|
||||
atom_ty = _cute_nvgpu_ir.UniversalFmaAtomType.get(
|
||||
shape_mnk_attr,
|
||||
self.abacc_dtype.mlir_type,
|
||||
self.abacc_dtype.mlir_type,
|
||||
self.abacc_dtype.mlir_type,
|
||||
)
|
||||
return MmaUniversalTrait(_cute_ir.atom(atom_ty, loc=loc, ip=ip))
|
||||
return MmaUniversalTrait(cute.make_atom(atom_ty, loc=loc, ip=ip))
|
||||
|
||||
def _verify_fragment_A(self, input, *, loc=None, ip=None):
|
||||
pass
|
||||
@@ -90,7 +91,8 @@ class MmaUniversalOp(core.MmaOp):
|
||||
def _verify_fragment_B(self, input, *, loc=None, ip=None):
|
||||
pass
|
||||
|
||||
class MmaUniversalTrait(core.Trait):
|
||||
|
||||
class MmaUniversalTrait(atom.Trait):
|
||||
pass
|
||||
|
||||
|
||||
@@ -137,8 +139,9 @@ class MemoryScope(enum.Enum):
|
||||
def _to_ir(self) -> _cute_ir.MemScopeKind:
|
||||
return self.value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyUniversalOp(core.CopyOp):
|
||||
class CopyUniversalOp(atom.CopyOp):
|
||||
"""
|
||||
The universal Copy Operation.
|
||||
|
||||
@@ -182,8 +185,8 @@ class CopyUniversalOp(core.CopyOp):
|
||||
memory_order._to_ir(),
|
||||
memory_scope._to_ir(),
|
||||
)
|
||||
return CopyUniversalTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return CopyUniversalTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class CopyUniversalTrait(core.Trait):
|
||||
class CopyUniversalTrait(atom.Trait):
|
||||
pass
|
||||
|
||||
@@ -36,4 +36,5 @@ __all__ = [
|
||||
"fence_tma_desc_acquire",
|
||||
"cp_fence_tma_desc_release",
|
||||
"fence_tma_desc_release",
|
||||
"group_bulk_copy_modes",
|
||||
]
|
||||
|
||||
@@ -13,13 +13,15 @@ import enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Type
|
||||
|
||||
from cutlass.cutlass_dsl import CuTeDSL, t
|
||||
from cutlass import cute
|
||||
from cutlass.base_dsl.arch import Arch
|
||||
from cutlass.cutlass_dsl import CuTeDSL
|
||||
|
||||
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, ReductionOp
|
||||
from ...atom import CopyOp, Trait
|
||||
from ...tensor import ReductionOp
|
||||
from ...typing import Int16, Pointer, Integer, Numeric
|
||||
from ..common import OpError
|
||||
from ..tcgen05.mma import CtaGroup
|
||||
@@ -73,19 +75,13 @@ class CopyG2SOp(CopyOp):
|
||||
|
||||
def _make_trait(
|
||||
self,
|
||||
copy_internal_type: Type[t.Numeric],
|
||||
copy_internal_type: Type[Numeric],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
**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 "
|
||||
@@ -100,7 +96,7 @@ class CopyG2SOp(CopyOp):
|
||||
ty = _cute_nvgpu_ir.CopyAtomSIMTAsyncCopyType.get(
|
||||
copy_internal_type.mlir_type, self.cache_mode._to_ir(), num_bits_per_copy
|
||||
)
|
||||
return CopyG2STrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return CopyG2STrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class CopyG2STrait(Trait):
|
||||
@@ -114,8 +110,18 @@ class CopyG2STrait(Trait):
|
||||
####################################################################################################
|
||||
|
||||
TMA_MBAR_PTR_FIELD_NAME = "tma_bar"
|
||||
TMA_MASK_FIELD_NAME = "mcast_mask"
|
||||
TMA_MCAST_MASK_FIELD_NAME = "mcast_mask"
|
||||
TMA_DESC_PTR_FIELD_NAME = "tma_descriptor_ptr"
|
||||
TMA_BYTE_MASK_FIELD_NAME = "byte_mask"
|
||||
|
||||
|
||||
class TmaCopyOp(CopyOp):
|
||||
"""
|
||||
Base class for all TMA copy operations.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
#
|
||||
# TMA GMEM -> SMEM copies
|
||||
@@ -123,7 +129,7 @@ TMA_DESC_PTR_FIELD_NAME = "tma_descriptor_ptr"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkTensorTileG2SOp(CopyOp):
|
||||
class CopyBulkTensorTileG2SOp(TmaCopyOp):
|
||||
"""
|
||||
Bulk tensor asynchrnous GMEM to SMEM Copy Operation using the TMA unit.
|
||||
|
||||
@@ -133,27 +139,20 @@ class CopyBulkTensorTileG2SOp(CopyOp):
|
||||
|
||||
cta_group: CtaGroup = CtaGroup.ONE
|
||||
|
||||
admissible_archs = [
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.cta_group, CtaGroup):
|
||||
raise OpError(
|
||||
self, "expects the 'cta_group' parameter to be a CtaGroup instance"
|
||||
)
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch >= Arch.sm_90:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
if (self.cta_group == CtaGroup.TWO) and arch[:5] == "sm_90":
|
||||
if (self.cta_group == CtaGroup.TWO) and arch.major == Arch.sm_90.major:
|
||||
raise OpError(
|
||||
self,
|
||||
f"CTA group of 2 is tcgen05-specific and is not and is not compatible with {arch}",
|
||||
@@ -163,7 +162,7 @@ class CopyBulkTensorTileG2SOp(CopyOp):
|
||||
def __str__(self) -> str:
|
||||
res = "cp.async GMEM -> SMEM bulk tensor copy Operation"
|
||||
if self.cta_group == CtaGroup.TWO:
|
||||
res += f"\n CTA group = 2"
|
||||
res += "\n CTA group = 2"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
@@ -225,7 +224,7 @@ class CopyBulkTensorTileG2SNonExecTrait(Trait):
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkTensorTileG2SMulticastOp(CopyOp):
|
||||
class CopyBulkTensorTileG2SMulticastOp(TmaCopyOp):
|
||||
"""
|
||||
Bulk tensor asynchrnous multicast GMEM to SMEM Copy Operation using the TMA unit.
|
||||
|
||||
@@ -235,27 +234,20 @@ class CopyBulkTensorTileG2SMulticastOp(CopyOp):
|
||||
|
||||
cta_group: CtaGroup = CtaGroup.ONE
|
||||
|
||||
admissible_archs = [
|
||||
"sm_90",
|
||||
"sm_90a",
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
]
|
||||
|
||||
def __post_init__(self):
|
||||
if not isinstance(self.cta_group, CtaGroup):
|
||||
raise OpError(
|
||||
self, "expects the 'cta_group' parameter to be a CtaGroup instance"
|
||||
)
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch >= Arch.sm_90:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
if (self.cta_group == CtaGroup.TWO) and arch[:5] == "sm_90":
|
||||
if (self.cta_group == CtaGroup.TWO) and arch.major == Arch.sm_90.major:
|
||||
raise OpError(
|
||||
self,
|
||||
f"CTA group of 2 is tcgen05-specific and is not and is not compatible with {arch}",
|
||||
@@ -265,7 +257,7 @@ class CopyBulkTensorTileG2SMulticastOp(CopyOp):
|
||||
def __str__(self) -> str:
|
||||
res = "cp.async GMEM -> SMEM bulk tensor multicast copy Operation"
|
||||
if self.cta_group == CtaGroup.TWO:
|
||||
res += f"\n CTA group = 2"
|
||||
res += "\n CTA group = 2"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
@@ -309,12 +301,7 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
|
||||
"expects a multicast mask to be provided via the mcast_mask kw argument"
|
||||
)
|
||||
exec_value = _cute_nvgpu_ir.atom_make_exec_tma(self.value, loc=loc, ip=ip)
|
||||
attr_str = f"#cute_nvgpu.atom_copy_field_tmaload<tma_bar>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
exec_value = _cute_nvgpu_ir.atom_set_value(
|
||||
exec_value, attr, tma_bar_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
attr_str = f"#cute_nvgpu.atom_copy_field_tmaload<mcast_mask>"
|
||||
attr_str = "#cute_nvgpu.atom_copy_field_tmaload<mcast_mask>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
exec_value = _cute_nvgpu_ir.atom_set_value(
|
||||
exec_value, attr, Int16(mcast_mask).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
@@ -325,6 +312,13 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
|
||||
exec_value = _cute_nvgpu_ir.atom_set_value(
|
||||
exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
# Set the tma_bar_ptr at last to ensure that the atom creation and setting
|
||||
# operations above can be moved outside the loop
|
||||
attr_str = "#cute_nvgpu.atom_copy_field_tmaload<tma_bar>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
exec_value = _cute_nvgpu_ir.atom_set_value(
|
||||
exec_value, attr, tma_bar_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
return exec_value
|
||||
|
||||
|
||||
@@ -334,7 +328,7 @@ class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkTensorTileS2GOp(CopyOp):
|
||||
class CopyBulkTensorTileS2GOp(TmaCopyOp):
|
||||
"""
|
||||
Bulk tensor asynchronous SMEM to GMEM Copy Operation using the TMA unit.
|
||||
|
||||
@@ -342,20 +336,13 @@ class CopyBulkTensorTileS2GOp(CopyOp):
|
||||
This Operation uses TMA in the ``.tile`` mode.
|
||||
"""
|
||||
|
||||
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:
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch >= Arch.sm_90:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
@@ -386,8 +373,9 @@ class CopyBulkTensorTileS2GTrait(Trait):
|
||||
)
|
||||
return exec_value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyReduceBulkTensorTileS2GOp(CopyOp):
|
||||
class CopyReduceBulkTensorTileS2GOp(TmaCopyOp):
|
||||
"""
|
||||
Bulk tensor asynchronous SMEM to GMEM Reduction Operation using the TMA unit.
|
||||
|
||||
@@ -397,20 +385,13 @@ class CopyReduceBulkTensorTileS2GOp(CopyOp):
|
||||
|
||||
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:
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch >= Arch.sm_90:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
@@ -461,11 +442,200 @@ class CopyReduceBulkTensorTileS2GTrait(Trait):
|
||||
)
|
||||
return exec_value
|
||||
|
||||
__all__ = [
|
||||
"LoadCacheMode",
|
||||
"CopyG2SOp",
|
||||
"CopyBulkTensorTileG2SOp",
|
||||
"CopyBulkTensorTileG2SMulticastOp",
|
||||
"CopyBulkTensorTileS2GOp",
|
||||
"CopyReduceBulkTensorTileS2GOp",
|
||||
]
|
||||
|
||||
#
|
||||
# Bulk GMEM -> SMEM copies
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkG2SOp(CopyOp):
|
||||
"""
|
||||
Bulk copy asynchrnous GMEM to SMEM Copy Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk>`__.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Arch verification
|
||||
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch >= Arch.sm_90:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
res = "cp.async GMEM -> SMEM bulk copy Operation"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "CopyBulkG2STrait":
|
||||
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
|
||||
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
|
||||
raise ValueError(
|
||||
"expects a 'num_bits_per_copy' kw argument of type int that is positive "
|
||||
f"when creating a copy Atom for {self.__class__.__name__}"
|
||||
)
|
||||
ty = _cute_nvgpu_ir.CopyAtomBulkCopyG2SType.get(
|
||||
copy_internal_type.mlir_type, num_bits_per_copy, False
|
||||
)
|
||||
return CopyBulkG2STrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class CopyBulkG2STrait(Trait):
|
||||
# We allow kw args to be dropped so that the user can write common code for non-multicast
|
||||
# and multicast loads.
|
||||
def unpack(
|
||||
self,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
mbar_ptr: Optional[Pointer] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Custom implementation of unpack for bulk copy load.
|
||||
|
||||
The non-multicast bulk load requires a `mbar_ptr` keyword argument to be provided when
|
||||
using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error.
|
||||
"""
|
||||
if not isinstance(mbar_ptr, Pointer):
|
||||
raise ValueError(
|
||||
"expects a pointer to an mbarrier to be provided via the mbar_ptr kw argument"
|
||||
)
|
||||
attr_str = f"#cute_nvgpu.atom_copy_field_bulkg2s<{TMA_MBAR_PTR_FIELD_NAME}>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
val = _cute_nvgpu_ir.atom_set_value(
|
||||
self.value, attr, mbar_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
return val
|
||||
|
||||
|
||||
#
|
||||
# Bulk GMEM -> SMEM Multicast copies
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkG2SMulticastOp(CopyOp):
|
||||
"""
|
||||
Bulk multicast copy asynchrnous GMEM to SMEM Copy Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk>`__.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Arch verification
|
||||
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch >= Arch.sm_90:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
res = "cp.async GMEM -> SMEM multicast bulk copy Operation"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "CopyBulkG2SMulticastTrait":
|
||||
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
|
||||
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
|
||||
raise ValueError(
|
||||
"expects a 'num_bits_per_copy' kw argument of type int that is positive "
|
||||
f"when creating a copy Atom for {self.__class__.__name__}"
|
||||
)
|
||||
ty = _cute_nvgpu_ir.CopyAtomBulkCopyG2SType.get(
|
||||
copy_internal_type.mlir_type, num_bits_per_copy, True
|
||||
)
|
||||
return CopyBulkG2SMulticastTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class CopyBulkG2SMulticastTrait(Trait):
|
||||
# We allow kw args to be dropped so that the user can write common code for non-multicast
|
||||
# and multicast loads.
|
||||
def unpack(
|
||||
self,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
mbar_ptr: Optional[Pointer] = None,
|
||||
mcast_mask: Optional[Integer] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Custom implementation of unpack for bulk copy load.
|
||||
|
||||
The non-multicast bulk load requires a `mbar_ptr` keyword argument to be provided when
|
||||
using `cute.copy`. Any other kw arguments will be ignored instead of triggering an error.
|
||||
"""
|
||||
if not isinstance(mbar_ptr, Pointer):
|
||||
raise ValueError(
|
||||
"expects a pointer to an mbarrier to be provided via the mbar_ptr kw argument"
|
||||
)
|
||||
if not isinstance(mcast_mask, Integer):
|
||||
raise ValueError(
|
||||
"expects a multicast mask to be provided via the mcast_mask kw argument"
|
||||
)
|
||||
attr_str = f"#cute_nvgpu.atom_copy_field_bulkg2s<{TMA_MBAR_PTR_FIELD_NAME}>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
val = _cute_nvgpu_ir.atom_set_value(
|
||||
self.value, attr, mbar_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
attr_str = f"#cute_nvgpu.atom_copy_field_bulkg2s<{TMA_MCAST_MASK_FIELD_NAME}>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
val = _cute_nvgpu_ir.atom_set_value(
|
||||
val, attr, Int16(mcast_mask).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
return val
|
||||
|
||||
|
||||
#
|
||||
# Bulk SMEM -> GMEM copies
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkS2GOp(CopyOp):
|
||||
"""
|
||||
Bulk copy asynchrnous SMEM to GMEM Copy Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk>`__.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Arch verification
|
||||
arch: Arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch >= Arch.sm_90:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be at least {Arch.sm_90.name}, but got {arch.name}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
res = "cp.async SMEM -> GMEM bulk copy Operation"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "CopyBulkS2GTrait":
|
||||
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
|
||||
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
|
||||
raise ValueError(
|
||||
"expects a 'num_bits_per_copy' kw argument of type int that is positive "
|
||||
f"when creating a copy Atom for {self.__class__.__name__}"
|
||||
)
|
||||
ty = _cute_nvgpu_ir.CopyAtomBulkCopyS2GType.get(
|
||||
copy_internal_type.mlir_type, num_bits_per_copy, False
|
||||
)
|
||||
return CopyBulkS2GTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class CopyBulkS2GTrait(Trait):
|
||||
pass
|
||||
|
||||
@@ -16,8 +16,18 @@ from cutlass.cutlass_dsl import dsl_user_op
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir.dialects import llvm
|
||||
|
||||
from ...typing import Coord, Layout, Tensor, Tiler, Pointer, Int16, Numeric, NumericMeta
|
||||
from ... import core
|
||||
from ...typing import (
|
||||
Coord,
|
||||
Layout,
|
||||
ComposedLayout,
|
||||
Tensor,
|
||||
Tiler,
|
||||
Pointer,
|
||||
Int16,
|
||||
Numeric,
|
||||
NumericMeta,
|
||||
)
|
||||
from ... import core, atom
|
||||
from .copy import (
|
||||
CopyBulkTensorTileG2SOp,
|
||||
CopyBulkTensorTileG2SMulticastOp,
|
||||
@@ -39,14 +49,14 @@ def make_tiled_tma_atom(
|
||||
CopyReduceBulkTensorTileS2GOp,
|
||||
],
|
||||
gmem_tensor: Tensor,
|
||||
smem_layout: Union[Layout, core.ComposedLayout],
|
||||
smem_layout: Union[Layout, ComposedLayout],
|
||||
cta_tiler: Tiler,
|
||||
num_multicast: int = 1,
|
||||
*,
|
||||
internal_type: Optional[Type[Numeric]] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[core.CopyAtom, Tensor]:
|
||||
) -> Tuple[atom.CopyAtom, Tensor]:
|
||||
"""
|
||||
Makes a TMA Copy Atom in the ``.tile`` mode to copy tiles of a GMEM tensor to/from SMEM
|
||||
buffer with the given Layout.
|
||||
@@ -74,7 +84,7 @@ def make_tiled_tma_atom(
|
||||
:param gmem_tensor: The GMEM tensor involved in the Copy
|
||||
:type gmem_tensor: Tensor
|
||||
:param smem_layout: The SMEM layout to construct the Copy Atom for
|
||||
:type smem_layout: Union[Layout, core.ComposedLayout]
|
||||
:type smem_layout: Union[Layout, ComposedLayout]
|
||||
:param cta_tiler: The CTA Tiler to use
|
||||
:type cta_tiler: Tiler
|
||||
:param num_multicast: The multicast factor
|
||||
@@ -82,7 +92,7 @@ def make_tiled_tma_atom(
|
||||
:param internal_type: An optional parameter for the internal data type to use when the actual data type is not supported by the TMA unit
|
||||
:type internal_type: Type[Numeric]
|
||||
:return: A Copy Atom for this Operation and the associated TMA tensor
|
||||
:rtype: Tuple[core.CopyAtom, Tensor]
|
||||
:rtype: Tuple[atom.CopyAtom, Tensor]
|
||||
"""
|
||||
|
||||
if internal_type is not None:
|
||||
@@ -97,6 +107,9 @@ def make_tiled_tma_atom(
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
if isinstance(smem_layout, core._ComposedLayout):
|
||||
smem_layout = smem_layout.value
|
||||
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
if num_multicast != 1:
|
||||
raise ValueError(
|
||||
@@ -113,7 +126,7 @@ def make_tiled_tma_atom(
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return core.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
return atom.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
elif isinstance(op, CopyBulkTensorTileG2SMulticastOp):
|
||||
if num_multicast < 1:
|
||||
raise ValueError(
|
||||
@@ -131,7 +144,7 @@ def make_tiled_tma_atom(
|
||||
ip=ip,
|
||||
)
|
||||
return (
|
||||
core.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
atom.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
res[1],
|
||||
)
|
||||
elif isinstance(op, CopyBulkTensorTileS2GOp):
|
||||
@@ -143,7 +156,7 @@ def make_tiled_tma_atom(
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return core.CopyAtom(op, CopyBulkTensorTileS2GTrait(res[0])), res[1]
|
||||
return atom.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,
|
||||
@@ -154,14 +167,14 @@ def make_tiled_tma_atom(
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return core.CopyAtom(op, CopyReduceBulkTensorTileS2GTrait(res[0])), res[1]
|
||||
return atom.CopyAtom(op, CopyReduceBulkTensorTileS2GTrait(res[0])), res[1]
|
||||
else:
|
||||
raise ValueError(f"expects a bulk tensor (TMA) Copy Op, but got {op}")
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def tma_partition(
|
||||
atom: core.CopyAtom,
|
||||
atom: atom.CopyAtom,
|
||||
cta_coord: Coord,
|
||||
cta_layout: Layout,
|
||||
smem_tensor: Tensor,
|
||||
@@ -221,7 +234,7 @@ def create_tma_multicast_mask(
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def prefetch_descriptor(tma_atom: core.CopyAtom, *, loc=None, ip=None) -> None:
|
||||
def prefetch_descriptor(tma_atom: atom.CopyAtom, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Prefetches the TMA descriptor associated with the TMA Atom.
|
||||
"""
|
||||
@@ -230,7 +243,7 @@ def prefetch_descriptor(tma_atom: core.CopyAtom, *, loc=None, ip=None) -> None:
|
||||
|
||||
@dsl_user_op
|
||||
def copy_tensormap(
|
||||
tma_atom: core.CopyAtom, tensormap_ptr: Pointer, *, loc=None, ip=None
|
||||
tma_atom: atom.CopyAtom, tensormap_ptr: Pointer, *, loc=None, ip=None
|
||||
) -> None:
|
||||
"""
|
||||
Copies the tensormap held by a TMA Copy Atom to the memory location pointed to by the provided
|
||||
@@ -248,7 +261,7 @@ def copy_tensormap(
|
||||
|
||||
@dsl_user_op
|
||||
def update_tma_descriptor(
|
||||
tma_atom: core.CopyAtom,
|
||||
tma_atom: atom.CopyAtom,
|
||||
gmem_tensor: Tensor,
|
||||
tma_desc_ptr: Pointer,
|
||||
*,
|
||||
@@ -289,7 +302,7 @@ def fence_tma_desc_acquire(
|
||||
"""
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar>`__.
|
||||
"""
|
||||
tma_desc_ptr_i64 = tma_desc_ptr.toint(loc=loc, ip=ip).ir_value()
|
||||
tma_desc_ptr_i64 = tma_desc_ptr.toint(loc=loc, ip=ip).ir_value(loc=loc, ip=ip)
|
||||
llvm.inline_asm(
|
||||
None,
|
||||
[tma_desc_ptr_i64],
|
||||
@@ -312,8 +325,12 @@ def cp_fence_tma_desc_release(
|
||||
"""
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-tensormap-cp-fenceproxy>`__.
|
||||
"""
|
||||
tma_desc_global_ptr_i64 = tma_desc_global_ptr.toint(loc=loc, ip=ip).ir_value()
|
||||
tma_desc_shared_ptr_i32 = tma_desc_shared_ptr.toint(loc=loc, ip=ip).ir_value()
|
||||
tma_desc_global_ptr_i64 = tma_desc_global_ptr.toint(loc=loc, ip=ip).ir_value(
|
||||
loc=loc, ip=ip
|
||||
)
|
||||
tma_desc_shared_ptr_i32 = tma_desc_shared_ptr.toint(loc=loc, ip=ip).ir_value(
|
||||
loc=loc, ip=ip
|
||||
)
|
||||
llvm.inline_asm(
|
||||
None,
|
||||
[tma_desc_global_ptr_i64, tma_desc_shared_ptr_i32],
|
||||
@@ -339,3 +356,13 @@ def fence_tma_desc_release(*, loc=None, ip=None) -> None:
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def group_bulk_copy_modes(src: Tensor, dst: Tensor, loc=None, ip=None) -> Tuple:
|
||||
"""
|
||||
Copy async bulk need group mode 0, acquiring whole tensor for bulk copy
|
||||
"""
|
||||
mSrc = core.group_modes(src, 0, core.rank(src))
|
||||
mDst = core.group_modes(dst, 0, core.rank(dst))
|
||||
return (mSrc, mDst)
|
||||
|
||||
@@ -15,8 +15,8 @@ from cutlass.cutlass_dsl import dsl_user_op
|
||||
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
|
||||
from .. import core
|
||||
from ..typing import Shape, Layout, Tensor, Numeric, NumericMeta
|
||||
from .. import core, atom
|
||||
from ..typing import Shape, Layout, ComposedLayout, Tensor, Numeric, NumericMeta
|
||||
from ...impl_utils import check_type_in
|
||||
from .cpasync.copy import (
|
||||
CopyBulkTensorTileG2SOp,
|
||||
@@ -37,15 +37,15 @@ from .cpasync.copy import (
|
||||
def make_tiled_tma_atom_A(
|
||||
op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
|
||||
gmem_tensor: Tensor,
|
||||
smem_layout: Union[Layout, core.ComposedLayout],
|
||||
smem_layout: Union[Layout, ComposedLayout],
|
||||
mma_tiler_mnk: Shape,
|
||||
tiled_mma: core.TiledMma,
|
||||
cluster_shape_vmnk: Shape,
|
||||
tiled_mma: atom.TiledMma,
|
||||
cluster_shape_vmnk: Union[Shape, None] = None,
|
||||
*,
|
||||
internal_type: Optional[Type[Numeric]] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[core.CopyAtom, Tensor]:
|
||||
) -> Tuple[atom.CopyAtom, Tensor]:
|
||||
"""
|
||||
Makes a TMA Copy atom mapping to ``.tile`` mode for ``cp.async.bulk.tensor`` PTX operation
|
||||
accounting for the MK projections of the TiledMMA for A tensor loads.
|
||||
@@ -76,18 +76,18 @@ def make_tiled_tma_atom_A(
|
||||
:param gmem_tensor: The GMEM tensor to be loaded by this copy atom
|
||||
:type gmem_tensor: Tensor
|
||||
:param smem_layout: Shared memory layout to load the tensor into (PDSL)
|
||||
:type smem_layout: Union[Layout, core.ComposedLayout]
|
||||
:type smem_layout: Union[Layout, ComposedLayout]
|
||||
:param mma_tiler_mnk: The MMA Tiler shape (TILE_M, TILE_N, TILE_K) in MNK dimensions
|
||||
:type mma_tiler_mnk: Shape
|
||||
:param tiled_mma: The TiledMMA that will consume the load as operands
|
||||
:type tiled_mma: core.TiledMma
|
||||
:type tiled_mma: atom.TiledMma
|
||||
:param cluster_shape_vmnk: The Cluster-level shape in VMNK dimensions
|
||||
:type cluster_shape_vmnk: Shape
|
||||
:param internal_type: An optional parameter for the internal data type to when element
|
||||
type does not match the copy type
|
||||
:type internal_type: Type[Numeric]
|
||||
:return: A copy atom for this operation and the associated TMA coord tensor
|
||||
:rtype: Tuple[core.CopyAtom, Tensor]
|
||||
:rtype: Tuple[atom.CopyAtom, Tensor]
|
||||
|
||||
"""
|
||||
|
||||
@@ -114,8 +114,15 @@ def make_tiled_tma_atom_A(
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
# multicast across the N-mode since those would share the same tile of A
|
||||
if cluster_shape_vmnk is None:
|
||||
raise ValueError(
|
||||
"cluster_shape_vmnk must be provided for multicast A tensor loads"
|
||||
)
|
||||
num_multicast = core.size(cluster_shape_vmnk, mode=[2])
|
||||
|
||||
if isinstance(smem_layout, core._ComposedLayout):
|
||||
smem_layout = smem_layout.value
|
||||
|
||||
# res[0] = the IR Value for the non-executable atom instance
|
||||
# res[1] = the IR Value for the associated TMA tensor
|
||||
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load(
|
||||
@@ -129,11 +136,11 @@ def make_tiled_tma_atom_A(
|
||||
ip=ip,
|
||||
)
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
return core.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
return atom.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
return (
|
||||
core.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
atom.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
res[1],
|
||||
)
|
||||
|
||||
@@ -142,15 +149,15 @@ def make_tiled_tma_atom_A(
|
||||
def make_tiled_tma_atom_B(
|
||||
op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
|
||||
gmem_tensor: Tensor,
|
||||
smem_layout: Union[Layout, core.ComposedLayout],
|
||||
smem_layout: Union[Layout, ComposedLayout],
|
||||
mma_tiler_mnk: Shape,
|
||||
tiled_mma: core.TiledMma,
|
||||
cluster_shape_vmnk: Shape,
|
||||
tiled_mma: atom.TiledMma,
|
||||
cluster_shape_vmnk: Union[Shape, None] = None,
|
||||
*,
|
||||
internal_type: Optional[Type[Numeric]] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[core.CopyAtom, Tensor]:
|
||||
) -> Tuple[atom.CopyAtom, Tensor]:
|
||||
"""
|
||||
Makes a TMA Copy atom mapping to ``.tile`` mode for ``cp.async.bulk.tensor`` PTX operation
|
||||
accounting for the NK projections of the TiledMMA for B tensor loads.
|
||||
@@ -181,7 +188,7 @@ def make_tiled_tma_atom_B(
|
||||
:param gmem_tensor: The GMEM tensor to be loaded by this copy atom
|
||||
:type gmem_tensor: Tensor
|
||||
:param smem_layout: Shared memory layout to load the tensor into (PDSL)
|
||||
:type smem_layout: Union[Layout, core.ComposedLayout]
|
||||
:type smem_layout: Union[Layout, ComposedLayout]
|
||||
:param mma_tiler_mnk: The MMA Tiler shape (TILE_M, TILE_N, TILE_K) in MNK dimensions
|
||||
:type mma_tiler_mnk: Shape
|
||||
:param tiled_mma: The TiledMMA that will consume the load as operands
|
||||
@@ -192,7 +199,7 @@ def make_tiled_tma_atom_B(
|
||||
type does not match the copy type
|
||||
:type internal_type: Type[Numeric]
|
||||
:return: A Copy Atom for this Operation and the associated TMA tensor
|
||||
:rtype: Tuple[core.CopyAtom, Tensor]
|
||||
:rtype: Tuple[atom.CopyAtom, Tensor]
|
||||
|
||||
"""
|
||||
|
||||
@@ -219,8 +226,15 @@ def make_tiled_tma_atom_B(
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
# multicast across the M-mode since those would share the same tile of B
|
||||
if cluster_shape_vmnk is None:
|
||||
raise ValueError(
|
||||
"cluster_shape_vmnk must be provided for multicast B tensor loads"
|
||||
)
|
||||
num_multicast = core.size(cluster_shape_vmnk, mode=[1])
|
||||
|
||||
if isinstance(smem_layout, core._ComposedLayout):
|
||||
smem_layout = smem_layout.value
|
||||
|
||||
# res[0] = the IR Value for the non-executable atom instance
|
||||
# res[1] = the IR Value for the associated TMA tensor
|
||||
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load(
|
||||
@@ -234,11 +248,11 @@ def make_tiled_tma_atom_B(
|
||||
ip=ip,
|
||||
)
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
return core.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
return atom.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
return (
|
||||
core.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
atom.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
res[1],
|
||||
)
|
||||
|
||||
|
||||
@@ -40,6 +40,7 @@ __all__ = [
|
||||
"Field",
|
||||
"MmaTF32Op",
|
||||
"MmaF16BF16Op",
|
||||
"MmaF16BF16SparseOp",
|
||||
"MmaI8Op",
|
||||
"MmaFP8Op",
|
||||
"MmaMXF8Op",
|
||||
|
||||
@@ -13,14 +13,15 @@ import enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
|
||||
from cutlass import cute
|
||||
from cutlass.base_dsl.arch import Arch
|
||||
from cutlass.cutlass_dsl import CuTeDSL
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..common import OpError
|
||||
from ...core import CopyOp, Trait
|
||||
from ...atom import CopyOp, Trait
|
||||
from ...typing import Numeric
|
||||
|
||||
from .mma import CtaGroup
|
||||
@@ -46,24 +47,6 @@ class Repetition(enum.Enum):
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
@classmethod
|
||||
def _missing_(cls, value):
|
||||
if isinstance(value, int):
|
||||
if value == 1:
|
||||
return Repetition.x1
|
||||
elif value == 2:
|
||||
return Repetition.x2
|
||||
elif value == 8:
|
||||
return Repetition.x8
|
||||
elif value == 16:
|
||||
return Repetition.x16
|
||||
elif value == 32:
|
||||
return Repetition.x32
|
||||
elif value == 64:
|
||||
return Repetition.x64
|
||||
elif value == 128:
|
||||
return Repetition.x128
|
||||
|
||||
|
||||
class Pack(enum.Enum):
|
||||
"""
|
||||
@@ -97,17 +80,40 @@ class Unpack(enum.Enum):
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _LdBase(CopyOp):
|
||||
"""
|
||||
Base class for TMEM load operations in the tcgen05 instruction set.
|
||||
|
||||
This abstract base class provides common functionality and validation for tensor memory (TMEM)
|
||||
load operations. It defines the fundamental parameters and architecture constraints that apply
|
||||
to all load operation variants.
|
||||
|
||||
:param repeat: Number of repetitions for the load operation, defaults to Repetition.x1
|
||||
:type repeat: Repetition, optional
|
||||
:param pack: Packing pattern for TMEM to RMEM copies, defaults to Pack.NONE
|
||||
:type pack: Pack, optional
|
||||
:raises OpError: If the current architecture is not supported or if invalid parameters are provided
|
||||
"""
|
||||
|
||||
repeat: Repetition = Repetition.x1
|
||||
pack: Pack = Pack.NONE
|
||||
|
||||
admissible_archs = [
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
]
|
||||
admissible_archs = Arch.filter(
|
||||
lambda arch: arch.is_family_of(Arch.sm_100f) or arch.is_family_of(Arch.sm_110f)
|
||||
)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""
|
||||
Post-initialization validation for TMEM load operations.
|
||||
|
||||
Performs comprehensive validation of operation parameters and architecture compatibility.
|
||||
This method is automatically called after object creation to ensure all constraints are met.
|
||||
|
||||
:raises OpError: If architecture is not supported
|
||||
:raises OpError: If repeat parameter is not a Repetition instance
|
||||
:raises OpError: If pack parameter is not a Pack instance
|
||||
"""
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
@@ -127,12 +133,21 @@ class _LdBase(CopyOp):
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
"""
|
||||
Generate a human-readable string representation of the load operation.
|
||||
|
||||
Creates a formatted description showing the operation type, repetition count,
|
||||
and any special packing configuration.
|
||||
|
||||
:return: Multi-line string describing the operation configuration
|
||||
:rtype: str
|
||||
"""
|
||||
res = (
|
||||
f"tcgen05 {self.__class__.__name__[:-2]} Copy Operation"
|
||||
+ f"\n number of repetitions = {self.repeat.value}"
|
||||
)
|
||||
if self.pack == Pack.PACK_16b_IN_32b:
|
||||
res += f"\n with 2x 16-bit to 32b packing"
|
||||
res += "\n with 2x 16-bit to 32b packing"
|
||||
return res
|
||||
|
||||
|
||||
@@ -148,6 +163,24 @@ class Ld16x64bOp(_LdBase):
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x64bTrait":
|
||||
"""
|
||||
Create a trait object for the 16x64b TMEM load operation.
|
||||
|
||||
Constructs an MLIR-based trait that encapsulates the specific parameters and
|
||||
characteristics of this load operation. The trait is used by the compiler
|
||||
infrastructure to generate the appropriate low-level code.
|
||||
|
||||
:param copy_internal_type: The data type for the copy operation
|
||||
:type copy_internal_type: Type[Numeric]
|
||||
:param loc: MLIR location information for debugging, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: MLIR insertion point for code generation, defaults to None
|
||||
:type ip: optional
|
||||
:param kwargs: Additional keyword arguments passed to the trait constructor
|
||||
:type kwargs: dict
|
||||
:return: A trait object that represents this specific load operation
|
||||
:rtype: Ld16x64bTrait
|
||||
"""
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
@@ -155,7 +188,7 @@ class Ld16x64bOp(_LdBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
|
||||
)
|
||||
return Ld16x64bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Ld16x64bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Ld16x64bTrait(Trait):
|
||||
@@ -172,6 +205,15 @@ class Ld16x128bOp(_LdBase):
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""
|
||||
Additional validation specific to 16x128b load operations.
|
||||
|
||||
Extends the base class validation with operation-specific constraints.
|
||||
The 16x128b operation has limitations on the maximum repetition count due to
|
||||
hardware register and bandwidth constraints.
|
||||
|
||||
:raises OpError: If x128 repetition is specified
|
||||
"""
|
||||
super().__post_init__()
|
||||
if self.repeat == Repetition.x128:
|
||||
raise OpError(
|
||||
@@ -183,6 +225,20 @@ class Ld16x128bOp(_LdBase):
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x128bTrait":
|
||||
"""
|
||||
Create a trait object for the 16x128b TMEM load operation.
|
||||
|
||||
:param copy_internal_type: The data type for the copy operation
|
||||
:type copy_internal_type: Type[Numeric]
|
||||
:param loc: MLIR location information for debugging, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: MLIR insertion point for code generation, defaults to None
|
||||
:type ip: optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
:return: A trait object for this load operation
|
||||
:rtype: Ld16x128bTrait
|
||||
"""
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
@@ -190,7 +246,7 @@ class Ld16x128bOp(_LdBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
|
||||
)
|
||||
return Ld16x128bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Ld16x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Ld16x128bTrait(Trait):
|
||||
@@ -207,6 +263,15 @@ class Ld16x256bOp(_LdBase):
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
"""
|
||||
Additional validation specific to 16x256b load operations.
|
||||
|
||||
Extends the base class validation with operation-specific constraints.
|
||||
The 16x256b operation has more restrictive limitations on repetition count due to
|
||||
the larger data size per operation requiring more hardware resources.
|
||||
|
||||
:raises OpError: If x64 or x128 repetition is specified
|
||||
"""
|
||||
super().__post_init__()
|
||||
if self.repeat in (Repetition.x128, Repetition.x64):
|
||||
raise OpError(
|
||||
@@ -218,6 +283,20 @@ class Ld16x256bOp(_LdBase):
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x256bTrait":
|
||||
"""
|
||||
Create a trait object for the 16x256b TMEM load operation.
|
||||
|
||||
:param copy_internal_type: The data type for the copy operation
|
||||
:type copy_internal_type: Type[Numeric]
|
||||
:param loc: MLIR location information for debugging, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: MLIR insertion point for code generation, defaults to None
|
||||
:type ip: optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
:return: A trait object for this load operation
|
||||
:rtype: Ld16x256bTrait
|
||||
"""
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
@@ -225,7 +304,7 @@ class Ld16x256bOp(_LdBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
|
||||
)
|
||||
return Ld16x256bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Ld16x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Ld16x256bTrait(Trait):
|
||||
@@ -244,6 +323,20 @@ class Ld16x32bx2Op(_LdBase):
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x32bx2Trait":
|
||||
"""
|
||||
Create a trait object for the 16x32bx2 TMEM load operation.
|
||||
|
||||
:param copy_internal_type: The data type for the copy operation
|
||||
:type copy_internal_type: Type[Numeric]
|
||||
:param loc: MLIR location information for debugging, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: MLIR insertion point for code generation, defaults to None
|
||||
:type ip: optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
:return: A trait object for this load operation
|
||||
:rtype: Ld16x32bx2Trait
|
||||
"""
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
@@ -251,7 +344,7 @@ class Ld16x32bx2Op(_LdBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
|
||||
)
|
||||
return Ld16x32bx2Trait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Ld16x32bx2Trait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Ld16x32bx2Trait(Trait):
|
||||
@@ -270,6 +363,20 @@ class Ld32x32bOp(_LdBase):
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld32x32bTrait":
|
||||
"""
|
||||
Create a trait object for the 32x32b TMEM load operation.
|
||||
|
||||
:param copy_internal_type: The data type for the copy operation
|
||||
:type copy_internal_type: Type[Numeric]
|
||||
:param loc: MLIR location information for debugging, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: MLIR insertion point for code generation, defaults to None
|
||||
:type ip: optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
:return: A trait object for this load operation
|
||||
:rtype: Ld32x32bTrait
|
||||
"""
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
32,
|
||||
@@ -277,7 +384,7 @@ class Ld32x32bOp(_LdBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.pack == Pack.PACK_16b_IN_32b else None,
|
||||
)
|
||||
return Ld32x32bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Ld32x32bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Ld32x32bTrait(Trait):
|
||||
@@ -286,17 +393,30 @@ class Ld32x32bTrait(Trait):
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _StBase(CopyOp):
|
||||
"""
|
||||
Base class for TMEM store operations in the tcgen05 instruction set.
|
||||
|
||||
This abstract base class provides common functionality and validation for tensor memory (TMEM)
|
||||
store operations. It defines the fundamental parameters and architecture constraints that apply
|
||||
to all store operation variants.
|
||||
|
||||
:param repeat: Number of repetitions for the store operation (required parameter)
|
||||
:type repeat: Repetition
|
||||
:param unpack: Unpacking pattern for RMEM to TMEM copies, defaults to Unpack.NONE
|
||||
:type unpack: Unpack, optional
|
||||
:raises OpError: If the current architecture is not supported or if invalid parameters are provided
|
||||
"""
|
||||
|
||||
repeat: Repetition
|
||||
unpack: Unpack = Unpack.NONE
|
||||
|
||||
admissible_archs = [
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
]
|
||||
admissible_archs = Arch.filter(
|
||||
lambda arch: arch.is_family_of(Arch.sm_100f) or arch.is_family_of(Arch.sm_110f)
|
||||
)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
@@ -312,7 +432,7 @@ class _StBase(CopyOp):
|
||||
if not isinstance(self.unpack, Unpack):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'pack' Op parameter to be a tcgen05.Unpack instance",
|
||||
"expects the 'unpack' Op parameter to be a tcgen05.Unpack instance",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
@@ -321,7 +441,7 @@ class _StBase(CopyOp):
|
||||
+ f"\n number of repetitions = {self.repeat.value}"
|
||||
)
|
||||
if self.unpack == Unpack.UNPACK_32b_IN_16b:
|
||||
res += f"\n with 32-bit to 2x 16b unpacking"
|
||||
res += "\n with 32-bit to 2x 16b unpacking"
|
||||
return res
|
||||
|
||||
|
||||
@@ -337,6 +457,20 @@ class St16x64bOp(_StBase):
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "St16x64bTrait":
|
||||
"""
|
||||
Create a trait object for the 16x64b TMEM store operation.
|
||||
|
||||
:param copy_internal_type: The data type for the copy operation
|
||||
:type copy_internal_type: Type[Numeric]
|
||||
:param loc: MLIR location information for debugging, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: MLIR insertion point for code generation, defaults to None
|
||||
:type ip: optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
:return: A trait object for this store operation
|
||||
:rtype: St16x64bTrait
|
||||
"""
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemStoreType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
@@ -344,7 +478,7 @@ class St16x64bOp(_StBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
|
||||
)
|
||||
return St16x64bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return St16x64bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class St16x64bTrait(Trait):
|
||||
@@ -379,7 +513,7 @@ class St16x128bOp(_StBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
|
||||
)
|
||||
return St16x128bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return St16x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class St16x128bTrait(Trait):
|
||||
@@ -414,7 +548,7 @@ class St16x256bOp(_StBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
|
||||
)
|
||||
return St16x256bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return St16x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class St16x256bTrait(Trait):
|
||||
@@ -440,7 +574,7 @@ class St16x32bx2Op(_StBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
|
||||
)
|
||||
return St16x32bx2Trait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return St16x32bx2Trait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class St16x32bx2Trait(Trait):
|
||||
@@ -466,7 +600,7 @@ class St32x32bOp(_StBase):
|
||||
self.repeat.value,
|
||||
ir.UnitAttr.get() if self.unpack == Unpack.UNPACK_32b_IN_16b else None,
|
||||
)
|
||||
return St32x32bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return St32x32bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class St32x32bTrait(Trait):
|
||||
@@ -475,20 +609,27 @@ class St32x32bTrait(Trait):
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _S2TCopyBase(CopyOp):
|
||||
cta_group: CtaGroup
|
||||
"""
|
||||
Base class for SMEM to TMEM copy operations in the tcgen05 instruction set.
|
||||
|
||||
admissible_archs = [
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
]
|
||||
This abstract base class provides common functionality and validation for shared memory (SMEM)
|
||||
to tensor memory (TMEM) copy operations. These operations are used for high-throughput data
|
||||
movement between different memory hierarchies in modern GPU architectures.
|
||||
|
||||
:param cta_group: Cooperative Thread Array (CTA) group configuration
|
||||
:type cta_group: CtaGroup
|
||||
:raises OpError: If the current architecture is not SM100f family or if invalid parameters are provided
|
||||
"""
|
||||
|
||||
cta_group: CtaGroup
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch.is_family_of(Arch.sm_100f):
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
f"expects arch to be one of {Arch.filter(lambda arch: arch.is_family_of(Arch.sm_100f))}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
# Verify that the user provided enum values
|
||||
@@ -519,6 +660,20 @@ class Cp128x256bOp(_S2TCopyBase):
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Cp128x256bTrait":
|
||||
"""
|
||||
Create a trait object for the 128x256b SMEM to TMEM copy operation.
|
||||
|
||||
:param copy_internal_type: The data type for the copy operation
|
||||
:type copy_internal_type: Type[Numeric]
|
||||
:param loc: MLIR location information for debugging, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: MLIR insertion point for code generation, defaults to None
|
||||
:type ip: optional
|
||||
:param kwargs: Additional keyword arguments
|
||||
:type kwargs: dict
|
||||
:return: A trait object for this S2T copy operation
|
||||
:rtype: Cp128x256bTrait
|
||||
"""
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100CopyS2TType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
128,
|
||||
@@ -526,7 +681,7 @@ class Cp128x256bOp(_S2TCopyBase):
|
||||
self.cta_group.value,
|
||||
_cute_nvgpu_ir.CopyS2TBroadcast.none,
|
||||
)
|
||||
return Cp128x256bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Cp128x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Cp128x256bTrait(Trait):
|
||||
@@ -552,7 +707,7 @@ class Cp128x128bOp(_S2TCopyBase):
|
||||
self.cta_group.value,
|
||||
_cute_nvgpu_ir.CopyS2TBroadcast.none,
|
||||
)
|
||||
return Cp128x128bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Cp128x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Cp128x128bTrait(Trait):
|
||||
@@ -578,7 +733,7 @@ class Cp4x256bOp(_S2TCopyBase):
|
||||
self.cta_group.value,
|
||||
_cute_nvgpu_ir.CopyS2TBroadcast.none,
|
||||
)
|
||||
return Cp4x256bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Cp4x256bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Cp4x256bTrait(Trait):
|
||||
@@ -604,7 +759,7 @@ class Cp4x32x128bOp(_S2TCopyBase):
|
||||
self.cta_group.value,
|
||||
_cute_nvgpu_ir.CopyS2TBroadcast.x4,
|
||||
)
|
||||
return Cp4x32x128bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Cp4x32x128bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Cp4x32x128bTrait(Trait):
|
||||
@@ -630,7 +785,7 @@ class Cp2x64x128b0213Op(_S2TCopyBase):
|
||||
self.cta_group.value,
|
||||
_cute_nvgpu_ir.CopyS2TBroadcast.lw_0213,
|
||||
)
|
||||
return Cp2x64x128b0213Trait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Cp2x64x128b0213Trait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Cp2x64x128b0213Trait(Trait):
|
||||
@@ -656,7 +811,7 @@ class Cp2x64x128b0123Op(_S2TCopyBase):
|
||||
self.cta_group.value,
|
||||
_cute_nvgpu_ir.CopyS2TBroadcast.lw_0123,
|
||||
)
|
||||
return Cp2x64x128b0123Trait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return Cp2x64x128b0123Trait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class Cp2x64x128b0123Trait(Trait):
|
||||
|
||||
@@ -13,7 +13,6 @@ from typing import overload, Type, Tuple, Union
|
||||
|
||||
from cutlass.cutlass_dsl import dsl_user_op
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir.dialects import nvvm
|
||||
|
||||
@@ -21,6 +20,7 @@ from ...typing import (
|
||||
Shape,
|
||||
IntTuple,
|
||||
Layout,
|
||||
ComposedLayout,
|
||||
Tensor,
|
||||
Int,
|
||||
Numeric,
|
||||
@@ -29,6 +29,8 @@ from ...typing import (
|
||||
Int32,
|
||||
)
|
||||
from ... import core
|
||||
from ...tensor import recast_tensor
|
||||
from ...atom import CopyAtom, TiledCopy
|
||||
from .mma import SmemLayoutAtomKind, CtaGroup
|
||||
from .copy import (
|
||||
Pack,
|
||||
@@ -56,7 +58,7 @@ from .copy import (
|
||||
@dsl_user_op
|
||||
def make_smem_layout_atom(
|
||||
kind: SmemLayoutAtomKind, element_type: Type[Numeric], *, loc=None, ip=None
|
||||
) -> core.ComposedLayout:
|
||||
) -> ComposedLayout:
|
||||
"""
|
||||
Makes a SMEM layout Atom.
|
||||
|
||||
@@ -68,7 +70,7 @@ def make_smem_layout_atom(
|
||||
:param element_type: The element data type to construct the layout for
|
||||
:type element_type: Type[Numeric]
|
||||
:return: The SMEM layout atom
|
||||
:rtype: core.ComposedLayout
|
||||
:rtype: ComposedLayout
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(f"element_type must be a Numeric, but got {element_type}")
|
||||
@@ -130,13 +132,13 @@ def tile_to_mma_shape(
|
||||
|
||||
@overload
|
||||
def tile_to_mma_shape(
|
||||
atom: core.ComposedLayout,
|
||||
atom: ComposedLayout,
|
||||
mma_tile_shape: Shape,
|
||||
order: IntTuple = None,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> core.ComposedLayout: ...
|
||||
) -> ComposedLayout: ...
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
@@ -152,7 +154,7 @@ def tile_to_mma_shape(
|
||||
if core.rank(order) != core.rank(mma_tile_shape) - 1:
|
||||
raise ValueError(
|
||||
f"rank(order)={core.rank(order)} must be equal to "
|
||||
f"rank(mma_tile_shape)-1={core.rank(mma_tile_shape)-1}"
|
||||
f"rank(mma_tile_shape)-1={core.rank(mma_tile_shape) - 1}"
|
||||
)
|
||||
order_val = core._pack_int_tuple(order, loc=loc, ip=ip)
|
||||
mma_tile_shape_val = core._pack_shape(mma_tile_shape, loc=loc, ip=ip)
|
||||
@@ -164,8 +166,12 @@ def tile_to_mma_shape(
|
||||
):
|
||||
raise ValueError("tile_to_mma_shape only supports static inputs")
|
||||
|
||||
if isinstance(atom, core._ComposedLayout):
|
||||
atom = atom.value
|
||||
|
||||
res_ty = _cute_nvgpu_ir.tile_to_mma_shape(atom, mma_tile_shape_val, order_val)
|
||||
return _cute_ir.static(res_ty, loc=loc, ip=ip)
|
||||
res_val = core.static(res_ty, loc=loc, ip=ip)
|
||||
return core.coalesce(res_val, target_profile=mma_tile_shape, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
@@ -209,7 +215,7 @@ def commit(
|
||||
####################################################################################################
|
||||
|
||||
|
||||
def is_tmem_load(atom: core.CopyAtom) -> bool:
|
||||
def is_tmem_load(atom: CopyAtom) -> bool:
|
||||
"""
|
||||
Returns whether a CopyAtom instance is a TMEM load.
|
||||
"""
|
||||
@@ -225,7 +231,7 @@ def is_tmem_load(atom: core.CopyAtom) -> bool:
|
||||
)
|
||||
|
||||
|
||||
def is_tmem_store(atom: core.CopyAtom) -> bool:
|
||||
def is_tmem_store(atom: CopyAtom) -> bool:
|
||||
"""
|
||||
Returns whether a CopyAtom instance is a TMEM store.
|
||||
"""
|
||||
@@ -242,7 +248,7 @@ def is_tmem_store(atom: core.CopyAtom) -> bool:
|
||||
|
||||
|
||||
def get_tmem_copy_properties(
|
||||
atom: core.CopyAtom,
|
||||
atom: CopyAtom,
|
||||
) -> Tuple[int, int, int, Union[Pack, Unpack]]:
|
||||
"""
|
||||
Returns the properties of a TMEM copy atom (number of data paths, bits, repetitions,
|
||||
@@ -279,7 +285,7 @@ def find_tmem_tensor_col_offset(tmem_tensor: Tensor, *, loc=None, ip=None) -> In
|
||||
"""
|
||||
tmem_col_mask = 0x0000FFFF
|
||||
offset = (
|
||||
core.cosize(core.recast_tensor(tmem_tensor, Int32).layout, loc=loc, ip=ip)
|
||||
core.cosize(recast_tensor(tmem_tensor, Int32).layout, loc=loc, ip=ip)
|
||||
& tmem_col_mask
|
||||
)
|
||||
if isinstance(offset, int):
|
||||
@@ -289,8 +295,8 @@ def find_tmem_tensor_col_offset(tmem_tensor: Tensor, *, loc=None, ip=None) -> In
|
||||
|
||||
@dsl_user_op
|
||||
def make_tmem_copy(
|
||||
atom: core.CopyAtom, tmem_tensor: Tensor, *, loc=None, ip=None
|
||||
) -> core.TiledCopy:
|
||||
atom: CopyAtom, tmem_tensor: Tensor, *, loc=None, ip=None
|
||||
) -> TiledCopy:
|
||||
"""
|
||||
Makes a Tiled Copy instance from a TMEM Copy Atom and a TMEM tensor.
|
||||
"""
|
||||
@@ -298,13 +304,13 @@ def make_tmem_copy(
|
||||
atom._trait.value, tmem_tensor.value, loc=loc, ip=ip
|
||||
)
|
||||
new_trait = type(atom._trait)(tiled_copy_val)
|
||||
return core.TiledCopy(atom.op, new_trait)
|
||||
return TiledCopy(atom.op, new_trait)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_s2t_copy(
|
||||
atom: core.CopyAtom, tmem_tensor: Tensor, *, loc=None, ip=None
|
||||
) -> core.TiledCopy:
|
||||
atom: CopyAtom, tmem_tensor: Tensor, *, loc=None, ip=None
|
||||
) -> TiledCopy:
|
||||
"""
|
||||
Makes a Tiled Copy instance from a TMEM Copy Atom and a TMEM tensor.
|
||||
"""
|
||||
@@ -312,12 +318,12 @@ def make_s2t_copy(
|
||||
atom._trait.value, tmem_tensor.value, loc=loc, ip=ip
|
||||
)
|
||||
new_trait = type(atom._trait)(tiled_copy_val)
|
||||
return core.TiledCopy(atom.op, new_trait)
|
||||
return TiledCopy(atom.op, new_trait)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def get_s2t_smem_desc_tensor(
|
||||
atom: core.CopyAtom, smem_tensor: Tensor, *, loc=None, ip=None
|
||||
atom: CopyAtom, smem_tensor: Tensor, *, loc=None, ip=None
|
||||
) -> Tensor:
|
||||
"""
|
||||
Returns the SMEM descriptor tensor from a S2T copy atom and a SMEM tensor.
|
||||
|
||||
@@ -11,8 +11,10 @@
|
||||
|
||||
import enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
from typing import Type, Any
|
||||
|
||||
from cutlass import cute
|
||||
from cutlass.base_dsl.arch import Arch
|
||||
from cutlass.cutlass_dsl import CuTeDSL, T
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
@@ -20,8 +22,9 @@ import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..common import OpError
|
||||
from ... import core
|
||||
from ...core import Trait, _pack_shape, rank, depth, _Tensor
|
||||
from ... import core, atom
|
||||
from ...core import _pack_shape, rank, depth
|
||||
from ...tensor import _Tensor
|
||||
from ...typing import (
|
||||
Shape,
|
||||
Float4E2M1FN,
|
||||
@@ -40,6 +43,9 @@ from ...typing import (
|
||||
AddressSpace,
|
||||
Pointer,
|
||||
)
|
||||
from ...atom import Trait
|
||||
|
||||
from ..warp.mma import SparseMetadataFormat
|
||||
|
||||
|
||||
####################################################################################################
|
||||
@@ -49,6 +55,14 @@ from ...typing import (
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class Tcgen05MmaOp(atom.MmaOp):
|
||||
"""
|
||||
Base class for all tcgen05 MMA operations.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class OperandMajorMode(enum.Enum):
|
||||
"""
|
||||
An enumeration for the majorness of the input operands of the MMA.
|
||||
@@ -108,6 +122,7 @@ class CtaGroup(enum.Enum):
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
|
||||
class Field(enum.Enum):
|
||||
"""
|
||||
An enumeration for the fields of the MMA Atom that can be modified at runtime.
|
||||
@@ -131,7 +146,7 @@ class Field(enum.Enum):
|
||||
|
||||
# Base class for all tcgen05 MMA Ops with syntax `tcgen05.mma.cta_group.kind` used to factor out some internal code
|
||||
@dataclass(frozen=True)
|
||||
class MmaOp(core.MmaOp):
|
||||
class MmaOp(Tcgen05MmaOp):
|
||||
a_dtype: Type[Numeric]
|
||||
b_dtype: Type[Numeric]
|
||||
acc_dtype: Type[Numeric]
|
||||
@@ -141,14 +156,13 @@ class MmaOp(core.MmaOp):
|
||||
a_major_mode: OperandMajorMode
|
||||
b_major_mode: OperandMajorMode
|
||||
|
||||
admissible_archs = [
|
||||
"sm_100a",
|
||||
"sm_100f",
|
||||
]
|
||||
admissible_archs = Arch.filter(
|
||||
lambda arch: arch.is_family_of(Arch.sm_100f) or arch.is_family_of(Arch.sm_110f)
|
||||
)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Verify arch
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
@@ -194,18 +208,18 @@ class MmaOp(core.MmaOp):
|
||||
f"expects the N-mode to satisfy 8 <= N <= 256 and N % 8 == 0, but got {n}",
|
||||
)
|
||||
elif m == 128:
|
||||
if (n < 16) or (n > 256) or (n % 16 != 0):
|
||||
if (n < 8) or (n > 256) or (n % 8 != 0):
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the N-mode to satisfy 8 <= N <= 256 and N % 16 == 0, but got {n}",
|
||||
f"expects the N-mode to satisfy 8 <= N <= 256 and N % 8 == 0, but got {n}",
|
||||
)
|
||||
else:
|
||||
if m not in [128, 256]:
|
||||
raise OpError(self, f"expects the M-mode to be 128 or 256, but got {m}")
|
||||
if (n < 32) or (n > 256) or (n % 32 != 0):
|
||||
if (n < 16) or (n > 256) or (n % 16 != 0):
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the N-mode to satisfy 32 <= N <= 256 and N % 32 == 0, but got {n}",
|
||||
f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
@@ -246,7 +260,7 @@ class MmaOp(core.MmaOp):
|
||||
return True
|
||||
|
||||
|
||||
class MmaTrait(Trait):
|
||||
class MmaTraits(Trait):
|
||||
admissible_fields = [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]
|
||||
|
||||
def set(self, field, value, *, loc=None, ip=None) -> None:
|
||||
@@ -260,10 +274,21 @@ class MmaTrait(Trait):
|
||||
self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
|
||||
def get(self, field, *, loc=None, ip=None) -> Any:
|
||||
if field not in self.admissible_fields:
|
||||
raise ValueError(
|
||||
f"expects field to be one of {self.admissible_fields}, but got {field}"
|
||||
)
|
||||
field_name = f"#cute_nvgpu.atom_mma_field_sm100<{field._to_ir_field_name()}>"
|
||||
attr = ir.Attribute.parse(field_name)
|
||||
return _cute_nvgpu_ir.atom_get_value(
|
||||
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
# Base class for all tcgen05 BlockScaled MMA Ops with syntax `tcgen05.mma.cta_group.kind.block_scale` used to factor out some internal code
|
||||
@dataclass(frozen=True)
|
||||
class BlockScaledMmaOp(core.MmaOp):
|
||||
class BlockScaledMmaOp(Tcgen05MmaOp):
|
||||
a_dtype: Type[Numeric]
|
||||
b_dtype: Type[Numeric]
|
||||
acc_dtype: Float32
|
||||
@@ -276,12 +301,12 @@ class BlockScaledMmaOp(core.MmaOp):
|
||||
b_major_mode: OperandMajorMode
|
||||
|
||||
admissible_archs = [
|
||||
"sm_100a",
|
||||
Arch.sm_100a,
|
||||
]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Verify arch
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
@@ -409,6 +434,170 @@ class BlockScaledMmaTraits(Trait):
|
||||
self.value, attr, value, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
def get(self, field, *, loc=None, ip=None) -> Any:
|
||||
if field not in [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]:
|
||||
raise ValueError(f"the get method for {field} is not supported")
|
||||
field_name = f"#cute_nvgpu.atom_mma_field_sm100_block_scaled<{field._to_ir_field_name()}>"
|
||||
attr = ir.Attribute.parse(field_name)
|
||||
return _cute_nvgpu_ir.atom_get_value(
|
||||
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
# Base class for all tcgen05 Sparse MMA Ops with syntax `tcgen05.mma.cta_group.kind.sparse` used to factor out some internal code
|
||||
@dataclass(frozen=True)
|
||||
class SparseMmaOp(Tcgen05MmaOp):
|
||||
a_dtype: Type[Numeric]
|
||||
b_dtype: Type[Numeric]
|
||||
acc_dtype: Type[Numeric]
|
||||
shape_mnk: Shape
|
||||
cta_group: CtaGroup
|
||||
a_src: OperandSource
|
||||
a_major_mode: OperandMajorMode
|
||||
b_major_mode: OperandMajorMode
|
||||
sparse_metadata_format: SparseMetadataFormat
|
||||
|
||||
admissible_archs = Arch.filter(
|
||||
lambda arch: arch.is_family_of(Arch.sm_100f) or arch.is_family_of(Arch.sm_110f)
|
||||
)
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Verify arch
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
# Verify that the user provided enum values
|
||||
if not isinstance(self.cta_group, CtaGroup):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'cta_group' Op parameter to be a tcgen05.CtaGroup instance",
|
||||
)
|
||||
if not isinstance(self.a_src, OperandSource):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'a_src' Op parameter to be a tcgen05.OperandSource instance",
|
||||
)
|
||||
if not isinstance(self.a_major_mode, OperandMajorMode):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'a_major_mode' Op parameter to be a tcgen05.OperandMajorMode instance",
|
||||
)
|
||||
if not isinstance(self.b_major_mode, OperandMajorMode):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'b_major_mode' Op parameter to be a tcgen05.OperandMajorMode instance",
|
||||
)
|
||||
if not isinstance(self.sparse_metadata_format, SparseMetadataFormat):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'sparse_metadata_format' Op parameter to be a tcgen05.SparseMetadataFormat instance",
|
||||
)
|
||||
# Verify the instruction shape
|
||||
if (rank(self.shape_mnk) not in [2, 3]) or (depth(self.shape_mnk) != 1):
|
||||
raise OpError(
|
||||
self,
|
||||
f"expected a flat rank 2 or 3 tuple for the 'shape_mnk' Op parameter, "
|
||||
f"but got {self.shape_mnk}",
|
||||
)
|
||||
m, n = self.shape_mnk[0], self.shape_mnk[1]
|
||||
# For sparse MMA, the shape validation follows the same rules as dense MMA
|
||||
# but the K dimension is typically doubled in the derived classes
|
||||
if self.cta_group == CtaGroup.ONE:
|
||||
if m not in [64, 128]:
|
||||
raise OpError(self, f"expects the M-mode to be 64 or 128, but got {m}")
|
||||
if m == 64:
|
||||
if (n < 8) or (n > 256) or (n % 8 != 0):
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the N-mode to satisfy 8 <= N <= 256 and N % 8 == 0, but got {n}",
|
||||
)
|
||||
elif m == 128:
|
||||
if (n < 16) or (n > 256) or (n % 16 != 0):
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the N-mode to satisfy 16 <= N <= 256 and N % 16 == 0, but got {n}",
|
||||
)
|
||||
else:
|
||||
if m not in [128, 256]:
|
||||
raise OpError(self, f"expects the M-mode to be 128 or 256, but got {m}")
|
||||
if (n < 32) or (n > 256) or (n % 32 != 0):
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the N-mode to satisfy 32 <= N <= 256 and N % 32 == 0, but got {n}",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
self.__class__.descriptive_name # type: ignore
|
||||
+ f"\n A data type = {self.a_dtype}"
|
||||
+ f"\n B data type = {self.b_dtype}"
|
||||
+ f"\n Accumulator data type = {self.acc_dtype}"
|
||||
+ f"\n CTA group = {self.cta_group}"
|
||||
+ f"\n A source location = {self.a_src}"
|
||||
+ f"\n A major mode = {self.a_major_mode}"
|
||||
+ f"\n B major mode = {self.b_major_mode}"
|
||||
+ f"\n Instruction shape MNK = {self.shape_mnk}"
|
||||
+ f"\n Sparse metadata format = {self.sparse_metadata_format}"
|
||||
)
|
||||
|
||||
def _verify_fragment_A(self, input: _Tensor, *, loc=None, ip=None):
|
||||
if input.memspace == AddressSpace.smem and isinstance(
|
||||
input.layout.type, _cute_ir.ComposedLayoutType
|
||||
):
|
||||
raise OpError(
|
||||
self,
|
||||
f"Expected affine layout for {self._make_trait()}'s operand A, "
|
||||
f"but got composed layout instead: {input.layout}"
|
||||
f"\nPlease use recast_ptr(ptr, {input.layout.inner}, element_type) operation to move swizzle to the ptr",
|
||||
)
|
||||
return True
|
||||
|
||||
def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None):
|
||||
if input.memspace == AddressSpace.smem and isinstance(
|
||||
input.layout.type, _cute_ir.ComposedLayoutType
|
||||
):
|
||||
raise OpError(
|
||||
self,
|
||||
f"Expected affine layout for {self._make_trait()}'s operand B, "
|
||||
f"but got composed layout instead: {input.layout}"
|
||||
f"\nPlease use recast_ptr(ptr, {input.layout.inner}, element_type) operation to move swizzle to the ptr",
|
||||
)
|
||||
return True
|
||||
|
||||
|
||||
class SparseMmaTraits(Trait):
|
||||
admissible_fields = [Field.ACCUMULATE, Field.NEGATE_A, Field.NEGATE_B]
|
||||
|
||||
def set(self, field, value, *, loc=None, ip=None) -> None:
|
||||
if field not in self.admissible_fields:
|
||||
raise ValueError(
|
||||
f"expects field to be one of {self.admissible_fields}, but got {field}"
|
||||
)
|
||||
field_name = (
|
||||
f"#cute_nvgpu.atom_mma_field_sm100_sparse<{field._to_ir_field_name()}>"
|
||||
)
|
||||
attr = ir.Attribute.parse(field_name)
|
||||
self.value = _cute_nvgpu_ir.atom_set_value(
|
||||
self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
|
||||
def get(self, field, *, loc=None, ip=None) -> Any:
|
||||
if field not in self.admissible_fields:
|
||||
raise ValueError(
|
||||
f"expects field to be one of {self.admissible_fields}, but got {field}"
|
||||
)
|
||||
field_name = (
|
||||
f"#cute_nvgpu.atom_mma_field_sm100_sparse<{field._to_ir_field_name()}>"
|
||||
)
|
||||
attr = ir.Attribute.parse(field_name)
|
||||
return _cute_nvgpu_ir.atom_get_value(
|
||||
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
#
|
||||
# TF32 MMA
|
||||
@@ -472,18 +661,20 @@ class MmaTF32Op(MmaOp):
|
||||
0,
|
||||
)
|
||||
return MmaTF32Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaTF32Trait(MmaTrait):
|
||||
class MmaTF32Trait(MmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
@@ -564,18 +755,123 @@ class MmaF16BF16Op(MmaOp):
|
||||
0,
|
||||
)
|
||||
return MmaF16BF16Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaF16BF16Trait(MmaTrait):
|
||||
class MmaF16BF16Trait(MmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaF16BF16SparseOp(SparseMmaOp):
|
||||
"""
|
||||
F16/BF16 tcgen05 Sparse MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma-sp>`__.
|
||||
This Operation corresponds to the ``.kind::f16`` qualifier with sparse support.
|
||||
"""
|
||||
|
||||
descriptive_name = "tcgen05 F16/BF16 Sparse MMA Operation"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ab_dtype: Type[Numeric],
|
||||
acc_dtype: Type[Numeric],
|
||||
instruction_shape: Shape,
|
||||
cta_group: CtaGroup,
|
||||
a_src: OperandSource,
|
||||
a_major_mode: OperandMajorMode,
|
||||
b_major_mode: OperandMajorMode,
|
||||
sparse_metadata_format: SparseMetadataFormat,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
ab_dtype,
|
||||
ab_dtype,
|
||||
acc_dtype,
|
||||
instruction_shape,
|
||||
cta_group,
|
||||
a_src,
|
||||
a_major_mode,
|
||||
b_major_mode,
|
||||
sparse_metadata_format,
|
||||
)
|
||||
self._verify()
|
||||
|
||||
def _verify(self) -> None:
|
||||
# Input data type verification
|
||||
if self.a_dtype not in [Float16, BFloat16]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'ab_dtype' Op parameter to be one of Float16 or BFloat16",
|
||||
)
|
||||
assert self.b_dtype == self.a_dtype, "a_dtype and b_dtype must be the same"
|
||||
# Accumulator data type verification
|
||||
if self.acc_dtype not in [Float16, Float32]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'acc_dtype' Op parameter to be one of Float16 or Float32",
|
||||
)
|
||||
# Instruction shape verification
|
||||
instruction_k = 32 # For sparse, K is doubled compared to dense F16/BF16
|
||||
if rank(self.shape_mnk) == 2:
|
||||
object.__setattr__(self, "shape_mnk", (*self.shape_mnk, instruction_k))
|
||||
if self.shape_mnk[2] != instruction_k:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the instruction extent in the K-mode to be {instruction_k}, "
|
||||
f"but got {self.shape_mnk[2]}",
|
||||
)
|
||||
|
||||
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaF16BF16SparseTrait":
|
||||
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.MmaAtomSM100UMMASparseType.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.cta_group.value,
|
||||
self.a_major_mode._to_ir(),
|
||||
self.b_major_mode._to_ir(),
|
||||
self.a_dtype.mlir_type,
|
||||
self.b_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
T.ui8(),
|
||||
self.sparse_metadata_format._to_ir(),
|
||||
self.a_src._to_ir(),
|
||||
0, # cScaleExp
|
||||
)
|
||||
|
||||
def get_e_ptr():
|
||||
ptr_type = _cute_ir.PtrType.get(T.ui8(), _cute_ir.AddressSpace.tmem, 8)
|
||||
address_value = Int32(0).ir_value(loc=loc, ip=ip)
|
||||
aligned_ty = _cute_ir.ConstrainedIntType.get(8, 32)
|
||||
aligned_intptr = _cute_ir.assume(aligned_ty, address_value, loc=loc, ip=ip)
|
||||
ui8_tmem_ptr = _cute_ir.inttoptr(ptr_type, aligned_intptr, loc=loc, ip=ip)
|
||||
return ui8_tmem_ptr
|
||||
|
||||
return MmaF16BF16SparseTrait(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
get_e_ptr().value,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaF16BF16SparseTrait(SparseMmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
@@ -649,18 +945,20 @@ class MmaI8Op(MmaOp):
|
||||
0,
|
||||
)
|
||||
return MmaI8Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaI8Trait(MmaTrait):
|
||||
class MmaI8Trait(MmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
@@ -689,7 +987,6 @@ class MmaFP8Op(MmaOp):
|
||||
a_major_mode: OperandMajorMode,
|
||||
b_major_mode: OperandMajorMode,
|
||||
) -> None:
|
||||
|
||||
super().__init__(
|
||||
ab_dtype,
|
||||
ab_dtype,
|
||||
@@ -741,18 +1038,20 @@ class MmaFP8Op(MmaOp):
|
||||
0,
|
||||
)
|
||||
return MmaFP8Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaFP8Trait(MmaTrait):
|
||||
class MmaFP8Trait(MmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
@@ -829,13 +1128,19 @@ class MmaMXF8Op(BlockScaledMmaOp):
|
||||
self.sf_vec_size,
|
||||
)
|
||||
return MmaMXF8Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma_bs(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
core.make_ptr(self.sf_dtype, 0, _cute_ir.AddressSpace.tmem).value,
|
||||
core.make_ptr(self.sf_dtype, 0, _cute_ir.AddressSpace.tmem).value,
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
core.make_ptr(
|
||||
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
|
||||
).value,
|
||||
core.make_ptr(
|
||||
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
|
||||
).value,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
@@ -909,13 +1214,19 @@ class MmaMXF4Op(BlockScaledMmaOp):
|
||||
self.sf_vec_size,
|
||||
)
|
||||
return MmaMXF4Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma_bs(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
core.make_ptr(self.sf_dtype, 0, _cute_ir.AddressSpace.tmem).value,
|
||||
core.make_ptr(self.sf_dtype, 0, _cute_ir.AddressSpace.tmem).value,
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
core.make_ptr(
|
||||
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
|
||||
).value,
|
||||
core.make_ptr(
|
||||
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
|
||||
).value,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
@@ -996,13 +1307,19 @@ class MmaMXF4NVF4Op(BlockScaledMmaOp):
|
||||
self.sf_vec_size,
|
||||
)
|
||||
return MmaMXF4NVF4Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma_bs(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
core.make_ptr(self.sf_dtype, 0, _cute_ir.AddressSpace.tmem).value,
|
||||
core.make_ptr(self.sf_dtype, 0, _cute_ir.AddressSpace.tmem).value,
|
||||
(
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
core.make_ptr(
|
||||
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
|
||||
).value,
|
||||
core.make_ptr(
|
||||
self.sf_dtype, 0, _cute_ir.AddressSpace.tmem, loc=loc, ip=ip
|
||||
).value,
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
@@ -1012,6 +1329,7 @@ class MmaMXF4NVF4Op(BlockScaledMmaOp):
|
||||
class MmaMXF4NVF4Trait(BlockScaledMmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# SMEM layout atoms
|
||||
|
||||
@@ -12,13 +12,14 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
from cutlass import cute
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..common import OpError
|
||||
from ...core import CopyOp, Trait, _pack_shape
|
||||
from ...core import _pack_shape
|
||||
from ...typing import Numeric
|
||||
from ...atom import CopyOp, Trait
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
@@ -39,7 +40,7 @@ class BaseOp(CopyOp):
|
||||
+ f"\n number of matrices = {self.num_matrices}"
|
||||
)
|
||||
if self.transpose:
|
||||
res += f"\n transposed"
|
||||
res += "\n transposed"
|
||||
return res
|
||||
|
||||
|
||||
@@ -71,7 +72,7 @@ class LdMatrix8x8x16bOp(BaseOp):
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get() if self.transpose else None,
|
||||
)
|
||||
return LdMatrix8x8x16bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return LdMatrix8x8x16bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class LdMatrix8x8x16bTrait(Trait):
|
||||
@@ -110,7 +111,7 @@ class LdMatrix16x16x8bOp(BaseOp):
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get(),
|
||||
)
|
||||
return LdMatrix16x16x8bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return LdMatrix16x16x8bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class LdMatrix16x16x8bTrait(Trait):
|
||||
@@ -144,7 +145,7 @@ class StMatrix8x8x16bOp(BaseOp):
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get() if self.transpose else None,
|
||||
)
|
||||
return StMatrix8x8x16bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return StMatrix8x8x16bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class StMatrix8x8x16bTrait(Trait):
|
||||
@@ -182,7 +183,7 @@ class StMatrix16x8x8bOp(BaseOp):
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get(),
|
||||
)
|
||||
return StMatrix16x8x8bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return StMatrix16x8x8bTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class StMatrix16x8x8bTrait(Trait):
|
||||
|
||||
@@ -12,16 +12,36 @@
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
import enum
|
||||
from cutlass import cute
|
||||
|
||||
from ..common import OpError
|
||||
from ...core import MmaOp, Trait, _pack_shape, _Tensor
|
||||
from ...typing import Shape, Float16, BFloat16, Float32, Numeric, AddressSpace
|
||||
from ...typing import Shape, Float16, BFloat16, Float32, Numeric
|
||||
from ...core import _pack_shape
|
||||
from ...tensor import _Tensor
|
||||
from ...atom import MmaOp, Trait
|
||||
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir.dialects.cute_nvgpu import SparseMetadataFormat
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# MMA Ops and Traits
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class WarpMmaOp(MmaOp):
|
||||
"""
|
||||
Base class for all warp-level MMA operations.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaF16BF16Op(MmaOp):
|
||||
class MmaF16BF16Op(WarpMmaOp):
|
||||
"""
|
||||
F16/BF16 tcgen05 MMA Operation.
|
||||
|
||||
@@ -63,7 +83,7 @@ class MmaF16BF16Op(MmaOp):
|
||||
self.ab_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
)
|
||||
return MmaF16BF16Trait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
return MmaF16BF16Trait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
@@ -79,5 +99,84 @@ class MmaF16BF16Op(MmaOp):
|
||||
def _verify_fragment_B(self, input: _Tensor, *, loc=None, ip=None):
|
||||
pass
|
||||
|
||||
|
||||
class MmaF16BF16Trait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
class SparseMetadataFormat(enum.Enum):
|
||||
"""
|
||||
An enumeration for the sparse metadata format of the MMA.
|
||||
"""
|
||||
|
||||
TID = SparseMetadataFormat.tid
|
||||
|
||||
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_nvgpu_ir.SparseMetadataFormat:
|
||||
return self.value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaF16BF16SparseOp(WarpMmaOp):
|
||||
ab_dtype: Type[Numeric]
|
||||
acc_dtype: Type[Numeric]
|
||||
shape_mnk: Shape
|
||||
sparse_metadata_format: SparseMetadataFormat
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# verify field after initialization
|
||||
if not isinstance(self.sparse_metadata_format, SparseMetadataFormat):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'sparse_metadata_format' Op parameter to be a SparseMetadataFormat instance",
|
||||
)
|
||||
# verify the instruction shape
|
||||
if self.ab_dtype not in [Float16, BFloat16]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'ab_dtype' Op parameter to be one of Float16 or BFloat16",
|
||||
)
|
||||
if self.acc_dtype not in [Float16, Float32]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'acc_dtype' Op parameter to be one of Float16 or Float32",
|
||||
)
|
||||
if (self.ab_dtype == BFloat16) and (self.acc_dtype != Float32):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'acc_dtype' Op parameter to be Float32 when 'ab_dtype' is BFloat16",
|
||||
)
|
||||
if self.shape_mnk not in [(16, 8, 16), (16, 8, 32)]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'shape_mnk' Op parameter to be one of (16,8,16) or (16,8,32)",
|
||||
)
|
||||
|
||||
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaF16BF16SparseTrait":
|
||||
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.MmaAtomSM80SparseType.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.ab_dtype.mlir_type,
|
||||
self.ab_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
self.sparse_metadata_format._to_ir(),
|
||||
)
|
||||
return MmaF16BF16SparseTrait(cute.make_atom(ty, loc=loc, ip=ip))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
"warp-level F16/BF16 Sparse MMA Operation"
|
||||
+ f"\n A/B data type = {self.ab_dtype}"
|
||||
+ f"\n Accumulator data type = {self.acc_dtype}"
|
||||
+ f"\n Instruction shape MNK = {self.shape_mnk}"
|
||||
+ f"\n Sparse metadata format = {self.sparse_metadata_format}"
|
||||
)
|
||||
|
||||
|
||||
class MmaF16BF16SparseTrait(Trait):
|
||||
pass
|
||||
|
||||
@@ -15,7 +15,7 @@ from cutlass.cutlass_dsl import dsl_user_op
|
||||
|
||||
from cutlass._mlir.dialects import nvvm
|
||||
|
||||
from ...typing import Numeric, NumericMeta
|
||||
from ...typing import Numeric, NumericMeta, ComposedLayout
|
||||
from ... import core
|
||||
from .mma import SmemLayoutAtomKind
|
||||
|
||||
@@ -23,7 +23,7 @@ from .mma import SmemLayoutAtomKind
|
||||
@dsl_user_op
|
||||
def make_smem_layout_atom(
|
||||
kind: SmemLayoutAtomKind, element_type: Type[Numeric], *, loc=None, ip=None
|
||||
) -> core.ComposedLayout:
|
||||
) -> ComposedLayout:
|
||||
"""
|
||||
Makes a SMEM layout Atom.
|
||||
|
||||
@@ -35,7 +35,7 @@ def make_smem_layout_atom(
|
||||
:param element_type: The element data type to construct the layout for
|
||||
:type element_type: Type[Numeric]
|
||||
:return: The SMEM layout atom
|
||||
:rtype: core.ComposedLayout
|
||||
:rtype: ComposedLayout
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(f"element_type must be a Numeric, but got {element_type}")
|
||||
|
||||
@@ -11,16 +11,19 @@
|
||||
|
||||
import enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
from typing import Type, Any
|
||||
|
||||
from cutlass.cutlass_dsl import CuTeDSL
|
||||
from cutlass import cute
|
||||
from cutlass.base_dsl.arch import Arch
|
||||
from cutlass.cutlass_dsl import CuTeDSL, T
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..common import OpError
|
||||
from ...core import MmaOp, Trait, _pack_shape, rank, depth, _Tensor
|
||||
from ...core import _pack_shape, rank, depth
|
||||
from ...tensor import _Tensor
|
||||
from ...typing import (
|
||||
Shape,
|
||||
Float16,
|
||||
@@ -29,9 +32,13 @@ from ...typing import (
|
||||
Boolean,
|
||||
Float8E5M2,
|
||||
Float8E4M3FN,
|
||||
Int32,
|
||||
Int8,
|
||||
Uint8,
|
||||
Numeric,
|
||||
AddressSpace,
|
||||
)
|
||||
from ...atom import MmaOp, Trait
|
||||
|
||||
|
||||
####################################################################################################
|
||||
@@ -41,6 +48,14 @@ from ...typing import (
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class WarpGroupMmaOp(MmaOp):
|
||||
"""
|
||||
Base class for all warpgroup-level MMA operations.
|
||||
"""
|
||||
|
||||
pass
|
||||
|
||||
|
||||
class OperandMajorMode(enum.Enum):
|
||||
"""
|
||||
An enumeration for the majorness of the input operands of the MMA.
|
||||
@@ -104,7 +119,7 @@ class Field(enum.Enum):
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaOp(MmaOp):
|
||||
class MmaOp(WarpGroupMmaOp):
|
||||
a_dtype: Type[Numeric]
|
||||
b_dtype: Type[Numeric]
|
||||
acc_dtype: Type[Numeric]
|
||||
@@ -113,15 +128,13 @@ class MmaOp(MmaOp):
|
||||
a_major_mode: OperandMajorMode
|
||||
b_major_mode: OperandMajorMode
|
||||
|
||||
admissible_archs = ["sm_90a"]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Verify arch
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
arch = CuTeDSL._get_dsl().get_arch_enum()
|
||||
if not arch == Arch.sm_90a:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
f"expects arch to be {Arch.sm_90a}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
# Verify that the user provided enum values
|
||||
@@ -193,7 +206,7 @@ class MmaOp(MmaOp):
|
||||
return True
|
||||
|
||||
|
||||
class MmaTrait(Trait):
|
||||
class MmaTraits(Trait):
|
||||
admissible_fields = [Field.ACCUMULATE]
|
||||
|
||||
def set(self, field, value, *, loc=None, ip=None) -> None:
|
||||
@@ -207,13 +220,24 @@ class MmaTrait(Trait):
|
||||
self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
|
||||
def get(self, field, *, loc=None, ip=None) -> Any:
|
||||
if field not in self.admissible_fields:
|
||||
raise ValueError(
|
||||
f"invalid field, must be {Field.ACCUMULATE}, but got {field}"
|
||||
)
|
||||
field_name = f"#cute_nvgpu.atom_mma_field_sm90<{field._to_ir_field_name()}>"
|
||||
attr = ir.Attribute.parse(field_name)
|
||||
return _cute_nvgpu_ir.atom_get_value(
|
||||
Boolean.mlir_type, self.value, attr, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaF16BF16Op(MmaOp):
|
||||
"""
|
||||
F16/BF16 warpgroup MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-multiply-and-accumulate-instruction-wgmma-mma-async>`__.
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions-wgmma-mma>`__.
|
||||
This Operation covers the instructions using the ``.f16`` or ``.bf16`` qualifiers for the input operands.
|
||||
"""
|
||||
|
||||
@@ -281,16 +305,16 @@ class MmaF16BF16Op(MmaOp):
|
||||
self.a_src._to_ir(),
|
||||
)
|
||||
return MmaF16BF16Trait(
|
||||
_cute_nvgpu_ir.make_sm90_mma(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
(Boolean(False).ir_value(loc=loc, ip=ip),),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaF16BF16Trait(MmaTrait):
|
||||
class MmaF16BF16Trait(MmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
@@ -299,7 +323,7 @@ class MmaF8Op(MmaOp):
|
||||
"""
|
||||
F16/BF16 warpgroup MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-multiply-and-accumulate-instruction-wgmma-mma-async>`__.
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions-wgmma-mma>`__.
|
||||
This Operation covers the instructions using the ``.e4m3`` or ``.e5m2`` qualifiers for the input operands.
|
||||
"""
|
||||
|
||||
@@ -367,13 +391,111 @@ class MmaF8Op(MmaOp):
|
||||
self.a_src._to_ir(),
|
||||
)
|
||||
return MmaF8Trait(
|
||||
_cute_nvgpu_ir.make_sm90_mma(
|
||||
ty, Boolean(False).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
cute.make_atom(
|
||||
ty,
|
||||
(Boolean(False).ir_value(loc=loc, ip=ip),),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaF8Trait(MmaTrait):
|
||||
class MmaF8Trait(MmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaI8Op(MmaOp):
|
||||
"""
|
||||
I8 warpgroup MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions-wgmma-mma>`__.
|
||||
This Operation covers the instructions using the ``.s8`` or ``.u8`` qualifiers for the input operands.
|
||||
"""
|
||||
|
||||
descriptive_name = "warpgroup I8 MMA Operation"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
a_dtype: Type[Numeric],
|
||||
b_dtype: Type[Numeric],
|
||||
acc_dtype: Type[Numeric],
|
||||
instruction_shape: Shape,
|
||||
a_src: OperandSource,
|
||||
a_major_mode: OperandMajorMode,
|
||||
b_major_mode: OperandMajorMode,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
a_dtype,
|
||||
b_dtype,
|
||||
acc_dtype,
|
||||
instruction_shape,
|
||||
a_src,
|
||||
a_major_mode,
|
||||
b_major_mode,
|
||||
)
|
||||
self._verify()
|
||||
|
||||
def _verify(self):
|
||||
# Input data type verification
|
||||
if self.a_dtype not in [Int8, Uint8]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'a_dtype' Op parameter to be one of Int8 or Uint8",
|
||||
)
|
||||
if self.b_dtype not in [Int8, Uint8]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'b_dtype' Op parameter to be one of Int8 or Uint8",
|
||||
)
|
||||
# Accumulator data type verification
|
||||
if self.acc_dtype != Int32:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'acc_dtype' Op parameter must be Int32",
|
||||
)
|
||||
|
||||
# Verify the instruction shape
|
||||
instruction_k = 32
|
||||
if rank(self.shape_mnk) == 2:
|
||||
object.__setattr__(self, "shape_mnk", (*self.shape_mnk, instruction_k))
|
||||
if self.shape_mnk[2] != instruction_k:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects the instruction extent in the K-mode to be {instruction_k}, "
|
||||
f"but got {self.shape_mnk[2]}",
|
||||
)
|
||||
|
||||
n = self.shape_mnk[1]
|
||||
if not (n >= 8 and n <= 256 and (n == 8 or n == 24 or n % 16 == 0)):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the N-mode to satisfy N=8*i where i={1,2,3,4} ",
|
||||
f"or N=16*i where i={{3,4,...,15,16}}. But got {n}",
|
||||
)
|
||||
|
||||
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaI8Trait":
|
||||
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.MmaAtomSM90Type.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.a_major_mode._to_ir(),
|
||||
self.b_major_mode._to_ir(),
|
||||
(T.si8() if self.a_dtype.signed else T.ui8()),
|
||||
(T.si8() if self.b_dtype.signed else T.ui8()),
|
||||
self.acc_dtype.mlir_type,
|
||||
self.a_src._to_ir(),
|
||||
)
|
||||
return MmaI8Trait(
|
||||
cute.make_atom(
|
||||
ty,
|
||||
(Boolean(False).ir_value(loc=loc, ip=ip),),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaI8Trait(MmaTraits):
|
||||
pass
|
||||
|
||||
|
||||
|
||||
@@ -13,41 +13,18 @@ import ctypes
|
||||
from functools import lru_cache
|
||||
import itertools
|
||||
import operator
|
||||
from time import time
|
||||
from typing import Union
|
||||
from typing import Union, Optional
|
||||
|
||||
# MLIR modules imports
|
||||
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 JitArgAdapterRegistry
|
||||
from cutlass.cutlass_dsl import JitArgAdapterRegistry, DSLRuntimeError
|
||||
|
||||
# Local modules imports
|
||||
from .typing import (
|
||||
AddressSpace,
|
||||
Tensor,
|
||||
Type,
|
||||
Pointer,
|
||||
Boolean,
|
||||
Numeric,
|
||||
Float4E2M1FN,
|
||||
Int64,
|
||||
Int32,
|
||||
Int16,
|
||||
Int8,
|
||||
Uint64,
|
||||
Uint32,
|
||||
Uint16,
|
||||
Uint8,
|
||||
Float64,
|
||||
Float32,
|
||||
Float16,
|
||||
BFloat16,
|
||||
Float8E5M2,
|
||||
)
|
||||
from .typing import AddressSpace, Tensor, Type, Pointer, Numeric
|
||||
from . import core
|
||||
from .core import _Tensor as CoreTensor
|
||||
from .tensor import _Tensor as CoreTensor
|
||||
|
||||
|
||||
class _Pointer(Pointer):
|
||||
@@ -88,9 +65,9 @@ class _Pointer(Pointer):
|
||||
self._assumed_align = assumed_align
|
||||
|
||||
self._c_pointer = None
|
||||
assert (
|
||||
int(self._pointer) % self._assumed_align == 0
|
||||
), f"pointer must be {self._assumed_align} bytes aligned"
|
||||
assert 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))
|
||||
@@ -109,9 +86,6 @@ class _Pointer(Pointer):
|
||||
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:
|
||||
@@ -146,11 +120,7 @@ class _Pointer(Pointer):
|
||||
|
||||
|
||||
class _Tensor(Tensor):
|
||||
def __init__(
|
||||
self,
|
||||
tensor,
|
||||
assumed_align=None,
|
||||
):
|
||||
def __init__(self, tensor, assumed_align=None, use_32bit_stride=False):
|
||||
# If tensor is already a DLPack object, use it directly
|
||||
if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"):
|
||||
self._dlpack_data = tensor
|
||||
@@ -161,13 +131,13 @@ class _Tensor(Tensor):
|
||||
self._is_dynamic = False
|
||||
self._memref_desc = None
|
||||
self._dtype = None
|
||||
self._use_32bit_stride = use_32bit_stride
|
||||
|
||||
@property
|
||||
def __class__(self) -> Type[Tensor]:
|
||||
# Cheat to let `type(_Tensor())` to return cute.Tensor
|
||||
return Tensor
|
||||
|
||||
@staticmethod
|
||||
def lazily_load_dltensor(func):
|
||||
"""Decorator to lazily load the DLTensorWrapper.
|
||||
|
||||
@@ -177,13 +147,15 @@ class _Tensor(Tensor):
|
||||
|
||||
def wrapper(self, *args, **kwargs):
|
||||
if self._dltensor_wrapper is None:
|
||||
self._dltensor_wrapper = _cute_ir.DLTensorWrapper(self._dlpack_data)
|
||||
self._dltensor_wrapper = _cute_ir.DLTensorWrapper(
|
||||
self._dlpack_data, self._use_32bit_stride
|
||||
)
|
||||
return func(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
@lazily_load_dltensor
|
||||
def mark_layout_dynamic(self, leading_dim: int | None = None):
|
||||
def mark_layout_dynamic(self, leading_dim: Optional[int] = None):
|
||||
"""Marks the tensor layout as dynamic based on the leading dimension.
|
||||
|
||||
:param leading_dim: The leading dimension of the layout, defaults to None
|
||||
@@ -209,7 +181,7 @@ class _Tensor(Tensor):
|
||||
def mark_compact_shape_dynamic(
|
||||
self,
|
||||
mode: int,
|
||||
stride_order: tuple[int, ...] | None = None,
|
||||
stride_order: Optional[tuple[int, ...]] = None,
|
||||
divisibility: int = 1,
|
||||
):
|
||||
"""Marks the tensor shape as dynamic and propagates dynamic and divisibility information to the corresponding strides.
|
||||
@@ -308,10 +280,10 @@ class _Tensor(Tensor):
|
||||
return self.__str__()
|
||||
|
||||
def __setitem__(self, crd, value):
|
||||
raise TypeError(f"runtime._Tensor is not indexable")
|
||||
raise TypeError("runtime._Tensor is not indexable")
|
||||
|
||||
def __getitem__(self, crd):
|
||||
raise TypeError(f"runtime._Tensor is not indexable")
|
||||
raise TypeError("runtime._Tensor is not indexable")
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
@@ -326,7 +298,7 @@ class _Tensor(Tensor):
|
||||
@property
|
||||
def layout(self):
|
||||
raise NotImplementedError(
|
||||
f"layout property is not supported in runtime, support in future"
|
||||
"layout property is not supported in runtime, support in future"
|
||||
)
|
||||
|
||||
@property
|
||||
@@ -363,7 +335,7 @@ class _Tensor(Tensor):
|
||||
return core.leading_dim(self.shape, self.stride)
|
||||
|
||||
def fill(self, value: Numeric):
|
||||
raise TypeError(f"fill function is not supported in runtime")
|
||||
raise TypeError("fill function is not supported in runtime")
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
@@ -389,6 +361,7 @@ class _Tensor(Tensor):
|
||||
def from_dlpack(
|
||||
tensor_dlpack,
|
||||
assumed_align=None,
|
||||
use_32bit_stride=False,
|
||||
) -> Tensor:
|
||||
"""Convert from tensor object supporting __dlpack__() to a CuTe Tensor.
|
||||
|
||||
@@ -397,6 +370,10 @@ def from_dlpack(
|
||||
:param assumed_align: Assumed alignment of the tensor (bytes), defaults to None,
|
||||
if None, will use the element size bytes as the assumed alignment.
|
||||
:type assumed_align: int, optional
|
||||
:param use_32bit_stride: Whether to use 32-bit stride, defaults to False. When True, the dynamic
|
||||
stride bitwidth will be set to 32 for small problem size (cosize(layout) <= Int32_max) for better performance.
|
||||
This is only applied when the dimension is dynamic.
|
||||
:type use_32bit_stride: bool, optional
|
||||
:return: A CuTe Tensor object
|
||||
:rtype: Tensor
|
||||
|
||||
@@ -415,6 +392,7 @@ def from_dlpack(
|
||||
return _Tensor(
|
||||
tensor_dlpack,
|
||||
assumed_align=assumed_align,
|
||||
use_32bit_stride=use_32bit_stride,
|
||||
)
|
||||
|
||||
|
||||
|
||||
1976
python/CuTeDSL/cutlass/cute/tensor.py
Normal file
1976
python/CuTeDSL/cutlass/cute/tensor.py
Normal file
File diff suppressed because it is too large
Load Diff
@@ -13,37 +13,37 @@ import functools
|
||||
import inspect
|
||||
import logging
|
||||
import os
|
||||
from enum import Enum
|
||||
from inspect import isclass
|
||||
from itertools import product
|
||||
from time import time
|
||||
from typing import Any, Callable, Dict, List, Optional, Type, Union
|
||||
from typing import Type, Union, Callable, Optional, Dict, List, Any
|
||||
|
||||
import cuda.bindings.driver as cuda_driver
|
||||
import cuda.bindings.runtime as cuda_runtime
|
||||
import numpy as np
|
||||
|
||||
import cutlass._mlir.ir as ir
|
||||
import cutlass.base_dsl.jit_executor
|
||||
from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, dsl_user_op
|
||||
|
||||
from .typing import Numeric, Int8, Boolean
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute import nvgpu
|
||||
|
||||
from cutlass._mlir.dialects import builtin, cf, nvvm, vector
|
||||
from cutlass.cute import core, nvgpu
|
||||
from cutlass.cutlass_dsl import Constexpr, CuTeDSL, T, t, dsl_user_op
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def assert_(cond, msg=None, *, loc=None, ip=None):
|
||||
cf.assert_(t.Boolean(cond).ir_value(), msg if msg else "", loc=loc, ip=ip)
|
||||
cf.assert_(Boolean(cond).ir_value(), msg if msg else "", loc=loc, ip=ip)
|
||||
|
||||
|
||||
def _maybe_recast_tensor_from_f4(src: core.Tensor, tv_layout: core.Layout):
|
||||
def _maybe_recast_tensor_from_f4(src: cute.Tensor, tv_layout: cute.Layout):
|
||||
if src.element_type.width == 4:
|
||||
tv_layout = core.recast_layout(8, 4, tv_layout)
|
||||
src = core.recast_tensor(src, dtype=t.Int8)
|
||||
tv_layout = cute.recast_layout(8, 4, tv_layout)
|
||||
src = cute.recast_tensor(src, dtype=Int8)
|
||||
return src, tv_layout
|
||||
|
||||
|
||||
def _maybe_recast_to_f4(input: core.TensorSSA, dtype: Type[core.Numeric]):
|
||||
def _maybe_recast_to_f4(input: cute.TensorSSA, dtype: Type[Numeric]):
|
||||
"""Conditionally recasts the tensor to 4-bit type if the destination type is 4-bit.
|
||||
|
||||
:param input: The input tensor to recast.
|
||||
@@ -51,22 +51,22 @@ def _maybe_recast_to_f4(input: core.TensorSSA, dtype: Type[core.Numeric]):
|
||||
:raises TypeError: If dtype is not a subclass of Numeric.
|
||||
:return: A new tensor recast to 4-bit if dtype is 4-bit, otherwise returns self unchanged.
|
||||
"""
|
||||
if not isclass(dtype) or not issubclass(dtype, core.Numeric):
|
||||
if not inspect.isclass(dtype) or not issubclass(dtype, Numeric):
|
||||
raise TypeError(f"dst_ty must be a type of Numeric, but got {dtype}")
|
||||
|
||||
if dtype.width == 4:
|
||||
recast_shape = core.recast_layout(4, 8, core.make_layout(input.shape)).shape
|
||||
recast_shape = cute.recast_layout(4, 8, cute.make_layout(input.shape)).shape
|
||||
i4_vec = vector.bitcast(
|
||||
T.vector(input.type.shape[0] * 2, T.i(4)), input.maybe_downcast()
|
||||
)
|
||||
res_vect = builtin.unrealized_conversion_cast(
|
||||
[T.vector(i4_vec.type.shape[0], dtype.mlir_type)], [i4_vec]
|
||||
)
|
||||
return core.TensorSSA(res_vect, recast_shape, dtype)
|
||||
return cute.TensorSSA(res_vect, recast_shape, dtype)
|
||||
return input
|
||||
|
||||
|
||||
def _maybe_recast_from_f4(input: core.TensorSSA, src_dtype: Type[core.Numeric]):
|
||||
def _maybe_recast_from_f4(input: cute.TensorSSA, src_dtype: Type[Numeric]):
|
||||
"""Conditionally recasts the tensor from 4-bit type if the source type is 4-bit.
|
||||
|
||||
:param input: The input tensor to recast.
|
||||
@@ -74,27 +74,27 @@ def _maybe_recast_from_f4(input: core.TensorSSA, src_dtype: Type[core.Numeric]):
|
||||
:raises TypeError: If src_dtype is not a subclass of Numeric.
|
||||
:return: A new tensor recast from 4-bit if src_dtype is 4-bit, otherwise returns self unchanged.
|
||||
"""
|
||||
if not isclass(src_dtype) or not issubclass(src_dtype, core.Numeric):
|
||||
if not inspect.isclass(src_dtype) or not issubclass(src_dtype, Numeric):
|
||||
raise TypeError(f"src_ty must be a type of Numeric, but got {src_dtype}")
|
||||
|
||||
if src_dtype.width == 4:
|
||||
recast_shape = core.recast_layout(8, 4, core.make_layout(input.shape)).shape
|
||||
recast_shape = cute.recast_layout(8, 4, cute.make_layout(input.shape)).shape
|
||||
i4_vec = builtin.unrealized_conversion_cast(
|
||||
[T.vector(input.type.shape[0], T.i(4))], [input.maybe_downcast()]
|
||||
)
|
||||
res_vect = vector.bitcast(T.vector(i4_vec.type.shape[0] // 2, T.i8()), i4_vec)
|
||||
return core.TensorSSA(res_vect, recast_shape, core.Int8)
|
||||
return cute.TensorSSA(res_vect, recast_shape, Int8)
|
||||
return input
|
||||
|
||||
|
||||
@CuTeDSL.kernel
|
||||
def _convert_kernel(
|
||||
gSrc: core.Tensor,
|
||||
gDst: core.Tensor,
|
||||
cSrc: core.Tensor,
|
||||
src_tv_layout: core.Layout,
|
||||
dst_tv_layout: core.Layout,
|
||||
src_shape: core.Shape,
|
||||
gSrc: cute.Tensor,
|
||||
gDst: cute.Tensor,
|
||||
cSrc: cute.Tensor,
|
||||
src_tv_layout: cute.Layout,
|
||||
dst_tv_layout: cute.Layout,
|
||||
src_shape: cute.Shape,
|
||||
src_ty,
|
||||
dst_ty,
|
||||
):
|
||||
@@ -110,9 +110,9 @@ def _convert_kernel(
|
||||
|
||||
# compose with CTA TV layout
|
||||
# tid, vid -> address
|
||||
tidfrgSrc = core.composition(ctaSrc, src_tv_layout) # (T,V)
|
||||
tidfrgDst = core.composition(ctaDst, dst_tv_layout) # (T,V)
|
||||
tidfrgCSrc = core.composition(ctaCSrc, src_tv_layout) # (T,V)
|
||||
tidfrgSrc = cute.composition(ctaSrc, src_tv_layout) # (T,V)
|
||||
tidfrgDst = cute.composition(ctaDst, dst_tv_layout) # (T,V)
|
||||
tidfrgCSrc = cute.composition(ctaCSrc, src_tv_layout) # (T,V)
|
||||
# print(f"tidfrgSrc = {tidfrgSrc.type}")
|
||||
|
||||
# slice for threads
|
||||
@@ -123,19 +123,19 @@ def _convert_kernel(
|
||||
# print(f"thrSrc = {thrSrc.type}")
|
||||
|
||||
# predicate
|
||||
if core.elem_less(thrCSrc[0], src_shape):
|
||||
if cute.elem_less(thrCSrc[0], src_shape):
|
||||
# allocate fragments for gmem->rmem
|
||||
frgSrc = core.make_fragment(
|
||||
core.get(src_tv_layout, mode=[1]), gSrc.element_type
|
||||
frgSrc = cute.make_rmem_tensor(
|
||||
cute.get(src_tv_layout, mode=[1]), gSrc.element_type
|
||||
) # (V)
|
||||
frgDst = core.make_fragment(
|
||||
core.get(dst_tv_layout, mode=[1]), gDst.element_type
|
||||
frgDst = cute.make_rmem_tensor(
|
||||
cute.get(dst_tv_layout, mode=[1]), gDst.element_type
|
||||
) # (V)
|
||||
# print(f"frgSrc = {frgSrc.type}")
|
||||
|
||||
# Move data to reg address space
|
||||
copy_atom_load = core.make_copy_atom(nvgpu.CopyUniversalOp(), gSrc.element_type)
|
||||
core.copy(copy_atom_load, thrSrc, frgSrc)
|
||||
copy_atom_load = cute.make_copy_atom(nvgpu.CopyUniversalOp(), gSrc.element_type)
|
||||
cute.copy(copy_atom_load, thrSrc, frgSrc)
|
||||
|
||||
vec_src = frgSrc.load()
|
||||
vec_src = _maybe_recast_to_f4(vec_src, src_ty)
|
||||
@@ -144,49 +144,48 @@ def _convert_kernel(
|
||||
frgDst.store(vec_dst)
|
||||
|
||||
# Copy the results back to c
|
||||
copy_atom_stg = core.make_copy_atom(nvgpu.CopyUniversalOp(), gDst.element_type)
|
||||
core.copy(copy_atom_stg, frgDst, thrDst)
|
||||
copy_atom_stg = cute.make_copy_atom(nvgpu.CopyUniversalOp(), gDst.element_type)
|
||||
cute.copy(copy_atom_stg, frgDst, thrDst)
|
||||
|
||||
|
||||
@CuTeDSL.jit(preprocess=False)
|
||||
def _convert(
|
||||
src: core.Tensor,
|
||||
dst: core.Tensor,
|
||||
src: cute.Tensor,
|
||||
dst: cute.Tensor,
|
||||
leading_mode: Constexpr,
|
||||
elem_per_copy: Constexpr,
|
||||
):
|
||||
|
||||
# Step 1. figure proper tv_layout
|
||||
src_ty = src.element_type
|
||||
dst_ty = dst.element_type
|
||||
|
||||
tv_layout = core.make_layout((128, elem_per_copy), stride=(elem_per_copy, 1))
|
||||
tv_layout = cute.make_layout((128, elem_per_copy), stride=(elem_per_copy, 1))
|
||||
|
||||
# Step 2. maybe recast from f4 tensor
|
||||
src, src_tv_layout = _maybe_recast_tensor_from_f4(src, tv_layout)
|
||||
dst, dst_tv_layout = _maybe_recast_tensor_from_f4(dst, tv_layout)
|
||||
src_shape = src.shape
|
||||
# predicate tensor
|
||||
idA = core.make_identity_tensor(src.shape)
|
||||
idA = cute.make_identity_tensor(src.shape)
|
||||
|
||||
# Step 3. select a proper tiling pattern as (...,TileV, ...)
|
||||
src_cta_tiler = [
|
||||
1,
|
||||
] * core.rank(src.layout)
|
||||
src_cta_tiler[leading_mode] = core.size(src_tv_layout) # (...,TileV,...)
|
||||
] * cute.rank(src.layout)
|
||||
src_cta_tiler[leading_mode] = cute.size(src_tv_layout) # (...,TileV,...)
|
||||
dst_cta_tiler = [
|
||||
1,
|
||||
] * core.rank(dst.layout)
|
||||
dst_cta_tiler[leading_mode] = core.size(dst_tv_layout) # (...,TileV,...)
|
||||
] * cute.rank(dst.layout)
|
||||
dst_cta_tiler[leading_mode] = cute.size(dst_tv_layout) # (...,TileV,...)
|
||||
|
||||
# Step 4. partition input and output tensor by cta tiler.
|
||||
gS = core.zipped_divide(
|
||||
gS = cute.zipped_divide(
|
||||
src, tuple(src_cta_tiler)
|
||||
) # ((...,TileV,...),(...,RestV,...))
|
||||
cS = core.zipped_divide(
|
||||
cS = cute.zipped_divide(
|
||||
idA, tuple(src_cta_tiler)
|
||||
) # ((...,TileV,...),(...,RestV,...))
|
||||
gD = core.zipped_divide(
|
||||
gD = cute.zipped_divide(
|
||||
dst, tuple(dst_cta_tiler)
|
||||
) # ((...,TileV,...),(...,RestV,...))
|
||||
# print(f"{gS.type=}")
|
||||
@@ -201,8 +200,8 @@ def _convert(
|
||||
src_ty,
|
||||
dst_ty,
|
||||
).launch(
|
||||
grid=[core.size(gS, mode=[1]), 1, 1],
|
||||
block=[core.size(src_tv_layout, mode=[0]), 1, 1],
|
||||
grid=[cute.size(gS, mode=[1]), 1, 1],
|
||||
block=[cute.size(src_tv_layout, mode=[0]), 1, 1],
|
||||
)
|
||||
|
||||
|
||||
@@ -210,10 +209,10 @@ def _convert(
|
||||
# And when src or dst dtype is narrow precision(Float4E2M1FN/Float8E8M0FNU/Float8E4M3FN), the shape of
|
||||
# their leading dimension should be 4(fp8)/8(fp4) element align. (nvgpu.cvt_fptrunc/cvt_fpext
|
||||
# needs 32-bits aligned input/output)
|
||||
def convert(src: core.Tensor, dst: core.Tensor):
|
||||
assert len(src.shape) == len(
|
||||
dst.shape
|
||||
), "Shape of src and dst tensors should be the same rank."
|
||||
def convert(src: cute.Tensor, dst: cute.Tensor):
|
||||
assert len(src.shape) == len(dst.shape), (
|
||||
"Shape of src and dst tensors should be the same rank."
|
||||
)
|
||||
# find leading mode
|
||||
leading_mode = [
|
||||
idx
|
||||
@@ -329,9 +328,9 @@ def _does_kernel_use_stream(
|
||||
:rtype: bool
|
||||
"""
|
||||
|
||||
assert int(stream) != int(
|
||||
cuda_driver.CUstream_flags.CU_STREAM_DEFAULT
|
||||
), "Stream must be a non-default stream"
|
||||
assert int(stream) != int(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT), (
|
||||
"Stream must be a non-default stream"
|
||||
)
|
||||
|
||||
err = cuda_runtime.cudaStreamBeginCapture(
|
||||
stream, cuda_runtime.cudaStreamCaptureMode.cudaStreamCaptureModeThreadLocal
|
||||
@@ -474,8 +473,13 @@ def benchmark(
|
||||
elapsed_time = float("nan")
|
||||
|
||||
if use_cuda_graphs:
|
||||
# Check if the callable is a JitExecutor
|
||||
if not isinstance(callable, cutlass.base_dsl.jit_executor.JitExecutor):
|
||||
# Check if the callable is a JitCompiledFunction or JitExecutor
|
||||
# These are functions that can be called to launch kernels
|
||||
compiled_types = (
|
||||
cutlass.base_dsl.jit_executor.JitCompiledFunction,
|
||||
cutlass.base_dsl.jit_executor.JitExecutor,
|
||||
)
|
||||
if not isinstance(callable, compiled_types):
|
||||
raise TypeError("Function must be precompiled to be used with CUDA Graphs")
|
||||
|
||||
# Check if the stream is a non-default stream
|
||||
@@ -502,7 +506,7 @@ def benchmark(
|
||||
# Assertion is >= since we may launch multiple kernels in one host function
|
||||
if num_nodes < warmup_iterations:
|
||||
raise ValueError(
|
||||
f"CUDA stream passed to benchmark does not match the stream the kernel was launched in"
|
||||
"CUDA stream passed to benchmark does not match the stream the kernel was launched in"
|
||||
)
|
||||
|
||||
# Capture profiling graph
|
||||
@@ -549,7 +553,6 @@ def benchmark(
|
||||
_cuda_success(err, "Error on destroying graph")
|
||||
|
||||
else:
|
||||
|
||||
if int(stream) != int(
|
||||
cuda_driver.CUstream_flags.CU_STREAM_DEFAULT
|
||||
) and not _does_kernel_use_stream(
|
||||
@@ -599,12 +602,474 @@ def get_workspace_count(
|
||||
:rtype: int
|
||||
"""
|
||||
num_l2_cache_bytes = cutlass.utils.HardwareInfo().get_l2_cache_size_in_bytes()
|
||||
return max(
|
||||
1,
|
||||
min(
|
||||
warmup_iterations + iterations, # Don't create more workspaces than needed
|
||||
(num_l2_cache_bytes + one_workspace_bytes - 1)
|
||||
// one_workspace_bytes, # Ceiling division
|
||||
),
|
||||
)
|
||||
num_workspaces = (num_l2_cache_bytes * 3) // one_workspace_bytes + 1
|
||||
num_iters = warmup_iterations + iterations
|
||||
return num_iters if num_iters < num_workspaces else num_workspaces
|
||||
|
||||
|
||||
#########################################
|
||||
# Autotuning/Tuning utilities
|
||||
#########################################
|
||||
|
||||
|
||||
def _benchmark_for_autotune(
|
||||
callable: Callable,
|
||||
*args,
|
||||
warmup_iterations: int,
|
||||
iterations: int,
|
||||
use_cold_l2: bool,
|
||||
print_verbose: bool,
|
||||
current_stream: Optional[cuda_driver.CUstream] = None,
|
||||
**kwargs,
|
||||
) -> float:
|
||||
"""Benchmarks a callable function with the specified parameters.
|
||||
|
||||
This function differs from the benchmark function in that it is used for autotuning. In this case we
|
||||
do not loop through workspaces to keep the L2 cache cold. Instead we rely on writing to an L2 cache sized address to keep the L2 cache cold.
|
||||
|
||||
The primary reason for doing this is that we do not have information on how to generate the workspaces for the kernel when autotuning.
|
||||
We also do not have information on how much memory the workspaces take up.
|
||||
|
||||
This benchmarking is done as a close approximation of the actual runtime of the kernel in an E2E system,
|
||||
where we may have clock throttling, a warm cache, or other factors that could affect the runtime of the kernel.
|
||||
|
||||
:param callable: The function to benchmark
|
||||
:type callable: Callable
|
||||
:param args: Arguments to pass to the callable function
|
||||
:param warmup_iterations: Number of warmup iterations, defaults to 10
|
||||
:type warmup_iterations: int, optional
|
||||
:param iterations: Number of benchmark iterations, defaults to 100
|
||||
:type iterations: int, optional
|
||||
:param use_cold_l2: Whether to clear L2 cache between runs, defaults to True
|
||||
:type use_cold_l2: bool, optional
|
||||
:param print_verbose: Whether to print verbose output, defaults to False
|
||||
:type print_verbose: bool, optional
|
||||
:param current_stream: Stream to benchmark in, defaults to CUDA stream default
|
||||
:type current_stream: CUstream, None
|
||||
:param kwargs: Additional keyword arguments to pass to the callable function
|
||||
|
||||
:return: The benchmark time in microseconds
|
||||
:rtype: float
|
||||
"""
|
||||
if current_stream is None:
|
||||
current_stream = cuda_driver.CUstream(
|
||||
cuda_driver.CUstream_flags.CU_STREAM_DEFAULT
|
||||
)
|
||||
|
||||
if int(current_stream) != int(
|
||||
cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT)
|
||||
) and not _does_kernel_use_stream(callable, current_stream, *args, **kwargs):
|
||||
raise ValueError(f"Incorrect stream passed to kernel: {current_stream}")
|
||||
|
||||
if use_cold_l2:
|
||||
from cutlass.utils import HardwareInfo
|
||||
|
||||
# use memset to clear L2 cache
|
||||
hardware_info = HardwareInfo()
|
||||
num_l2_cache_bytes = hardware_info.get_l2_cache_size_in_bytes()
|
||||
err, cache_ptr = cuda_driver.cuMemAlloc(int(num_l2_cache_bytes))
|
||||
_cuda_success(err, "Error on allocating memory")
|
||||
|
||||
# Create CUDA events for timing
|
||||
err, start_event = cuda_driver.cuEventCreate(
|
||||
cuda_driver.CUevent_flags.CU_EVENT_DEFAULT
|
||||
)
|
||||
_cuda_success(err, "Error on creating event")
|
||||
err, end_event = cuda_driver.cuEventCreate(
|
||||
cuda_driver.CUevent_flags.CU_EVENT_DEFAULT
|
||||
)
|
||||
_cuda_success(err, "Error on creating event")
|
||||
try:
|
||||
# warmup
|
||||
for _ in range(warmup_iterations):
|
||||
callable(*args, **kwargs)
|
||||
|
||||
time = 0
|
||||
execution_time_ms = []
|
||||
for _ in range(iterations):
|
||||
if use_cold_l2:
|
||||
# clear L2 cache by memset to zero for every run
|
||||
err = cuda_driver.cuMemsetD32Async(
|
||||
cache_ptr, 0, int(num_l2_cache_bytes // 4), current_stream
|
||||
)
|
||||
_cuda_success(err, "Error on memset")
|
||||
err = cuda_driver.cuEventRecord(start_event, current_stream)
|
||||
_cuda_success(err, "Error on recording event")
|
||||
callable(*args, **kwargs)
|
||||
err = cuda_driver.cuEventRecord(end_event, current_stream)
|
||||
_cuda_success(err, "Error on recording event")
|
||||
err = cuda_driver.cuEventSynchronize(end_event)
|
||||
_cuda_success(err, "Error on synchronizing event")
|
||||
err, elapsed_time = cuda_driver.cuEventElapsedTime(start_event, end_event)
|
||||
_cuda_success(err, "Error on querying event")
|
||||
execution_time_ms.append(elapsed_time)
|
||||
# unit: us
|
||||
time_us = sum(execution_time_ms) / len(execution_time_ms)
|
||||
except Exception as e:
|
||||
print(f"This config execution error: {e}")
|
||||
time_us = float("inf")
|
||||
if print_verbose:
|
||||
print(f"Execution time: {time_us:.4f} us")
|
||||
|
||||
if use_cold_l2:
|
||||
err = cuda_driver.cuMemFree(cache_ptr)
|
||||
_cuda_success(err, "Error on freeing memory")
|
||||
err = cuda_driver.cuEventDestroy(start_event)
|
||||
_cuda_success(err, "Error on destroying event")
|
||||
err = cuda_driver.cuEventDestroy(end_event)
|
||||
_cuda_success(err, "Error on destroying event")
|
||||
return time_us
|
||||
|
||||
|
||||
class autotune_jit:
|
||||
"""Auto-tuning tool supporting both dictionary and parameterized decorator styles.
|
||||
The autotune_jit class can be used as a decorator or a function.
|
||||
When used as a decorator, it will automatically tune the function based on the parameters.
|
||||
When used as a function, it will return a decorator that can be used to decorate a function.
|
||||
For example:
|
||||
.. code-block:: python
|
||||
|
||||
@autotune_jit(params_dict={'param1': [1, 2, 3], 'param2': [4, 5, 6]}, update_on_change=['param3'])
|
||||
@cute.jit
|
||||
def user_function(param1=1, param2=2, param3=3):
|
||||
# contents of the function
|
||||
pass
|
||||
|
||||
The function will be automatically tuned over all combinations of param1 and param2 whenever param3 changes .
|
||||
For non-specified parameters, the default value in user_function will be used (e.g., `param3` in `user_function`).
|
||||
.. code-block:: python
|
||||
user_function(a, b, c) # Autotunes code
|
||||
user_function(a, b, c) # This call pulls the best kernel from cache
|
||||
|
||||
Known Limitations:
|
||||
- Only supports functions that are decorated with cute.jit
|
||||
- If the function which is decorated with cute.jit is call method of a class, and the class has internal state that
|
||||
is used as constexpr arguments in the function, the autotuner will not be able to find the best configuration.
|
||||
|
||||
Note: The autotuner has the same semantics as cute.compile. If the function is compiled, but global variables are changed,
|
||||
the autotuner will not recompile the kernel.
|
||||
"""
|
||||
|
||||
logger = None
|
||||
|
||||
@classmethod
|
||||
def _initialize_logger(cls):
|
||||
"""Ensure the logger is initialized"""
|
||||
if cls.logger is None:
|
||||
cls.logger = logging.getLogger(__name__ + "_Autotune")
|
||||
if not cls.logger.handlers:
|
||||
handler = logging.StreamHandler()
|
||||
formatter = logging.Formatter(
|
||||
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
handler.setFormatter(formatter)
|
||||
cls.logger.addHandler(handler)
|
||||
if (
|
||||
os.environ.get("CUTE_DSL_LOG_AUTOTUNE") is not None
|
||||
and os.environ.get("CUTE_DSL_LOG_AUTOTUNE") != "0"
|
||||
):
|
||||
cls.logger.setLevel(logging.INFO)
|
||||
|
||||
@classmethod
|
||||
def _create_tuning_wrapper(
|
||||
cls, func, warmup_iterations, iterations, autotune_update_params
|
||||
):
|
||||
"""Create a wrapper function that performs auto-tuning
|
||||
|
||||
Args:
|
||||
func: Original function
|
||||
|
||||
Returns:
|
||||
Decorated wrapper function
|
||||
"""
|
||||
|
||||
# Initialize autotune parameters
|
||||
if not hasattr(func, "_autotune_params"):
|
||||
func._original_func = func
|
||||
func._autotune_params = {}
|
||||
func._autotune_update_params = autotune_update_params
|
||||
func._best_kernel = dict()
|
||||
func._best_config = dict()
|
||||
|
||||
# Create wrapper function for auto-tuning
|
||||
@functools.wraps(func)
|
||||
def tuning_wrapper(*args, **kwargs):
|
||||
parameters = inspect.signature(func._original_func).parameters.keys()
|
||||
tuning_key = list()
|
||||
for param_name in func._autotune_update_params:
|
||||
if param_name in kwargs.keys():
|
||||
tuning_key.append(kwargs[param_name])
|
||||
else:
|
||||
index = list(parameters).index(param_name)
|
||||
if index < len(args):
|
||||
tuning_key.append(args[index])
|
||||
tuning_key = tuple(tuning_key)
|
||||
if tuning_key in func._best_kernel.keys():
|
||||
cls.logger.info(
|
||||
f"Using cached best configuration: {func._best_config[tuning_key]}"
|
||||
)
|
||||
return func._best_kernel[tuning_key](*args, **kwargs)
|
||||
|
||||
# Get all parameter configurations
|
||||
params_dict = func._autotune_params
|
||||
keys = list(params_dict.keys())
|
||||
values = list(params_dict.values())
|
||||
|
||||
min_time = float("inf")
|
||||
|
||||
best_kernel = None
|
||||
# Record start time
|
||||
start = time()
|
||||
|
||||
# Iterate through all possible configuration combinations
|
||||
for config_values in product(*values):
|
||||
# Build current configuration
|
||||
current_config = dict(zip(keys, config_values))
|
||||
cls.logger.info(f"Tuning configuration: {current_config}")
|
||||
|
||||
try:
|
||||
# Call the original function, using current configuration to replace default parameters
|
||||
# For example, if current_config contains "cluster_shape_mn": (2, 1)
|
||||
# It will override func's default parameter value
|
||||
merged_kwargs = {**kwargs, **current_config}
|
||||
compiled_func = cute.compile(
|
||||
func._original_func, *args, **merged_kwargs
|
||||
)
|
||||
|
||||
# Detect which constexpr arguments we need to remove from args and merged_kwargs
|
||||
# This is done because after compiling our function signature will change, removing all constexpr arguments.
|
||||
indexes_to_remove = list()
|
||||
for arg in compiled_func.args_spec.get_constexpr_args():
|
||||
if arg["argument_name"] in merged_kwargs:
|
||||
del merged_kwargs[arg["argument_name"]]
|
||||
elif arg["argument_index"] is not None:
|
||||
indexes_to_remove.append(arg["argument_index"])
|
||||
if arg["argument_name"] not in func._autotune_update_params:
|
||||
# Handle the case where the programmer avoided autotuning over constexpr values, and
|
||||
# recompile in that case
|
||||
func._autotune_update_params.append(
|
||||
arg["argument_name"]
|
||||
)
|
||||
|
||||
# Remove constexpr arguments from args
|
||||
args_no_constexpr = list(args)
|
||||
for index in sorted(indexes_to_remove, reverse=True):
|
||||
del args_no_constexpr[index]
|
||||
|
||||
# Benchmark the compiled function
|
||||
cur_time = _benchmark_for_autotune(
|
||||
compiled_func,
|
||||
*args_no_constexpr,
|
||||
warmup_iterations=warmup_iterations,
|
||||
iterations=iterations,
|
||||
use_cold_l2=True,
|
||||
print_verbose=False,
|
||||
**merged_kwargs,
|
||||
)
|
||||
|
||||
cls.logger.info(f" Execution time: {cur_time} us")
|
||||
|
||||
# Update best results
|
||||
if cur_time < min_time:
|
||||
min_time = cur_time
|
||||
best_kernel = compiled_func
|
||||
best_config = current_config
|
||||
|
||||
except NotImplementedError as e:
|
||||
cls.logger.info(
|
||||
f" Encountered unimplemented error, abort execution: {e}"
|
||||
)
|
||||
raise e
|
||||
except (ValueError, TypeError) as e:
|
||||
cls.logger.info(f" Configuration parameter skipping: {e}")
|
||||
raise e
|
||||
continue
|
||||
except Exception as e:
|
||||
cls.logger.info(f" Execution error skipping: {e}")
|
||||
raise e
|
||||
continue
|
||||
|
||||
end = time()
|
||||
tuning_time = end - start
|
||||
|
||||
if best_kernel is None:
|
||||
raise ValueError("No best kernel found")
|
||||
|
||||
cls.logger.info(
|
||||
f"Best configuration: {best_config}, execution time: {min_time} us"
|
||||
)
|
||||
cls.logger.info(f"Total tuning time: {tuning_time} s")
|
||||
func._best_kernel[tuning_key] = best_kernel
|
||||
func._best_config[tuning_key] = best_config
|
||||
return best_kernel(*args, **kwargs)
|
||||
|
||||
# Append autotune wrapper to not conflict with the jit kernel names
|
||||
tuning_wrapper.__name__ = func.__name__ + "_autotune_wrapper"
|
||||
tuning_wrapper.__qualname__ = func.__qualname__ + "_autotune_wrapper"
|
||||
|
||||
return tuning_wrapper
|
||||
|
||||
return func # If already has a wrapper, return the original function
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params_dict: Dict[str, List[Any]] = None,
|
||||
update_on_change: List[str] = None,
|
||||
warmup_iterations=10,
|
||||
iterations=100,
|
||||
):
|
||||
"""Initialize the autotune_jit decorator.
|
||||
|
||||
:param params_dict: Dictionary containing parameter names and their possible values
|
||||
:type params_dict: Dict[str, List[Any]], optional
|
||||
:param update_on_change: Whether to retune when the parameters changes, defaults to None
|
||||
:type update_on_change: bool, optional
|
||||
:param warmup_iterations: Number of warmup iterations, defaults to 100
|
||||
:type warmup_iterations: int, optional
|
||||
:param iterations: Number of benchmark iterations, defaults to 100
|
||||
:type iterations: int, optional
|
||||
"""
|
||||
# Initialize logger
|
||||
self._initialize_logger()
|
||||
|
||||
# Save parameter dictionary
|
||||
self.params_dict = params_dict or {}
|
||||
self.update_on_change = update_on_change or list()
|
||||
|
||||
# Save iterations
|
||||
self.warmup_iterations = warmup_iterations
|
||||
self.iterations = iterations
|
||||
|
||||
def __call__(self, func):
|
||||
"""Called when class instance is used as a decorator.
|
||||
|
||||
:param func: Function to be decorated
|
||||
:type func: Callable
|
||||
:return: Decorated function
|
||||
:rtype: Callable
|
||||
"""
|
||||
# Create wrapper function
|
||||
decorated_func = self._create_tuning_wrapper(
|
||||
func, self.warmup_iterations, self.iterations, self.update_on_change
|
||||
)
|
||||
|
||||
# Use the wrapper if it exists, otherwise use the original function
|
||||
result_func = (
|
||||
decorated_func if hasattr(decorated_func, "_autotune_params") else func
|
||||
)
|
||||
|
||||
# Add parameters from the dictionary to the function's autotune parameters
|
||||
for param_name, param_values in self.params_dict.items():
|
||||
result_func._autotune_params[param_name] = param_values
|
||||
|
||||
return result_func
|
||||
|
||||
|
||||
def tune(
|
||||
func: Callable[[Any], Callable[[], Any]],
|
||||
params_dict: Dict[str, List[Any]] = None,
|
||||
kernel_arguments: JitArguments = JitArguments(),
|
||||
warmup_iterations=10,
|
||||
iterations=100,
|
||||
stream: Optional[cuda_driver.CUstream] = None,
|
||||
) -> Dict[str, Any]:
|
||||
"""Tuning tool to suport arbitrary functions. The user must provide a function that returns a callable, which
|
||||
takes no arguments to be tuned over.
|
||||
Best practice is to return a jit function that is compiled with cute.compile for optimal performance.
|
||||
For example:
|
||||
.. code-block:: python
|
||||
|
||||
def user_function(param1=1, param2=2, param3=3) -> Callable[[], Any]:
|
||||
# contents of the function
|
||||
return lambda : compiled_func(param1, param2, param3)
|
||||
|
||||
config = tune(user_function, params_dict={'param1': [1, 2, 3], 'param2': [4, 5, 6]}, update_on_change=['param3'])
|
||||
|
||||
:param func: Function to be tuned, note that errors raised in the function will be ignored and the next configuration will be tried.
|
||||
:type func: Callable[[Any], Callable[[], Any]]
|
||||
:param params_dict: Dictionary containing parameter names and their possible values
|
||||
:type params_dict: Dict[str, List[Any]], optional
|
||||
:param kernel_arguments: Kernel arguments to launch callable with, defaults to JitArguments()
|
||||
:type kernel_arguments: JitArguments, optional
|
||||
:param warmup_iterations: Number of warmup iterations, defaults to 10
|
||||
:type warmup_iterations: int, optional
|
||||
:param iterations: Number of benchmark iterations, defaults to 100
|
||||
:type iterations: int, optional
|
||||
:param stream: Stream kernel is launched in, defaults to CUDA stream default
|
||||
:type stream: CUstream, None
|
||||
:return: Best configuration
|
||||
:rtype: Dict[str, Any]
|
||||
"""
|
||||
logger = logging.getLogger(__name__ + "_Autotune")
|
||||
if not logger.handlers:
|
||||
handler = logging.StreamHandler()
|
||||
formatter = logging.Formatter(
|
||||
"%(asctime)s - %(name)s - %(levelname)s - %(message)s"
|
||||
)
|
||||
handler.setFormatter(formatter)
|
||||
logger.addHandler(handler)
|
||||
if (
|
||||
os.environ.get("CUTE_DSL_LOG_AUTOTUNE") is not None
|
||||
and os.environ.get("CUTE_DSL_LOG_AUTOTUNE") != "0"
|
||||
):
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
if stream is None:
|
||||
stream = cuda_driver.CUstream(cuda_driver.CUstream_flags.CU_STREAM_DEFAULT)
|
||||
|
||||
# Get all parameter configurations
|
||||
keys = list(params_dict.keys())
|
||||
values = list(params_dict.values())
|
||||
|
||||
min_time = float("inf")
|
||||
|
||||
best_config = None
|
||||
# Record start time
|
||||
start = time()
|
||||
|
||||
# Iterate through all possible configuration combinations
|
||||
for config_values in product(*values):
|
||||
# Build current configuration
|
||||
current_config = dict(zip(keys, config_values))
|
||||
logger.info(f"Tuning configuration: {current_config}")
|
||||
|
||||
try:
|
||||
merged_kwargs = {**kernel_arguments.kwargs, **current_config}
|
||||
|
||||
compiled_func = func(*kernel_arguments.args, **merged_kwargs)
|
||||
# Benchmark the compiled function
|
||||
cur_time = _benchmark_for_autotune(
|
||||
compiled_func,
|
||||
warmup_iterations=warmup_iterations,
|
||||
iterations=iterations,
|
||||
use_cold_l2=True,
|
||||
print_verbose=False,
|
||||
current_stream=stream,
|
||||
)
|
||||
|
||||
logger.info(f" Execution time: {cur_time} us")
|
||||
|
||||
# Update best results
|
||||
if cur_time < min_time:
|
||||
min_time = cur_time
|
||||
best_config = current_config
|
||||
|
||||
except NotImplementedError as e:
|
||||
logger.info(f" Encountered unimplemented error, abort execution: {e}")
|
||||
raise e
|
||||
except (ValueError, TypeError) as e:
|
||||
logger.info(f" Configuration parameter skipping: {e}")
|
||||
continue
|
||||
except Exception as e:
|
||||
logger.info(f" Execution error skipping: {e}")
|
||||
continue
|
||||
|
||||
end = time()
|
||||
tuning_time = end - start
|
||||
|
||||
if best_config is None:
|
||||
raise ValueError("No best kernel found")
|
||||
|
||||
logger.info(f"Best configuration: {best_config}, execution time: {min_time} us")
|
||||
logger.info(f"Total tuning time: {tuning_time} s")
|
||||
return best_config
|
||||
|
||||
331
python/CuTeDSL/cutlass/cute/tuple.py
Normal file
331
python/CuTeDSL/cutlass/cute/tuple.py
Normal file
@@ -0,0 +1,331 @@
|
||||
# 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 inspect import signature
|
||||
from itertools import chain
|
||||
from typing import Any, Callable, Union, Tuple, List, Iterable
|
||||
|
||||
from cutlass.cutlass_dsl import is_dynamic_expression, dsl_user_op
|
||||
from cutlass._mlir import ir
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
|
||||
from .typing import XTuple, IntTuple, Shape, Coord, Boolean, is_integer
|
||||
|
||||
|
||||
def wrap(x) -> Tuple[Any, ...]:
|
||||
"""
|
||||
Wraps the input into a tuple if not a tuple.
|
||||
"""
|
||||
if isinstance(x, tuple):
|
||||
return x
|
||||
return (x,)
|
||||
|
||||
|
||||
def flatten_to_tuple(a: XTuple) -> Tuple[Any, ...]:
|
||||
"""Flattens a potentially nested tuple structure into a flat tuple.
|
||||
|
||||
This function recursively traverses the input structure and flattens it into
|
||||
a single-level tuple, preserving the order of elements.
|
||||
|
||||
:param a: The structure to flatten
|
||||
:type a: Union[IntTuple, Coord, Shape, Stride]
|
||||
:return: A flattened tuple containing all elements from the input
|
||||
:rtype: tuple
|
||||
|
||||
**Examples:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
flatten_to_tuple((1, 2, 3)) # Returns (1, 2, 3)
|
||||
flatten_to_tuple(((1, 2), 3)) # Returns (1, 2, 3)
|
||||
flatten_to_tuple((1, (2, (3,)))) # Returns (1, 2, 3)
|
||||
"""
|
||||
if not isinstance(a, tuple):
|
||||
return wrap(a)
|
||||
else:
|
||||
return tuple(chain.from_iterable(tuple(flatten_to_tuple(x) for x in a)))
|
||||
|
||||
|
||||
def unflatten(
|
||||
sequence: Union[Tuple[Any, ...], List[Any], Iterable[Any]], profile: XTuple
|
||||
) -> XTuple:
|
||||
"""Unflatten a flat tuple into a nested tuple structure according to a profile.
|
||||
|
||||
This function transforms a flat sequence of elements into a nested tuple structure
|
||||
that matches the structure defined by the profile parameter. It traverses the profile
|
||||
structure and populates it with elements from the sequence.
|
||||
|
||||
sequence must be long enough to fill the profile. Raises RuntimeError if it is not.
|
||||
|
||||
:param sequence: A flat sequence of elements to be restructured
|
||||
:type sequence: Union[Tuple[Any, ...], List[Any], Iterable[Any]]
|
||||
:param profile: A nested tuple structure that defines the shape of the output
|
||||
:type profile: XTuple
|
||||
:return: A nested tuple with the same structure as profile but containing elements from sequence
|
||||
:rtype: XTuple
|
||||
|
||||
**Examples:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
unflatten([1, 2, 3, 4], ((0, 0), (0, 0))) # Returns ((1, 2), (3, 4))
|
||||
"""
|
||||
|
||||
def _make_generator():
|
||||
for element in sequence:
|
||||
yield element
|
||||
|
||||
xs = _make_generator()
|
||||
return transform_leaf(lambda _: next(xs), profile)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def product(a: Union[IntTuple, Shape], *, loc=None, ip=None):
|
||||
# Local import to avoid circular dependency
|
||||
from .core import _pack_int_tuple, _unpack_x_tuple
|
||||
|
||||
"""Return product of the given IntTuple or Shape.
|
||||
|
||||
Computes the product of all elements in the input tuple or shape.
|
||||
|
||||
Returns static value if type is static otherwise dynamic value.
|
||||
|
||||
:param a: The input tuple or shape
|
||||
:type a: IntTuple or Shape
|
||||
:param loc: Source location for MLIR, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: Insertion point, defaults to None
|
||||
:type ip: optional
|
||||
:return: Static product of IntTuple or Shape if static, otherwise a Value
|
||||
:rtype: int or Value
|
||||
:raises TypeError: If input is not an IntTuple or Shape
|
||||
"""
|
||||
if is_integer(a):
|
||||
return a
|
||||
if isinstance(a, tuple):
|
||||
a_val = _pack_int_tuple(a, loc=loc, ip=ip)
|
||||
res = _cute_ir.tuple_product(a_val, loc=loc, ip=ip)
|
||||
return _unpack_x_tuple(res, loc=loc, ip=ip)
|
||||
else:
|
||||
raise TypeError(f"expects IntTuple or Shape, but got {type(a)}")
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def product_like(a: IntTuple, target_profile: XTuple, *, loc=None, ip=None) -> IntTuple:
|
||||
"""Return product of the given IntTuple or Shape at leaves of `target_profile`.
|
||||
|
||||
This function computes products according to the structure defined by target_profile.
|
||||
|
||||
:param a: The input tuple or shape
|
||||
:type a: IntTuple or Shape
|
||||
:param target_profile: The profile that guides how products are computed
|
||||
:type target_profile: XTuple
|
||||
:param loc: Source location for MLIR, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: Insertion point, defaults to None
|
||||
:type ip: optional
|
||||
:return: The resulting tuple with products computed according to target_profile
|
||||
:rtype: IntTuple or Shape
|
||||
:raises TypeError: If inputs have incompatible types
|
||||
:raises ValueError: If inputs have incompatible shapes
|
||||
"""
|
||||
# Perform product at leaf of `target_profile`
|
||||
if not isinstance(target_profile, tuple):
|
||||
return product(a, loc=loc, ip=ip)
|
||||
|
||||
if not isinstance(a, tuple):
|
||||
raise TypeError(f"expects `a` tuple but got {a}")
|
||||
|
||||
if len(a) != len(target_profile):
|
||||
raise ValueError("expects `a` and `guide` have the same rank")
|
||||
|
||||
return tuple(product_like(x, g, loc=loc, ip=ip) for x, g in zip(a, target_profile))
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def product_each(a: IntTuple, *, loc=None, ip=None) -> IntTuple:
|
||||
from .core import _pack_int_tuple, _unpack_x_tuple
|
||||
|
||||
"""Compute products for each component of the input.
|
||||
|
||||
Returns a rank(a) tuple result such that ``get(result, mode=[i]) == product(get(a, mode=[i]))``
|
||||
|
||||
:param a: The input IntTuple or Shape
|
||||
:type a: IntTuple or Shape
|
||||
:param loc: Source location for MLIR, defaults to None
|
||||
:type loc: optional
|
||||
:param ip: Insertion point, defaults to None
|
||||
:type ip: optional
|
||||
:return: A tuple containing products for each component
|
||||
:rtype: tuple
|
||||
:raises TypeError: If input is not an IntTuple or Shape
|
||||
"""
|
||||
if is_integer(a):
|
||||
return a
|
||||
|
||||
if not isinstance(a, tuple):
|
||||
raise TypeError(f"expects IntTuple or Shape, but got {type(a)}")
|
||||
|
||||
if a == ():
|
||||
return 1
|
||||
|
||||
a_val = _pack_int_tuple(a, loc=loc, ip=ip)
|
||||
res = _cute_ir.tuple_product_each(a_val, loc=loc, ip=ip)
|
||||
return _unpack_x_tuple(res, loc=loc, ip=ip)
|
||||
|
||||
|
||||
def find_if(
|
||||
t: Union[tuple, ir.Value, int],
|
||||
pred_fn: Callable[[int, Tuple[int, ...]], bool],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Union[int, Tuple[int, ...], None]:
|
||||
from .core import rank, get
|
||||
|
||||
"""Find the first position in t where pred_fn(val, pos) returns True.
|
||||
|
||||
:param t: The search space
|
||||
:type t: Union[tuple, ir.Value, int]
|
||||
:param pred_fn: A callable object (lambda, function, etc.) that predicates the value and position in t.
|
||||
It takes the current leaf value and position, returns True if the value or position is satisfied.
|
||||
:type pred_fn: Callable[[int, Tuple[int, ...]], bool]
|
||||
:return: Index if found at top level, tuple of indices showing nested position, or None if not found
|
||||
:rtype: Union[int, Tuple[int, ...], None]
|
||||
|
||||
**Examples:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Find the first position of x in t
|
||||
t = (3, 4)
|
||||
find_if(t, pred_fn=lambda val, pos: val == x)
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# find the leading dimension
|
||||
shape = (3, 4)
|
||||
stride = (4, 1)
|
||||
# Find value 1 in stride where the corresponding shape is not 1
|
||||
def pred_fn(val, pos):
|
||||
mode = [pos] if isinstance(pos, int) else list(pos)
|
||||
return val == 1 and get(shape, mode) != 1
|
||||
find_if(stride, pred_fn=pred_fn)
|
||||
"""
|
||||
|
||||
def _find_if_impl(curr, pos, *, loc=None, ip=None):
|
||||
if isinstance(curr, tuple):
|
||||
# Recursively search nested tuple
|
||||
for i in range(rank(curr)):
|
||||
sub_curr = get(curr, mode=[i], loc=loc, ip=ip)
|
||||
sub_pos = (pos, i) if isinstance(pos, int) else pos + (i,)
|
||||
res_pos = _find_if_impl(sub_curr, sub_pos, loc=loc, ip=ip)
|
||||
if res_pos is not None:
|
||||
return res_pos
|
||||
else:
|
||||
# For leaf values, check if it matches x
|
||||
if pred_fn(curr, pos):
|
||||
return pos
|
||||
return None
|
||||
|
||||
def _check_pred_fn():
|
||||
if not callable(pred_fn):
|
||||
raise TypeError(f"pred_fn must be callable, but got {type(pred_fn)}")
|
||||
|
||||
sig = signature(pred_fn)
|
||||
if len(sig.parameters) != 2:
|
||||
raise ValueError(
|
||||
f"pred_fn must have two parameters (value, pos), but got {len(sig.parameters)}"
|
||||
)
|
||||
|
||||
_check_pred_fn()
|
||||
|
||||
for i in range(rank(t)):
|
||||
curr = get(t, mode=[i], loc=loc, ip=ip)
|
||||
res_pos = _find_if_impl(curr, i, loc=loc, ip=ip)
|
||||
if res_pos is not None:
|
||||
return res_pos
|
||||
return None
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def find(
|
||||
t: Union[tuple, ir.Value, int], x: int, *, loc=None, ip=None
|
||||
) -> Union[int, Tuple[int, ...], None]:
|
||||
"""Find the first position of a value ``x`` in a hierarchical structure ``t``.
|
||||
|
||||
Searches for the first occurrence of x in t, optionally excluding positions
|
||||
where a comparison value matches. The search can traverse nested structures
|
||||
and returns either a single index or a tuple of indices for nested positions.
|
||||
|
||||
:param t: The search space
|
||||
:type t: Union[tuple, ir.Value, int]
|
||||
:param x: The static integer x to search for
|
||||
:type x: int
|
||||
:return: Index if found at top level, tuple of indices showing nested position, or None if not found
|
||||
:rtype: Union[int, Tuple[int, ...], None]
|
||||
"""
|
||||
if not isinstance(x, int):
|
||||
raise TypeError(f"find() requires a static x to search for, but got {x}")
|
||||
|
||||
def pred_fn(val, pos):
|
||||
# Skip dynamic values which can't be compared
|
||||
return not is_dynamic_expression(val) and val == x
|
||||
|
||||
return find_if(t, pred_fn=pred_fn, loc=loc, ip=ip)
|
||||
|
||||
|
||||
def transform_leaf(f, *args):
|
||||
"""
|
||||
Apply a function to the leaf nodes of nested tuple structures.
|
||||
|
||||
This function traverses nested tuple structures in parallel and applies the function f
|
||||
to corresponding leaf nodes. All input tuples must have the same nested structure.
|
||||
|
||||
:param f: Function to apply to leaf nodes
|
||||
:type f: Callable
|
||||
:param args: One or more nested tuple structures with matching profiles
|
||||
:return: A new nested tuple with the same structure as the inputs, but with leaf values transformed by f
|
||||
:raises TypeError: If the input tuples have different nested structures
|
||||
|
||||
**Example:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
>>> transform_leaf(lambda x: x + 1, (1, 2))
|
||||
(2, 3)
|
||||
>>> transform_leaf(lambda x, y: x + y, (1, 2), (3, 4))
|
||||
(4, 6)
|
||||
>>> transform_leaf(lambda x: x * 2, ((1, 2), (3, 4)))
|
||||
((2, 4), (6, 8))
|
||||
"""
|
||||
if all(isinstance(t, tuple) for t in args):
|
||||
return tuple(transform_leaf(f, *_args) for _args in zip(*args))
|
||||
elif all(not isinstance(t, tuple) for t in args):
|
||||
return f(*args)
|
||||
else:
|
||||
raise TypeError(f"profile of input tuples doesn't match: {args}")
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def elem_less(
|
||||
lhs: Union[Shape, IntTuple, Coord],
|
||||
rhs: Union[Shape, IntTuple, Coord],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Boolean:
|
||||
from .core import _pack_coord
|
||||
|
||||
# Coord is super set of IntTuple and Shape
|
||||
lhs_val = _pack_coord(lhs, loc=loc, ip=ip)
|
||||
rhs_val = _pack_coord(rhs, loc=loc, ip=ip)
|
||||
return Boolean(_cute_ir.elem_less(lhs_val, rhs_val, loc=loc, ip=ip))
|
||||
@@ -10,13 +10,12 @@
|
||||
# is strictly prohibited.
|
||||
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import ForwardRef, Tuple, Union, Any, Type, List
|
||||
from typing import ForwardRef, Tuple, Union, Any, Type, List, Optional
|
||||
|
||||
from cutlass.base_dsl.typing import *
|
||||
|
||||
from cutlass._mlir import ir
|
||||
import cutlass._mlir.extras.types as T
|
||||
from cutlass._mlir.dialects.cute import AddressSpace
|
||||
from cutlass._mlir.dialects.cute import AddressSpace, ConstrainedIntType
|
||||
|
||||
|
||||
Int = Union[int, Integer]
|
||||
@@ -24,7 +23,6 @@ Int = Union[int, Integer]
|
||||
|
||||
ScaledBasis = ForwardRef("ScaledBasis")
|
||||
|
||||
|
||||
IntTuple = Union[Int, Tuple["IntTuple", ...]]
|
||||
Shape = Union[Int, Tuple["Shape", ...]]
|
||||
Stride = Union[Int, ScaledBasis, Tuple["Stride", ...]]
|
||||
@@ -35,7 +33,7 @@ class Layout(ir.Value):
|
||||
def __init__(self, op_result):
|
||||
super().__init__(op_result)
|
||||
|
||||
def __str__(self): ...
|
||||
def __str__(self) -> str: ...
|
||||
|
||||
def get_hier_coord(self, idx) -> Coord:
|
||||
"""Return the (hierarchical) ND logical coordinate corresponding to the linear index"""
|
||||
@@ -48,12 +46,95 @@ class Layout(ir.Value):
|
||||
def stride(self, *, loc=None, ip=None) -> Stride: ...
|
||||
|
||||
|
||||
class ComposedLayout(ABC):
|
||||
r"""ComposedLayout represents the functional composition of layouts in CuTe.
|
||||
|
||||
**Formally:**
|
||||
|
||||
.. math::
|
||||
|
||||
R(c) := (inner \circ offset \circ outer)(c) := inner(offset + outer(c))
|
||||
|
||||
where:
|
||||
|
||||
- inner: The inner layout or swizzle that is applied last
|
||||
- offset: An integer tuple representing a coordinate offset
|
||||
- outer: The outer layout that is applied first
|
||||
|
||||
This composition allows for complex transformations of coordinates and indices,
|
||||
enabling operations like tiling, partitioning, and reshaping of data.
|
||||
|
||||
:ivar inner: The inner layout or swizzle component
|
||||
:ivar offset: The coordinate offset applied between inner and outer layouts
|
||||
:ivar outer: The outer layout component
|
||||
:ivar max_alignment: The maximum alignment of the composed layout
|
||||
|
||||
**Examples:**
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Create a composed layout with inner layout, offset, and outer layout
|
||||
|
||||
# inner layout: (4, 8):(1, 4)
|
||||
inner_layout = make_layout((4, 8))
|
||||
|
||||
offset = (0, 0)
|
||||
|
||||
# outer layout: (2, 2):(1@0, 1@1)
|
||||
outer_layout = make_layout((2, 2), stride=(1 * E(0), 1 * E(1)))
|
||||
|
||||
# composed layout: (inner o offset o outer)
|
||||
composed = make_composed_layout(inner_layout, offset, outer_layout)
|
||||
|
||||
# Accessing components of the composed layout
|
||||
inner = composed.inner
|
||||
offset = composed.offset
|
||||
outer = composed.outer
|
||||
|
||||
# map coordinate (0, 1) to linear index
|
||||
# - outer(0, 1) = (0, 1)
|
||||
# - offset + outer(0, 1) = (0, 1)
|
||||
# - inner(0, 1) = 0 * 1 + 1 * 4 = 4
|
||||
idx = crd2idx((0, 1), composed)
|
||||
|
||||
# Composition is used in many tiling operations
|
||||
# For example, in logical_product, raked_product, and blocked_product
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def type(self) -> ir.Type: ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def is_normal(self) -> bool: ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def inner(self, *, loc=None, ip=None): ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def offset(self, *, loc=None, ip=None) -> IntTuple: ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def outer(self, *, loc=None, ip=None) -> Layout: ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def shape(self, *, loc=None, ip=None): ...
|
||||
|
||||
@abstractmethod
|
||||
def __call__(self, coord: Coord, loc=None, ip=None) -> IntTuple: ...
|
||||
|
||||
|
||||
Tile = Union[Int, None, Layout, Tuple["Tile", ...]]
|
||||
Tiler = Union[Shape, Layout, Tile]
|
||||
|
||||
|
||||
# XTuple is super set of above types
|
||||
XTuple = Union[IntTuple, Shape, Stride, Coord, Tile]
|
||||
|
||||
Tiler = Union[Shape, Layout, Tile]
|
||||
XTuple = Union[Any, Tuple["XTuple", ...]]
|
||||
|
||||
|
||||
class Pointer(ABC):
|
||||
@@ -70,6 +151,8 @@ class Pointer(ABC):
|
||||
|
||||
def align(self, min_align: int) -> "Pointer": ...
|
||||
|
||||
def __add__(self, other: int, *, loc=None, ip=None) -> "Pointer": ...
|
||||
|
||||
def __get_mlir_types__(self) -> List[ir.Type]: ...
|
||||
|
||||
def __extract_mlir_values__(self) -> List[ir.Value]: ...
|
||||
@@ -78,51 +161,69 @@ class Pointer(ABC):
|
||||
|
||||
|
||||
class Tensor(ABC):
|
||||
"""
|
||||
Abstract base class for CuTe jit function and runtime _Tensor
|
||||
r"""Abstract base class for Tensor representations in CuTe DSL.
|
||||
|
||||
A CuTe Tensor is iterator with layout
|
||||
A CuTe Tensor is iterator with layout. A tensor evaluates the layout by mapping a
|
||||
coordinate to the codomain, offsets the iterator accordingly, and dereferences
|
||||
the result to obtain the tensor's value.
|
||||
|
||||
:Examples:
|
||||
**Formally:**
|
||||
|
||||
.. math::
|
||||
|
||||
T(c) = (E \circ L)(c) = *(E + L(c))
|
||||
|
||||
where
|
||||
|
||||
- :math:`E` is the iterator/engine
|
||||
- :math:`L` is the layout
|
||||
|
||||
**Notes:**
|
||||
|
||||
- The tensor supports both direct element access via coordinates and slicing operations
|
||||
- Load/store operations are only supported for specific memory spaces (rmem, smem, gmem, generic)
|
||||
- For composed layouts, stride information is not directly accessible
|
||||
- Dynamic layouts do not support vector load/store operations
|
||||
|
||||
**Examples:**
|
||||
|
||||
Create tensor from torch.tensor with Host Runtime:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
>>> import torch
|
||||
>>> from cutlass.cute.runtime import from_dlpack
|
||||
>>> mA = from_dlpack(torch.tensor([1, 3, 5], dtype=torch.int32))
|
||||
>>> mA.shape
|
||||
(3,)
|
||||
>>> mA.stride
|
||||
(1,)
|
||||
>>> mA.layout
|
||||
(3,):(1,)
|
||||
import torch
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
|
||||
mA = from_dlpack(torch.tensor([1, 3, 5], dtype=torch.int32))
|
||||
print(mA.shape) # (3,)
|
||||
print(mA.stride) # (1,)
|
||||
print(mA.layout) # (3,):(1,)
|
||||
|
||||
Define JIT function:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def add(a: Tensor, b: Tensor, res: Tensor): ...
|
||||
def add(a: Tensor, b: Tensor, res: Tensor):
|
||||
res.store(a.load() + b.load())
|
||||
|
||||
Call JIT function from python:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
>>> import torch
|
||||
>>> a = torch.tensor([1, 3, 5], dtype=torch.int32)
|
||||
>>> b = torch.tensor([2, 4, 6], dtype=torch.int32)
|
||||
>>> c = torch.zeros([3], dtype=torch.int32)
|
||||
>>> mA = from_dlpack(a)
|
||||
>>> mB = from_dlpack(b)
|
||||
>>> mC = from_dlpack(c)
|
||||
>>> add(mA, mB, mC)
|
||||
>>> c
|
||||
tensor([3, 7, 11], dtype=torch.int32)
|
||||
import torch
|
||||
a = torch.tensor([1, 3, 5], dtype=torch.int32)
|
||||
b = torch.tensor([2, 4, 6], dtype=torch.int32)
|
||||
c = torch.zeros([3], dtype=torch.int32)
|
||||
mA = from_dlpack(a)
|
||||
mB = from_dlpack(b)
|
||||
mC = from_dlpack(c)
|
||||
add(mA, mB, mC)
|
||||
print(c) # tensor([3, 7, 11], dtype=torch.int32)
|
||||
"""
|
||||
|
||||
def __str__(self): ...
|
||||
@abstractmethod
|
||||
def __str__(self) -> str: ...
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, idx) -> Union["Tensor", ir.Value, IntTuple]: ...
|
||||
@@ -143,7 +244,7 @@ class Tensor(ABC):
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def iterator(self): ...
|
||||
def iterator(self) -> Union[Pointer, IntTuple]: ...
|
||||
|
||||
@property
|
||||
def layout(self) -> Union[Layout, "ComposedLayout"]: ...
|
||||
@@ -151,16 +252,19 @@ class Tensor(ABC):
|
||||
@property
|
||||
def shape(self) -> Shape: ...
|
||||
|
||||
@property
|
||||
def stride(self) -> Stride: ...
|
||||
|
||||
def load(self, *, loc=None, ip=None) -> "TensorSSA": ...
|
||||
|
||||
def store(self, data: "TensorSSA", *, loc=None, ip=None): ...
|
||||
|
||||
def mark_layout_dynamic(self, leading_dim: int | None = None) -> "Tensor": ...
|
||||
def mark_layout_dynamic(self, leading_dim: Optional[int] = None) -> "Tensor": ...
|
||||
|
||||
def mark_compact_shape_dynamic(
|
||||
self,
|
||||
mode: int,
|
||||
stride_order: tuple[int, ...] | None = None,
|
||||
stride_order: Optional[tuple[int, ...]] = None,
|
||||
divisibility: int = 1,
|
||||
) -> "Tensor": ...
|
||||
|
||||
@@ -168,11 +272,27 @@ class Tensor(ABC):
|
||||
def fill(self, value: Numeric) -> None: ...
|
||||
|
||||
|
||||
def is_integer(a) -> bool:
|
||||
"""Check if an object is static integer or dynamic integer"""
|
||||
return isinstance(a, (int, Integer)) or (
|
||||
isinstance(a, ir.Value)
|
||||
and isinstance(a.type, (ir.IntegerType, ConstrainedIntType))
|
||||
)
|
||||
|
||||
|
||||
def is_int_tuple(a) -> bool:
|
||||
if isinstance(a, tuple):
|
||||
return all([is_int_tuple(x) for x in a])
|
||||
else:
|
||||
return is_integer(a)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Coord",
|
||||
"Numeric",
|
||||
"Integer",
|
||||
"Boolean",
|
||||
"Int4",
|
||||
"Int8",
|
||||
"Int16",
|
||||
"Int32",
|
||||
@@ -204,4 +324,6 @@ __all__ = [
|
||||
"Tile",
|
||||
"Tiler",
|
||||
"XTuple",
|
||||
"is_integer",
|
||||
"is_int_tuple",
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user