Release v4.0.0 (#2294)
This commit is contained in:
Kihiro Bando
@@ -0,0 +1,310 @@
|
||||
# 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.
|
||||
|
||||
# Use the auto-generated enum AddressSpace
|
||||
from cutlass._mlir.dialects.cute import AddressSpace
|
||||
|
||||
# Explicitly import types that might be directly used by other modules.
|
||||
# This is a fix for using Sphinx to generate documentation
|
||||
# Because Sphinx processes each module in isolation, it won't be able to rely
|
||||
# on re-exported symbols via wildcard imports (from .typing import *) in the
|
||||
# same way that Python does at runtime.
|
||||
from .typing import (
|
||||
Shape,
|
||||
Stride,
|
||||
IntTuple,
|
||||
Coord,
|
||||
Tile,
|
||||
XTuple,
|
||||
Tiler,
|
||||
Layout,
|
||||
Pointer,
|
||||
Tensor,
|
||||
)
|
||||
|
||||
# Import everything else
|
||||
from .typing import *
|
||||
|
||||
from .core import (
|
||||
assume,
|
||||
is_integer,
|
||||
is_int_tuple,
|
||||
is_static,
|
||||
size,
|
||||
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_composed_layout,
|
||||
make_layout_tv,
|
||||
make_swizzle,
|
||||
recast_ptr,
|
||||
make_tensor,
|
||||
make_identity_tensor,
|
||||
make_fragment,
|
||||
recast_tensor,
|
||||
get,
|
||||
select,
|
||||
front,
|
||||
is_major,
|
||||
find,
|
||||
coalesce,
|
||||
group_modes,
|
||||
cosize,
|
||||
dice,
|
||||
product_each,
|
||||
prepend,
|
||||
append,
|
||||
prepend_ones,
|
||||
append_ones,
|
||||
ceil_div,
|
||||
slice_and_offset,
|
||||
crd2idx,
|
||||
domain_offset,
|
||||
elem_less,
|
||||
transform_leaf,
|
||||
filter_zeros,
|
||||
filter,
|
||||
tile_to_shape,
|
||||
shape_div,
|
||||
composition,
|
||||
complement,
|
||||
right_inverse,
|
||||
left_inverse,
|
||||
max_common_layout,
|
||||
max_common_vector,
|
||||
logical_product,
|
||||
zipped_product,
|
||||
tiled_product,
|
||||
flat_product,
|
||||
raked_product,
|
||||
blocked_product,
|
||||
flat_divide,
|
||||
logical_divide,
|
||||
zipped_divide,
|
||||
tiled_divide,
|
||||
local_partition,
|
||||
local_tile,
|
||||
printf,
|
||||
print_tensor,
|
||||
# tiled mma/tiled copy
|
||||
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_C_atom,
|
||||
basic_copy,
|
||||
basic_copy_if,
|
||||
autovec_copy,
|
||||
copy,
|
||||
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 . import arch
|
||||
from . import nvgpu
|
||||
from . import testing
|
||||
from . import runtime
|
||||
|
||||
# Export all math ops without "math."
|
||||
from .math import *
|
||||
|
||||
# Used as internal symbol
|
||||
from .. import cutlass_dsl as _dsl
|
||||
|
||||
# Aliases
|
||||
jit = _dsl.CuTeDSL.jit
|
||||
kernel = _dsl.CuTeDSL.kernel
|
||||
register_jit_arg_adapter = _dsl.JitArgAdapterRegistry.register_jit_arg_adapter
|
||||
compile = _dsl.compile
|
||||
|
||||
# Explicitly export all symbols for documentation generation
|
||||
__all__ = [
|
||||
# Core types
|
||||
"AddressSpace",
|
||||
"Tensor",
|
||||
"Layout",
|
||||
"ComposedLayout",
|
||||
"Swizzle",
|
||||
"E",
|
||||
"Atom",
|
||||
"MmaAtom",
|
||||
"CopyAtom",
|
||||
"TiledCopy",
|
||||
"TiledMma",
|
||||
"TensorSSA",
|
||||
# Basic utility functions
|
||||
"assume",
|
||||
"is_integer",
|
||||
"is_int_tuple",
|
||||
"is_static",
|
||||
"size",
|
||||
"has_underscore",
|
||||
"slice_",
|
||||
"depth",
|
||||
"rank",
|
||||
"shape",
|
||||
"printf",
|
||||
"print_tensor",
|
||||
"pretty_str",
|
||||
# Layout functions
|
||||
"make_layout",
|
||||
"recast_layout",
|
||||
"make_identity_layout",
|
||||
"make_ordered_layout",
|
||||
"make_composed_layout",
|
||||
"make_layout_tv",
|
||||
"make_layout_image_mask",
|
||||
# Tensor functions
|
||||
"make_ptr",
|
||||
"make_tensor",
|
||||
"make_identity_tensor",
|
||||
"make_fragment",
|
||||
"make_fragment_like",
|
||||
"recast_ptr",
|
||||
"recast_tensor",
|
||||
# Tensor manipulation
|
||||
"get",
|
||||
"select",
|
||||
"front",
|
||||
"is_major",
|
||||
"find",
|
||||
"coalesce",
|
||||
"group_modes",
|
||||
"cosize",
|
||||
"size_in_bytes",
|
||||
# Tuple operations
|
||||
"flatten_to_tuple",
|
||||
"flatten",
|
||||
"product",
|
||||
"product_like",
|
||||
"product_each",
|
||||
"prepend",
|
||||
"append",
|
||||
"prepend_ones",
|
||||
"append_ones",
|
||||
# Math operations
|
||||
"ceil_div",
|
||||
"round_up",
|
||||
# Layout operations
|
||||
"slice_and_offset",
|
||||
"crd2idx",
|
||||
"domain_offset",
|
||||
"elem_less",
|
||||
"filter_zeros",
|
||||
"filter",
|
||||
"tile_to_shape",
|
||||
"shape_div",
|
||||
"dice",
|
||||
# Layout algebra
|
||||
"composition",
|
||||
"complement",
|
||||
"right_inverse",
|
||||
"left_inverse",
|
||||
"max_common_layout",
|
||||
"max_common_vector",
|
||||
"is_congruent",
|
||||
"is_weakly_congruent",
|
||||
# Product operations
|
||||
"logical_product",
|
||||
"zipped_product",
|
||||
"tiled_product",
|
||||
"flat_product",
|
||||
"raked_product",
|
||||
"blocked_product",
|
||||
# Division operations
|
||||
"flat_divide",
|
||||
"logical_divide",
|
||||
"zipped_divide",
|
||||
"tiled_divide",
|
||||
"local_partition",
|
||||
"local_tile",
|
||||
# MMA and Copy operations
|
||||
"make_mma_atom",
|
||||
"make_tiled_mma",
|
||||
"make_copy_atom",
|
||||
"make_tiled_copy_tv",
|
||||
"make_tiled_copy",
|
||||
"make_tiled_copy_C_atom",
|
||||
"basic_copy",
|
||||
"basic_copy_if",
|
||||
"autovec_copy",
|
||||
"copy",
|
||||
"gemm",
|
||||
# Tensor creation
|
||||
"full",
|
||||
"full_like",
|
||||
"empty_like",
|
||||
"ones_like",
|
||||
"zeros_like",
|
||||
"where",
|
||||
"any_",
|
||||
"all_",
|
||||
"repeat_like",
|
||||
"ScaledBasis",
|
||||
# User defined struct
|
||||
"struct",
|
||||
# Modules
|
||||
"arch",
|
||||
"nvgpu",
|
||||
"testing",
|
||||
"runtime",
|
||||
# Decorators and code generation
|
||||
"jit",
|
||||
"kernel",
|
||||
"register_jit_arg_adapter",
|
||||
"compile",
|
||||
]
|
||||
@@ -0,0 +1,98 @@
|
||||
# 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 .elect import *
|
||||
from .mbar import *
|
||||
from .nvvm_wrappers import *
|
||||
from .smem import *
|
||||
from .tmem import *
|
||||
|
||||
# __all__ is required here for documentation generation
|
||||
__all__ = [
|
||||
#
|
||||
# elect.py
|
||||
#
|
||||
"make_warp_uniform",
|
||||
"elect_one",
|
||||
#
|
||||
# mbar.py
|
||||
#
|
||||
"mbarrier_init_arrive_cnt",
|
||||
"mbarrier_init_fence",
|
||||
"mbarrier_init_tx_bytes",
|
||||
"mbarrier_wait",
|
||||
"mbarrier_try_wait",
|
||||
"conditional_mbarrier_try_wait",
|
||||
"mbarrier_arrive",
|
||||
#
|
||||
# nvvm_wrappers.py
|
||||
#
|
||||
"lane_idx",
|
||||
"warp_idx",
|
||||
"thread_idx",
|
||||
"block_dim",
|
||||
"block_idx",
|
||||
"grid_dim",
|
||||
"cluster_idx",
|
||||
"cluster_dim",
|
||||
"block_in_cluster_idx",
|
||||
"block_in_cluster_dim",
|
||||
"block_idx_in_cluster",
|
||||
"shuffle_sync",
|
||||
"shuffle_sync_up",
|
||||
"shuffle_sync_down",
|
||||
"shuffle_sync_bfly",
|
||||
"barrier",
|
||||
"sync_threads",
|
||||
"sync_warp",
|
||||
"fence_acq_rel_cta",
|
||||
"fence_acq_rel_cluster",
|
||||
"fence_acq_rel_gpu",
|
||||
"fence_acq_rel_sys",
|
||||
"cp_async_commit_group",
|
||||
"cp_async_wait_group",
|
||||
"cp_async_bulk_commit_group",
|
||||
"cp_async_bulk_wait_group",
|
||||
"cluster_wait",
|
||||
"cluster_arrive",
|
||||
"cluster_arrive_relaxed",
|
||||
"fence_proxy",
|
||||
"vote_ballot_sync",
|
||||
"popc",
|
||||
"fence_view_async_tmem_load",
|
||||
"fence_view_async_tmem_store",
|
||||
"warpgroup_reg_alloc",
|
||||
"warpgroup_reg_dealloc",
|
||||
"fma_packed_f32x2",
|
||||
"mul_packed_f32x2",
|
||||
"add_packed_f32x2",
|
||||
"fmax",
|
||||
"rcp_approx",
|
||||
"exp2",
|
||||
# Constants
|
||||
"WARP_SIZE",
|
||||
# Forward from auto-generated nvvm python
|
||||
"ProxyKind",
|
||||
"SharedSpace",
|
||||
"RoundingModeKind",
|
||||
#
|
||||
# smem.py
|
||||
#
|
||||
"alloc_smem",
|
||||
"get_dyn_smem",
|
||||
#
|
||||
# tmem.py
|
||||
#
|
||||
"retrieve_tmem_ptr",
|
||||
"alloc_tmem",
|
||||
"relinquish_tmem_alloc_permit",
|
||||
"dealloc_tmem",
|
||||
]
|
||||
@@ -0,0 +1,75 @@
|
||||
# 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 CuTeDSL, T, dsl_user_op
|
||||
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
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.
|
||||
|
||||
:param value: The integer to make warp uniform.
|
||||
:type value: Int
|
||||
:return: The warp-uniform value equal to the input.
|
||||
:rtype: Int32
|
||||
"""
|
||||
return Int32(
|
||||
_cute_nvgpu_ir.arch_make_warp_uniform(
|
||||
Int32(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class IfOpRegion:
|
||||
"""
|
||||
A context manager for if Op.
|
||||
Automatically inserts `scf.yield([])` when exiting the context.
|
||||
"""
|
||||
|
||||
def __init__(self, block, *, loc=None, ip=None):
|
||||
self.block = block
|
||||
self.insert_point = ir.InsertionPoint(self.block)
|
||||
self.loc = loc
|
||||
self.ip = ip
|
||||
|
||||
def __enter__(self):
|
||||
self.insert_point.__enter__()
|
||||
return self.block.arguments
|
||||
|
||||
def __exit__(self, exc_type, exc_value, traceback):
|
||||
scf.yield_([], loc=self.loc, ip=self.ip)
|
||||
self.insert_point.__exit__(exc_type, exc_value, traceback)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def elect_one(*, loc=None, ip=None) -> IfOpRegion:
|
||||
"""
|
||||
Elects one thread within a warp.
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
with elect_one():
|
||||
# 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"], "arch")
|
||||
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)
|
||||
@@ -0,0 +1,208 @@
|
||||
# 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 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
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# Mbarrier management utilities
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def mbarrier_init_arrive_cnt(mbar_ptr: Pointer, cnt: Int, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Initializes a mbarrier with the specified thread arrival count.
|
||||
|
||||
:param mbar_ptr: A pointer to the mbarrier in SMEM
|
||||
:type mbar_ptr: Pointer
|
||||
:param cnt: The arrival count of the mbarrier
|
||||
:type cnt: Int
|
||||
"""
|
||||
nvvm.mbarrier_init_shared(
|
||||
mbar_ptr.llvm_ptr, Int32(cnt).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
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"], "arch")
|
||||
nvvm.fence_mbarrier_init(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def mbarrier_init_tx_bytes(
|
||||
mbar_ptr: Pointer, bytes: Int, peer_cta_rank_in_cluster=None, *, loc=None, ip=None
|
||||
) -> None:
|
||||
"""
|
||||
Initializes a mbarrier with the specified number of transaction bytes.
|
||||
|
||||
:param mbar_ptr: A pointer to the mbarrier in SMEM
|
||||
:type mbar_ptr: Pointer
|
||||
:param bytes: The number of transaction bytes
|
||||
:type bytes: Int
|
||||
:param peer_cta_rank_in_cluster: An optional CTA rank in cluster. If provided, the pointer to
|
||||
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"], "arch")
|
||||
|
||||
mbar_llvm_ptr = mbar_ptr.llvm_ptr
|
||||
if peer_cta_rank_in_cluster is not None:
|
||||
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
|
||||
mbar_llvm_ptr.type,
|
||||
mbar_llvm_ptr,
|
||||
Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
space = nvvm.MBarrierSpaceKind.CLUSTER
|
||||
else:
|
||||
space = nvvm.MBarrierSpaceKind.CTA
|
||||
|
||||
nvvm.mbarrier_txn(
|
||||
mbar_llvm_ptr,
|
||||
Int32(bytes).ir_value(loc=loc, ip=ip),
|
||||
kind=nvvm.MBarrierTxnKind.ARRIVE_EXPECT_TX,
|
||||
space=space,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def mbarrier_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Waits on a mbarrier with a specified phase.
|
||||
|
||||
:param mbar_ptr: A pointer to the mbarrier in SMEM
|
||||
:type mbar_ptr: Pointer
|
||||
: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"], "arch")
|
||||
|
||||
timeout_ns = 10000000
|
||||
# This NVVM Op is a spin-loop wrapping the mbarrier.try_wait.parity.shared.b64 PTX
|
||||
# The timeout in ns only applies to the latter and this call is truly blocking
|
||||
nvvm.mbarrier_try_wait_parity_shared(
|
||||
mbar_ptr.llvm_ptr,
|
||||
Int32(phase).ir_value(loc=loc, ip=ip),
|
||||
Int32(timeout_ns).ir_value(loc=loc, ip=ip),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def mbarrier_try_wait(mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None) -> Boolean:
|
||||
"""
|
||||
Attempts to wait on a mbarrier with a specified phase in a non-blocking fashion.
|
||||
|
||||
:param mbar_ptr: A pointer to the mbarrier in SMEM
|
||||
:type mbar_ptr: Pointer
|
||||
:param phase: The phase to wait for (either 0 or 1)
|
||||
:type phase: Int
|
||||
: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"], "arch")
|
||||
|
||||
return Boolean(
|
||||
nvvm.mbarrier_wait_parity(
|
||||
T.bool(),
|
||||
mbar_ptr.llvm_ptr,
|
||||
Int32(phase).ir_value(loc=loc, ip=ip),
|
||||
nvvm.MBarrierWaitKind.TRY,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def conditional_mbarrier_try_wait(
|
||||
cond, mbar_ptr: Pointer, phase: Int, *, loc=None, ip=None
|
||||
) -> Boolean:
|
||||
"""
|
||||
Conditionally attempts to wait on a mbarrier with a specified phase in a non-blocking fashion.
|
||||
|
||||
:param cond: A boolean predicate
|
||||
:param mbar_ptr: A pointer to the mbarrier in SMEM
|
||||
:type mbar_ptr: Pointer
|
||||
:param phase: The phase to wait for (either 0 or 1)
|
||||
:type phase: Int
|
||||
: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"], "arch")
|
||||
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],
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def mbarrier_arrive(
|
||||
mbar_ptr: Pointer, peer_cta_rank_in_cluster: Int = None, *, loc=None, ip=None
|
||||
) -> None:
|
||||
"""
|
||||
Arrives on an mbarrier.
|
||||
|
||||
:param mbar_ptr: A pointer to the mbarrier in SMEM
|
||||
:type mbar_ptr: Pointer
|
||||
:param peer_cta_rank_in_cluster: An optional CTA rank in cluster. If provided, the pointer to
|
||||
the mbarrier is converted to a remote address in the peer CTA's
|
||||
SMEM.
|
||||
"""
|
||||
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"], "arch")
|
||||
|
||||
mbar_llvm_ptr = nvvm.mapa_shared_cluster(
|
||||
mbar_llvm_ptr.type,
|
||||
mbar_llvm_ptr,
|
||||
Int32(peer_cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
space = nvvm.MBarrierSpaceKind.CLUSTER
|
||||
else:
|
||||
space = nvvm.MBarrierSpaceKind.CTA
|
||||
|
||||
nvvm.mbarrier_txn(
|
||||
mbar_llvm_ptr,
|
||||
Int32(1).ir_value(loc=loc, ip=ip),
|
||||
kind=nvvm.MBarrierTxnKind.ARRIVE,
|
||||
space=space,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
@@ -0,0 +1,547 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# Use of this software is governed by the terms and conditions of the
|
||||
# NVIDIA End User License Agreement (EULA), available at:
|
||||
# https://docs.nvidia.com/cutlass/media/docs/pythonDSL/license.html
|
||||
#
|
||||
# Any use, reproduction, disclosure, or distribution of this software
|
||||
# and related documentation outside the scope permitted by the EULA
|
||||
# is strictly prohibited.
|
||||
|
||||
from functools import partial
|
||||
from typing import Optional, Tuple, Union, Callable
|
||||
|
||||
from cutlass.cutlass_dsl import T, dsl_user_op
|
||||
|
||||
from cutlass._mlir import ir
|
||||
from cutlass._mlir.dialects import llvm, nvvm, vector
|
||||
|
||||
# Forward nvvm enums
|
||||
from cutlass._mlir.dialects.nvvm import (
|
||||
ProxyKind,
|
||||
SharedSpace,
|
||||
Tcgen05WaitKind,
|
||||
SetMaxRegisterAction,
|
||||
RoundingModeKind,
|
||||
)
|
||||
|
||||
from ..typing import Int, Boolean, Int32, Float32, Numeric, as_numeric
|
||||
|
||||
WARP_SIZE = 32
|
||||
FULL_MASK = 0xFFFFFFFF
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def lane_idx(*, loc=None, ip=None) -> Int32:
|
||||
"""
|
||||
Returns the lane index of the current thread within the warp.
|
||||
"""
|
||||
return Int32(nvvm.read_ptx_sreg_laneid(T.i32(), loc=loc, ip=ip))
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def warp_idx(*, loc=None, ip=None) -> Int32:
|
||||
"""
|
||||
Returns the warp index within a CTA.
|
||||
"""
|
||||
warp_size = 32
|
||||
tid_x = Int32(nvvm.read_ptx_sreg_tid_x(T.i32(), loc=loc, ip=ip))
|
||||
tid_y = Int32(nvvm.read_ptx_sreg_tid_y(T.i32(), loc=loc, ip=ip))
|
||||
tid_z = Int32(nvvm.read_ptx_sreg_tid_z(T.i32(), loc=loc, ip=ip))
|
||||
ntid_x = Int32(nvvm.read_ptx_sreg_ntid_x(T.i32(), loc=loc, ip=ip))
|
||||
ntid_y = Int32(nvvm.read_ptx_sreg_ntid_y(T.i32(), loc=loc, ip=ip))
|
||||
tid = tid_x + tid_y * ntid_x + tid_z * ntid_x * ntid_y
|
||||
return tid // warp_size
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def thread_idx(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the thread index within a CTA.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_tid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_tid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_tid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def block_dim(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the number of threads in each dimension of the CTA.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_ntid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_ntid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_ntid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def block_idx(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the CTA identifier within a grid.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_ctaid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_ctaid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_ctaid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def grid_dim(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the number of CTAs in each dimension of the grid.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_nctaid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_nctaid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_nctaid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cluster_idx(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the cluster identifier within a grid.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_clusterid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_clusterid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_clusterid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cluster_dim(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the number of clusters in each dimension of the grid.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_nclusterid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_nclusterid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_nclusterid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def block_in_cluster_idx(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the CTA index within a cluster across all dimensions.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_cluster_ctaid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_cluster_ctaid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_cluster_ctaid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def block_in_cluster_dim(*, loc=None, ip=None) -> Tuple[Int32, Int32, Int32]:
|
||||
"""
|
||||
Returns the dimensions of the cluster.
|
||||
"""
|
||||
return (
|
||||
Int32(nvvm.read_ptx_sreg_cluster_nctaid_x(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_cluster_nctaid_y(T.i32(), loc=loc, ip=ip)),
|
||||
Int32(nvvm.read_ptx_sreg_cluster_nctaid_z(T.i32(), loc=loc, ip=ip)),
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def block_idx_in_cluster(*, loc=None, ip=None) -> Int32:
|
||||
"""
|
||||
Returns the linearized identifier of the CTA within the cluster.
|
||||
"""
|
||||
return Int32(nvvm.read_ptx_sreg_cluster_ctarank(T.i32(), loc=loc, ip=ip))
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def shuffle_sync_op(
|
||||
value: Numeric,
|
||||
offset: Int,
|
||||
mask: Int = FULL_MASK,
|
||||
mask_and_clamp: Int = WARP_SIZE - 1,
|
||||
kind: nvvm.ShflKind = nvvm.ShflKind.idx,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Numeric:
|
||||
"""
|
||||
Shuffles a value within the threads of a warp.
|
||||
|
||||
:param value: The value to shuffle
|
||||
:type value: Numeric
|
||||
:param mask: A mask describing the threads participating in this operation
|
||||
:type mask: Int
|
||||
:param offset: A source lane or a source lane offset depending on kind
|
||||
:type offset: Int
|
||||
:param mask_and_clamp: An integer containing two packed values specifying a mask for logically
|
||||
splitting warps into sub-segments and an upper bound for clamping the
|
||||
source lane index.
|
||||
:type mask_and_clamp: Int
|
||||
:param kind: The kind of shuffle, can be idx, up, down, or bfly
|
||||
:type kind: ShflKind
|
||||
:return: The shuffled value
|
||||
:rtype: Numeric
|
||||
"""
|
||||
if not isinstance(value, Numeric):
|
||||
value = as_numeric(value)
|
||||
return type(value)(
|
||||
nvvm.shfl_sync(
|
||||
type(value).mlir_type,
|
||||
Int32(mask).ir_value(loc=loc, ip=ip),
|
||||
value.ir_value(loc=loc, ip=ip),
|
||||
Int32(offset).ir_value(loc=loc, ip=ip),
|
||||
Int32(mask_and_clamp).ir_value(loc=loc, ip=ip),
|
||||
kind,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
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 barrier(*, barrier_id=None, number_of_threads=None, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Creates a barrier, optionally named.
|
||||
"""
|
||||
if barrier_id is not None:
|
||||
barrier_id = Int32(barrier_id).ir_value(loc=loc, ip=ip)
|
||||
|
||||
if number_of_threads is not None:
|
||||
number_of_threads = Int32(number_of_threads).ir_value(loc=loc, ip=ip)
|
||||
|
||||
nvvm.barrier(
|
||||
barrier_id=barrier_id, number_of_threads=number_of_threads, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
@dsl_user_op
|
||||
def sync_threads(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Synchronizes all threads within a CTA.
|
||||
"""
|
||||
nvvm.barrier(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def sync_warp(mask: Int = FULL_MASK, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Performs a warp-wide sync with an optional mask.
|
||||
"""
|
||||
nvvm.bar_warp_sync(Int32(mask).ir_value(loc=loc, ip=ip), loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_acq_rel_cta(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Fence operation with acquire-release semantics.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar>`__.
|
||||
"""
|
||||
nvvm.fence_acq_rel_cta(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_acq_rel_cluster(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Fence operation with acquire-release semantics.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar>`__.
|
||||
"""
|
||||
nvvm.fence_acq_rel_cluster(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_acq_rel_gpu(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Fence operation with acquire-release semantics.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar>`__.
|
||||
"""
|
||||
nvvm.fence_acq_rel_gpu(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_acq_rel_sys(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Fence operation with acquire-release semantics.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar>`__.
|
||||
"""
|
||||
nvvm.fence_acq_rel_sys(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cp_async_commit_group(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Commits all prior initiated but uncommitted cp.async instructions.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-commit-group>`__.
|
||||
"""
|
||||
nvvm.cp_async_commit_group(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cp_async_wait_group(n, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Waits till only a specified numbers of cp.async groups are pending.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-wait-group-cp-async-wait-all>`__.
|
||||
"""
|
||||
nvvm.cp_async_wait_group(n, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cp_async_bulk_commit_group(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Commits all prior initiated but uncommitted cp.async.bulk instructions.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-commit-group>`__.
|
||||
"""
|
||||
nvvm.cp_async_bulk_commit_group(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cp_async_bulk_wait_group(group, *, read=None, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Waits till only a specified numbers of cp.async.bulk groups are pending.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-wait-group>`__.
|
||||
"""
|
||||
nvvm.cp_async_bulk_wait_group(group, read=read, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cluster_wait(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
A cluster-wide wait operation.
|
||||
"""
|
||||
nvvm.cluster_wait(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cluster_arrive(*, aligned=None, loc=None, ip=None) -> None:
|
||||
"""
|
||||
A cluster-wide arrive operation.
|
||||
"""
|
||||
nvvm.cluster_arrive(aligned=aligned, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cluster_arrive_relaxed(*, aligned=None, loc=None, ip=None) -> None:
|
||||
"""
|
||||
A cluster-wide arrive operation with relaxed semantics.
|
||||
"""
|
||||
nvvm.cluster_arrive_relaxed(aligned=aligned, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_proxy(
|
||||
kind: ProxyKind,
|
||||
*,
|
||||
space: Optional[SharedSpace] = None,
|
||||
use_intrinsic=None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
nvvm.fence_proxy(
|
||||
kind=kind, space=space, use_intrinsic=use_intrinsic, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def vote_ballot_sync(
|
||||
pred: Boolean, mask: Int = FULL_MASK, *, loc=None, ip=None
|
||||
) -> Int32:
|
||||
"""
|
||||
Performs a ballot operation across the warp.
|
||||
"""
|
||||
return Int32(
|
||||
nvvm.vote_ballot_sync(
|
||||
T.i32(),
|
||||
Int32(mask).ir_value(loc=loc, ip=ip),
|
||||
Boolean(pred).ir_value(loc=loc, ip=ip),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def popc(value: Numeric, *, loc=None, ip=None) -> Numeric:
|
||||
"""
|
||||
Performs a population count operation.
|
||||
"""
|
||||
if not isinstance(value, Numeric):
|
||||
value = as_numeric(value)
|
||||
return type(value)(llvm.intr_ctpop(value.ir_value(), loc=loc, ip=ip))
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_view_async_tmem_op(
|
||||
kind: Tcgen05WaitKind,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
Perform a fence operation on the async TMEM load or store.
|
||||
|
||||
.. note::
|
||||
This function is only available on sm_100a and above.
|
||||
The fence is required to synchronize the TMEM load/store
|
||||
and let the pipeline release or commit the buffer.
|
||||
|
||||
Take a mma2acc pipeline as an example of LOAD fence, the ACC tensor is from TMEM.
|
||||
```
|
||||
# Start to copy ACC from TMEM to register
|
||||
cute.copy(tmem_load, tACC, rACC)
|
||||
fence_view_async_tmem_load()
|
||||
# After fence, we can ensure the TMEM buffer is consumed totally.
|
||||
# Release the buffer to let the MMA know it can overwrite the buffer.
|
||||
mma2accum_pipeline.consumer_release(curr_consumer_state)
|
||||
```
|
||||
Take a TS GEMM kernel as an example of STORE fence, the A tensor is from TMEM.
|
||||
```
|
||||
# Start to copy A from register to TMEM
|
||||
cute.copy(tmem_store, rA, tA)
|
||||
fence_view_async_tmem_store()
|
||||
# After fence, we can ensure the TMEM buffer is ready.
|
||||
# Commit the buffer to let the MMA know it can start to load A.
|
||||
tmem_mma_pipeline.producer_commit(curr_producer_state)
|
||||
```
|
||||
|
||||
|
||||
:param kind: The kind of fence operation to perform including LOAD and STORE.
|
||||
:type kind: Tcgen05WaitKind
|
||||
"""
|
||||
nvvm.tcgen05_wait(kind, loc=loc, ip=ip)
|
||||
|
||||
|
||||
fence_view_async_tmem_load = partial(
|
||||
fence_view_async_tmem_op, kind=Tcgen05WaitKind.LOAD
|
||||
)
|
||||
fence_view_async_tmem_store = partial(
|
||||
fence_view_async_tmem_op, kind=Tcgen05WaitKind.STORE
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def warpgroup_reg_realloc_op(
|
||||
reg_count: int,
|
||||
kind: SetMaxRegisterAction,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
nvvm.setmaxregister(reg_count, kind, loc=loc, ip=ip)
|
||||
|
||||
|
||||
warpgroup_reg_alloc = partial(
|
||||
warpgroup_reg_realloc_op, kind=SetMaxRegisterAction.increase
|
||||
)
|
||||
warpgroup_reg_dealloc = partial(
|
||||
warpgroup_reg_realloc_op, kind=SetMaxRegisterAction.decrease
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def calc_packed_f32x2_op(
|
||||
src_a: Tuple[Float32, Float32],
|
||||
src_b: Tuple[Float32, Float32],
|
||||
src_c: Tuple[Float32, Float32] | None,
|
||||
calc_func: Callable,
|
||||
*,
|
||||
rnd=RoundingModeKind.RZ,
|
||||
ftz=True,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> 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_src_b = vector.from_elements(
|
||||
vec_type, tuple(as_numeric(b).ir_value() 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_res = calc_func(
|
||||
vec_type, vec_src_a, vec_src_b, vec_src_c, rnd=rnd, ftz=ftz, loc=loc, ip=ip
|
||||
)
|
||||
else:
|
||||
vec_res = calc_func(
|
||||
vec_type, vec_src_a, vec_src_b, rnd=rnd, ftz=ftz, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
res0 = Float32(
|
||||
vector.extract(
|
||||
vec_res, dynamic_position=[], static_position=[0], loc=loc, ip=ip
|
||||
)
|
||||
)
|
||||
res1 = Float32(
|
||||
vector.extract(
|
||||
vec_res, dynamic_position=[], static_position=[1], loc=loc, ip=ip
|
||||
)
|
||||
)
|
||||
return res0, res1
|
||||
|
||||
|
||||
fma_packed_f32x2 = partial(calc_packed_f32x2_op, calc_func=nvvm.fma_packed_f32x2)
|
||||
mul_packed_f32x2 = partial(
|
||||
calc_packed_f32x2_op, src_c=None, calc_func=nvvm.mul_packed_f32x2
|
||||
)
|
||||
add_packed_f32x2 = partial(
|
||||
calc_packed_f32x2_op, src_c=None, calc_func=nvvm.add_packed_f32x2
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fmax(
|
||||
a: Union[float, Float32], b: Union[float, Float32], *, loc=None, ip=None
|
||||
) -> Float32:
|
||||
return Float32(
|
||||
nvvm.fmax(
|
||||
T.f32(),
|
||||
Float32(a).ir_value(loc=loc, ip=ip),
|
||||
Float32(b).ir_value(loc=loc, ip=ip),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def rcp_approx(a: Union[float, Float32], *, loc=None, ip=None):
|
||||
return Float32(
|
||||
nvvm.rcp_approx_ftz_f(
|
||||
T.f32(), Float32(a).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def exp2(a: Union[float, Float32], *, loc=None, ip=None) -> Float32:
|
||||
return Float32(
|
||||
llvm.inline_asm(
|
||||
T.f32(),
|
||||
[Float32(a).ir_value(loc=loc, ip=ip)],
|
||||
"ex2.approx.ftz.f32 $0, $1;",
|
||||
"=f,f",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
)
|
||||
@@ -0,0 +1,96 @@
|
||||
# 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 typing import Optional, Type
|
||||
|
||||
from cutlass.cutlass_dsl import T, dsl_user_op
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..typing import Pointer, Numeric, NumericMeta
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def alloc_smem(
|
||||
element_type: Type[Numeric],
|
||||
size_in_elems: int,
|
||||
alignment: Optional[int] = None,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Pointer:
|
||||
"""
|
||||
Statically allocates SMEM.
|
||||
|
||||
:param element_type: The pointee type of the pointer.
|
||||
:type element_type: Type[Numeric]
|
||||
:param size_in_elems: The size of the allocation in terms of number of elements of the
|
||||
pointee type
|
||||
:type size_in_elems: int
|
||||
:param alignment: An optional pointer alignment for the allocation
|
||||
:type alignment: int
|
||||
:return: A pointer to the start of the allocation
|
||||
:rtype: Pointer
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(
|
||||
f"element_type must be a type of Numeric, but got {element_type}"
|
||||
)
|
||||
|
||||
if alignment is None:
|
||||
# Default alignment based on the element type's width
|
||||
alignment = element_type.width // 8
|
||||
ptr_ty = _cute_ir.PtrType.get(
|
||||
element_type.mlir_type, _cute_ir.AddressSpace.smem, alignment
|
||||
)
|
||||
return _cute_nvgpu_ir.arch_alloc_smem(
|
||||
ptr=ptr_ty,
|
||||
input=ir.IntegerAttr.get(T.i32(), size_in_elems),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def get_dyn_smem(
|
||||
element_type: Type[Numeric],
|
||||
alignment: Optional[int] = None,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Pointer:
|
||||
"""
|
||||
Retrieves a pointer to a dynamic SMEM allocation.
|
||||
|
||||
:param element_type: The pointee type of the pointer.
|
||||
:type element_type: Type[Numeric]
|
||||
:param alignment: An optional pointer alignment, the result pointer is offset appropriately
|
||||
:type alignment: int
|
||||
:return: A pointer to the start of the dynamic SMEM allocation with a correct
|
||||
alignement
|
||||
:rtype: Pointer
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(
|
||||
f"element_type must be a type of Numeric, but got {element_type}"
|
||||
)
|
||||
|
||||
if alignment is None:
|
||||
# Default alignment based on the element type's width
|
||||
alignment = element_type.width // 8
|
||||
ptr_ty = _cute_ir.PtrType.get(
|
||||
element_type.mlir_type,
|
||||
_cute_ir.AddressSpace.smem,
|
||||
alignment,
|
||||
)
|
||||
return _cute_nvgpu_ir.arch_get_dyn_smem(ptr=ptr_ty, loc=loc, ip=ip)
|
||||
@@ -0,0 +1,142 @@
|
||||
# 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 typing import Type
|
||||
|
||||
from cutlass.cutlass_dsl import dsl_user_op
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
|
||||
from ..typing import Pointer, Int, Int32, Numeric, NumericMeta
|
||||
|
||||
|
||||
SM100_TMEM_CAPACITY_COLUMNS = 512
|
||||
SM100_TMEM_MIN_ALLOC_COLUMNS = 32
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def retrieve_tmem_ptr(
|
||||
element_type: Type[Numeric],
|
||||
alignment: int,
|
||||
ptr_to_buffer_holding_addr: Pointer,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Pointer:
|
||||
"""
|
||||
Retrieves a pointer to TMEM with the provided element type and alignment.
|
||||
|
||||
:param element_type: The pointee type of the pointer.
|
||||
:type element_type: Type[Numeric]
|
||||
:param alignment: The alignment of the result pointer
|
||||
:type alignment: int
|
||||
:param ptr_to_buffer_holding_addr: A pointer to a SMEM buffer holding the TMEM address of the
|
||||
start of the allocation allocation
|
||||
:type ptr_to_buffer_holding_addr: Pointer
|
||||
:return: A pointer to TMEM
|
||||
:rtype: Pointer
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(
|
||||
f"element_type must be a type of Numeric, but got {element_type}"
|
||||
)
|
||||
|
||||
res_ty = _cute_ir.PtrType.get(
|
||||
element_type.mlir_type, _cute_ir.AddressSpace.tmem, alignment
|
||||
)
|
||||
return _cute_nvgpu_ir.arch_sm100_retrieve_tmem_ptr(
|
||||
res_ty, ptr_to_buffer_holding_addr.value, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def alloc_tmem(
|
||||
num_columns: Int,
|
||||
smem_ptr_to_write_address: Pointer,
|
||||
is_two_cta=None,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
Allocates TMEM.
|
||||
|
||||
:param num_columns: The number of TMEM columns to allocate
|
||||
:type num_columns: Int
|
||||
:param smem_ptr_to_write_address: A pointer to a SMEM buffer where the TMEM address is written
|
||||
to
|
||||
:type smem_ptr_to_write_address: Pointer
|
||||
:param is_two_cta: Optional boolean parameter for 2-CTA MMAs
|
||||
"""
|
||||
if isinstance(num_columns, int):
|
||||
if (
|
||||
num_columns < SM100_TMEM_MIN_ALLOC_COLUMNS
|
||||
or num_columns > SM100_TMEM_CAPACITY_COLUMNS
|
||||
or not (num_columns & (num_columns - 1) == 0)
|
||||
):
|
||||
raise ValueError(
|
||||
f"num_columns must be between 32 and 512, and must be pow of 2, but got {num_columns}"
|
||||
)
|
||||
_cute_nvgpu_ir.arch_sm100_alloc_tmem(
|
||||
Int32(num_columns).ir_value(loc=loc, ip=ip),
|
||||
smem_ptr_to_write_address.value,
|
||||
is_two_cta=is_two_cta,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def relinquish_tmem_alloc_permit(is_two_cta=None, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Relinquishes the right to allocate TMEM so that other CTAs potentially in a different grid can
|
||||
allocate.
|
||||
"""
|
||||
_cute_nvgpu_ir.arch_sm100_relinquish_tmem_alloc_permit(
|
||||
is_two_cta=is_two_cta, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def dealloc_tmem(
|
||||
tmem_ptr: Pointer,
|
||||
num_columns: Int,
|
||||
is_two_cta=None,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
Deallocates TMEM using the provided pointer and number of columns.
|
||||
|
||||
:param tmem_ptr: A pointer to the TMEM allocation to de-allocate
|
||||
:type tmem_ptr: Pointer
|
||||
:param num_columns: The number of columns in the TMEM allocation
|
||||
:type num_columns: Int
|
||||
:param is_two_cta: Optional boolean parameter for 2-CTA MMAs
|
||||
"""
|
||||
if isinstance(num_columns, int):
|
||||
if (
|
||||
num_columns < SM100_TMEM_MIN_ALLOC_COLUMNS
|
||||
or num_columns > SM100_TMEM_CAPACITY_COLUMNS
|
||||
or not (num_columns & (num_columns - 1) == 0)
|
||||
):
|
||||
raise ValueError(
|
||||
f"num_columns must be between 32 and 512, and must be pow of 2, but got {num_columns}"
|
||||
)
|
||||
_cute_nvgpu_ir.arch_sm100_dealloc_tmem(
|
||||
tmem_ptr.value,
|
||||
Int32(num_columns).ir_value(loc=loc, ip=ip),
|
||||
is_two_cta=is_two_cta,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,354 @@
|
||||
# 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 .core import TensorSSA
|
||||
from cutlass._mlir.dialects import math, arith
|
||||
|
||||
|
||||
def acos(a: TensorSSA) -> TensorSSA:
|
||||
"""Compute element-wise arc cosine of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:return: Tensor containing the arc cosine of each element in input tensor
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = acos(y) # Compute arc cosine
|
||||
"""
|
||||
return TensorSSA(math.acos(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def asin(a: TensorSSA) -> TensorSSA:
|
||||
"""Compute element-wise arc sine of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:return: Tensor containing the arc sine of each element in input tensor
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = asin(y) # Compute arc sine
|
||||
"""
|
||||
return TensorSSA(math.asin(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def atan(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise arc tangent of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the arc tangent of each element in input tensor
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = atan(y) # Compute arc tangent
|
||||
"""
|
||||
raise NotImplementedError("atan is not implemented")
|
||||
return TensorSSA(math.atan(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def atan2(a: TensorSSA, b: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise arc tangent of two tensors.
|
||||
|
||||
Computes atan2(a, b) element-wise. The function atan2(a, b) is the angle in radians
|
||||
between the positive x-axis and the point given by the coordinates (b, a).
|
||||
|
||||
:param a: First input tensor (y-coordinates)
|
||||
:type a: TensorSSA
|
||||
:param b: Second input tensor (x-coordinates)
|
||||
:type b: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the arc tangent of a/b element-wise
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
y = cute.make_fragment(ptr1, layout).load() # y coordinates
|
||||
x = cute.make_fragment(ptr2, layout).load() # x coordinates
|
||||
theta = atan2(y, x) # Compute angles
|
||||
"""
|
||||
return TensorSSA(
|
||||
math.atan2(a, b, fastmath=arith.FastMathFlags.none), a.shape, a.dtype
|
||||
)
|
||||
|
||||
|
||||
def cos(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise cosine of the input tensor.
|
||||
|
||||
:param a: Input tensor (in radians)
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the cosine of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = cos(y) # Compute cosine
|
||||
"""
|
||||
return TensorSSA(math.cos(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def erf(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise error function of the input tensor.
|
||||
|
||||
The error function is defined as:
|
||||
erf(x) = 2/√π ∫[0 to x] exp(-t²) dt
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the error function value for each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = erf(y) # Compute error function
|
||||
"""
|
||||
return TensorSSA(math.erf(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def exp2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise base-2 exponential of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing 2 raised to the power of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = exp2(y) # Compute 2^x
|
||||
"""
|
||||
return TensorSSA(math.exp2(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def log(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise natural logarithm of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the natural logarithm of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = log(y) # Compute natural logarithm
|
||||
"""
|
||||
return TensorSSA(math.log(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def log2(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise base-2 logarithm of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the base-2 logarithm of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = log2(y) # Compute log base 2
|
||||
"""
|
||||
return TensorSSA(math.log2(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def log10(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise base-10 logarithm of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the base-10 logarithm of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = log10(y) # Compute log base 10
|
||||
"""
|
||||
return TensorSSA(math.log10(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def rsqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise reciprocal square root of the input tensor.
|
||||
|
||||
Computes 1/√x element-wise.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the reciprocal square root of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = rsqrt(y) # Compute 1/√x
|
||||
"""
|
||||
return TensorSSA(math.rsqrt(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def sin(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise sine of the input tensor.
|
||||
|
||||
:param a: Input tensor (in radians)
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the sine of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = sin(y) # Compute sine
|
||||
"""
|
||||
return TensorSSA(math.sin(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def sqrt(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise square root of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the square root of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = sqrt(y) # Compute square root
|
||||
"""
|
||||
return TensorSSA(math.sqrt(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def tan(a: TensorSSA) -> TensorSSA:
|
||||
"""Compute element-wise tangent of the input tensor.
|
||||
|
||||
:param a: Input tensor (in radians)
|
||||
:type a: TensorSSA
|
||||
:return: Tensor containing the tangent of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = tan(y) # Compute tangent
|
||||
"""
|
||||
return TensorSSA(math.tan(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
def tanh(a: TensorSSA, fastmath: bool = False) -> TensorSSA:
|
||||
"""Compute element-wise hyperbolic tangent of the input tensor.
|
||||
|
||||
:param a: Input tensor
|
||||
:type a: TensorSSA
|
||||
:param fastmath: Enable fast math optimizations, defaults to False
|
||||
:type fastmath: bool, optional
|
||||
:return: Tensor containing the hyperbolic tangent of each element
|
||||
:rtype: TensorSSA
|
||||
|
||||
Example:
|
||||
|
||||
.. code-block::
|
||||
|
||||
x = cute.make_fragment(layout) # Create tensor
|
||||
y = x.load() # Load values
|
||||
z = tanh(y) # Compute hyperbolic tangent
|
||||
"""
|
||||
return TensorSSA(math.tanh(a, fastmath=arith.FastMathFlags.none), a.shape, a.dtype)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"acos",
|
||||
"asin",
|
||||
"atan",
|
||||
"atan2",
|
||||
"cos",
|
||||
"erf",
|
||||
"exp2",
|
||||
"log",
|
||||
"log10",
|
||||
"log2",
|
||||
"rsqrt",
|
||||
"sin",
|
||||
"sqrt",
|
||||
"tan",
|
||||
"tanh",
|
||||
]
|
||||
@@ -0,0 +1,26 @@
|
||||
# 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 . import warp
|
||||
from . import cpasync
|
||||
from . import warpgroup
|
||||
from . import tcgen05
|
||||
|
||||
from .common import *
|
||||
from .helpers import *
|
||||
|
||||
|
||||
# __all__ is required here for documentation generation
|
||||
__all__ = [
|
||||
"OpError",
|
||||
"MmaUniversalOp",
|
||||
"CopyUniversalOp",
|
||||
]
|
||||
@@ -0,0 +1,143 @@
|
||||
# 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 dataclasses import dataclass
|
||||
from typing import Type, Optional
|
||||
|
||||
from cutlass.cutlass_dsl import DSLBaseError
|
||||
|
||||
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 ..typing import Float16, Float32, Float64, Numeric
|
||||
|
||||
|
||||
class OpError(DSLBaseError):
|
||||
"""
|
||||
An exception class for Op construction errors.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self, op: core.Op, message: str, suggestion: Optional[str] = None
|
||||
) -> None:
|
||||
if suggestion is None:
|
||||
# Default suggestion
|
||||
suggestion = "Check your Op construction code"
|
||||
super().__init__(
|
||||
message,
|
||||
error_code=f"{op.__class__.__name__} error",
|
||||
suggestion=suggestion,
|
||||
)
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# MMA Ops and Traits
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaUniversalOp(core.MmaOp):
|
||||
"""
|
||||
The universal MMA Operation.
|
||||
|
||||
This Operation currently expects the A/B operands as well as the accumulator to share the same
|
||||
data types.
|
||||
|
||||
:param abacc_dtype: The data type for the A/B operands and the accumulator
|
||||
:type abacc_dtype: Type[Numeric]
|
||||
"""
|
||||
|
||||
abacc_dtype: Type[Numeric]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
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",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
"universal MMA Operation using FMA"
|
||||
f"\n A/B/Accumulator data type = {self.abacc_dtype}"
|
||||
)
|
||||
|
||||
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaUniversalTrait":
|
||||
shape_mnk_attr = ir.Attribute.parse(f'#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))
|
||||
|
||||
|
||||
class MmaUniversalTrait(core.Trait):
|
||||
pass
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# Copy Ops and Traits
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyUniversalOp(core.CopyOp):
|
||||
"""
|
||||
The universal Copy Operation.
|
||||
|
||||
When creating a Copy Atom out of this operation, the expected usage pattern is
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
op = cute.nvgpu.CopyUniversalOp()
|
||||
atom = cute.make_copy_atom(op, tensor_dtype, num_bits_per_copy=64)
|
||||
|
||||
- ``tensor_dtype`` is the data type used to build the reference TV Layout (either the source \
|
||||
or the destination TV Layout) in unit of tensor elements and is used for partitioning by \
|
||||
``TiledCopy`` for example
|
||||
- ``num_bits_per_copy`` is a kw argument specifying the number of bits to copy per Atom \
|
||||
execution. This can be larger than the width of the above data type. When not provided, \
|
||||
the compiler will do a best effort at auto-vectorizing.
|
||||
"""
|
||||
|
||||
def __str__(self) -> str:
|
||||
return "universal Copy Operation"
|
||||
|
||||
def _make_trait(
|
||||
self,
|
||||
copy_internal_type: Type[Numeric],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
**kwargs,
|
||||
) -> "CopyUniversalTrait":
|
||||
num_bits_per_copy = kwargs.get("num_bits_per_copy", 0)
|
||||
if not isinstance(num_bits_per_copy, int) or (num_bits_per_copy < 0):
|
||||
raise ValueError(
|
||||
"expects a 'num_bits_per_copy' kw argument of type int that is non-negative "
|
||||
f"when creating a copy Atom for {self.__class__.__name__}"
|
||||
)
|
||||
ty = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
|
||||
copy_internal_type.mlir_type, num_bits_per_copy
|
||||
)
|
||||
return CopyUniversalTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class CopyUniversalTrait(core.Trait):
|
||||
pass
|
||||
@@ -0,0 +1,38 @@
|
||||
# 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 .copy import *
|
||||
from .helpers import *
|
||||
|
||||
|
||||
# __all__ is required here for documentation generation
|
||||
__all__ = [
|
||||
#
|
||||
# copy.py
|
||||
#
|
||||
"LoadCacheMode",
|
||||
"CopyG2SOp",
|
||||
"CopyBulkTensorTileG2SOp",
|
||||
"CopyBulkTensorTileG2SMulticastOp",
|
||||
"CopyBulkTensorTileS2GOp",
|
||||
#
|
||||
# helpers.py
|
||||
#
|
||||
"make_tma_tile_atom",
|
||||
"tma_partition",
|
||||
"create_tma_multicast_mask",
|
||||
"prefetch_descriptor",
|
||||
"copy_tensormap",
|
||||
"update_tma_descriptor",
|
||||
"fence_tma_desc_acquire",
|
||||
"cp_fence_tma_desc_release",
|
||||
"fence_tma_desc_release",
|
||||
]
|
||||
@@ -0,0 +1,366 @@
|
||||
# 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 enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional, Type
|
||||
|
||||
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 ...core import CopyOp, Trait
|
||||
from ...typing import Int16, Pointer, Integer, Numeric
|
||||
from ..common import OpError
|
||||
from ..tcgen05.mma import CtaGroup
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# Aynchronous copies
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class LoadCacheMode(enum.Enum):
|
||||
"""
|
||||
An enumeration for the possible cache modes of a non-bulk ``cp.async`` instruction.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#cache-operators>`__.
|
||||
"""
|
||||
|
||||
ALWAYS = _cute_nvgpu_ir.LoadCacheMode.always
|
||||
GLOBAL = _cute_nvgpu_ir.LoadCacheMode.global_
|
||||
STREAMING = _cute_nvgpu_ir.LoadCacheMode.streaming
|
||||
LAST_USE = _cute_nvgpu_ir.LoadCacheMode.last_use
|
||||
NONE = _cute_nvgpu_ir.LoadCacheMode.none
|
||||
|
||||
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.LoadCacheMode:
|
||||
return self.value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyG2SOp(CopyOp):
|
||||
"""
|
||||
Non-bulk asynchronous GMEM to SMEM Copy Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-non-bulk-copy>`__.
|
||||
"""
|
||||
|
||||
cache_mode: LoadCacheMode = LoadCacheMode.ALWAYS
|
||||
|
||||
def __str__(self) -> str:
|
||||
res = "cp.async GMEM -> SMEM copy Operation"
|
||||
if self.cache_mode != LoadCacheMode.ALWAYS:
|
||||
res += f"\n with cache mode = {self.cache_mode}"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
self,
|
||||
copy_internal_type: Type[t.Numeric],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
**kwargs,
|
||||
) -> "CopyG2STrait":
|
||||
num_bits_per_copy = kwargs.get("num_bits_per_copy", None)
|
||||
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__}"
|
||||
)
|
||||
# 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",
|
||||
)
|
||||
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))
|
||||
|
||||
|
||||
class CopyG2STrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# Bulk tensor copies a.k.a TMA copies
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
TMA_MBAR_PTR_FIELD_NAME = "tma_bar"
|
||||
TMA_MASK_FIELD_NAME = "mcast_mask"
|
||||
TMA_DESC_PTR_FIELD_NAME = "tma_descriptor_ptr"
|
||||
|
||||
#
|
||||
# TMA GMEM -> SMEM copies
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkTensorTileG2SOp(CopyOp):
|
||||
"""
|
||||
Bulk tensor asynchrnous GMEM to SMEM Copy Operation using the TMA unit.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor>`__.
|
||||
This Operation uses TMA in the ``.tile`` mode.
|
||||
"""
|
||||
|
||||
cta_group: CtaGroup = CtaGroup.ONE
|
||||
|
||||
admissible_archs = ["sm_90", "sm_90a", "sm_100a"]
|
||||
|
||||
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:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
if (self.cta_group == CtaGroup.TWO) and arch[:5] == "sm_90":
|
||||
raise OpError(
|
||||
self,
|
||||
f"CTA group of 2 is tcgen05-specific and is not and is not compatible with {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
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"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "CopyBulkTensorTileG2SNonExecTrait":
|
||||
raise NotImplementedError(
|
||||
"Use cpasync.make_tma_tile_atom to obtain a copy Atom for TMA"
|
||||
)
|
||||
|
||||
def _to_ir(self) -> _cute_nvgpu_ir.TiledTmaLoadEnum:
|
||||
if self.cta_group == CtaGroup.ONE:
|
||||
return _cute_nvgpu_ir.TiledTmaLoadEnum.sm_90
|
||||
elif self.cta_group == CtaGroup.TWO:
|
||||
return _cute_nvgpu_ir.TiledTmaLoadEnum.sm_100_2sm
|
||||
else:
|
||||
assert False, "unrecognized self.cta_group"
|
||||
|
||||
|
||||
class CopyBulkTensorTileG2SNonExecTrait(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,
|
||||
tma_bar_ptr: Optional[Pointer] = None,
|
||||
tma_desc_ptr: Optional[Pointer] = None,
|
||||
**kwargs,
|
||||
):
|
||||
"""
|
||||
Custom implementation of unpack for non-executable TMAs.
|
||||
|
||||
The non-multicast TMA load requires a `tma_bar_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(tma_bar_ptr, Pointer):
|
||||
raise ValueError(
|
||||
"expects a pointer to an mbarrier to be provided via the tma_bar_ptr 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_MBAR_PTR_FIELD_NAME}>"
|
||||
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
|
||||
)
|
||||
if isinstance(tma_desc_ptr, Pointer):
|
||||
attr_str = f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_DESC_PTR_FIELD_NAME}>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
exec_value = _cute_nvgpu_ir.atom_set_value(
|
||||
exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
return exec_value
|
||||
|
||||
|
||||
#
|
||||
# TMA GMEM -> SMEM multicast copies
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkTensorTileG2SMulticastOp(CopyOp):
|
||||
"""
|
||||
Bulk tensor asynchrnous multicast GMEM to SMEM Copy Operation using the TMA unit.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor>`__.
|
||||
This Operation uses TMA in the ``.tile`` mode.
|
||||
"""
|
||||
|
||||
cta_group: CtaGroup = CtaGroup.ONE
|
||||
|
||||
admissible_archs = ["sm_90", "sm_90a", "sm_100a"]
|
||||
|
||||
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:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
if (self.cta_group == CtaGroup.TWO) and arch[:5] == "sm_90":
|
||||
raise OpError(
|
||||
self,
|
||||
f"CTA group of 2 is tcgen05-specific and is not and is not compatible with {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
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"
|
||||
return res
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "CopyBulkTensorTileG2SMulticastNonExecTrait":
|
||||
raise NotImplementedError(
|
||||
"Use cpasync.make_tma_tile_atom to obtain a copy Atom for TMA"
|
||||
)
|
||||
|
||||
def _to_ir(self) -> _cute_nvgpu_ir.TiledTmaLoadEnum:
|
||||
if self.cta_group == CtaGroup.ONE:
|
||||
return _cute_nvgpu_ir.TiledTmaLoadEnum.sm_90_multicast
|
||||
elif self.cta_group == CtaGroup.TWO:
|
||||
return _cute_nvgpu_ir.TiledTmaLoadEnum.sm_100_2sm_multicast
|
||||
else:
|
||||
assert False, "unrecognized self.cta_group"
|
||||
|
||||
|
||||
class CopyBulkTensorTileG2SMulticastNonExecTrait(Trait):
|
||||
def unpack(
|
||||
self,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
tma_bar_ptr: Optional[Pointer] = None,
|
||||
mcast_mask=None,
|
||||
tma_desc_ptr=None,
|
||||
):
|
||||
"""
|
||||
Custom implementation of unpack for non-executable TMAs.
|
||||
|
||||
The multicast TMA load requires a `tma_bar_ptr` and a `mcast_mask` keyword arguments to be
|
||||
provided when using `cute.copy`.
|
||||
"""
|
||||
if not isinstance(tma_bar_ptr, Pointer):
|
||||
raise ValueError(
|
||||
"expects a pointer to an mbarrier to be provided via the tma_bar_ptr kw argument"
|
||||
)
|
||||
if not isinstance(mcast_mask, Integer):
|
||||
raise ValueError(
|
||||
"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 = 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
|
||||
)
|
||||
if isinstance(tma_desc_ptr, Pointer):
|
||||
attr_str = f"#cute_nvgpu.atom_copy_field_tmaload<{TMA_DESC_PTR_FIELD_NAME}>"
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
exec_value = _cute_nvgpu_ir.atom_set_value(
|
||||
exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
return exec_value
|
||||
|
||||
|
||||
#
|
||||
# TMA SMEM -> GMEM copies
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class CopyBulkTensorTileS2GOp(CopyOp):
|
||||
"""
|
||||
Bulk tensor asynchrnous SMEM to GMEM Copy Operation using the TMA unit.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#data-movement-and-conversion-instructions-cp-async-bulk-tensor>`__.
|
||||
This Operation uses TMA in the ``.tile`` mode.
|
||||
"""
|
||||
|
||||
admissible_archs = ["sm_90", "sm_90a", "sm_100a"]
|
||||
|
||||
def __post_init__(self):
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
return "cp.async SMEM -> GMEM bulk tensor copy Operation"
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "CopyBulkTensorTileS2GTrait":
|
||||
raise NotImplementedError(
|
||||
"Use cpasync.make_tma_tile_atom to obtain a copy Atom for TMA"
|
||||
)
|
||||
|
||||
|
||||
class CopyBulkTensorTileS2GTrait(Trait):
|
||||
def unpack(self, *, loc=None, ip=None, tma_desc_ptr: Optional[Pointer] = None):
|
||||
"""
|
||||
Custom implementation of unpack for non-executable TMAs.
|
||||
"""
|
||||
exec_value = _cute_nvgpu_ir.atom_make_exec_tma(self.value, loc=loc, ip=ip)
|
||||
if isinstance(tma_desc_ptr, Pointer):
|
||||
attr_str = (
|
||||
f"#cute_nvgpu.atom_copy_field_tmastore<{TMA_DESC_PTR_FIELD_NAME}>"
|
||||
)
|
||||
attr = ir.Attribute.parse(attr_str)
|
||||
exec_value = _cute_nvgpu_ir.atom_set_value(
|
||||
exec_value, attr, tma_desc_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
return exec_value
|
||||
@@ -0,0 +1,327 @@
|
||||
# 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 typing import Optional, Tuple, Type, Union
|
||||
|
||||
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 .copy import (
|
||||
CopyBulkTensorTileG2SOp,
|
||||
CopyBulkTensorTileG2SMulticastOp,
|
||||
CopyBulkTensorTileS2GOp,
|
||||
CopyBulkTensorTileG2SNonExecTrait,
|
||||
CopyBulkTensorTileG2SMulticastNonExecTrait,
|
||||
CopyBulkTensorTileS2GTrait,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_tma_tile_atom(
|
||||
op: Union[
|
||||
CopyBulkTensorTileG2SOp,
|
||||
CopyBulkTensorTileG2SMulticastOp,
|
||||
CopyBulkTensorTileS2GOp,
|
||||
],
|
||||
gmem_tensor: Tensor,
|
||||
smem_layout: Layout,
|
||||
cta_tiler: Tiler,
|
||||
num_multicast: int = 1,
|
||||
*,
|
||||
internal_type: Optional[Type[Numeric]] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[core.CopyAtom, Tensor]:
|
||||
"""
|
||||
Makes a TMA Copy Atom in the ``.tile`` mode to copy tiles of a GMEM tensor to/from and SMEM
|
||||
buffer with the given Layout.
|
||||
|
||||
Given
|
||||
|
||||
- a GMEM tensor
|
||||
- a SMEM layout
|
||||
- a CTA-level Tiler
|
||||
|
||||
this function figures out the bulk tensor asynchronous copy instruction to use with the maximum
|
||||
"TMA vector length" to copy tiles of the GMEM tensor to/from an SMEM buffer with the provided
|
||||
layout and consistent with the provided Tiler.
|
||||
|
||||
This function returns two results:
|
||||
|
||||
1. the Copy Atom
|
||||
2. the so-called TMA tensor used to map logical coordinates of the GMEM tensor to coordinates \
|
||||
that the TMA unit can consume. TMA tensors have so-called basis stride elements so that the \
|
||||
associated layout can output coordinates. Otherwise, TMA tensors can be partitioned \
|
||||
similarly to any other CuTe tensors using the algebra.
|
||||
|
||||
:param op: The Copy Operation to construct an Atom for
|
||||
:type op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp, CopyBulkTensorTileS2GOp]
|
||||
: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: Layout
|
||||
:param cta_tiler: The CTA Tiler to use
|
||||
:type cta_tiler: Tiler
|
||||
:param num_multicast: The multicast factor
|
||||
:type num_multicast: int
|
||||
:param internal_type: An optional parameter for the internal data type to use when the actual data type is not supported by the TMA unit
|
||||
:type internal_type: Type[Numeric]
|
||||
:return: A Copy Atom for this Operation and the associated TMA tensor
|
||||
:rtype: Tuple[core.CopyAtom, Tensor]
|
||||
"""
|
||||
|
||||
if internal_type is not None:
|
||||
if not isinstance(internal_type, NumericMeta):
|
||||
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
|
||||
internal_type = internal_type.mlir_type
|
||||
|
||||
cta_v_map = core.composition(
|
||||
core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip),
|
||||
cta_tiler,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
if num_multicast != 1:
|
||||
raise ValueError(
|
||||
f"expects num_multicast to be 1 for non multicast G2S copies, "
|
||||
f"but got {num_multicast}"
|
||||
)
|
||||
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load(
|
||||
gmem_tensor.value,
|
||||
smem_layout,
|
||||
cta_v_map,
|
||||
op._to_ir(),
|
||||
num_multicast=num_multicast,
|
||||
internal_type=internal_type,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return core.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
elif isinstance(op, CopyBulkTensorTileG2SMulticastOp):
|
||||
if num_multicast < 1:
|
||||
raise ValueError(
|
||||
f"expects num_multicast to be >= 1 for multicast G2S copies, "
|
||||
f"but got {num_multicast}"
|
||||
)
|
||||
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_load(
|
||||
gmem_tensor.value,
|
||||
smem_layout,
|
||||
cta_v_map,
|
||||
op._to_ir(),
|
||||
num_multicast=num_multicast,
|
||||
internal_type=internal_type,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return (
|
||||
core.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
res[1],
|
||||
)
|
||||
elif isinstance(op, CopyBulkTensorTileS2GOp):
|
||||
res = _cute_nvgpu_ir.atom_make_non_exec_tiled_tma_store(
|
||||
gmem_tensor.value,
|
||||
smem_layout,
|
||||
cta_v_map,
|
||||
internal_type=internal_type,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return core.CopyAtom(op, CopyBulkTensorTileS2GTrait(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,
|
||||
cta_coord: Coord,
|
||||
cta_layout: Layout,
|
||||
smem_tensor: Tensor,
|
||||
gmem_tensor: Tensor,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[Tensor, Tensor]:
|
||||
"""
|
||||
Tiles the GMEM and SMEM tensors for the provided TMA Copy Atom.
|
||||
"""
|
||||
cta_coord_val = core._pack_coord(cta_coord, loc=loc, ip=ip)
|
||||
s, d = _cute_nvgpu_ir.atom_tma_partition(
|
||||
atom._trait.value,
|
||||
cta_coord=cta_coord_val,
|
||||
cta_layout=cta_layout,
|
||||
smem_tensor=smem_tensor.value,
|
||||
gmem_tensor=gmem_tensor.value,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return s, d
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def create_tma_multicast_mask(
|
||||
cta_layout_vmnk: Layout,
|
||||
cta_coord_vmnk: Coord,
|
||||
mcast_mode: int,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Int16:
|
||||
"""
|
||||
Computes a multicast mask for a TMA load Copy.
|
||||
|
||||
:param cta_layout_vmnk: The VMNK layout of the cluster
|
||||
:type cta_layout_vmnk: Layout
|
||||
:param cta_coord_vmnk: The VMNK coordinate of the current CTA
|
||||
:type cta_coord_vmnk: Coord
|
||||
:param mcast_mode: The tensor mode in which to multicast
|
||||
:type mcast_mode: int
|
||||
:return: The resulting mask
|
||||
:rtype: Int16
|
||||
"""
|
||||
if core.rank(cta_layout_vmnk) != 4:
|
||||
raise ValueError(
|
||||
f"cta_layout_vmnk must be rank 4, but got {core.pretty_str(cta_layout_vmnk)}"
|
||||
)
|
||||
if core.rank(cta_coord_vmnk) != 4:
|
||||
raise ValueError(
|
||||
f"cta_coord_vmnk must be rank 4, but got {core.pretty_str(cta_coord_vmnk)}"
|
||||
)
|
||||
return core.make_layout_image_mask(
|
||||
cta_layout_vmnk, cta_coord_vmnk, mcast_mode, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def prefetch_descriptor(tma_atom: core.CopyAtom, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
Prefetches the TMA descriptor associated with the TMA Atom.
|
||||
"""
|
||||
_cute_nvgpu_ir.prefetch_tma_desc(tma_atom._trait.value, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def copy_tensormap(
|
||||
tma_atom: core.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
|
||||
pointer.
|
||||
|
||||
:param tma_atom: The TMA Copy Atom
|
||||
:type tma_atom: CopyAtom
|
||||
:param tensormap_ptr: The pointer to the memory location to copy the tensormap to
|
||||
:type tensormap_ptr: Pointer
|
||||
"""
|
||||
_cute_nvgpu_ir.copy_tma_desc(
|
||||
tma_atom._trait.value, tensormap_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def update_tma_descriptor(
|
||||
tma_atom: core.CopyAtom,
|
||||
gmem_tensor: Tensor,
|
||||
tma_desc_ptr: Pointer,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
Updates the TMA descriptor in the memory location pointed to by the provided pointer using
|
||||
information from a TMA Copy Atom and the provided GMEM tensor.
|
||||
|
||||
Specifically, the following fields of the TMA descriptor will be updated:
|
||||
|
||||
1. the GMEM tensor base address
|
||||
2. the GMEM tensor shape
|
||||
3. the GMEM tensor stride
|
||||
|
||||
Other fields of the TMA descriptor are left unchanged.
|
||||
|
||||
:param tma_atom: The TMA Copy Atom
|
||||
:type tma_atom: CopyAtom
|
||||
:param gmem_tensor: The GMEM tensor
|
||||
:type gmem_tensor: Tensor
|
||||
:param tensormap_ptr: The pointer to the memory location of the descriptor to udpate
|
||||
:type tensormap_ptr: Pointer
|
||||
"""
|
||||
_cute_nvgpu_ir.update_tma_desc(
|
||||
tma_atom._trait.value, gmem_tensor.value, tma_desc_ptr.value, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_tma_desc_acquire(
|
||||
tma_desc_ptr: Pointer,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
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()
|
||||
llvm.inline_asm(
|
||||
None,
|
||||
[tma_desc_ptr_i64],
|
||||
"fence.proxy.tensormap::generic.acquire.gpu [$0], 128;",
|
||||
"l",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def cp_fence_tma_desc_release(
|
||||
tma_desc_global_ptr: Pointer,
|
||||
tma_desc_shared_ptr: Pointer,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
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()
|
||||
llvm.inline_asm(
|
||||
None,
|
||||
[tma_desc_global_ptr_i64, tma_desc_shared_ptr_i32],
|
||||
"tensormap.cp_fenceproxy.global.shared::cta.tensormap::generic.release.gpu.sync.aligned [$0], [$1], 128;",
|
||||
"l,r",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence_tma_desc_release(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#parallel-synchronization-and-communication-instructions-membar>`__.
|
||||
"""
|
||||
llvm.inline_asm(
|
||||
None,
|
||||
[],
|
||||
"fence.proxy.tensormap::generic.release.gpu;",
|
||||
"",
|
||||
has_side_effects=True,
|
||||
is_align_stack=False,
|
||||
asm_dialect=llvm.AsmDialect.AD_ATT,
|
||||
)
|
||||
@@ -0,0 +1,159 @@
|
||||
# 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 typing import Optional, Tuple, Type, Union
|
||||
|
||||
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 ...impl_utils import check_type_in
|
||||
from .cpasync.copy import (
|
||||
CopyBulkTensorTileG2SOp,
|
||||
CopyBulkTensorTileG2SNonExecTrait,
|
||||
CopyBulkTensorTileG2SMulticastOp,
|
||||
CopyBulkTensorTileG2SMulticastNonExecTrait,
|
||||
)
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# TMA creation helpers for tcgen05 MMAs
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_tma_tile_atom_A(
|
||||
op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
|
||||
gmem_tensor: Tensor,
|
||||
smem_layout: Layout,
|
||||
mma_tiler_mnk: Shape,
|
||||
tiled_mma: core.TiledMma,
|
||||
cluster_shape_vmnk: Shape,
|
||||
*,
|
||||
internal_type: Optional[Type[Numeric]] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[core.CopyAtom, Tensor]:
|
||||
if internal_type is not None:
|
||||
if not isinstance(internal_type, NumericMeta):
|
||||
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
|
||||
internal_type = internal_type.mlir_type
|
||||
check_type_in(
|
||||
op,
|
||||
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
|
||||
"op",
|
||||
"make_tma_tile_atom_A",
|
||||
)
|
||||
|
||||
ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
|
||||
mma_tiler_mk = (mma_tiler_mnk[0], *mma_tiler_mnk[2:])
|
||||
g_tile = core.composition(ident, mma_tiler_mk, loc=loc, ip=ip)
|
||||
cta_v_map = tiled_mma._thrfrg_A(g_tile)
|
||||
cta_v_map = core.get(cta_v_map, mode=[1])
|
||||
cta_v_map = core.dice(cta_v_map, (1, (1,) * core.rank(g_tile)))
|
||||
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
num_multicast = 1
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
# multicast across the N-mode since those would share the same tile of A
|
||||
num_multicast = core.size(cluster_shape_vmnk, mode=[2])
|
||||
|
||||
# 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(
|
||||
gmem_tensor.value,
|
||||
smem_layout,
|
||||
cta_v_map,
|
||||
op._to_ir(),
|
||||
num_multicast=num_multicast,
|
||||
internal_type=internal_type,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
return core.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
return (
|
||||
core.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
res[1],
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_tma_tile_atom_B(
|
||||
op: Union[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
|
||||
gmem_tensor: Tensor,
|
||||
smem_layout: Layout,
|
||||
mma_tiler_mnk: Shape,
|
||||
tiled_mma: core.TiledMma,
|
||||
cluster_shape_vmnk: Shape,
|
||||
*,
|
||||
internal_type: Optional[Type[Numeric]] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[core.CopyAtom, Tensor]:
|
||||
if internal_type is not None:
|
||||
if not isinstance(internal_type, NumericMeta):
|
||||
raise TypeError(f"internal_type must be a Numeric, but got {internal_type}")
|
||||
internal_type = internal_type.mlir_type
|
||||
check_type_in(
|
||||
op,
|
||||
[CopyBulkTensorTileG2SOp, CopyBulkTensorTileG2SMulticastOp],
|
||||
"op",
|
||||
"make_tma_tile_atom_B",
|
||||
)
|
||||
|
||||
ident = core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
|
||||
mma_tiler_nk = (mma_tiler_mnk[1], *mma_tiler_mnk[2:])
|
||||
g_tile = core.composition(ident, mma_tiler_nk, loc=loc, ip=ip)
|
||||
cta_v_map = tiled_mma._thrfrg_B(g_tile)
|
||||
cta_v_map = core.get(cta_v_map, mode=[1])
|
||||
cta_v_map = core.dice(cta_v_map, (1, (1,) * core.rank(g_tile)))
|
||||
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
num_multicast = 1
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
# multicast across the M-mode since those would share the same tile of B
|
||||
num_multicast = core.size(cluster_shape_vmnk, mode=[1])
|
||||
|
||||
# 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(
|
||||
gmem_tensor.value,
|
||||
smem_layout,
|
||||
cta_v_map,
|
||||
op._to_ir(),
|
||||
num_multicast=num_multicast,
|
||||
internal_type=internal_type,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
if isinstance(op, CopyBulkTensorTileG2SOp):
|
||||
return core.CopyAtom(op, CopyBulkTensorTileG2SNonExecTrait(res[0])), res[1]
|
||||
else:
|
||||
assert isinstance(op, CopyBulkTensorTileG2SMulticastOp)
|
||||
return (
|
||||
core.CopyAtom(op, CopyBulkTensorTileG2SMulticastNonExecTrait(res[0])),
|
||||
res[1],
|
||||
)
|
||||
|
||||
|
||||
__all__ = [
|
||||
"make_tma_tile_atom_A",
|
||||
"make_tma_tile_atom_B",
|
||||
]
|
||||
@@ -0,0 +1,57 @@
|
||||
# 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 .copy import *
|
||||
from .mma import *
|
||||
from .helpers import *
|
||||
|
||||
# __all__ is required here for documentation generation
|
||||
__all__ = [
|
||||
#
|
||||
# copy.py
|
||||
#
|
||||
"Repetition",
|
||||
"Pack",
|
||||
"Unpack",
|
||||
"Ld16x64bOp",
|
||||
"Ld16x128bOp",
|
||||
"Ld16x256bOp",
|
||||
"Ld16x32bx2Op",
|
||||
"Ld32x32bOp",
|
||||
"St16x64bOp",
|
||||
"St16x128bOp",
|
||||
"St16x256bOp",
|
||||
"St16x32bx2Op",
|
||||
"St32x32bOp",
|
||||
#
|
||||
# mma.py
|
||||
#
|
||||
"OperandMajorMode",
|
||||
"OperandSource",
|
||||
"CtaGroup",
|
||||
"Field",
|
||||
"MmaTF32Op",
|
||||
"MmaF16BF16Op",
|
||||
"MmaI8Op",
|
||||
"MmaFP8Op",
|
||||
"SmemLayoutAtomKind",
|
||||
#
|
||||
# helpers.py
|
||||
#
|
||||
"make_smem_layout_atom",
|
||||
"tile_to_mma_shape",
|
||||
"commit",
|
||||
"is_tmem_load",
|
||||
"is_tmem_store",
|
||||
"get_tmem_copy_properties",
|
||||
"find_tmem_tensor_col_offset",
|
||||
"make_tmem_copy",
|
||||
]
|
||||
@@ -0,0 +1,465 @@
|
||||
# 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 enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
|
||||
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 ...typing import Numeric
|
||||
|
||||
|
||||
class Repetition(enum.Enum):
|
||||
"""
|
||||
An enumeration for the number of repetitions of a given TMEM copy within the instruction.
|
||||
"""
|
||||
|
||||
x1 = 1
|
||||
x2 = 2
|
||||
x4 = 4
|
||||
x8 = 8
|
||||
x16 = 16
|
||||
x32 = 32
|
||||
x64 = 64
|
||||
x128 = 128
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
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):
|
||||
"""
|
||||
An enumeration for the possible packing patterns for TMEM to RMEM copies.
|
||||
"""
|
||||
|
||||
NONE = enum.auto()
|
||||
PACK_16b_IN_32b = enum.auto()
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
|
||||
class Unpack(enum.Enum):
|
||||
"""
|
||||
An enumeration for the possible unpacking patterns for RMEM to TMEM copies.
|
||||
"""
|
||||
|
||||
NONE = enum.auto()
|
||||
UNPACK_32b_IN_16b = enum.auto()
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _LdBase(CopyOp):
|
||||
repeat: Repetition = Repetition.x1
|
||||
pack: Pack = Pack.NONE
|
||||
|
||||
admissible_archs = ["sm_100a"]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
if not isinstance(self.repeat, Repetition):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'repeat' Op parameter to be a tcgen05.Repetition instance",
|
||||
)
|
||||
if not isinstance(self.pack, Pack):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'pack' Op parameter to be a tcgen05.Pack instance",
|
||||
)
|
||||
|
||||
def __str__(self) -> 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"
|
||||
return res
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Ld16x64bOp(_LdBase):
|
||||
"""
|
||||
16x64b TMEM load Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
|
||||
This Operation corresponds to the ``.16x64b`` qualifier.
|
||||
"""
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x64bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
64,
|
||||
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))
|
||||
|
||||
|
||||
class Ld16x64bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Ld16x128bOp(_LdBase):
|
||||
"""
|
||||
16x128b TMEM load Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
|
||||
This Operation corresponds to the ``.16x128b`` qualifier.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.repeat == Repetition.x128:
|
||||
raise OpError(
|
||||
self,
|
||||
"x128 repetition is not supported",
|
||||
suggestion="choose one of x1, x2, x4, x8, x16, x32, x64",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x128bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
128,
|
||||
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))
|
||||
|
||||
|
||||
class Ld16x128bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Ld16x256bOp(_LdBase):
|
||||
"""
|
||||
16x256b TMEM load Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
|
||||
This Operation corresponds to the ``.16x256b`` qualifier.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.repeat in (Repetition.x128, Repetition.x64):
|
||||
raise OpError(
|
||||
self,
|
||||
"x64 and x128 repetition is not supported",
|
||||
suggestion="choose one of x1, x2, x4, x8, x16, x32",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x256bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
256,
|
||||
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))
|
||||
|
||||
|
||||
class Ld16x256bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Ld16x32bx2Op(_LdBase):
|
||||
"""
|
||||
16x32bx2 TMEM load Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
|
||||
This Operation corresponds to the ``.16x32bx2`` qualifier.
|
||||
"""
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld16x32bx2Trait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
32,
|
||||
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))
|
||||
|
||||
|
||||
class Ld16x32bx2Trait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class Ld32x32bOp(_LdBase):
|
||||
"""
|
||||
32x32b TMEM load Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
|
||||
This Operation corresponds to the ``.32x32`` qualifier.
|
||||
"""
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "Ld32x32bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemLoadType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
32,
|
||||
32,
|
||||
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))
|
||||
|
||||
|
||||
class Ld32x32bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class _StBase(CopyOp):
|
||||
repeat: Repetition
|
||||
unpack: Unpack = Unpack.NONE
|
||||
|
||||
admissible_archs = ["sm_100a"]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Arch verification
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
|
||||
if not isinstance(self.repeat, Repetition):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'repeat' Op parameter to be a tcgen05.Repetition instance",
|
||||
)
|
||||
if not isinstance(self.unpack, Unpack):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'pack' Op parameter to be a tcgen05.Unpack instance",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
res = (
|
||||
f"tcgen05 {self.__class__.__name__[:-2]} Copy Operation"
|
||||
+ 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"
|
||||
return res
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class St16x64bOp(_StBase):
|
||||
"""
|
||||
16x64b TMEM store Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st>`__.
|
||||
This Operation corresponds to the ``.16x64`` qualifier.
|
||||
"""
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "St16x64bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemStoreType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
64,
|
||||
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))
|
||||
|
||||
|
||||
class St16x64bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class St16x128bOp(_StBase):
|
||||
"""
|
||||
16x128b TMEM store Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st>`__.
|
||||
This Operation corresponds to the ``.16x128`` qualifier.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.repeat == Repetition.x128:
|
||||
raise OpError(
|
||||
self,
|
||||
"x128 repetition is not supported",
|
||||
suggestion="choose one of x1, x2, x4, x8, x16, x32, x64",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "St16x128bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemStoreType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
128,
|
||||
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))
|
||||
|
||||
|
||||
class St16x128bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class St16x256bOp(_StBase):
|
||||
"""
|
||||
16x256b TMEM store Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st>`__.
|
||||
This Operation corresponds to the ``.16x256`` qualifier.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.repeat in (Repetition.x128, Repetition.x64):
|
||||
raise OpError(
|
||||
self,
|
||||
"x64 and x128 repetition is not supported",
|
||||
suggestion="choose one of x1, x2, x4, x8, x16, x32",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "St16x256bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemStoreType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
256,
|
||||
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))
|
||||
|
||||
|
||||
class St16x256bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class St16x32bx2Op(_StBase):
|
||||
"""
|
||||
16x32x2b TMEM store Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st>`__.
|
||||
This Operation corresponds to the ``.16x32x2`` qualifier.
|
||||
"""
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "St16x32bx2Trait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemStoreType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
16,
|
||||
32,
|
||||
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))
|
||||
|
||||
|
||||
class St16x32bx2Trait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class St32x32bOp(_StBase):
|
||||
"""
|
||||
32x32b TMEM store Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-st>`__.
|
||||
This Operation corresponds to the ``.32x32`` qualifier.
|
||||
"""
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "St32x32bTrait":
|
||||
ty = _cute_nvgpu_ir.CopyAtomSM100TmemStoreType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
32,
|
||||
32,
|
||||
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))
|
||||
|
||||
|
||||
class St32x32bTrait(Trait):
|
||||
pass
|
||||
@@ -0,0 +1,301 @@
|
||||
# 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 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
|
||||
|
||||
from ...typing import (
|
||||
Shape,
|
||||
IntTuple,
|
||||
Layout,
|
||||
Tensor,
|
||||
Int,
|
||||
Numeric,
|
||||
NumericMeta,
|
||||
Int16,
|
||||
Int32,
|
||||
)
|
||||
from ... import core
|
||||
from .mma import SmemLayoutAtomKind, CtaGroup
|
||||
from .copy import (
|
||||
Pack,
|
||||
Unpack,
|
||||
Ld16x64bOp,
|
||||
Ld16x128bOp,
|
||||
Ld16x256bOp,
|
||||
Ld16x32bx2Op,
|
||||
Ld32x32bOp,
|
||||
St16x64bOp,
|
||||
St16x128bOp,
|
||||
St16x256bOp,
|
||||
St16x32bx2Op,
|
||||
St32x32bOp,
|
||||
)
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# Helper functions for MMA
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_smem_layout_atom(
|
||||
kind: SmemLayoutAtomKind, element_type: Type[Numeric], *, loc=None, ip=None
|
||||
) -> core.ComposedLayout:
|
||||
"""
|
||||
Makes a SMEM layout Atom.
|
||||
|
||||
This function creates a composed layout in unit of elements consistent with the requested layout
|
||||
Atom kind and element data type.
|
||||
|
||||
:param kind: The kind of layout Atom
|
||||
:type kind: SmemLayoutAtomKind
|
||||
: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
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(f"element_type must be a Numeric, but got {element_type}")
|
||||
|
||||
if kind in (SmemLayoutAtomKind.MN_INTER, SmemLayoutAtomKind.K_INTER):
|
||||
num_contiguous_bits = 128
|
||||
sw = core.make_swizzle(0, 4, 3)
|
||||
elif kind in (SmemLayoutAtomKind.MN_SW32, SmemLayoutAtomKind.K_SW32):
|
||||
num_contiguous_bits = 256
|
||||
sw = core.make_swizzle(1, 4, 3)
|
||||
elif kind in (SmemLayoutAtomKind.MN_SW64, SmemLayoutAtomKind.K_SW64):
|
||||
num_contiguous_bits = 512
|
||||
sw = core.make_swizzle(2, 4, 3)
|
||||
elif kind in (SmemLayoutAtomKind.MN_SW128, SmemLayoutAtomKind.K_SW128):
|
||||
num_contiguous_bits = 1024
|
||||
sw = core.make_swizzle(3, 4, 3)
|
||||
elif kind == SmemLayoutAtomKind.MN_SW128_32B:
|
||||
num_contiguous_bits = 1024
|
||||
sw = core.make_swizzle(2, 5, 2)
|
||||
else:
|
||||
raise ValueError("unrecognized SMEM layout atom kind")
|
||||
num_contiguous_elems = num_contiguous_bits // element_type.width
|
||||
|
||||
if kind in (
|
||||
SmemLayoutAtomKind.MN_INTER,
|
||||
SmemLayoutAtomKind.MN_SW32,
|
||||
SmemLayoutAtomKind.MN_SW64,
|
||||
SmemLayoutAtomKind.MN_SW128,
|
||||
SmemLayoutAtomKind.MN_SW128_32B,
|
||||
):
|
||||
# M/N-major layout
|
||||
return core.make_composed_layout(
|
||||
sw,
|
||||
0,
|
||||
core.make_layout(
|
||||
(num_contiguous_elems, 8), stride=(1, num_contiguous_elems)
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
else:
|
||||
# K-major layout
|
||||
return core.make_composed_layout(
|
||||
sw,
|
||||
0,
|
||||
core.make_layout(
|
||||
(8, num_contiguous_elems), stride=(num_contiguous_elems, 1)
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@overload
|
||||
def tile_to_mma_shape(
|
||||
atom: Layout, mma_tile_shape: Shape, order: IntTuple = None, *, loc=None, ip=None
|
||||
) -> Layout: ...
|
||||
|
||||
|
||||
@overload
|
||||
def tile_to_mma_shape(
|
||||
atom: core.ComposedLayout,
|
||||
mma_tile_shape: Shape,
|
||||
order: IntTuple = None,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> core.ComposedLayout: ...
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def tile_to_mma_shape(
|
||||
atom, mma_tile_shape: Shape, order: IntTuple = None, *, loc=None, ip=None
|
||||
):
|
||||
"""
|
||||
Tiles a layout to an MMA shape.
|
||||
"""
|
||||
# Default order is colexicographical
|
||||
if order is None:
|
||||
order = tuple(range(core.rank(mma_tile_shape) - 1))
|
||||
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}"
|
||||
)
|
||||
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)
|
||||
|
||||
if not (
|
||||
core.is_static(atom)
|
||||
and core.is_static(mma_tile_shape_val)
|
||||
and core.is_static(order_val)
|
||||
):
|
||||
raise ValueError("tile_to_mma_shape only supports static inputs")
|
||||
|
||||
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)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def commit(
|
||||
mbar_ptr: core.Pointer,
|
||||
mask=None,
|
||||
cta_group: CtaGroup = CtaGroup.ONE,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> None:
|
||||
"""
|
||||
Perform an arrive operation on a mbarrier upon completion of previous MMA operations.
|
||||
|
||||
:param mbar_ptr: A pointer to the mbarrier in SMEM
|
||||
:type mbar_ptr: Pointer
|
||||
:param mask: An optional multicast mask for the CTAs in the cluster to signal arrival to
|
||||
:type mask: Int
|
||||
"""
|
||||
if cta_group == CtaGroup.ONE:
|
||||
group = nvvm.Tcgen05GroupKind.CTA_1
|
||||
else:
|
||||
assert cta_group == CtaGroup.TWO
|
||||
group = nvvm.Tcgen05GroupKind.CTA_2
|
||||
|
||||
mbar_ptr = mbar_ptr.llvm_ptr
|
||||
if mask is not None:
|
||||
mask = Int16(mask).ir_value(loc=loc, ip=ip)
|
||||
nvvm.tcgen05_commit_arrive(
|
||||
mbar_ptr, multicast_mask=mask, group=group, loc=loc, ip=ip
|
||||
)
|
||||
else:
|
||||
nvvm.tcgen05_commit_arrive(mbar_ptr, group=group, loc=loc, ip=ip)
|
||||
return
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# Helper functions for Copies
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
def is_tmem_load(atom: core.CopyAtom) -> bool:
|
||||
"""
|
||||
Returns whether a CopyAtom instance is a TMEM load.
|
||||
"""
|
||||
return isinstance(
|
||||
atom.op,
|
||||
(
|
||||
Ld16x64bOp,
|
||||
Ld16x128bOp,
|
||||
Ld16x256bOp,
|
||||
Ld16x32bx2Op,
|
||||
Ld32x32bOp,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def is_tmem_store(atom: core.CopyAtom) -> bool:
|
||||
"""
|
||||
Returns whether a CopyAtom instance is a TMEM store.
|
||||
"""
|
||||
return isinstance(
|
||||
atom.op,
|
||||
(
|
||||
St16x64bOp,
|
||||
St16x128bOp,
|
||||
St16x256bOp,
|
||||
St16x32bx2Op,
|
||||
St32x32bOp,
|
||||
),
|
||||
)
|
||||
|
||||
|
||||
def get_tmem_copy_properties(
|
||||
atom: core.CopyAtom,
|
||||
) -> Tuple[int, int, int, Union[Pack, Unpack]]:
|
||||
"""
|
||||
Returns the properties of a TMEM copy atom (number of data paths, bits, repetitions,
|
||||
and whether packing/unpacking is used).
|
||||
"""
|
||||
if isinstance(atom.op, (Ld16x64bOp, St16x64bOp)):
|
||||
num_dp, num_bits = 16, 64
|
||||
elif isinstance(atom.op, (Ld16x128bOp, St16x128bOp)):
|
||||
num_dp, num_bits = 16, 128
|
||||
elif isinstance(atom.op, (Ld16x256bOp, St16x256bOp)):
|
||||
num_dp, num_bits = 16, 256
|
||||
elif isinstance(atom.op, (Ld16x32bx2Op, St16x32bx2Op)):
|
||||
num_dp, num_bits = 16, 32
|
||||
elif isinstance(atom.op, (Ld32x32bOp, St32x32bOp)):
|
||||
num_dp, num_bits = 32, 32
|
||||
else:
|
||||
raise ValueError(f"expects 'atom' to be a TMEM copy, but got {atom}")
|
||||
if is_tmem_load(atom):
|
||||
return num_dp, num_bits, atom.op.repeat.value, atom.op.pack
|
||||
else:
|
||||
assert is_tmem_store(atom), "atom must be a TMEM store"
|
||||
return num_dp, num_bits, atom.op.repeat.value, atom.op.unpack
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def find_tmem_tensor_col_offset(tmem_tensor: Tensor, *, loc=None, ip=None) -> Int:
|
||||
"""
|
||||
Computes the TMEM column offset given a TMEM tensor.
|
||||
|
||||
:param tmem_tensor: The TMEM tensor to use to compute the columns offset
|
||||
:type tmem_tensor: Tensor
|
||||
:return: The columns offset
|
||||
:rtype: Int
|
||||
"""
|
||||
tmem_col_mask = 0x0000FFFF
|
||||
offset = (
|
||||
core.cosize(core.recast_tensor(tmem_tensor, Int32).layout, loc=loc, ip=ip)
|
||||
& tmem_col_mask
|
||||
)
|
||||
if isinstance(offset, int):
|
||||
return offset
|
||||
return Int32(offset, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_tmem_copy(
|
||||
atom: core.CopyAtom, tmem_tensor: Tensor, *, loc=None, ip=None
|
||||
) -> core.TiledCopy:
|
||||
"""
|
||||
Makes a Tiled Copy instance from a TMEM Copy Atom and a TMEM tensor.
|
||||
"""
|
||||
tiled_copy_val = _cute_nvgpu_ir.atom_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)
|
||||
@@ -0,0 +1,603 @@
|
||||
# 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 enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
|
||||
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
|
||||
from ...typing import (
|
||||
Shape,
|
||||
Float8E5M2,
|
||||
Float8E4M3FN,
|
||||
Float16,
|
||||
BFloat16,
|
||||
Float32,
|
||||
TFloat32,
|
||||
Boolean,
|
||||
Int8,
|
||||
Uint8,
|
||||
Int32,
|
||||
Numeric,
|
||||
)
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# MMA Ops and Traits
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class OperandMajorMode(enum.Enum):
|
||||
"""
|
||||
An enumeration for the majorness of the input operands of the MMA.
|
||||
"""
|
||||
|
||||
MN = _cute_ir.MajorMode.mn
|
||||
K = _cute_ir.MajorMode.k
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
@classmethod
|
||||
def _missing_(cls, value):
|
||||
if isinstance(value, str):
|
||||
value = value.upper()
|
||||
if value == "MN":
|
||||
return OperandMajorMode.MN
|
||||
elif value == "K":
|
||||
return OperandMajorMode.K
|
||||
|
||||
def _to_ir(self) -> _cute_ir.MajorMode:
|
||||
return self.value
|
||||
|
||||
|
||||
class OperandSource(enum.Enum):
|
||||
"""
|
||||
An enumeration for the source memory location of the A input operand of the MMA.
|
||||
"""
|
||||
|
||||
TMEM = _cute_ir.MmaFragKind.tmem
|
||||
SMEM = _cute_ir.MmaFragKind.smem_desc
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
def _to_ir(self) -> _cute_ir.MmaFragKind:
|
||||
return self.value
|
||||
|
||||
|
||||
class CtaGroup(enum.Enum):
|
||||
"""
|
||||
An enumeration for the ``cta_group`` qualifier of the MMA.
|
||||
"""
|
||||
|
||||
ONE = 1
|
||||
TWO = 2
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
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.
|
||||
"""
|
||||
|
||||
NEGATE_A = "neg_a"
|
||||
NEGATE_B = "neg_b"
|
||||
ACCUMULATE = "accum_c"
|
||||
|
||||
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_field_name(self) -> str:
|
||||
return self.value
|
||||
|
||||
|
||||
# Base class for all tcgen05 MMA Ops used to factor out some internal code
|
||||
@dataclass(frozen=True)
|
||||
class MmaOp(MmaOp):
|
||||
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
|
||||
|
||||
admissible_archs = ["sm_100a"]
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
# Verify arch
|
||||
arch = CuTeDSL._get_dsl().envar.arch
|
||||
if arch not in self.admissible_archs:
|
||||
raise OpError(
|
||||
self,
|
||||
f"expects arch to be one of {self.admissible_archs}, but got {arch}",
|
||||
suggestion="Ensure env CUTE_DSL_ARCH matches your GPU architecture",
|
||||
)
|
||||
# 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",
|
||||
)
|
||||
# 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]
|
||||
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 8 <= 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}"
|
||||
)
|
||||
|
||||
|
||||
class MmaTrait(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<{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
|
||||
)
|
||||
|
||||
|
||||
#
|
||||
# TF32 MMA
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaTF32Op(MmaOp):
|
||||
"""
|
||||
TF32 tcgen05 MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma>`__.
|
||||
This Operation corresponds to the ``.kind::tf32`` qualifier.
|
||||
"""
|
||||
|
||||
descriptive_name = "tcgen05 TF32 MMA Operation"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
instruction_shape: Shape,
|
||||
cta_group: CtaGroup,
|
||||
a_src: OperandSource,
|
||||
a_major_mode: OperandMajorMode,
|
||||
b_major_mode: OperandMajorMode,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
TFloat32,
|
||||
TFloat32,
|
||||
Float32,
|
||||
instruction_shape,
|
||||
cta_group,
|
||||
a_src,
|
||||
a_major_mode,
|
||||
b_major_mode,
|
||||
)
|
||||
self._verify()
|
||||
|
||||
def _verify(self) -> None:
|
||||
# Verify the instruction shape
|
||||
instruction_k = 8
|
||||
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) -> "MmaTF32Trait":
|
||||
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.MmaAtomSM100UMMAType.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.cta_group.value,
|
||||
self.a_major_mode._to_ir(),
|
||||
self.b_major_mode._to_ir(),
|
||||
self.a_dtype.mlir_type,
|
||||
self.b_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
self.a_src._to_ir(),
|
||||
0,
|
||||
)
|
||||
return MmaTF32Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
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),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaTF32Trait(MmaTrait):
|
||||
pass
|
||||
|
||||
|
||||
#
|
||||
# F16/BF16 MMA
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaF16BF16Op(MmaOp):
|
||||
"""
|
||||
F16/BF16 tcgen05 MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma>`__.
|
||||
This Operation corresponds to the ``.kind::f16`` qualifier.
|
||||
"""
|
||||
|
||||
descriptive_name = "tcgen05 F16/BF16 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,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
ab_dtype,
|
||||
ab_dtype,
|
||||
acc_dtype,
|
||||
instruction_shape,
|
||||
cta_group,
|
||||
a_src,
|
||||
a_major_mode,
|
||||
b_major_mode,
|
||||
)
|
||||
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 = 16
|
||||
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) -> "MmaF16BF16Trait":
|
||||
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.MmaAtomSM100UMMAType.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.cta_group.value,
|
||||
self.a_major_mode._to_ir(),
|
||||
self.b_major_mode._to_ir(),
|
||||
self.a_dtype.mlir_type,
|
||||
self.b_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
self.a_src._to_ir(),
|
||||
0,
|
||||
)
|
||||
return MmaF16BF16Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
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),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaF16BF16Trait(MmaTrait):
|
||||
pass
|
||||
|
||||
|
||||
#
|
||||
# I8 MMA
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaI8Op(MmaOp):
|
||||
"""
|
||||
I8 tcgen05 MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma>`__.
|
||||
This Operation corresponds to the ``.kind::i8`` qualifier.
|
||||
"""
|
||||
|
||||
descriptive_name = "tcgen05 I8 MMA Operation"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ab_dtype: Type[Numeric],
|
||||
instruction_shape: Shape,
|
||||
cta_group: CtaGroup,
|
||||
a_src: OperandSource,
|
||||
a_major_mode: OperandMajorMode,
|
||||
b_major_mode: OperandMajorMode,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
ab_dtype,
|
||||
ab_dtype,
|
||||
Int32,
|
||||
instruction_shape,
|
||||
cta_group,
|
||||
a_src,
|
||||
a_major_mode,
|
||||
b_major_mode,
|
||||
)
|
||||
self._verify()
|
||||
|
||||
def _verify(self) -> None:
|
||||
# Input data type verification
|
||||
if self.a_dtype not in [Int8, Uint8]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'ab_dtype' Op parameter to be one of Int8 or Uint8",
|
||||
)
|
||||
assert self.b_dtype == self.a_dtype, "a_dtype and b_dtype must be the same"
|
||||
# Instruction shape verification
|
||||
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]}",
|
||||
)
|
||||
|
||||
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.MmaAtomSM100UMMAType.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.cta_group.value,
|
||||
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()),
|
||||
T.si32(),
|
||||
self.a_src._to_ir(),
|
||||
0,
|
||||
)
|
||||
return MmaI8Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
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),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaI8Trait(MmaTrait):
|
||||
pass
|
||||
|
||||
|
||||
#
|
||||
# F8F6F4 MMA
|
||||
#
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaFP8Op(MmaOp):
|
||||
"""
|
||||
F8 tcgen05 MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-mma-instructions-mma>`__.
|
||||
"""
|
||||
|
||||
descriptive_name = "tcgen05 F8 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,
|
||||
) -> None:
|
||||
|
||||
super().__init__(
|
||||
ab_dtype,
|
||||
ab_dtype,
|
||||
acc_dtype,
|
||||
instruction_shape,
|
||||
cta_group,
|
||||
a_src,
|
||||
a_major_mode,
|
||||
b_major_mode,
|
||||
)
|
||||
self._verify()
|
||||
|
||||
def _verify(self) -> None:
|
||||
# Input data type verification
|
||||
if self.a_dtype not in [Float8E5M2, Float8E4M3FN]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'ab_dtype' Op parameter to be one of Float8E5M2 or Float8E4M3FN",
|
||||
)
|
||||
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
|
||||
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) -> "MmaFP8Trait":
|
||||
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.MmaAtomSM100UMMAType.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.cta_group.value,
|
||||
self.a_major_mode._to_ir(),
|
||||
self.b_major_mode._to_ir(),
|
||||
self.a_dtype.mlir_type,
|
||||
self.b_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
self.a_src._to_ir(),
|
||||
0,
|
||||
)
|
||||
return MmaFP8Trait(
|
||||
_cute_nvgpu_ir.make_sm100_mma(
|
||||
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),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaFP8Trait(MmaTrait):
|
||||
pass
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# SMEM layout atoms
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class SmemLayoutAtomKind(enum.Enum):
|
||||
"""
|
||||
Enum class for the kinds of SMEM layout atoms for SM100.
|
||||
|
||||
Given a swizzle kind, an SMEM layout atom is the compact layout of smallest size that can be
|
||||
used to construct an SMEM layout using blocked product for operand A or B such that the
|
||||
resulting layout is legal for both TMA and UMMA.
|
||||
|
||||
Note that there are other ways of creating legal layouts for operand A and B.
|
||||
"""
|
||||
|
||||
MN_INTER = enum.auto()
|
||||
MN_SW32 = enum.auto()
|
||||
MN_SW64 = enum.auto()
|
||||
MN_SW128 = enum.auto()
|
||||
MN_SW128_32B = enum.auto()
|
||||
K_INTER = enum.auto()
|
||||
K_SW32 = enum.auto()
|
||||
K_SW64 = enum.auto()
|
||||
K_SW128 = enum.auto()
|
||||
@@ -0,0 +1,25 @@
|
||||
# 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 .copy import *
|
||||
from .mma import *
|
||||
|
||||
|
||||
# __all__ is required here for documentation generation
|
||||
__all__ = [
|
||||
# mma.py
|
||||
"MmaF16BF16Op",
|
||||
# copy.py
|
||||
"LdMatrix8x8x16bOp",
|
||||
"LdMatrix16x16x8bOp",
|
||||
"StMatrix8x8x16bOp",
|
||||
"StMatrix16x8x8bOp",
|
||||
]
|
||||
@@ -0,0 +1,189 @@
|
||||
# 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 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
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from ..common import OpError
|
||||
from ...core import CopyOp, Trait, _pack_shape
|
||||
from ...typing import Numeric
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class BaseOp(CopyOp):
|
||||
transpose: bool = False
|
||||
num_matrices: int = 1
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
if not isinstance(self.transpose, bool):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'transpose' Op parameter to be a bool instance",
|
||||
)
|
||||
|
||||
def __str__(self) -> str:
|
||||
res = (
|
||||
f"{self.__class__.__name__[:-2]} Copy Operation"
|
||||
+ f"\n number of matrices = {self.num_matrices}"
|
||||
)
|
||||
if self.transpose:
|
||||
res += f"\n transposed"
|
||||
return res
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LdMatrix8x8x16bOp(BaseOp):
|
||||
"""
|
||||
8x8 ``ldmatrix`` Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-load-instruction-ldmatrix>`__.
|
||||
This operation corresponds to the ``.m8n8`` qualifier.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.num_matrices not in [1, 2, 4]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "LdMatrix8x8x16bTrait":
|
||||
mode = _pack_shape((8, 8), loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.CopyAtomLdsmType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
mode.type.attribute,
|
||||
_cute_nvgpu_ir.LdsmSzPattern.u16,
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get() if self.transpose else None,
|
||||
)
|
||||
return LdMatrix8x8x16bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class LdMatrix8x8x16bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class LdMatrix16x16x8bOp(BaseOp):
|
||||
"""
|
||||
16x16 8-bit ``ldmatrix`` Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-load-instruction-ldmatrix>`__.
|
||||
This operation corresponds to the ``.m16n16`` and the ``.b16`` qualifiers.
|
||||
"""
|
||||
|
||||
def __init__(self, num_matrices: int) -> None:
|
||||
super().__init__(transpose=True, num_matrices=num_matrices)
|
||||
self._verify()
|
||||
|
||||
def _verify(self):
|
||||
assert self.transpose, "transpose must be True"
|
||||
if self.num_matrices not in [1, 2]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'num_matrices' Op parameter to be one of [1,2]",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "LdMatrix16x16x8bTrait":
|
||||
mode = _pack_shape((16, 16), loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.CopyAtomLdsmType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
mode.type.attribute,
|
||||
_cute_nvgpu_ir.LdsmSzPattern.u8,
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get(),
|
||||
)
|
||||
return LdMatrix16x16x8bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class LdMatrix16x16x8bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StMatrix8x8x16bOp(BaseOp):
|
||||
"""
|
||||
8x8 ``stmatrix`` Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-stmatrix>`__.
|
||||
This operation corresponds to the ``m8n8`` qualifier.
|
||||
"""
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
super().__post_init__()
|
||||
if self.num_matrices not in [1, 2, 4]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "StMatrix8x8x16bTrait":
|
||||
mode = _pack_shape((8, 8), loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.CopyAtomStsmType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
mode.type.attribute,
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get() if self.transpose else None,
|
||||
)
|
||||
return StMatrix8x8x16bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class StMatrix8x8x16bTrait(Trait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class StMatrix16x8x8bOp(BaseOp):
|
||||
"""
|
||||
16x8 ``stmatrix`` Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-stmatrix>`__.
|
||||
This operation corresponds to the ``m16n8`` qualifier.
|
||||
"""
|
||||
|
||||
def __init__(self, num_matrices: int) -> None:
|
||||
super().__init__(transpose=True, num_matrices=num_matrices)
|
||||
self._verify()
|
||||
|
||||
def _verify(self):
|
||||
if self.num_matrices not in [1, 2, 4]:
|
||||
assert self.transpose, "transpose must be True"
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
|
||||
)
|
||||
|
||||
def _make_trait(
|
||||
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
|
||||
) -> "StMatrix16x8x8bTrait":
|
||||
mode = _pack_shape((16, 8), loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.CopyAtomStsmType.get(
|
||||
copy_internal_type.mlir_type,
|
||||
mode.type.attribute,
|
||||
self.num_matrices,
|
||||
ir.UnitAttr.get(),
|
||||
)
|
||||
return StMatrix16x8x8bTrait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
|
||||
|
||||
class StMatrix16x8x8bTrait(Trait):
|
||||
pass
|
||||
@@ -0,0 +1,78 @@
|
||||
# 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 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
|
||||
|
||||
from ..common import OpError
|
||||
from ...core import MmaOp, Trait, _pack_shape
|
||||
from ...typing import Shape, Float16, BFloat16, Float32, Numeric
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaF16BF16Op(MmaOp):
|
||||
"""
|
||||
F16/BF16 tcgen05 MMA Operation.
|
||||
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#warp-level-matrix-instructions-mma>`__.
|
||||
This Operation covers the instructions using the ``.f16`` or ``.bf16`` qualifiers for the input operands.
|
||||
"""
|
||||
|
||||
ab_dtype: Type[Numeric]
|
||||
acc_dtype: Type[Numeric]
|
||||
shape_mnk: Shape
|
||||
|
||||
def __post_init__(self) -> None:
|
||||
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, 8), (16, 8, 16)]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'shape_mnk' Op parameter to be one of (16,8,8) or (16,8,16)",
|
||||
)
|
||||
|
||||
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaF16BF16Trait":
|
||||
shape_mnk = _pack_shape(self.shape_mnk, loc=loc, ip=ip)
|
||||
ty = _cute_nvgpu_ir.MmaAtomSM80Type.get(
|
||||
shape_mnk.type.attribute,
|
||||
self.ab_dtype.mlir_type,
|
||||
self.ab_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
)
|
||||
return MmaF16BF16Trait(_cute_ir.atom(ty, loc=loc, ip=ip))
|
||||
|
||||
def __str__(self) -> str:
|
||||
return (
|
||||
"warp-level F16/BF16 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}"
|
||||
)
|
||||
|
||||
|
||||
class MmaF16BF16Trait(Trait):
|
||||
pass
|
||||
@@ -0,0 +1,29 @@
|
||||
# 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 .mma import *
|
||||
from .helpers import *
|
||||
|
||||
# __all__ is required here for documentation generation
|
||||
__all__ = [
|
||||
# mma.py
|
||||
"OperandMajorMode",
|
||||
"OperandSource",
|
||||
"Field",
|
||||
"MmaF16BF16Op",
|
||||
"MmaF8Op",
|
||||
"SmemLayoutAtomKind",
|
||||
# helpers.py
|
||||
"make_smem_layout_atom",
|
||||
"fence",
|
||||
"commit_group",
|
||||
"wait_group",
|
||||
]
|
||||
@@ -0,0 +1,109 @@
|
||||
# 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 typing import Type
|
||||
|
||||
from cutlass.cutlass_dsl import dsl_user_op
|
||||
|
||||
from cutlass._mlir.dialects import nvvm
|
||||
|
||||
from ...typing import Numeric, NumericMeta
|
||||
from ... import core
|
||||
from .mma import SmemLayoutAtomKind
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_smem_layout_atom(
|
||||
kind: SmemLayoutAtomKind, element_type: Type[Numeric], *, loc=None, ip=None
|
||||
) -> core.ComposedLayout:
|
||||
"""
|
||||
Makes a SMEM layout Atom.
|
||||
|
||||
This function creates a composed layout in unit of elements consistent with the requested layout
|
||||
Atom kind and element data type.
|
||||
|
||||
:param kind: The kind of layout Atom
|
||||
:type kind: SmemLayoutAtomKind
|
||||
: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
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(f"element_type must be a Numeric, but got {element_type}")
|
||||
|
||||
if kind in (SmemLayoutAtomKind.MN_INTER, SmemLayoutAtomKind.K_INTER):
|
||||
num_contiguous_bits = 128
|
||||
sw = core.make_swizzle(0, 4, 3)
|
||||
elif kind in (SmemLayoutAtomKind.MN_SW32, SmemLayoutAtomKind.K_SW32):
|
||||
num_contiguous_bits = 256
|
||||
sw = core.make_swizzle(1, 4, 3)
|
||||
elif kind in (SmemLayoutAtomKind.MN_SW64, SmemLayoutAtomKind.K_SW64):
|
||||
num_contiguous_bits = 512
|
||||
sw = core.make_swizzle(2, 4, 3)
|
||||
elif kind in (SmemLayoutAtomKind.MN_SW128, SmemLayoutAtomKind.K_SW128):
|
||||
num_contiguous_bits = 1024
|
||||
sw = core.make_swizzle(3, 4, 3)
|
||||
else:
|
||||
raise ValueError("unrecognized SMEM layout atom kind")
|
||||
num_contiguous_elems = num_contiguous_bits // element_type.width
|
||||
|
||||
if kind in (
|
||||
SmemLayoutAtomKind.MN_INTER,
|
||||
SmemLayoutAtomKind.MN_SW32,
|
||||
SmemLayoutAtomKind.MN_SW64,
|
||||
SmemLayoutAtomKind.MN_SW128,
|
||||
):
|
||||
# M/N-major layout
|
||||
return core.make_composed_layout(
|
||||
sw,
|
||||
0,
|
||||
core.make_layout(
|
||||
(num_contiguous_elems, 8), stride=(1, num_contiguous_elems)
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
else:
|
||||
# K-major layout
|
||||
return core.make_composed_layout(
|
||||
sw,
|
||||
0,
|
||||
core.make_layout(
|
||||
(8, num_contiguous_elems), stride=(num_contiguous_elems, 1)
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def fence(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-multiply-and-accumulate-instruction-wgmma-fence>`__.
|
||||
"""
|
||||
nvvm.wgmma_fence_aligned(loc=None, ip=None)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def commit_group(*, loc=None, ip=None) -> None:
|
||||
"""
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-warpgroup-level-matrix-instructions-wgmma-commit-group>`__.
|
||||
"""
|
||||
nvvm.wgmma_commit_group_sync_aligned(loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def wait_group(group, *, loc=None, ip=None) -> None:
|
||||
"""
|
||||
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#asynchronous-multiply-and-accumulate-instruction-wgmma-wait-group>`__.
|
||||
"""
|
||||
nvvm.wgmma_wait_group_sync_aligned(group, loc=loc, ip=ip)
|
||||
@@ -0,0 +1,380 @@
|
||||
# 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 enum
|
||||
from dataclasses import dataclass
|
||||
from typing import Type
|
||||
|
||||
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 MmaOp, Trait, _pack_shape, rank, depth
|
||||
from ...typing import (
|
||||
Shape,
|
||||
Float16,
|
||||
BFloat16,
|
||||
Float32,
|
||||
Boolean,
|
||||
Float8E5M2,
|
||||
Float8E4M3FN,
|
||||
Numeric,
|
||||
)
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# MMA Ops and Traits
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class OperandMajorMode(enum.Enum):
|
||||
"""
|
||||
An enumeration for the majorness of the input operands of the MMA.
|
||||
"""
|
||||
|
||||
MN = _cute_ir.MajorMode.mn
|
||||
K = _cute_ir.MajorMode.k
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
@classmethod
|
||||
def _missing_(cls, value):
|
||||
if isinstance(value, str):
|
||||
value = value.upper()
|
||||
if value == "MN":
|
||||
return OperandMajorMode.MN
|
||||
elif value == "K":
|
||||
return OperandMajorMode.K
|
||||
|
||||
def _to_ir(self) -> _cute_ir.MajorMode:
|
||||
return self.value
|
||||
|
||||
|
||||
class OperandSource(enum.Enum):
|
||||
"""
|
||||
An enumeration for the source memory location of the A input operand of the MMA.
|
||||
"""
|
||||
|
||||
RMEM = _cute_ir.MmaFragKind.rmem
|
||||
SMEM = _cute_ir.MmaFragKind.smem_desc
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"{self.__class__.__name__}.{self.name}"
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"<{self.__class__.__name__}.{self.name}>"
|
||||
|
||||
def _to_ir(self) -> _cute_ir.MmaFragKind:
|
||||
return self.value
|
||||
|
||||
|
||||
class Field(enum.Enum):
|
||||
"""
|
||||
An enumeration for the fields of the MMA Atom that can be modified at runtime.
|
||||
"""
|
||||
|
||||
ACCUMULATE = "accum_c"
|
||||
|
||||
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_field_name(self) -> str:
|
||||
return self.value
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class MmaOp(MmaOp):
|
||||
a_dtype: Type[Numeric]
|
||||
b_dtype: Type[Numeric]
|
||||
acc_dtype: Type[Numeric]
|
||||
shape_mnk: Shape
|
||||
a_src: OperandSource
|
||||
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:
|
||||
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.a_src, OperandSource):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'a_src' Op parameter to be a warpgroup.OperandSource instance",
|
||||
)
|
||||
if not isinstance(self.a_major_mode, OperandMajorMode):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'a_major_mode' Op parameter to be a warpgroup.OperandMajorMode instance",
|
||||
)
|
||||
if not isinstance(self.b_major_mode, OperandMajorMode):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'b_major_mode' Op parameter to be a warpgroup.OperandMajorMode instance",
|
||||
)
|
||||
# Verify 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]
|
||||
if m != 64:
|
||||
raise OpError(self, f"expects the M-mode to be 64, but got {m}")
|
||||
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}",
|
||||
)
|
||||
|
||||
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 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}"
|
||||
)
|
||||
|
||||
|
||||
class MmaTrait(Trait):
|
||||
admissible_fields = [Field.ACCUMULATE]
|
||||
|
||||
def set(self, field, value, *, loc=None, ip=None) -> None:
|
||||
if field not in self.admissible_fields:
|
||||
raise ValueError(
|
||||
f"invalid field, must be {Field.ACCUMULATE}, but got {field}"
|
||||
)
|
||||
field_name = f"#cute_nvgpu.atom_mma_field_sm90<{field._to_ir_field_name()}>"
|
||||
attr = ir.Attribute.parse(field_name)
|
||||
self.value = _cute_nvgpu_ir.atom_set_value(
|
||||
self.value, attr, Boolean(value).ir_value(loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@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>`__.
|
||||
This Operation covers the instructions using the ``.f16`` or ``.bf16`` qualifiers for the input operands.
|
||||
"""
|
||||
|
||||
descriptive_name = "warpgroup F16/BF16 MMA Operation"
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
ab_dtype: Type[Numeric],
|
||||
acc_dtype: Type[Numeric],
|
||||
instruction_shape: Shape,
|
||||
a_src: OperandSource,
|
||||
a_major_mode: OperandMajorMode,
|
||||
b_major_mode: OperandMajorMode,
|
||||
) -> None:
|
||||
super().__init__(
|
||||
ab_dtype,
|
||||
ab_dtype,
|
||||
acc_dtype,
|
||||
instruction_shape,
|
||||
a_src,
|
||||
a_major_mode,
|
||||
b_major_mode,
|
||||
)
|
||||
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",
|
||||
)
|
||||
if (self.a_dtype == BFloat16) and (self.acc_dtype != Float32):
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'acc_dtype' Op parameter to be Float32 when 'ab_dtype' is BFloat16",
|
||||
)
|
||||
# Verify the instruction shape
|
||||
instruction_k = 16
|
||||
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) -> "MmaF16BF16Trait":
|
||||
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(),
|
||||
self.a_dtype.mlir_type,
|
||||
self.b_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
self.a_src._to_ir(),
|
||||
)
|
||||
return MmaF16BF16Trait(
|
||||
_cute_nvgpu_ir.make_sm90_mma(
|
||||
ty,
|
||||
Boolean(False).ir_value(loc=loc, ip=ip),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaF16BF16Trait(MmaTrait):
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
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>`__.
|
||||
This Operation covers the instructions using the ``.e4m3`` or ``.e5m2`` qualifiers for the input operands.
|
||||
"""
|
||||
|
||||
descriptive_name = "warpgroup F8 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 [Float8E5M2, Float8E4M3FN]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'a_dtype' Op parameter to be one of Float8E5M2 or Float8E4M3FN",
|
||||
)
|
||||
if self.b_dtype not in [Float8E5M2, Float8E4M3FN]:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'b_dtype' Op parameter to be one of Float8E5M2 or Float8E4M3FN",
|
||||
)
|
||||
# Accumulator data type verification
|
||||
if self.acc_dtype != Float32:
|
||||
raise OpError(
|
||||
self,
|
||||
"expects the 'acc_dtype' Op parameter to be Float32",
|
||||
)
|
||||
# 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]}",
|
||||
)
|
||||
|
||||
def _make_trait(self, *, loc=None, ip=None, **kwargs) -> "MmaF8Trait":
|
||||
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(),
|
||||
self.a_dtype.mlir_type,
|
||||
self.b_dtype.mlir_type,
|
||||
self.acc_dtype.mlir_type,
|
||||
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
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
class MmaF8Trait(MmaTrait):
|
||||
pass
|
||||
|
||||
|
||||
####################################################################################################
|
||||
#
|
||||
# SMEM layout atoms
|
||||
#
|
||||
####################################################################################################
|
||||
|
||||
|
||||
class SmemLayoutAtomKind(enum.Enum):
|
||||
"""
|
||||
Enum class for the kinds of SMEM layout atoms for SM90.
|
||||
|
||||
Given a swizzle kind, an SMEM layout atom is the compact layout of smallest size that can
|
||||
be used to construct an SMEM layout using blocked product for operand A or B such that the
|
||||
resulting layout is legal for both TMA and UMMA.
|
||||
|
||||
Note that there are other ways of creating legal layouts for operand A and B.
|
||||
"""
|
||||
|
||||
MN_INTER = enum.auto()
|
||||
MN_SW32 = enum.auto()
|
||||
MN_SW64 = enum.auto()
|
||||
MN_SW128 = enum.auto()
|
||||
K_INTER = enum.auto()
|
||||
K_SW32 = enum.auto()
|
||||
K_SW64 = enum.auto()
|
||||
K_SW128 = enum.auto()
|
||||
@@ -0,0 +1,515 @@
|
||||
# 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 ctypes
|
||||
from functools import lru_cache
|
||||
import itertools
|
||||
import operator
|
||||
from time import time
|
||||
from typing import Union
|
||||
|
||||
# MLIR modules imports
|
||||
from cutlass._mlir import ir
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
|
||||
from cutlass.cutlass_dsl import TensorFormat, JitArgAdapterRegistry
|
||||
|
||||
# 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 .core import find, _Tensor as CoreTensor
|
||||
|
||||
|
||||
class _Pointer(Pointer):
|
||||
"""Runtime representation of a pointer that can inter-operate with various data structures,
|
||||
including numpy arrays and device memory.
|
||||
|
||||
:param pointer: The pointer to the data
|
||||
:type pointer: int or pointer-like object
|
||||
:param dtype: Data type of the elements pointed to
|
||||
:type dtype: Type
|
||||
:param mem_space: Memory space where the pointer resides, defaults to generic
|
||||
:type mem_space: _cute_ir.AddressSpace, optional
|
||||
:param assumed_align: Assumed alignment of input pointer in bytes, defaults to None
|
||||
:type assumed_align: int, optional
|
||||
|
||||
:ivar _pointer: The underlying pointer
|
||||
:ivar _dtype: Data type of the elements
|
||||
:ivar _addr_space: Memory space of the pointer
|
||||
:ivar _assumed_align: Alignment of the pointer in bytes
|
||||
:ivar _desc: C-type descriptor for the pointer
|
||||
:ivar _c_pointer: C-compatible pointer representation
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
pointer,
|
||||
dtype,
|
||||
mem_space: _cute_ir.AddressSpace = _cute_ir.AddressSpace.generic,
|
||||
assumed_align=None,
|
||||
):
|
||||
self._pointer = pointer
|
||||
self._dtype = dtype
|
||||
self._addr_space = mem_space
|
||||
|
||||
is_in_device = mem_space == _cute_ir.AddressSpace.gmem
|
||||
if assumed_align is None:
|
||||
if is_in_device:
|
||||
self._assumed_align = 32
|
||||
else:
|
||||
self._assumed_align = dtype.width // 8
|
||||
else:
|
||||
self._assumed_align = assumed_align
|
||||
|
||||
class PtrDescriptor(ctypes.Structure):
|
||||
"""A ctype descriptor for CuTe memref ptr"""
|
||||
|
||||
_fields_ = [("ptr", ctypes.c_void_p)]
|
||||
|
||||
def __str__(self):
|
||||
return f"0x{self.ptr:016x}"
|
||||
|
||||
self._desc = PtrDescriptor(int(self._pointer))
|
||||
self._c_pointer = ctypes.cast(ctypes.pointer(self._desc), ctypes.c_void_p)
|
||||
assert (
|
||||
self._desc.ptr % self._assumed_align == 0
|
||||
), f"pointer must be {self._assumed_align} bytes aligned"
|
||||
|
||||
def size_in_bytes(self) -> int:
|
||||
return ctypes.sizeof(self._desc)
|
||||
|
||||
def __get_mlir_types__(self):
|
||||
return [self.mlir_type]
|
||||
|
||||
def __c_pointers__(self):
|
||||
return [self._c_pointer]
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
assert len(values) == 1
|
||||
return values[0]
|
||||
|
||||
# Move mlir Type out of __init__ to decouple with mlir Context
|
||||
@property
|
||||
def mlir_type(self) -> ir.Type:
|
||||
return _cute_ir.PtrType.get(
|
||||
self._dtype.mlir_type, self._addr_space, self._assumed_align
|
||||
)
|
||||
|
||||
@property
|
||||
def element_type(self) -> Type[Numeric]:
|
||||
return self._dtype
|
||||
|
||||
@property
|
||||
def memspace(self):
|
||||
return self._addr_space
|
||||
|
||||
def verify(self, expected_py_type):
|
||||
if expected_py_type is Pointer:
|
||||
return True
|
||||
elif isinstance(expected_py_type, ir.Value) and expected_py_type.ty is Pointer:
|
||||
return True
|
||||
|
||||
return False
|
||||
|
||||
def __str__(self) -> str:
|
||||
return f"Ptr<0x{self._desc.ptr:016x}@{self._addr_space}>"
|
||||
|
||||
def __repr__(self):
|
||||
return self.__str__()
|
||||
|
||||
|
||||
class _Tensor(Tensor):
|
||||
def __init__(
|
||||
self,
|
||||
tensor,
|
||||
assumed_align=None,
|
||||
):
|
||||
# If tensor is already a DLPack object, use it directly
|
||||
if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"):
|
||||
self._dlpack_data = tensor
|
||||
else:
|
||||
self._dlpack_data = tensor.__dlpack__()
|
||||
self._dltensor_wrapper = None
|
||||
self._assumed_align = assumed_align
|
||||
self._is_dynamic = False
|
||||
self._memref_desc = None
|
||||
self._dtype = None
|
||||
|
||||
@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.
|
||||
|
||||
This decorator loads the DLTensorWrapper when needed,
|
||||
avoiding overhead in the critical path of calling JIT functions.
|
||||
"""
|
||||
|
||||
def wrapper(self, *args, **kwargs):
|
||||
if self._dltensor_wrapper is None:
|
||||
self._dltensor_wrapper = _cute_ir.DLTensorWrapper(self._dlpack_data)
|
||||
return func(self, *args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
@lazily_load_dltensor
|
||||
def mark_layout_dynamic(self, leading_dim: int | None = None):
|
||||
"""Marks the tensor layout as dynamic based on the leading dimension.
|
||||
|
||||
:param leading_dim: The leading dimension of the layout, defaults to None
|
||||
:type leading_dim: int, optional
|
||||
|
||||
When ``leading_dim`` is None, automatically deduces the leading dimension from the tensor layout.
|
||||
The layout can be deduced only when exactly one dimension has a stride of 1. Raises an error
|
||||
if the layout cannot be automatically deduced.
|
||||
|
||||
When ``leading_dim`` is explicitly specified, marks the layout as dynamic while setting the
|
||||
stride at ``leading_dim`` to 1. Also validates that the specified ``leading_dim`` is consistent
|
||||
with the existing layout by checking that the corresponding stride of that dimension is 1.
|
||||
|
||||
Limitation: only support flat layout for now. Will work on supporting nested layout in the future.
|
||||
|
||||
:return: The tensor with dynamic layout
|
||||
:rtype: _Tensor
|
||||
"""
|
||||
self._dltensor_wrapper.mark_layout_dynamic(leading_dim)
|
||||
return self
|
||||
|
||||
@lazily_load_dltensor
|
||||
def mark_compact_shape_dynamic(
|
||||
self,
|
||||
mode: int,
|
||||
stride_order: tuple[int, ...] | None = None,
|
||||
divisibility: int = 1,
|
||||
):
|
||||
"""Marks the tensor shape as dynamic and propagates dynamic and divisibility information to the corresponding strides.
|
||||
|
||||
:param mode: The mode of the compact shape, defaults to 0
|
||||
:type mode: int
|
||||
:param stride_order: Consistent with `torch.Tensor.dim_order`. Defaults to None.
|
||||
Indicates the order of the modes (dimensions) if the current layout were converted to row-major order.
|
||||
It starts from the outermost to the innermost dimension.
|
||||
:type stride_order: tuple[int, ...], optional
|
||||
:param divisibility: The divisibility constraint for the compact shape, defaults to 1
|
||||
:type divisibility: int, optional
|
||||
:return: The tensor with dynamic compact shape
|
||||
:rtype: _Tensor
|
||||
|
||||
If ``stride_order`` is not provided, the stride ordering will be automatically deduced from the layout.
|
||||
Automatic deduction is only possible when exactly one dimension has a stride of 1 (compact layout).
|
||||
An error is raised if automatic deduction fails.
|
||||
|
||||
If ``stride_order`` is explicitly specified, it does the consistency check with the layout.
|
||||
|
||||
For example:
|
||||
- Layout: (4,2):(1,4) has stride_order: (1,0) indicates the innermost dimension is 0(`4:1`), the outermost dimension is 1(`2:4`)
|
||||
- Layout: (5,3,2,4):(3,1,15,30) has stride_order: (3,2,0,1) indicates the innermost dimension is 1(`3:1`), the outermost dimension is 3(`4:30`).
|
||||
|
||||
Using `torch.Tensor.dim_order()` to get the stride order of the torch tensor.
|
||||
.. code-block:: python
|
||||
a = torch.empty(3, 4)
|
||||
t = cute.runtime.from_dlpack(a)
|
||||
t = t.mark_compact_shape_dynamic(mode=0, stride_order=a.dim_order())
|
||||
"""
|
||||
self._dltensor_wrapper.mark_compact_shape_dynamic(
|
||||
mode, stride_order, divisibility
|
||||
)
|
||||
return self
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def element_type(self) -> Type[Numeric]:
|
||||
if self._dtype is None:
|
||||
self._dtype = self._dltensor_wrapper.dtype
|
||||
return self._dtype
|
||||
|
||||
@element_type.setter
|
||||
def element_type(self, new_type):
|
||||
"""Set the element type of the tensor.
|
||||
|
||||
:warning: This API is added for narrow precision before we have a clean `recast_tensor` story.
|
||||
|
||||
:note: It is only used for the case that frameworks don't natively support narrow precision but we get tensor
|
||||
from frameworks with storage type like uint8.
|
||||
|
||||
**Example**:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
# Create a tensor from a numpy array
|
||||
import numpy as np
|
||||
from cutlass.cute import from_dlpack
|
||||
|
||||
# Create a tensor with Float32 elements
|
||||
a = np.zeros(shape, dtype=np.uint8)
|
||||
tensor = from_dlpack(a)
|
||||
|
||||
# Change the element type to Float4E2M1FN even storage type is uint8
|
||||
tensor.element_type = cutlass.Float4E2M1FN
|
||||
|
||||
src = from_dlpack(... data tensor ...)
|
||||
# convert and initialize narrow precision tensor
|
||||
cute.testing.convert(src, tensor)
|
||||
"""
|
||||
self._dtype = new_type
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def memspace(self):
|
||||
return self._dltensor_wrapper.address_space
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def size_in_bytes(self) -> int:
|
||||
return self._dltensor_wrapper.size_in_bytes()
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def mlir_type(self) -> ir.Type:
|
||||
return self._dltensor_wrapper.get_type(
|
||||
self.element_type.mlir_type, self._assumed_align
|
||||
)
|
||||
|
||||
@lazily_load_dltensor
|
||||
def __str__(self) -> str:
|
||||
return f"Tensor<0x{self._dltensor_wrapper.str}>"
|
||||
|
||||
def __repr__(self):
|
||||
return self.__str__()
|
||||
|
||||
def __setitem__(self, crd, value):
|
||||
raise TypeError(f"runtime._Tensor is not indexable")
|
||||
|
||||
def __getitem__(self, crd):
|
||||
raise TypeError(f"runtime._Tensor is not indexable")
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def iterator(self):
|
||||
return _Pointer(
|
||||
self._dltensor_wrapper.data_ptr,
|
||||
self.element_type,
|
||||
self.memspace,
|
||||
self._assumed_align,
|
||||
)
|
||||
|
||||
@property
|
||||
def layout(self):
|
||||
raise NotImplementedError(
|
||||
f"layout property is not supported in runtime, support in future"
|
||||
)
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def shape(self):
|
||||
return self._dltensor_wrapper.shape
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def stride(self):
|
||||
strides = self._dltensor_wrapper.stride
|
||||
if strides is None:
|
||||
strides = itertools.accumulate(
|
||||
reversed(self.shape), func=operator.mul, initial=1
|
||||
)
|
||||
strides = tuple(reversed(list(strides)[:-1]))
|
||||
|
||||
return strides
|
||||
|
||||
@property
|
||||
@lru_cache(maxsize=128, typed=True)
|
||||
def leading_dim(self):
|
||||
"""Get the leading dimension of this Tensor.
|
||||
|
||||
:return: The leading dimension index or indices
|
||||
:rtype: int or tuple or None
|
||||
|
||||
The return value depends on the tensor's stride pattern:
|
||||
|
||||
* If a single leading dimension is found, returns an integer index
|
||||
* If nested leading dimensions are found, returns a tuple of indices
|
||||
* If no leading dimension is found, returns None
|
||||
"""
|
||||
return find(1, self.stride, exclude_when=(1, self.shape))
|
||||
|
||||
def fill(self, value: Numeric):
|
||||
raise TypeError(f"fill function is not supported in runtime")
|
||||
|
||||
@property
|
||||
@lazily_load_dltensor
|
||||
def data_ptr(self):
|
||||
return self._dltensor_wrapper.data_ptr
|
||||
|
||||
@lazily_load_dltensor
|
||||
def __c_pointers__(self):
|
||||
self._memref_desc = self._dltensor_wrapper.build_memref_desc(
|
||||
self._assumed_align
|
||||
)
|
||||
return [_cute_ir.pycapsule_get_pointer(self._memref_desc)]
|
||||
|
||||
def __get_mlir_types__(self):
|
||||
return [self.mlir_type]
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
assert len(values) == 1
|
||||
assert isinstance(values[0], CoreTensor)
|
||||
return CoreTensor(values[0].value, self._dtype)
|
||||
|
||||
|
||||
def from_dlpack(
|
||||
tensor_dlpack,
|
||||
assumed_align=None,
|
||||
) -> Tensor:
|
||||
"""Convert from tensor object supporting __dlpack__() to a CuTe Tensor.
|
||||
|
||||
:param tensor_dlpack: Tensor object that supports the DLPack protocol
|
||||
:type tensor_dlpack: object
|
||||
: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
|
||||
:return: A CuTe Tensor object
|
||||
:rtype: Tensor
|
||||
|
||||
Examples:
|
||||
.. code-block:: python
|
||||
|
||||
import torch
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
x = torch.randn(100, 100)
|
||||
y = from_dlpack(x)
|
||||
y.shape
|
||||
# (100, 100)
|
||||
type(y)
|
||||
# <class 'cutlass.cute.Tensor'>
|
||||
"""
|
||||
return _Tensor(
|
||||
tensor_dlpack,
|
||||
assumed_align=assumed_align,
|
||||
)
|
||||
|
||||
|
||||
def make_ptr(
|
||||
dtype: Type[Numeric],
|
||||
value: Union[int, ctypes._Pointer],
|
||||
mem_space: AddressSpace = AddressSpace.generic,
|
||||
assumed_align=None,
|
||||
) -> Pointer:
|
||||
"""Create a pointer from a memory address
|
||||
|
||||
:param dtype: Data type of the pointer elements
|
||||
:type dtype: Type[Numeric]
|
||||
:param value: Memory address as integer or ctypes pointer
|
||||
:type value: Union[int, ctypes._Pointer]
|
||||
:param mem_space: Memory address space, defaults to AddressSpace.generic
|
||||
:type mem_space: AddressSpace, optional
|
||||
:param align_bytes: Alignment in bytes, defaults to None
|
||||
:type align_bytes: int, optional
|
||||
:return: A pointer object
|
||||
:rtype: Pointer
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
import numpy as np
|
||||
import ctypes
|
||||
|
||||
from cutlass import Float32
|
||||
from cutlass.cute.runtime import make_ptr
|
||||
|
||||
# Create a numpy array
|
||||
a = np.random.randn(16, 32).astype(np.float32)
|
||||
|
||||
# Get pointer address as integer
|
||||
ptr_address = a.ctypes.data_as(ctypes.POINTER(ctypes.c_float))
|
||||
|
||||
# Create pointer from address
|
||||
y = make_ptr(cutlass.Float32, ptr_address)
|
||||
|
||||
# Check properties
|
||||
print(y.element_type)
|
||||
print(type(y)) # <class 'cutlass.cute.Pointer'>
|
||||
"""
|
||||
# check if value is int or ctypes.POINTER
|
||||
if isinstance(value, int):
|
||||
address_value = value
|
||||
elif isinstance(value, ctypes._Pointer):
|
||||
# get address value
|
||||
address_value = ctypes.cast(value, ctypes.c_void_p).value
|
||||
assert address_value is not None, "Pointer address is None"
|
||||
else:
|
||||
raise TypeError(
|
||||
f"Expect int or ctypes.POINTER for value but got {type(value)=}"
|
||||
)
|
||||
|
||||
return _Pointer(address_value, dtype, mem_space, assumed_align=assumed_align)
|
||||
|
||||
|
||||
class TensorAdapter:
|
||||
"""
|
||||
Convert a DLPack protocol supported tensor/array to a cute tensor.
|
||||
"""
|
||||
|
||||
# Need reference these capsules to avoid being garbage collected
|
||||
tensor_capsules = []
|
||||
|
||||
def __init__(self, arg):
|
||||
self._arg = from_dlpack(arg).mark_layout_dynamic()
|
||||
self.tensor_capsules.append(self._arg)
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
return self._arg.__new_from_mlir_values__(values)
|
||||
|
||||
def __c_pointers__(self):
|
||||
return self._arg.__c_pointers__()
|
||||
|
||||
def __get_mlir_types__(self):
|
||||
return self._arg.__get_mlir_types__()
|
||||
|
||||
|
||||
# -------------------------------------------------------------------------
|
||||
# Try to register_jit_arg_adapter for TensorAdapter
|
||||
# -------------------------------------------------------------------------
|
||||
|
||||
try: # Register for numpy.ndarray
|
||||
import numpy
|
||||
|
||||
JitArgAdapterRegistry.register_jit_arg_adapter(numpy.ndarray)(TensorAdapter)
|
||||
except ImportError:
|
||||
pass # silent attempt, suppress error
|
||||
|
||||
try: # Register for torch.Tensor
|
||||
import torch
|
||||
|
||||
JitArgAdapterRegistry.register_jit_arg_adapter(torch.Tensor)(TensorAdapter)
|
||||
except ImportError:
|
||||
pass # silent attempt, suppress error
|
||||
@@ -0,0 +1,285 @@
|
||||
# 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 random
|
||||
import numpy as np
|
||||
import functools
|
||||
import hashlib
|
||||
|
||||
from cutlass.cutlass_dsl import (
|
||||
const,
|
||||
T,
|
||||
CuTeDSL,
|
||||
BaseDSL,
|
||||
t,
|
||||
Constexpr,
|
||||
detect_gpu_arch,
|
||||
)
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.ir as ir
|
||||
from cutlass._mlir.dialects import nvvm, cf, vector, builtin
|
||||
|
||||
from cutlass.cute import core
|
||||
from cutlass.cute import nvgpu
|
||||
from typing import Type
|
||||
from inspect import isclass
|
||||
|
||||
|
||||
def assert_(cond, msg=None):
|
||||
if isinstance(cond, ir.Value):
|
||||
if ir.VectorType.isinstance(cond.type):
|
||||
assert (
|
||||
cond.type.element_type == T.bool()
|
||||
), f"only expects vector type with boolean elements, but got {cond.type}"
|
||||
cond_val = vector.multi_reduction(
|
||||
vector.CombiningKind.AND, cond, const(True), range(cond.type.rank)
|
||||
)
|
||||
else:
|
||||
cond_val = cond
|
||||
else:
|
||||
cond_val = const(cond, t.Boolean)
|
||||
|
||||
cf.assert_(cond_val, msg if msg else "")
|
||||
|
||||
|
||||
def _maybe_recast_tensor_from_f4(src: core.Tensor, tv_layout: core.Layout):
|
||||
if src.element_type.width == 4:
|
||||
tv_layout = core.recast_layout(8, 4, tv_layout)
|
||||
src = core.recast_tensor(src, dtype=t.Int8)
|
||||
return src, tv_layout
|
||||
|
||||
|
||||
def _maybe_recast_to_f4(input: core.TensorSSA, dtype: Type[core.Numeric]):
|
||||
"""Conditionally recasts the tensor to 4-bit type if the destination type is 4-bit.
|
||||
|
||||
:param input: The input tensor to recast.
|
||||
:param dtype: The target numeric type to potentially recast to.
|
||||
: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):
|
||||
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
|
||||
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 input
|
||||
|
||||
|
||||
def _maybe_recast_from_f4(input: core.TensorSSA, src_dtype: Type[core.Numeric]):
|
||||
"""Conditionally recasts the tensor from 4-bit type if the source type is 4-bit.
|
||||
|
||||
:param input: The input tensor to recast.
|
||||
:param src_dtype: The source numeric type to potentially recast from.
|
||||
: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):
|
||||
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
|
||||
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 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,
|
||||
src_ty,
|
||||
dst_ty,
|
||||
):
|
||||
tidx = nvvm.read_ptx_sreg_tid_x(T.i32())
|
||||
bidx = nvvm.read_ptx_sreg_ctaid_x(T.i32())
|
||||
|
||||
cta_coord = (None, bidx)
|
||||
# logical idx -> address
|
||||
ctaSrc = gSrc[cta_coord] # (...,TileV,...)
|
||||
ctaDst = gDst[cta_coord] # (...,TileV,...)
|
||||
ctaCSrc = cSrc[cta_coord] # (...,TileV,...)
|
||||
# print(f"ctaSrc = {ctaSrc.type}")
|
||||
|
||||
# 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)
|
||||
# print(f"tidfrgSrc = {tidfrgSrc.type}")
|
||||
|
||||
# slice for threads
|
||||
thr_coord = (tidx, None)
|
||||
thrSrc = tidfrgSrc[thr_coord] # (V)
|
||||
thrDst = tidfrgDst[thr_coord] # (V)
|
||||
thrCSrc = tidfrgCSrc[thr_coord] # (V)
|
||||
# print(f"thrSrc = {thrSrc.type}")
|
||||
|
||||
# predicate
|
||||
if core.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
|
||||
) # (V)
|
||||
frgDst = core.make_fragment(
|
||||
core.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)
|
||||
|
||||
vec_src = frgSrc.load()
|
||||
vec_src = _maybe_recast_to_f4(vec_src, src_ty)
|
||||
vec_dst = vec_src.to(dst_ty)
|
||||
vec_dst = _maybe_recast_from_f4(vec_dst, dst_ty)
|
||||
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)
|
||||
|
||||
|
||||
@CuTeDSL.jit(preprocess=False)
|
||||
def _convert(
|
||||
src: core.Tensor,
|
||||
dst: core.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))
|
||||
|
||||
# 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)
|
||||
|
||||
# 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,...)
|
||||
dst_cta_tiler = [
|
||||
1,
|
||||
] * core.rank(dst.layout)
|
||||
dst_cta_tiler[leading_mode] = core.size(dst_tv_layout) # (...,TileV,...)
|
||||
|
||||
# Step 4. partition input and output tensor by cta tiler.
|
||||
gS = core.zipped_divide(
|
||||
src, tuple(src_cta_tiler)
|
||||
) # ((...,TileV,...),(...,RestV,...))
|
||||
cS = core.zipped_divide(
|
||||
idA, tuple(src_cta_tiler)
|
||||
) # ((...,TileV,...),(...,RestV,...))
|
||||
gD = core.zipped_divide(
|
||||
dst, tuple(dst_cta_tiler)
|
||||
) # ((...,TileV,...),(...,RestV,...))
|
||||
# print(f"{gS.type=}")
|
||||
|
||||
_convert_kernel(
|
||||
gS,
|
||||
gD,
|
||||
cS,
|
||||
src_tv_layout,
|
||||
dst_tv_layout,
|
||||
src_shape,
|
||||
src_ty,
|
||||
dst_ty,
|
||||
).launch(
|
||||
grid=[core.size(gS, mode=[1]), 1, 1],
|
||||
block=[core.size(src_tv_layout, mode=[0]), 1, 1],
|
||||
)
|
||||
|
||||
|
||||
# Converts from src tensor to dst tensor, their logical shape are required to be the same.
|
||||
# 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."
|
||||
# find leading mode
|
||||
leading_mode = np.argmin([np.min(s) for s in src.stride])
|
||||
|
||||
elem_per_copy = 2
|
||||
|
||||
if src.element_type.width == 4 or dst.element_type.width == 4:
|
||||
elem_per_copy = 8
|
||||
elif src.element_type.width == 8 or dst.element_type.width == 8:
|
||||
elem_per_copy = 4
|
||||
assert (
|
||||
src.shape[leading_mode] % elem_per_copy == 0
|
||||
and dst.shape[leading_mode] % elem_per_copy == 0
|
||||
)
|
||||
_convert(src, dst, leading_mode, elem_per_copy)
|
||||
|
||||
|
||||
#########################################
|
||||
# Testing utilities
|
||||
#########################################
|
||||
|
||||
|
||||
def sample_pytest(rand_cfg=None):
|
||||
"""
|
||||
Decorator to randomly sample pytest parametrized tests.
|
||||
rand_cfg: Tuple[int, float] - (random_seed, sample_ratio)
|
||||
Sampling is disabled when:
|
||||
- A specific test is selected (via -k or direct test path)
|
||||
- Not running under pytest
|
||||
"""
|
||||
import functools
|
||||
import os
|
||||
import random
|
||||
import pytest
|
||||
import sys
|
||||
|
||||
seed, sample_ratio = rand_cfg
|
||||
random.seed(seed)
|
||||
|
||||
def decorator(func):
|
||||
@functools.wraps(func)
|
||||
def wrapper(*args, **kwargs):
|
||||
if rand_cfg is not None and "PYTEST_CURRENT_TEST" in os.environ:
|
||||
# Check if test was explicitly selected like ::test_name[param1-param2-...]
|
||||
if "-k" in sys.argv or any(".py::" in arg for arg in sys.argv):
|
||||
# Test was explicitly selected, don't skip
|
||||
return func(*args, **kwargs)
|
||||
|
||||
if random.uniform(0.0, 1.0) > sample_ratio:
|
||||
pytest.skip(f"Randomly skipped (sampling ratio: {sample_ratio})")
|
||||
return func(*args, **kwargs)
|
||||
|
||||
return wrapper
|
||||
|
||||
return decorator
|
||||
@@ -0,0 +1,193 @@
|
||||
# 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 abc import ABC, abstractmethod
|
||||
from typing import ForwardRef, Tuple, Union, Any, Type, List
|
||||
|
||||
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
|
||||
|
||||
|
||||
Int = Union[int, Integer]
|
||||
|
||||
|
||||
ScaledBasis = ForwardRef("ScaledBasis")
|
||||
|
||||
|
||||
IntTuple = Union[Int, Tuple["IntTuple", ...]]
|
||||
Shape = Union[Int, Tuple["Shape", ...]]
|
||||
Stride = Union[Int, ScaledBasis, Tuple["Stride", ...]]
|
||||
Coord = Union[Int, None, Tuple["Coord", ...]]
|
||||
|
||||
|
||||
class Layout(ir.Value):
|
||||
def __init__(self, op_result):
|
||||
super().__init__(op_result)
|
||||
|
||||
def __str__(self): ...
|
||||
|
||||
def get_hier_coord(self, idx) -> Coord:
|
||||
"""Return the (hierarchical) ND logical coordinate corresponding to the linear index"""
|
||||
...
|
||||
|
||||
@property
|
||||
def shape(self, *, loc=None, ip=None) -> Shape: ...
|
||||
|
||||
@property
|
||||
def stride(self, *, loc=None, ip=None) -> Stride: ...
|
||||
|
||||
|
||||
Tile = Union[Int, None, Layout, Tuple["Tile", ...]]
|
||||
|
||||
# XTuple is super set of above types
|
||||
XTuple = Union[IntTuple, Shape, Stride, Coord, Tile]
|
||||
|
||||
Tiler = Union[Shape, Layout, Tile]
|
||||
|
||||
|
||||
class Pointer:
|
||||
"""
|
||||
Abstract base class for CuTe jit function and runtime _Pointer
|
||||
"""
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
# Doesn't matter just return a value
|
||||
return [self]
|
||||
|
||||
|
||||
class Tensor(ABC):
|
||||
"""
|
||||
Abstract base class for CuTe jit function and runtime _Tensor
|
||||
|
||||
A CuTe Tensor is iterator with layout
|
||||
|
||||
: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,)
|
||||
|
||||
Define JIT function:
|
||||
|
||||
.. code-block:: python
|
||||
|
||||
@cute.jit
|
||||
def add(a: Tensor, b: Tensor, res: Tensor): ...
|
||||
|
||||
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)
|
||||
"""
|
||||
|
||||
def __str__(self): ...
|
||||
|
||||
@abstractmethod
|
||||
def __getitem__(self, idx) -> Union["Tensor", ir.Value, IntTuple]: ...
|
||||
|
||||
@abstractmethod
|
||||
def __setitem__(self, idx, value): ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def element_type(self) -> Union[Type[Numeric], Type[IntTuple]]: ...
|
||||
|
||||
@element_type.setter
|
||||
def element_type(self, new_type): ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def memspace(self) -> AddressSpace: ...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def iterator(self): ...
|
||||
|
||||
@property
|
||||
def layout(self) -> Union[Layout, "ComposedLayout"]: ...
|
||||
|
||||
@property
|
||||
def shape(self) -> Shape: ...
|
||||
|
||||
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_compact_shape_dynamic(
|
||||
self, mode: int, stride_order: tuple[int, ...]|None = None, divisibility: int = 1
|
||||
) -> "Tensor": ...
|
||||
|
||||
@abstractmethod
|
||||
def fill(self, value: Numeric) -> None: ...
|
||||
|
||||
|
||||
__all__ = [
|
||||
"Coord",
|
||||
"Numeric",
|
||||
"Integer",
|
||||
"Boolean",
|
||||
"Int8",
|
||||
"Int16",
|
||||
"Int32",
|
||||
"Int64",
|
||||
"Uint8",
|
||||
"Uint16",
|
||||
"Uint32",
|
||||
"Uint64",
|
||||
"Float",
|
||||
"Float16",
|
||||
"BFloat16",
|
||||
"TFloat32",
|
||||
"Float32",
|
||||
"Float64",
|
||||
"Float8E5M2",
|
||||
"Float8E4M3FN",
|
||||
"Float8E4M3B11FNUZ",
|
||||
"Float8E4M3",
|
||||
"Float8E8M0FNU",
|
||||
"Float4E2M1FN",
|
||||
"Float6E2M3FN",
|
||||
"Float6E3M2FN",
|
||||
"IntTuple",
|
||||
"Layout",
|
||||
"Pointer",
|
||||
"Shape",
|
||||
"Stride",
|
||||
"Tensor",
|
||||
"Tile",
|
||||
"Tiler",
|
||||
"XTuple",
|
||||
]
|
||||
Reference in New Issue
Block a user