Release v4.0.0 (#2294)
This commit is contained in:
Kihiro Bando
@@ -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 .cutlass_dsl import (
|
||||
Constexpr,
|
||||
as_numeric,
|
||||
min,
|
||||
max,
|
||||
and_,
|
||||
or_,
|
||||
all_,
|
||||
any_,
|
||||
not_,
|
||||
all_,
|
||||
any_,
|
||||
select_,
|
||||
# Control-flow without AST pre-processor
|
||||
if_generate,
|
||||
for_generate,
|
||||
LoopUnroll,
|
||||
while_generate,
|
||||
yield_out,
|
||||
# Control-flow with AST pre-processor
|
||||
range_constexpr,
|
||||
range_dynamic,
|
||||
const_expr,
|
||||
dynamic_expr,
|
||||
# Data types
|
||||
dtype, # Provides conversions to types inheriting from NumericType
|
||||
DSLRuntimeError,
|
||||
JitArgAdapterRegistry,
|
||||
# Construction utilities for user-defined classes
|
||||
extract_mlir_values,
|
||||
new_from_mlir_values,
|
||||
)
|
||||
|
||||
from .cute.typing import *
|
||||
|
||||
# Utilities not belonging to CuTe
|
||||
from . import utils as utils
|
||||
|
||||
# Used as internal symbol
|
||||
from . import cutlass_dsl as _dsl
|
||||
|
||||
# Aliases
|
||||
LaunchConfig = _dsl.BaseDSL.LaunchConfig
|
||||
register_jit_arg_adapter = _dsl.JitArgAdapterRegistry.register_jit_arg_adapter
|
||||
gpu = _dsl.cutlass_gpu
|
||||
cuda = _dsl.cuda_helpers
|
||||
@@ -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",
|
||||
]
|
||||
@@ -0,0 +1,32 @@
|
||||
# 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.
|
||||
|
||||
|
||||
def check_value_in(
|
||||
value, possible_values: list, value_description: str, prefix=""
|
||||
) -> None:
|
||||
if value not in possible_values:
|
||||
err_msg = prefix
|
||||
if err_msg != "":
|
||||
err_msg += ": "
|
||||
err_msg += f"invalid {value_description}, got {value}, must be one of {possible_values}"
|
||||
raise ValueError(err_msg)
|
||||
|
||||
|
||||
def check_type_in(ty, possible_types: list, type_description: str, prefix="") -> None:
|
||||
if not isinstance(ty, type):
|
||||
ty = type(ty)
|
||||
if ty not in possible_types:
|
||||
err_msg = prefix
|
||||
if err_msg != "":
|
||||
err_msg += ": "
|
||||
err_msg += f"invalid type for {type_description}, got {ty}, must be one of {possible_types}"
|
||||
raise TypeError(err_msg)
|
||||
@@ -0,0 +1,169 @@
|
||||
# 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 enum import Enum
|
||||
from typing import Optional, Type, Union
|
||||
|
||||
from cutlass.cute.typing import (
|
||||
Numeric,
|
||||
Boolean,
|
||||
Float,
|
||||
Integer,
|
||||
TFloat32,
|
||||
Float8E4M3B11FNUZ,
|
||||
Float8E4M3FN,
|
||||
Float8E5M2,
|
||||
Float8E8M0FNU,
|
||||
Float4E2M1FN,
|
||||
Tensor,
|
||||
)
|
||||
from cutlass.cute.runtime import from_dlpack
|
||||
import cutlass.cute as cute
|
||||
import torch
|
||||
|
||||
|
||||
def dtype(ty: Type[Numeric]):
|
||||
"""
|
||||
Return the corresponding torch.dtype per the given DSL type
|
||||
"""
|
||||
torch_dtype = getattr(torch, ty.__name__.lower(), None)
|
||||
|
||||
torch_type_map = {
|
||||
Boolean: torch.bool,
|
||||
# TFloat32 is just alias of float32
|
||||
TFloat32: torch.float32,
|
||||
Float8E5M2: torch.float8_e5m2,
|
||||
Float8E4M3FN: torch.float8_e4m3fn,
|
||||
Float8E4M3B11FNUZ: torch.float8_e4m3fnuz,
|
||||
}
|
||||
if torch_dtype is None:
|
||||
torch_dtype = torch_type_map.get(ty)
|
||||
|
||||
if torch_dtype is None:
|
||||
raise TypeError(f"{ty} is not supported by torch")
|
||||
return torch_dtype
|
||||
|
||||
|
||||
@dataclass
|
||||
class ScalarInitConfig:
|
||||
"""Configuration for scalar initialization"""
|
||||
|
||||
value: float = 0.0
|
||||
|
||||
|
||||
@dataclass
|
||||
class RandomInitConfig:
|
||||
"""Configuration for random initialization"""
|
||||
|
||||
min_val: int = -2
|
||||
max_val: int = 2
|
||||
|
||||
|
||||
@dataclass
|
||||
class GaussianInitConfig:
|
||||
"""Configuration for Gaussian initialization"""
|
||||
|
||||
mean: float = 0.0
|
||||
std: float = 1.0
|
||||
scale: float = 1.0
|
||||
|
||||
|
||||
class TensorInitType(Enum):
|
||||
"""Enumeration of tensor initialization types"""
|
||||
|
||||
SKIP = "skip"
|
||||
SCALAR = "scalar"
|
||||
RANDOM = "random"
|
||||
GAUSSIAN = "gaussian"
|
||||
|
||||
|
||||
def create_and_permute_torch_tensor(
|
||||
shape,
|
||||
dtype: "torch.dtype",
|
||||
permute_order=None,
|
||||
init_type: TensorInitType = TensorInitType.RANDOM,
|
||||
init_config: Optional[
|
||||
Union[RandomInitConfig, ScalarInitConfig, GaussianInitConfig]
|
||||
] = None,
|
||||
) -> "torch.Tensor":
|
||||
"""
|
||||
Create a torch tensor with specified shape and dtype. Optionally permute it and initialize it with specified init type and config
|
||||
"""
|
||||
init_dtype = torch.int32 if init_type == TensorInitType.RANDOM else torch.float32
|
||||
init_torch_tensor = torch.empty(*shape, dtype=init_dtype)
|
||||
if init_type == TensorInitType.SKIP:
|
||||
assert init_config is None
|
||||
f32_torch_tensor = init_torch_tensor
|
||||
elif init_type == TensorInitType.SCALAR:
|
||||
if init_config is None:
|
||||
init_config = ScalarInitConfig()
|
||||
else:
|
||||
if not isinstance(init_config, ScalarInitConfig):
|
||||
raise ValueError("init_config must be ScalarInitConfig()")
|
||||
f32_torch_tensor = init_torch_tensor.fill_(init_config.value)
|
||||
elif init_type == TensorInitType.RANDOM:
|
||||
if init_config is None:
|
||||
init_config = RandomInitConfig()
|
||||
else:
|
||||
if not isinstance(init_config, RandomInitConfig):
|
||||
raise ValueError("init_config must be RandomInitConfig()")
|
||||
f32_torch_tensor = init_torch_tensor.random_(
|
||||
init_config.min_val, init_config.max_val
|
||||
).to(dtype=torch.float32)
|
||||
elif init_type == TensorInitType.GAUSSIAN:
|
||||
if init_config is None:
|
||||
init_config = GaussianInitConfig()
|
||||
else:
|
||||
if not isinstance(init_config, GaussianInitConfig):
|
||||
raise ValueError("init_config must be GaussianInitConfig()")
|
||||
f32_torch_tensor = init_torch_tensor.normal_(init_config.mean, init_config.std)
|
||||
f32_torch_tensor = f32_torch_tensor * (1 << init_config.scale)
|
||||
else:
|
||||
raise ValueError(f"Invalid init type: {init_type}")
|
||||
|
||||
if permute_order is not None:
|
||||
f32_torch_tensor = f32_torch_tensor.permute(permute_order)
|
||||
|
||||
dtype_torch_tensor = f32_torch_tensor.to(dtype=dtype)
|
||||
|
||||
return dtype_torch_tensor
|
||||
|
||||
|
||||
def convert_cute_tensor(
|
||||
f32_torch_tensor: "torch.Tensor",
|
||||
cute_tensor: Tensor,
|
||||
dtype: Type[Numeric],
|
||||
is_dynamic_layout: bool = True,
|
||||
) -> Tensor:
|
||||
"""
|
||||
Change the value of the cute tensor to make its value converted from a fp32 torch tensor.
|
||||
Used for fp8 types tensor creatation now.
|
||||
"""
|
||||
# if torch_tensor is on cpu, create a gpu copy
|
||||
if f32_torch_tensor.device.type == "cpu":
|
||||
f32_torch_tensor = f32_torch_tensor.cuda()
|
||||
|
||||
# Fp8 type need explicit type conversion
|
||||
if dtype in {
|
||||
Float8E5M2,
|
||||
Float8E4M3FN,
|
||||
Float8E8M0FNU,
|
||||
Float4E2M1FN,
|
||||
}:
|
||||
fp32_cute_tensor = from_dlpack(f32_torch_tensor)
|
||||
if is_dynamic_layout:
|
||||
fp32_cute_tensor = fp32_cute_tensor.mark_layout_dynamic(
|
||||
f32_torch_tensor.dim_order()[-1]
|
||||
)
|
||||
# Copy and convert from f32 cute tensor to dtype cute tensor
|
||||
cute.testing.convert(fp32_cute_tensor, cute_tensor)
|
||||
return cute_tensor
|
||||
@@ -0,0 +1,9 @@
|
||||
# Utilities
|
||||
|
||||
This folder contains various utilties for kernel authoring. Specifically, the implementation of the
|
||||
followings can be considered experimental and subject to breaking changes:
|
||||
|
||||
- static persistent tile scheduler defined in [`static_persistent_tile_scheduler.py`](./static_persistent_tile_scheduler.py)
|
||||
- pipeline abstractions defined in [`pipeline.py`](./pipeline.py)
|
||||
- grouped GEMM utilties defined [`grouped_gemm_tile_scheduler_helper.py`](./grouped_gemm_tile_scheduler_helper.py)
|
||||
and [`tensormap_manager.py`](./tensormap_manager.py)
|
||||
@@ -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 .static_persistent_tile_scheduler import (
|
||||
WorkTileInfo,
|
||||
PersistentTileSchedulerParams,
|
||||
StaticPersistentTileScheduler,
|
||||
)
|
||||
|
||||
from .pipeline import (
|
||||
Agent,
|
||||
CooperativeGroup,
|
||||
PipelineUserType,
|
||||
PipelineState,
|
||||
make_pipeline_state,
|
||||
PipelineAsync,
|
||||
PipelineTmaAsync,
|
||||
PipelineTmaUmma,
|
||||
PipelineUmmaAsync,
|
||||
PipelineTmaStore,
|
||||
pipeline_init_wait,
|
||||
)
|
||||
|
||||
from .hardware_info import (
|
||||
HardwareInfo,
|
||||
)
|
||||
|
||||
from .blackwell_helpers import (
|
||||
compute_epilogue_tile_shape,
|
||||
get_smem_store_op,
|
||||
get_tmem_load_op,
|
||||
get_num_tmem_alloc_cols,
|
||||
make_smem_layout_a,
|
||||
make_smem_layout_b,
|
||||
make_smem_layout_epi,
|
||||
make_trivial_tiled_mma,
|
||||
)
|
||||
|
||||
from .hopper_helpers import (
|
||||
sm90_get_smem_store_op,
|
||||
)
|
||||
|
||||
from .grouped_gemm_tile_scheduler_helper import (
|
||||
GroupSearchResult,
|
||||
GroupedGemmGroupSearchState,
|
||||
GroupedGemmTileSchedulerHelper,
|
||||
create_initial_search_state,
|
||||
)
|
||||
|
||||
from .tensormap_manager import (
|
||||
TensorMapUpdateMode,
|
||||
TensorMapManager,
|
||||
)
|
||||
|
||||
from .smem_allocator import SmemAllocator
|
||||
|
||||
from .layout import LayoutEnum
|
||||
|
||||
__all__ = [
|
||||
"WorkTileInfo",
|
||||
"PersistentTileSchedulerParams",
|
||||
"StaticPersistentTileScheduler",
|
||||
"TensorMapUpdateMode",
|
||||
"TensorMapManager",
|
||||
"GroupSearchResult",
|
||||
"GroupedGemmGroupSearchState",
|
||||
"create_initial_search_state",
|
||||
"GroupedGemmTileSchedulerHelper",
|
||||
"HardwareInfo",
|
||||
]
|
||||
@@ -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 enum import Enum
|
||||
|
||||
|
||||
class SmemCapacity(Enum):
|
||||
SM80_SMEM_CAPACITY_BYTES = (164 - 1) * 1024
|
||||
SM86_SMEM_CAPACITY_BYTES = (100 - 1) * 1024
|
||||
SM89_SMEM_CAPACITY_BYTES = (100 - 1) * 1024
|
||||
|
||||
|
||||
# Dictionary to map compute capability to SMEM capacity
|
||||
SMEM_CAPACITY = {
|
||||
"sm80": SmemCapacity.SM80_SMEM_CAPACITY_BYTES.value,
|
||||
"sm86": SmemCapacity.SM86_SMEM_CAPACITY_BYTES.value,
|
||||
"sm89": SmemCapacity.SM89_SMEM_CAPACITY_BYTES.value,
|
||||
}
|
||||
@@ -0,0 +1,910 @@
|
||||
# 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 enum import Enum
|
||||
from math import log2, ceil
|
||||
from typing import List, Type, Union, Tuple
|
||||
|
||||
from cutlass.cutlass_dsl import (
|
||||
Float16,
|
||||
BFloat16,
|
||||
TFloat32,
|
||||
Float32,
|
||||
Uint8,
|
||||
Int8,
|
||||
Float8E4M3FN,
|
||||
Float8E5M2,
|
||||
Numeric,
|
||||
NumericMeta,
|
||||
dsl_user_op,
|
||||
)
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.nvgpu.common import CopyUniversalOp
|
||||
from cutlass.cute.nvgpu.warp import StMatrix8x8x16bOp, StMatrix16x8x8bOp
|
||||
from cutlass.cute.nvgpu.tcgen05 import (
|
||||
MmaF16BF16Op,
|
||||
MmaTF32Op,
|
||||
MmaI8Op,
|
||||
MmaFP8Op,
|
||||
OperandSource,
|
||||
OperandMajorMode,
|
||||
CtaGroup,
|
||||
Ld16x64bOp,
|
||||
Ld16x128bOp,
|
||||
Ld16x256bOp,
|
||||
Ld16x32bx2Op,
|
||||
Ld32x32bOp,
|
||||
Repetition,
|
||||
Pack,
|
||||
find_tmem_tensor_col_offset,
|
||||
SmemLayoutAtomKind,
|
||||
make_smem_layout_atom,
|
||||
tile_to_mma_shape,
|
||||
is_tmem_load,
|
||||
get_tmem_copy_properties,
|
||||
)
|
||||
from cutlass.utils.layout import LayoutEnum
|
||||
|
||||
@dsl_user_op
|
||||
def compute_epilogue_tile_shape(
|
||||
cta_tile_shape: cute.Shape,
|
||||
use_2cta_instrs: bool,
|
||||
layout_d: LayoutEnum,
|
||||
elem_ty_d: Type[Numeric],
|
||||
*,
|
||||
layout_c: LayoutEnum = None,
|
||||
elem_ty_c: Union[Type[Numeric], None] = None,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> cute.Tile:
|
||||
"""Attempts to compute a reasonable epilogue tile based on block tile shape or allows the user to provide one.
|
||||
|
||||
:param cta_tile_shape: A tuple or list representing the dimensions of the CTA tile, where
|
||||
cta_tile_shape[0] corresponds to the height (M) and cta_tile_shape[1]
|
||||
corresponds to the width (N) of the tile.
|
||||
:type cta_tile_shape: cute.Shape
|
||||
:param use_2cta_instrs: A flag indicating whether the configuration is for a 2SM setup.
|
||||
:type use_2cta_instrs: bool
|
||||
:param layout_d: The layout enum of the output tensor D.
|
||||
:type layout_d: LayoutEnum
|
||||
:param elem_ty_d: The element type of output tensor D.
|
||||
:type elem_ty_d: Type[Numeric]
|
||||
:param layout_c: The layout enum of the input tensor C. Defaults to None.
|
||||
:type layout_c: LayoutEnum, optional
|
||||
:param elem_ty_c: The element type for input tensor C. Defaults to None.
|
||||
:type elem_ty_c: Union[Type[Numeric], None], optional
|
||||
|
||||
:return: Returns epilog tiler, which is used in subsequent epilog partitions.
|
||||
:rtype: cute.Tile
|
||||
|
||||
:raises ValueError: If the computed tile cute.size does not meet minimum requirements based on CTA dimensions.
|
||||
"""
|
||||
|
||||
def validate_type(ty, ty_name):
|
||||
if not isinstance(ty, NumericMeta):
|
||||
raise TypeError(f"{ty_name} must be Numeric, but got {ty}")
|
||||
|
||||
validate_type(elem_ty_d, "elem_ty_d")
|
||||
if elem_ty_c is not None:
|
||||
validate_type(elem_ty_c, "elem_ty_c")
|
||||
|
||||
cta_m, cta_n = cta_tile_shape[:2]
|
||||
(warp_m, warp_n) = (2, 2) if (cta_m == 64 and use_2cta_instrs) else (4, 1)
|
||||
disable_source = elem_ty_c == None
|
||||
max_bits = (
|
||||
elem_ty_d.width if disable_source else max(elem_ty_c.width, elem_ty_d.width)
|
||||
)
|
||||
|
||||
dp_full = 32
|
||||
tile_m = min(cta_m, dp_full * warp_m)
|
||||
n_perf = 0
|
||||
if disable_source:
|
||||
if max_bits == 4:
|
||||
compute_elts = 8192
|
||||
else:
|
||||
compute_elts = 4096
|
||||
n_perf = compute_elts // tile_m
|
||||
else:
|
||||
if max_bits == 32:
|
||||
n_perf = 16 if (cta_m > 64 and cta_n <= 128) else 32
|
||||
elif max_bits == 16:
|
||||
n_perf = 32 if cta_n <= 128 else 64
|
||||
else:
|
||||
n_perf = 64
|
||||
|
||||
d_is_m_major = layout_d.is_m_major_c()
|
||||
c_is_m_major = True if layout_c is None else layout_c.is_m_major_c()
|
||||
|
||||
n_min_d = (
|
||||
8 * warp_n
|
||||
if d_is_m_major
|
||||
else (128 * warp_n if elem_ty_d.width == 6 else 128 // elem_ty_d.width * warp_n)
|
||||
)
|
||||
n_min_c = (
|
||||
8 * warp_n
|
||||
if (c_is_m_major or disable_source)
|
||||
else (128 * warp_n if elem_ty_c.width == 6 else 128 // elem_ty_c.width * warp_n)
|
||||
)
|
||||
tile_n = min(cta_n, max(n_perf, n_min_c, n_min_d))
|
||||
|
||||
if cta_n < n_min_c or cta_n < n_min_d:
|
||||
raise ValueError(f"CTA tile too small: {cta_tile_shape=}")
|
||||
|
||||
# stride by tmem warp layout and return a by-mode tiler
|
||||
tile_m_layout = cute.make_layout(tile_m, loc=loc, ip=ip)
|
||||
tile_n_layout = cute.make_layout(
|
||||
(tile_n // warp_n, warp_n), stride=(1, cta_n // warp_n), loc=loc, ip=ip
|
||||
)
|
||||
return (tile_m_layout, cute.coalesce(tile_n_layout, loc=loc, ip=ip))
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def get_smem_store_op(
|
||||
layout_d: LayoutEnum,
|
||||
elem_ty_d: Type[Numeric],
|
||||
elem_ty_acc: Type[Numeric],
|
||||
tiled_tmem_load: cute.TiledCopy,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> cute.CopyAtom:
|
||||
"""Selects the largest vectorized smem store atom available subject to
|
||||
constraint of gmem layout and chosen TMEM_LOAD's thread-value ownership.
|
||||
|
||||
:param layout_d: The layout enum of the output tensor D.
|
||||
:type layout_d: LayoutEnum
|
||||
:param elem_ty_d: The element type for output tensor D.
|
||||
:type elem_ty_d: Type[Numeric]
|
||||
:param elem_ty_acc: The element type for accumulator.
|
||||
:type elem_ty_acc: Type[Numeric]
|
||||
:param tiled_tmem_load: An instance of TiledCopy that represents the tmem load operation.
|
||||
:type tiled_tmem_load: cute.TiledCopy
|
||||
|
||||
:return: Either SmemStoreMatrix or SimtSyncCopy, based on the input parameters.
|
||||
:rtype: cute.CopyAtom
|
||||
"""
|
||||
|
||||
def validate_type(ty, ty_name):
|
||||
if not isinstance(ty, NumericMeta):
|
||||
raise TypeError(f"{ty_name} must be a Numeric, but got {ty}")
|
||||
|
||||
validate_type(elem_ty_d, "elem_ty_d")
|
||||
validate_type(elem_ty_acc, "elem_ty_acc")
|
||||
|
||||
is_m_major = layout_d.is_m_major_c()
|
||||
is_n_major = layout_d.is_n_major_c()
|
||||
|
||||
if not is_tmem_load(tiled_tmem_load):
|
||||
return cute.make_copy_atom(CopyUniversalOp(), elem_ty_d, loc=loc, ip=ip)
|
||||
|
||||
num_dp, num_bits, num_rep, pack = get_tmem_copy_properties(tiled_tmem_load)
|
||||
|
||||
use_stmatrix_m8n8_4x = (
|
||||
all(
|
||||
[
|
||||
elem_ty_acc.width == 32,
|
||||
elem_ty_d.width == 32,
|
||||
is_n_major,
|
||||
num_dp == 16,
|
||||
num_bits == 128,
|
||||
num_rep in (2, 4, 8, 16, 32, 64),
|
||||
pack == Pack.NONE,
|
||||
]
|
||||
)
|
||||
or all(
|
||||
[
|
||||
elem_ty_acc.width == 32,
|
||||
elem_ty_d.width == 16,
|
||||
num_dp == 16,
|
||||
num_bits == 256,
|
||||
num_rep in (2, 4, 8, 16, 32),
|
||||
pack == Pack.NONE,
|
||||
]
|
||||
)
|
||||
or all(
|
||||
[
|
||||
elem_ty_acc.width == 16,
|
||||
elem_ty_d.width == 16,
|
||||
num_dp == 16,
|
||||
num_bits == 128,
|
||||
num_rep in (2, 4, 8, 16, 32, 64),
|
||||
pack == Pack.PACK_16b_IN_32b,
|
||||
]
|
||||
)
|
||||
)
|
||||
use_stmatrix_m16n8_4x = all(
|
||||
[
|
||||
elem_ty_acc.width == 32,
|
||||
elem_ty_d.width == 8,
|
||||
is_m_major,
|
||||
num_dp == 16,
|
||||
num_bits == 256,
|
||||
num_rep in (4, 8, 16, 32),
|
||||
pack == Pack.NONE,
|
||||
]
|
||||
)
|
||||
use_stmatrix_m8n8_2x = (
|
||||
all(
|
||||
[
|
||||
elem_ty_acc.width == 32,
|
||||
elem_ty_d.width == 32,
|
||||
is_n_major,
|
||||
num_dp == 16,
|
||||
num_bits == 128,
|
||||
num_rep == 1,
|
||||
pack == Pack.NONE,
|
||||
]
|
||||
)
|
||||
or all(
|
||||
[
|
||||
elem_ty_acc.width == 32,
|
||||
elem_ty_d.width == 16,
|
||||
num_dp == 16,
|
||||
num_bits == 256,
|
||||
num_rep == 1,
|
||||
pack == Pack.NONE,
|
||||
]
|
||||
)
|
||||
or all(
|
||||
[
|
||||
elem_ty_acc.width == 16,
|
||||
elem_ty_d.width == 16,
|
||||
num_dp == 16,
|
||||
num_bits == 128,
|
||||
num_rep == 1,
|
||||
pack == Pack.PACK_16b_IN_32b,
|
||||
]
|
||||
)
|
||||
)
|
||||
use_stmatrix_m16n8_2x = all(
|
||||
[
|
||||
elem_ty_acc.width == 32,
|
||||
elem_ty_d.width == 8,
|
||||
is_m_major,
|
||||
num_dp == 16,
|
||||
num_bits == 256,
|
||||
num_rep == 2,
|
||||
pack == Pack.NONE,
|
||||
]
|
||||
)
|
||||
use_stmatrix_m16n8_1x = all(
|
||||
[
|
||||
elem_ty_acc.width == 32,
|
||||
elem_ty_d.width == 8,
|
||||
is_m_major,
|
||||
num_dp == 16,
|
||||
num_bits == 256,
|
||||
num_rep == 1,
|
||||
pack == Pack.NONE,
|
||||
]
|
||||
)
|
||||
|
||||
if use_stmatrix_m8n8_4x:
|
||||
op = StMatrix8x8x16bOp(is_m_major, 4)
|
||||
return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip)
|
||||
elif use_stmatrix_m8n8_2x:
|
||||
op = StMatrix8x8x16bOp(is_m_major, 2)
|
||||
return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip)
|
||||
elif use_stmatrix_m16n8_4x:
|
||||
op = StMatrix16x8x8bOp(4)
|
||||
return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip)
|
||||
elif use_stmatrix_m16n8_2x:
|
||||
op = StMatrix16x8x8bOp(2)
|
||||
return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip)
|
||||
elif use_stmatrix_m16n8_1x:
|
||||
op = StMatrix16x8x8bOp(1)
|
||||
return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip)
|
||||
else:
|
||||
op = CopyUniversalOp()
|
||||
return cute.make_copy_atom(op, elem_ty_d, loc=loc, ip=ip)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def get_tmem_load_op(
|
||||
cta_tile_shape: cute.Shape,
|
||||
layout_d: LayoutEnum,
|
||||
elem_ty_d: Type[Numeric],
|
||||
elem_ty_acc: Type[Numeric],
|
||||
epi_tile: cute.Tile,
|
||||
use_2cta_instrs: bool,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> cute.CopyAtom:
|
||||
"""Finds a performant TMEM_LOAD copy op for the selected epilogue
|
||||
tile (epi_tile), element types, and tcgen05.mma instruction used.
|
||||
|
||||
:param cta_tile_shape: A tuple or list representing the dimensions of the CTA tile.
|
||||
:type cta_tile_shape: cute.Shape
|
||||
:param layout_d: The layout enum of the output tensor D.
|
||||
:type layout_d: LayoutEnum
|
||||
:param elem_ty_d: The element type for output tensor D.
|
||||
:type elem_ty_d: Type[Numeric]
|
||||
:param elem_ty_acc: The element type for accumulation.
|
||||
:type elem_ty_acc: Type[Numeric]
|
||||
:param epi_tile: The epilogue tile configuration.
|
||||
:type epi_tile: cute.Tile
|
||||
:param use_2cta_instrs: A flag indicating whether the configuration is for 2 SMs.
|
||||
:type use_2cta_instrs: bool
|
||||
|
||||
:return: An instance of Sm100TmemLoad with the computed configuration.
|
||||
:rtype: cute.CopyAtom
|
||||
|
||||
:raises ValueError: If the function cannot handle the given combination of accumulation
|
||||
and dimension types, or if it cannot determine the appropriate configuration based on
|
||||
the input parameters.
|
||||
"""
|
||||
is_m_major = layout_d.is_m_major_c()
|
||||
|
||||
acc_bits = elem_ty_acc.width
|
||||
d_bits = elem_ty_d.width
|
||||
|
||||
tmem_warp_shape_mn = (
|
||||
(2, 2) if (cta_tile_shape[0] == 64 and use_2cta_instrs) else (4, 1)
|
||||
)
|
||||
epilog_tile_shape_mn = cute.product_each(
|
||||
cute.shape(epi_tile, loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
epilog_warp_tile_shape_mn = cute.shape_div(
|
||||
epilog_tile_shape_mn, tmem_warp_shape_mn, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
num_dp = cute.size(epilog_warp_tile_shape_mn[0], loc=loc, ip=ip)
|
||||
if num_dp not in {16, 32}:
|
||||
raise ValueError("Cta tile and 2sm config does not generate correct num dp.")
|
||||
|
||||
num_col_bits = cute.size(epilog_warp_tile_shape_mn[1], loc=loc, ip=ip) * acc_bits
|
||||
|
||||
tmem_dp = 0
|
||||
tmem_bit = 0
|
||||
tmem_rep = 0
|
||||
tmem_pack16b = False
|
||||
if acc_bits == 32 and d_bits == 32:
|
||||
if num_dp == 16:
|
||||
if is_m_major:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 256
|
||||
else:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 128
|
||||
else:
|
||||
tmem_dp = 32
|
||||
tmem_bit = 32
|
||||
elif acc_bits == 32 and d_bits == 16:
|
||||
if num_dp == 16:
|
||||
if is_m_major:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 256
|
||||
else:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 256
|
||||
else:
|
||||
if is_m_major:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 256
|
||||
else:
|
||||
tmem_dp = 32
|
||||
tmem_bit = 32
|
||||
elif acc_bits == 32 and d_bits == 8:
|
||||
if num_dp == 16:
|
||||
if is_m_major:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 256
|
||||
else:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 32
|
||||
else:
|
||||
if is_m_major:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 256
|
||||
else:
|
||||
tmem_dp = 32
|
||||
tmem_bit = 32
|
||||
elif acc_bits == 16 and d_bits == 16:
|
||||
tmem_pack16b = True
|
||||
if num_dp == 16:
|
||||
if is_m_major:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 128
|
||||
else:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 128
|
||||
else:
|
||||
if is_m_major:
|
||||
tmem_dp = 16
|
||||
tmem_bit = 128
|
||||
else:
|
||||
tmem_dp = 32
|
||||
tmem_bit = 32
|
||||
elif acc_bits == 32 and d_bits == 6:
|
||||
if not num_dp == 32:
|
||||
raise ValueError("Num dp must be 32.")
|
||||
tmem_dp = 32
|
||||
tmem_bit = 32
|
||||
elif acc_bits == 32 and d_bits == 4:
|
||||
if not num_dp == 32:
|
||||
raise ValueError("Num dp must be 32.")
|
||||
tmem_dp = 32
|
||||
tmem_bit = 32
|
||||
else:
|
||||
raise ValueError(
|
||||
f"Can not handle acc/d type combination: {elem_ty_acc=}, {elem_ty_d=}"
|
||||
)
|
||||
|
||||
num_bit_div = tmem_bit
|
||||
if tmem_dp == 16 and tmem_bit == 32:
|
||||
num_bit_div = 64
|
||||
|
||||
if (num_col_bits % (num_bit_div * 128) == 0) and (
|
||||
(tmem_dp == 16 and tmem_bit == 64)
|
||||
or (tmem_dp == 16 and tmem_bit == 32)
|
||||
or (tmem_dp == 32 and tmem_bit == 32)
|
||||
):
|
||||
tmem_rep = 128
|
||||
elif (num_col_bits % (num_bit_div * 64) == 0) and (
|
||||
(tmem_dp == 16 and tmem_bit == 128)
|
||||
or (tmem_dp == 16 and tmem_bit == 64)
|
||||
or (tmem_dp == 16 and tmem_bit == 32)
|
||||
or (tmem_dp == 32 and tmem_bit == 32)
|
||||
):
|
||||
tmem_rep = 64
|
||||
elif num_col_bits % (num_bit_div * 32) == 0:
|
||||
tmem_rep = 32
|
||||
elif num_col_bits % (num_bit_div * 16) == 0:
|
||||
tmem_rep = 16
|
||||
elif num_col_bits % (num_bit_div * 8) == 0:
|
||||
tmem_rep = 8
|
||||
elif num_col_bits % (num_bit_div * 4) == 0:
|
||||
tmem_rep = 4
|
||||
elif num_col_bits % (num_bit_div * 2) == 0:
|
||||
tmem_rep = 2
|
||||
elif num_col_bits % (num_bit_div * 1) == 0:
|
||||
tmem_rep = 1
|
||||
else:
|
||||
raise ValueError("Can not pick tmem_rep based on cta tile shape and tmem atom.")
|
||||
|
||||
if tmem_dp == 16 and tmem_bit == 64:
|
||||
op = Ld16x64bOp(
|
||||
Repetition(tmem_rep), Pack.PACK_16b_IN_32b if tmem_pack16b else Pack.NONE
|
||||
)
|
||||
return cute.make_copy_atom(op, elem_ty_acc, loc=loc, ip=ip)
|
||||
elif tmem_dp == 16 and tmem_bit == 128:
|
||||
op = Ld16x128bOp(
|
||||
Repetition(tmem_rep), Pack.PACK_16b_IN_32b if tmem_pack16b else Pack.NONE
|
||||
)
|
||||
return cute.make_copy_atom(op, elem_ty_acc, loc=loc, ip=ip)
|
||||
elif tmem_dp == 16 and tmem_bit == 256:
|
||||
op = Ld16x256bOp(
|
||||
Repetition(tmem_rep), Pack.PACK_16b_IN_32b if tmem_pack16b else Pack.NONE
|
||||
)
|
||||
return cute.make_copy_atom(op, elem_ty_acc, loc=loc, ip=ip)
|
||||
elif tmem_dp == 16 and tmem_bit == 32:
|
||||
op = Ld16x32bx2Op(
|
||||
Repetition(tmem_rep), Pack.PACK_16b_IN_32b if tmem_pack16b else Pack.NONE
|
||||
)
|
||||
return cute.make_copy_atom(op, elem_ty_acc, loc=loc, ip=ip)
|
||||
|
||||
elif tmem_dp == 32 and tmem_bit == 32:
|
||||
op = Ld32x32bOp(
|
||||
Repetition(tmem_rep), Pack.PACK_16b_IN_32b if tmem_pack16b else Pack.NONE
|
||||
)
|
||||
return cute.make_copy_atom(op, elem_ty_acc, loc=loc, ip=ip)
|
||||
else:
|
||||
raise ValueError()
|
||||
|
||||
|
||||
def get_num_tmem_alloc_cols(
|
||||
tmem_tensors: Union[cute.Tensor, List[cute.Tensor]], rounding=True
|
||||
) -> int:
|
||||
"""Get the total number of TMEM allocation columns for the given TMEM tensors.
|
||||
|
||||
:param tmem_tensors: The TMEM tensors to get the number of allocation columns for.
|
||||
:type tmem_tensors: Union[cute.Tensor, List[cute.Tensor]]
|
||||
:param rounding: Whether to round up the number of allocation columns to the nearest power of 2.
|
||||
:type rounding: bool
|
||||
|
||||
:return: The total number of TMEM allocation columns.
|
||||
:rtype: int
|
||||
|
||||
:raises ValueError: If the number of TMEM allocation columns exceeds the maximum capacity of 512 or is less than 32.
|
||||
"""
|
||||
# Turn tmem_tensors into a list
|
||||
if isinstance(tmem_tensors, cute.Tensor):
|
||||
tmem_tensors = [tmem_tensors]
|
||||
|
||||
# For each tensor in tmem_tensors, find the tmem_tensor_col_offset
|
||||
num_tmem_alloc_cols_per_tensor = [
|
||||
find_tmem_tensor_col_offset(t) for t in tmem_tensors
|
||||
]
|
||||
|
||||
# Sum up the num_tmem_alloc_cols_per_tensor
|
||||
num_tmem_alloc_cols = sum(num_tmem_alloc_cols_per_tensor)
|
||||
|
||||
# Round up num_tmem_cols_total to the nearest power of 2
|
||||
if rounding:
|
||||
num_tmem_alloc_cols = 1 << ceil(log2(num_tmem_alloc_cols))
|
||||
|
||||
# Validate the number of TMEM allocation columns
|
||||
SM100_TMEM_CAPACITY_COLUMNS = 512
|
||||
SM100_TMEM_MIN_ALLOC_COLUMNS = 32
|
||||
if (
|
||||
num_tmem_alloc_cols > SM100_TMEM_CAPACITY_COLUMNS
|
||||
or num_tmem_alloc_cols < SM100_TMEM_MIN_ALLOC_COLUMNS
|
||||
):
|
||||
raise ValueError(
|
||||
f"TMEM allocation columns {num_tmem_alloc_cols} exceeds the maximum capacity of {SM100_TMEM_CAPACITY_COLUMNS} or less than {SM100_TMEM_MIN_ALLOC_COLUMNS}"
|
||||
)
|
||||
return num_tmem_alloc_cols
|
||||
|
||||
|
||||
def get_smem_layout_atom_ab(
|
||||
major_mode: OperandMajorMode,
|
||||
element_type: Type[Numeric],
|
||||
smem_shape_mn_k: Tuple[int, int],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> SmemLayoutAtomKind:
|
||||
"""Simple heuristics to select the optimal SMEM layout atom based on the
|
||||
majorness, the data type, and the major mode size.
|
||||
|
||||
:param major_mode: The major mode for the SMEM tensor is K major.
|
||||
:type major_mode: OperandMajorMode
|
||||
:param element_type: The element type for the SMEM tensor.
|
||||
:type element_type: Type[Numeric]
|
||||
:param smem_shape_mn_k: The shape of the SMEM tensor.
|
||||
:type smem_shape_mn_k: Tuple[int, int]
|
||||
|
||||
:return: The SMEM layout atom kind
|
||||
:rtype: SmemLayoutAtomKind
|
||||
"""
|
||||
is_k_major = major_mode == OperandMajorMode.K
|
||||
major_mode_size = smem_shape_mn_k[1] if is_k_major else smem_shape_mn_k[0]
|
||||
|
||||
assert major_mode_size % 8 == 0
|
||||
sw128_num_contiguous_bits = 1024
|
||||
sw64_num_contiguous_bits = 512
|
||||
sw32_num_contiguous_bits = 256
|
||||
inter_num_contiguous_bits = 128
|
||||
major_mode_size_bits = major_mode_size * element_type.width
|
||||
assert major_mode_size_bits % inter_num_contiguous_bits == 0
|
||||
|
||||
if not is_k_major:
|
||||
if (element_type.width == 32) and (
|
||||
major_mode_size_bits % sw128_num_contiguous_bits == 0
|
||||
):
|
||||
return SmemLayoutAtomKind.MN_SW128_32B
|
||||
if major_mode_size_bits % sw128_num_contiguous_bits == 0:
|
||||
return SmemLayoutAtomKind.MN_SW128
|
||||
if major_mode_size_bits % sw64_num_contiguous_bits == 0:
|
||||
return SmemLayoutAtomKind.MN_SW64
|
||||
if major_mode_size_bits % sw32_num_contiguous_bits == 0:
|
||||
return SmemLayoutAtomKind.MN_SW32
|
||||
return SmemLayoutAtomKind.MN_INTER
|
||||
if major_mode_size_bits % sw128_num_contiguous_bits == 0:
|
||||
return SmemLayoutAtomKind.K_SW128
|
||||
if major_mode_size_bits % sw64_num_contiguous_bits == 0:
|
||||
return SmemLayoutAtomKind.K_SW64
|
||||
if major_mode_size_bits % sw32_num_contiguous_bits == 0:
|
||||
return SmemLayoutAtomKind.K_SW32
|
||||
return SmemLayoutAtomKind.K_INTER
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_smem_layout_a(
|
||||
tiled_mma: cute.TiledMma,
|
||||
mma_tiler_mnk: cute.Tile,
|
||||
a_dtype: Type[Numeric],
|
||||
num_stages: int,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Union[cute.Layout, cute.ComposedLayout]:
|
||||
"""This function helps with:
|
||||
1. Get the partitioned shape of the A tensor based on the tiled_mma & MMA tiler.
|
||||
2. Select the heuristic SMEM layout atom based on the A tensor's majorness, the data type, and the major mode size.
|
||||
3. cute.Tile the SMEM layout atom to the MMA tile shape.
|
||||
4. Stage the SMEM layout based on the number of stages.
|
||||
|
||||
:param tiled_mma: The tiled MMA used to partition tensor A
|
||||
:type tiled_mma: cute.TiledMma
|
||||
:param mma_tiler_mnk: The MMA tile shape
|
||||
:type mma_tiler_mnk: cute.cute.Tile
|
||||
:param a_dtype: The element type for tensor A
|
||||
:type a_dtype: Type[Numeric]
|
||||
:param num_stages: The number of pipeline stages for tensor A
|
||||
:type num_stages: int
|
||||
|
||||
:return: SMEM layout for tensor A
|
||||
:rtype: Union[cute.Layout, cute.ComposedLayout]
|
||||
"""
|
||||
|
||||
is_k_major = tiled_mma.op.a_major_mode == OperandMajorMode.K
|
||||
a_smem_shape = tiled_mma.partition_shape_A(
|
||||
cute.dice(mma_tiler_mnk, (1, None, 1), loc=loc, ip=ip)
|
||||
)
|
||||
a_smem_shape_mn_k = (
|
||||
cute.size(a_smem_shape[0][0], loc=loc, ip=ip) * a_smem_shape[1],
|
||||
cute.size(a_smem_shape[0][1], loc=loc, ip=ip) * a_smem_shape[2],
|
||||
)
|
||||
a_smem_layout_atom = make_smem_layout_atom(
|
||||
get_smem_layout_atom_ab(
|
||||
tiled_mma.op.a_major_mode,
|
||||
a_dtype,
|
||||
a_smem_shape_mn_k,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
),
|
||||
a_dtype,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
a_smem_layout_staged = tile_to_mma_shape(
|
||||
a_smem_layout_atom,
|
||||
cute.append(a_smem_shape, num_stages, loc=loc, ip=ip),
|
||||
order=((1, 0, 2) if not is_k_major else (0, 1, 2)),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
return a_smem_layout_staged
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_smem_layout_b(
|
||||
tiled_mma: cute.TiledMma,
|
||||
mma_tiler_mnk: cute.Tile,
|
||||
b_dtype: Type[Numeric],
|
||||
num_stages: int,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Union[cute.Layout, cute.ComposedLayout]:
|
||||
"""This function helps:
|
||||
1. Get the partitioned shape of the B tensor based on the tiled_mma & MMA tiler.
|
||||
2. Select the heuristic SMEM layout atom based on the B tensor's majorness, the data type, and the major mode size.
|
||||
3. cute.Tile the SMEM layout atom to the MMA tile shape.
|
||||
4. Stage the SMEM layout based on the number of stages.
|
||||
|
||||
:param tiled_mma: The tiled MMA which is used to partition the B tensor.
|
||||
:type tiled_mma: cute.TiledMma
|
||||
:param mma_tiler_mnk: The MMA tile shape.
|
||||
:type mma_tiler_mnk: cute.cute.Tile
|
||||
:param b_dtype: The element type for the B tensor.
|
||||
:type b_dtype: Type[Numeric]
|
||||
:param num_stages: The stage of the B tensor.
|
||||
:type num_stages: int
|
||||
|
||||
:return: SMEM layout for the B tensor.
|
||||
:rtype: Union[cute.Layout, cute.ComposedLayout]
|
||||
"""
|
||||
|
||||
is_k_major = tiled_mma.op.b_major_mode == OperandMajorMode.K
|
||||
b_smem_shape = tiled_mma.partition_shape_B(
|
||||
cute.dice(mma_tiler_mnk, (None, 1, 1), loc=loc, ip=ip)
|
||||
)
|
||||
b_smem_shape_nk = (
|
||||
cute.size(b_smem_shape[0][0], loc=loc, ip=ip) * b_smem_shape[1],
|
||||
cute.size(b_smem_shape[0][1], loc=loc, ip=ip) * b_smem_shape[2],
|
||||
)
|
||||
b_smem_layout_atom = make_smem_layout_atom(
|
||||
get_smem_layout_atom_ab(
|
||||
tiled_mma.op.b_major_mode,
|
||||
b_dtype,
|
||||
b_smem_shape_nk,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
),
|
||||
b_dtype,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
b_smem_layout_staged = tile_to_mma_shape(
|
||||
b_smem_layout_atom,
|
||||
cute.append(b_smem_shape, num_stages, loc=loc, ip=ip),
|
||||
order=((1, 0, 2) if not is_k_major else (0, 1, 2)),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
return b_smem_layout_staged
|
||||
|
||||
@dsl_user_op
|
||||
def get_smem_layout_atom_epi(
|
||||
layout: LayoutEnum,
|
||||
element_type: Type[Numeric],
|
||||
epi_tile: cute.Tile,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> SmemLayoutAtomKind:
|
||||
"""Simple heuristics to select the optimal SMEM layout atom for epilog tensors.
|
||||
|
||||
:param layout: The layout enum for the SMEM tensor.
|
||||
:type layout: LayoutEnum
|
||||
:param element_type: The element type for the SMEM tensor.
|
||||
:type element_type: Type[Numeric]
|
||||
:param epi_tile: The epilogue tile shape.
|
||||
:type epi_tile: cute.Tile
|
||||
|
||||
:return: The SMEM layout atom kind
|
||||
:rtype: SmemLayoutAtomKind
|
||||
"""
|
||||
# Get the max contiguous tile usable by TMA
|
||||
tma_shape = tuple(
|
||||
(
|
||||
# assumes get<0>(epi_tile) is coalesced and unit stride
|
||||
cute.coalesce(cute.right_inverse(x, loc=loc, ip=ip), loc=loc, ip=ip).shape
|
||||
if isinstance(x, cute.Layout)
|
||||
else x
|
||||
)
|
||||
for x in epi_tile
|
||||
)
|
||||
|
||||
if layout.is_m_major_c():
|
||||
# ColMajor C/D (M-major)
|
||||
return get_smem_layout_atom_ab(
|
||||
OperandMajorMode.MN, element_type, tma_shape, loc=loc, ip=ip
|
||||
)
|
||||
else:
|
||||
# RowMajor C/D (N-major)
|
||||
return get_smem_layout_atom_ab(
|
||||
OperandMajorMode.K, element_type, tma_shape, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
|
||||
@dsl_user_op
|
||||
def make_smem_layout_epi(
|
||||
epi_dtype: Type[Numeric],
|
||||
epi_layout: LayoutEnum,
|
||||
epi_tile: cute.Tile,
|
||||
epi_stage: int,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Union[cute.Layout, cute.ComposedLayout]:
|
||||
"""This function helps:
|
||||
1. Select the heuristic SMEM layout atom based on the epilog tile shape,
|
||||
the epilog tensor's majorness, and the element type.
|
||||
2. cute.Tile the SMEM layout atom to the epilog tile shape.
|
||||
3. Stage the SMEM layout based on the number of stages.
|
||||
|
||||
:param epi_dtype: The element type for the epilog tensor.
|
||||
:type epi_dtype: Type[Numeric]
|
||||
:param epi_layout: The layout enum for the epilog tensor.
|
||||
:type epi_layout: LayoutEnum
|
||||
:param epi_tile: The epilogue tile shape.
|
||||
:type epi_tile: cute.cute.Tile
|
||||
:param epi_stage: The stage of the epilog tensor.
|
||||
:type epi_stage: int
|
||||
|
||||
:return: SMEM layout for epilog tensors (usually C & D which are processed in the epilog)
|
||||
:rtype: Union[cute.Layout, cute.ComposedLayout]
|
||||
"""
|
||||
|
||||
epilog_shape = cute.product_each(
|
||||
cute.shape(epi_tile, loc=loc, ip=ip), loc=loc, ip=ip
|
||||
)
|
||||
|
||||
c_smem_layout_atom = make_smem_layout_atom(
|
||||
get_smem_layout_atom_epi(
|
||||
epi_layout,
|
||||
epi_dtype,
|
||||
epi_tile,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
),
|
||||
epi_dtype,
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
epi_smem_layout_staged = cute.tile_to_shape(
|
||||
c_smem_layout_atom,
|
||||
cute.append(epilog_shape, epi_stage, loc=loc, ip=ip),
|
||||
order=((1, 0, 2) if not epi_layout.is_n_major_c() else (0, 1, 2)),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
return epi_smem_layout_staged
|
||||
|
||||
|
||||
class SmemCapacity(Enum):
|
||||
SM100_SMEM_CAPACITY_BYTES = (228 - 1) * 1024
|
||||
SM120_SMEM_CAPACITY_BYTES = (100 - 1) * 1024
|
||||
|
||||
|
||||
# Dictionary to map compute capability to SMEM capacity
|
||||
SMEM_CAPACITY = {
|
||||
"sm100": SmemCapacity.SM100_SMEM_CAPACITY_BYTES.value,
|
||||
"sm120": SmemCapacity.SM120_SMEM_CAPACITY_BYTES.value,
|
||||
}
|
||||
|
||||
@dsl_user_op
|
||||
def make_trivial_tiled_mma(
|
||||
ab_dtype: Type[Numeric],
|
||||
a_leading_mode: OperandMajorMode,
|
||||
b_leading_mode: OperandMajorMode,
|
||||
acc_dtype: Type[Numeric],
|
||||
cta_group: CtaGroup,
|
||||
mma_tiler_mn: Tuple[int, int],
|
||||
a_source: OperandSource = OperandSource.SMEM,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> cute.TiledMma:
|
||||
"""Make a tiled MMA atom with given data type, leading dimension, cta group and mma tile shape.
|
||||
By default, the MMA atom is created with SMEM operand source for A.
|
||||
|
||||
:param ab_dtype: Data type of operands A and B.
|
||||
:type ab_dtype: type[Numeric]
|
||||
:param a_leading_mode: Leading dimension of operand A (1 for K, 0 for M/N).
|
||||
:type a_leading_mode: tcgen05.OperandMajorMode
|
||||
:param b_leading_mode: Leading dimension of operand B (1 for K, 0 for M/N).
|
||||
:type b_leading_mode: tcgen05.OperandMajorMode
|
||||
:param acc_dtype: Data type of the accumulator.
|
||||
:type acc_dtype: type[Numeric]
|
||||
:param cta_group: The CTA group to use.
|
||||
:type cta_group: tcgen05.CtaGroup
|
||||
:param mma_tiler_mn: The shape (M, N, K) of the MMA tiler.
|
||||
:type mma_tiler_mn: Tuple[int, int]
|
||||
:param a_source: The source of operand A (SMEM by default or TMEM).
|
||||
:type a_source: OperandSource
|
||||
|
||||
:return: A tiled MMA atom.
|
||||
:rtype: cute.TiledMma
|
||||
|
||||
:raises TypeError: If the data type is not supported.
|
||||
"""
|
||||
|
||||
if ab_dtype in {Float16, BFloat16}:
|
||||
mma_op = MmaF16BF16Op(
|
||||
ab_dtype,
|
||||
acc_dtype,
|
||||
(*mma_tiler_mn, 16),
|
||||
cta_group,
|
||||
a_source,
|
||||
a_leading_mode,
|
||||
b_leading_mode,
|
||||
)
|
||||
elif ab_dtype in {TFloat32, Float32}:
|
||||
mma_op = MmaTF32Op(
|
||||
(*mma_tiler_mn, 8),
|
||||
cta_group,
|
||||
a_source,
|
||||
a_leading_mode,
|
||||
b_leading_mode,
|
||||
)
|
||||
elif ab_dtype in {
|
||||
Uint8,
|
||||
Int8,
|
||||
}:
|
||||
mma_op = MmaI8Op(
|
||||
ab_dtype,
|
||||
(*mma_tiler_mn, 32),
|
||||
cta_group,
|
||||
a_source,
|
||||
a_leading_mode,
|
||||
b_leading_mode,
|
||||
)
|
||||
elif ab_dtype in {Float8E4M3FN, Float8E5M2}:
|
||||
mma_op = MmaFP8Op(
|
||||
ab_dtype,
|
||||
acc_dtype,
|
||||
(*mma_tiler_mn, 32),
|
||||
cta_group,
|
||||
a_source,
|
||||
a_leading_mode,
|
||||
b_leading_mode,
|
||||
)
|
||||
else:
|
||||
raise TypeError(f"unsupported ab_dtype, got {ab_dtype}")
|
||||
|
||||
return cute.make_tiled_mma(cute.make_mma_atom(mma_op))
|
||||
@@ -0,0 +1,466 @@
|
||||
# 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 List, Tuple
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cutlass_dsl import Int32, extract_mlir_values, new_from_mlir_values
|
||||
from cutlass._mlir import ir
|
||||
|
||||
from cutlass.utils.static_persistent_tile_scheduler import PersistentTileSchedulerParams
|
||||
|
||||
|
||||
class GroupSearchResult:
|
||||
"""
|
||||
The result of the group search for grouped gemm.
|
||||
|
||||
:param group_idx: The result group index
|
||||
:type group_idx: Int32
|
||||
:param cta_tile_idx_m: CTA tile index along M dimension after rasterization
|
||||
:type cta_tile_idx_m: Int32
|
||||
:param cta_tile_idx_n: CTA tile index along N dimension after rasterization
|
||||
:type cta_tile_idx_n: Int32
|
||||
:param problem_shape_m: The M dimension of the gemm problem
|
||||
:type problem_shape_m: Int32
|
||||
:param problem_shape_n: The N dimension of the gemm problem
|
||||
:type problem_shape_n: Int32
|
||||
:param problem_shape_k: The K dimension of the gemm problem
|
||||
:type problem_shape_k: Int32
|
||||
:param cta_tile_count_k: Number of tiles along K dimension
|
||||
:type cta_tile_count_k: Int32
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
group_idx: Int32,
|
||||
cta_tile_idx_m: Int32,
|
||||
cta_tile_idx_n: Int32,
|
||||
problem_shape_m: Int32,
|
||||
problem_shape_n: Int32,
|
||||
problem_shape_k: Int32,
|
||||
cta_tile_count_k: Int32,
|
||||
) -> None:
|
||||
self.group_idx = group_idx
|
||||
self.cta_tile_idx_m = cta_tile_idx_m
|
||||
self.cta_tile_idx_n = cta_tile_idx_n
|
||||
self.problem_shape_m = problem_shape_m
|
||||
self.problem_shape_n = problem_shape_n
|
||||
self.problem_shape_k = problem_shape_k
|
||||
self.cta_tile_count_k = cta_tile_count_k
|
||||
|
||||
def __extract_mlir_values__(self) -> List[ir.Value]:
|
||||
values = extract_mlir_values(self.group_idx)
|
||||
values.extend(extract_mlir_values(self.cta_tile_idx_m))
|
||||
values.extend(extract_mlir_values(self.cta_tile_idx_n))
|
||||
values.extend(extract_mlir_values(self.problem_shape_m))
|
||||
values.extend(extract_mlir_values(self.problem_shape_n))
|
||||
values.extend(extract_mlir_values(self.problem_shape_k))
|
||||
values.extend(extract_mlir_values(self.cta_tile_count_k))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(self, values: List[ir.Value]) -> "GroupSearchResult":
|
||||
assert len(values) == 7
|
||||
return GroupSearchResult(*tuple(values))
|
||||
|
||||
|
||||
class GroupedGemmGroupSearchState:
|
||||
"""
|
||||
The state of group index search for grouped gemm.
|
||||
|
||||
The state will be initialized once and updated in every round of group index search.
|
||||
|
||||
:param start_group_idx: The group idx to start the search with
|
||||
:type start_group_idx: Int32
|
||||
:param tile_count_prev_group: Number of tiles before the matched group
|
||||
:type tile_count_prev_group: Int32
|
||||
:param tile_count_searched: Number of tiles we have searched. When the matched group is found,
|
||||
it records the number of tiles including the matched group
|
||||
:type tile_count_searched: Int32
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
start_group_idx: Int32,
|
||||
tile_count_prev_group: Int32,
|
||||
tile_count_searched: Int32,
|
||||
) -> None:
|
||||
self.start_group_idx = start_group_idx
|
||||
self.tile_count_prev_group = tile_count_prev_group
|
||||
self.tile_count_searched = tile_count_searched
|
||||
|
||||
def __extract_mlir_values__(self) -> List[ir.Value]:
|
||||
values = extract_mlir_values(self.start_group_idx)
|
||||
values.extend(extract_mlir_values(self.tile_count_prev_group))
|
||||
values.extend(extract_mlir_values(self.tile_count_searched))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(
|
||||
self, values: List[ir.Value]
|
||||
) -> "GroupedGemmGroupSearchState":
|
||||
start_group_idx = new_from_mlir_values(self.start_group_idx, [values[0]])
|
||||
tile_count_prev_group = new_from_mlir_values(
|
||||
self.tile_count_prev_group, [values[1]]
|
||||
)
|
||||
tile_count_searched = new_from_mlir_values(
|
||||
self.tile_count_searched, [values[2]]
|
||||
)
|
||||
return GroupedGemmGroupSearchState(
|
||||
start_group_idx, tile_count_prev_group, tile_count_searched
|
||||
)
|
||||
|
||||
|
||||
def create_initial_search_state() -> GroupedGemmGroupSearchState:
|
||||
"""
|
||||
Create an initial search state for grouped gemm.
|
||||
|
||||
:return: A new search state with initial values
|
||||
:rtype: GroupedGemmGroupSearchState
|
||||
"""
|
||||
return GroupedGemmGroupSearchState(
|
||||
start_group_idx=Int32(0),
|
||||
tile_count_prev_group=Int32(0),
|
||||
tile_count_searched=Int32(0),
|
||||
)
|
||||
|
||||
|
||||
class GroupedGemmTileSchedulerHelper:
|
||||
"""
|
||||
A helper to translate the raw block index (x, y, z) from tile scheduler to real CTA tile index for grouped gemm.
|
||||
|
||||
:param group_count: Number of groups in current grouped gemm problem
|
||||
:type group_count: int
|
||||
:param tile_sched_params: Parameter used to create the tile scheduler this helper works with
|
||||
:type tile_sched_params: PersistentTileSchedulerParams
|
||||
:param cluster_tile_shape_mnk: The shape of cluster tile as (m, n, k)
|
||||
:type cluster_tile_shape_mnk: tuple[int, int, int]
|
||||
:param search_state: The initial search state
|
||||
:type search_state: GroupedGemmGroupSearchState
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
group_count: int,
|
||||
tile_sched_params: PersistentTileSchedulerParams,
|
||||
cluster_tile_shape_mnk: tuple[int, int, int],
|
||||
search_state: GroupedGemmGroupSearchState,
|
||||
) -> None:
|
||||
self.tile_sched_params = tile_sched_params
|
||||
self.group_count = group_count
|
||||
self.lane_idx = cute.arch.lane_idx()
|
||||
self.cluster_tile_shape_mnk = cluster_tile_shape_mnk
|
||||
self.search_state = search_state
|
||||
|
||||
def __extract_mlir_values__(self) -> List[ir.Value]:
|
||||
values = extract_mlir_values(self.tile_sched_params)
|
||||
values.extend(extract_mlir_values(self.search_state))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(
|
||||
self, values: List[ir.Value]
|
||||
) -> "GroupedGemmTileSchedulerHelper":
|
||||
tile_sched_params = new_from_mlir_values(self.tile_sched_params, values)
|
||||
search_state = new_from_mlir_values(self.search_state, values[1:])
|
||||
return GroupedGemmTileSchedulerHelper(
|
||||
self.group_count,
|
||||
tile_sched_params,
|
||||
self.cluster_tile_shape_mnk,
|
||||
search_state,
|
||||
)
|
||||
|
||||
def delinearize_z(
|
||||
self,
|
||||
cta_tile_coord: tuple,
|
||||
problem_shape_mnkl: cute.Tensor,
|
||||
) -> GroupSearchResult:
|
||||
"""
|
||||
Delinearize the linear z index and return GroupSearchResult.
|
||||
|
||||
This function should be used by warps that need to know the CTA tile index on M and N dimensions.
|
||||
|
||||
:param cta_tile_coord: The raw CTA coordinate from tile scheduler
|
||||
:type cta_tile_coord: tuple of Int32
|
||||
:param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for each group
|
||||
:type problem_shape_mnkl: cute.Tensor
|
||||
:return: The search result containing group index and tile coordinates
|
||||
:rtype: GroupSearchResult
|
||||
"""
|
||||
# delinear the z coord
|
||||
linear_idx = cta_tile_coord[2]
|
||||
group_idx, problem_mnkl = self._group_search_and_load_problem_shape(
|
||||
linear_idx,
|
||||
problem_shape_mnkl,
|
||||
self.search_state.start_group_idx,
|
||||
self.search_state.tile_count_prev_group,
|
||||
)
|
||||
# linear index local to current group
|
||||
cluster_tile_idx_in_current_group = (
|
||||
linear_idx - self.search_state.tile_count_prev_group
|
||||
)
|
||||
cluster_count_m, cluster_count_n, cluster_count_k = cute.ceil_div(
|
||||
(problem_mnkl[0], problem_mnkl[1], problem_mnkl[2]),
|
||||
(
|
||||
self.cluster_tile_shape_mnk[0],
|
||||
self.cluster_tile_shape_mnk[1],
|
||||
self.cluster_tile_shape_mnk[2],
|
||||
),
|
||||
)
|
||||
# decompose to get indices on M and N
|
||||
cta_tile_idx_m, cta_tile_idx_n = self._compute_cta_tile_coord(
|
||||
cluster_tile_idx_in_current_group,
|
||||
cta_tile_coord,
|
||||
cluster_count_m,
|
||||
cluster_count_n,
|
||||
)
|
||||
return GroupSearchResult(
|
||||
group_idx,
|
||||
cta_tile_idx_m,
|
||||
cta_tile_idx_n,
|
||||
problem_mnkl[0],
|
||||
problem_mnkl[1],
|
||||
problem_mnkl[2],
|
||||
cluster_count_k,
|
||||
)
|
||||
|
||||
def search_cluster_tile_count_k(
|
||||
self,
|
||||
cta_tile_coord: tuple,
|
||||
problem_shape_mnkl: cute.Tensor,
|
||||
) -> Tuple[Int32, Int32]:
|
||||
"""
|
||||
Search the matched group for given linear index and compute the number of tiles along K dimension for the matched group.
|
||||
|
||||
This function should be used by warps that are only interested in the number of tiles along K dimension.
|
||||
|
||||
:param cta_tile_coord: The raw CTA coordinate from tile scheduler
|
||||
:type cta_tile_coord: tuple of Int32
|
||||
:param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups
|
||||
:type problem_shape_mnkl: cute.Tensor
|
||||
:return: A tuple containing cluster count along K dimension and the group index
|
||||
:rtype: Tuple[Int32, Int32]
|
||||
"""
|
||||
group_idx, problem_mnk = self._group_search_and_load_problem_shape(
|
||||
cta_tile_coord[2],
|
||||
problem_shape_mnkl,
|
||||
self.search_state.start_group_idx,
|
||||
self.search_state.tile_count_prev_group,
|
||||
)
|
||||
cluster_count_k = (
|
||||
problem_mnk[2] + self.cluster_tile_shape_mnk[2] - 1
|
||||
) // self.cluster_tile_shape_mnk[2]
|
||||
return cluster_count_k, group_idx
|
||||
|
||||
@cute.jit
|
||||
def _prefix_sum(self, value_per_thread: Int32) -> Int32:
|
||||
"""
|
||||
Perform prefix sum within a full warp.
|
||||
|
||||
:param value_per_thread: The value for this thread to contribute to the prefix sum
|
||||
:type value_per_thread: Int32
|
||||
:return: The prefix sum result for this thread
|
||||
:rtype: Int32
|
||||
"""
|
||||
clamp_value = 0
|
||||
idx = 1
|
||||
sum_per_thread = value_per_thread
|
||||
while idx < cute.arch.WARP_SIZE:
|
||||
value = cute.arch.shuffle_sync_up(
|
||||
sum_per_thread, idx, mask_and_clamp=clamp_value
|
||||
)
|
||||
if self.lane_idx >= idx:
|
||||
sum_per_thread += value
|
||||
idx = idx << 1
|
||||
return sum_per_thread
|
||||
|
||||
def _get_problem_for_group(
|
||||
self, problem_shape_mnkl: cute.Tensor, group_idx: Int32
|
||||
) -> cute.Tensor:
|
||||
"""
|
||||
Load gemm problem (m,n,k,l) for the specified group from global memory to register.
|
||||
|
||||
:param problem_shape_mnkl: Tensor in global memory with layout (group_count, 4):(4, 1)
|
||||
:type problem_shape_mnkl: cute.Tensor
|
||||
:param group_idx: The index of the group to load
|
||||
:type group_idx: Int32
|
||||
:return: The problem shape tensor for the specified group
|
||||
:rtype: cute.Tensor
|
||||
"""
|
||||
cur_problem_mnkl = cute.make_fragment(
|
||||
cute.make_layout(4), problem_shape_mnkl.element_type
|
||||
)
|
||||
cute.autovec_copy(problem_shape_mnkl[(group_idx, None)], cur_problem_mnkl)
|
||||
return cur_problem_mnkl
|
||||
|
||||
def _get_cluster_tile_count_mn(self, problem_shape: cute.Tensor) -> Int32:
|
||||
"""
|
||||
Compute total cluster count.
|
||||
|
||||
:param problem_shape: Tensor containing problem shape (m, n, k, l)
|
||||
:type problem_shape: cute.Tensor
|
||||
:return: The total cluster tile count for M and N dimensions
|
||||
:rtype: Int32
|
||||
"""
|
||||
cur_ntile_m = (
|
||||
problem_shape[0] + self.cluster_tile_shape_mnk[0] - 1
|
||||
) // self.cluster_tile_shape_mnk[0]
|
||||
cur_ntile_n = (
|
||||
problem_shape[1] + self.cluster_tile_shape_mnk[1] - 1
|
||||
) // self.cluster_tile_shape_mnk[1]
|
||||
cur_ntile_mn = cur_ntile_m * cur_ntile_n
|
||||
return cur_ntile_mn
|
||||
|
||||
def _compute_cta_tile_coord(
|
||||
self,
|
||||
cluster_tile_idx: Int32,
|
||||
cta_tile_coord_in_cluster: tuple,
|
||||
cluster_tile_count_m: Int32,
|
||||
cluster_tile_count_n: Int32,
|
||||
) -> tuple:
|
||||
"""
|
||||
Compute CTA tile indices along M and N dimensions based on the linear index within a group.
|
||||
|
||||
It uses the AlongM mode to decompose the linear index onto M and N dimensions.
|
||||
|
||||
:param cluster_tile_idx: The linear index within a group
|
||||
:type cluster_tile_idx: Int32
|
||||
:param cta_tile_coord_in_cluster: CTA indices along M and N dimensions within a cluster
|
||||
:type cta_tile_coord_in_cluster: tuple of Int32
|
||||
:param cluster_tile_count_m: The number of clusters along M dimension of the matched group
|
||||
:type cluster_tile_count_m: Int32
|
||||
:param cluster_tile_count_n: The number of clusters along N dimension of the matched group
|
||||
:type cluster_tile_count_n: Int32
|
||||
:return: A tuple containing CTA tile indices along M and N dimensions
|
||||
:rtype: tuple of (Int32, Int32)
|
||||
"""
|
||||
cluster_layout_mn = cute.make_layout(
|
||||
(cluster_tile_count_m, cluster_tile_count_n)
|
||||
)
|
||||
(mi, ni) = cluster_layout_mn.get_hier_coord(cluster_tile_idx)
|
||||
cta_tile_idx_m = (
|
||||
mi * self.tile_sched_params.cluster_shape_mn[0]
|
||||
+ cta_tile_coord_in_cluster[0]
|
||||
)
|
||||
cta_tile_idx_n = (
|
||||
ni * self.tile_sched_params.cluster_shape_mn[1]
|
||||
+ cta_tile_coord_in_cluster[1]
|
||||
)
|
||||
return (cta_tile_idx_m, cta_tile_idx_n)
|
||||
|
||||
@cute.jit
|
||||
def _group_search(
|
||||
self,
|
||||
linear_idx: Int32,
|
||||
problem_shape_mnkl: cute.Tensor,
|
||||
init_group_idx: Int32,
|
||||
init_tile_count_searched: Int32,
|
||||
) -> GroupedGemmGroupSearchState:
|
||||
"""
|
||||
Search which group the linear index belongs to.
|
||||
|
||||
:param linear_idx: The linear index to be decomposed
|
||||
:type linear_idx: Int32
|
||||
:param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups
|
||||
:type problem_shape_mnkl: cute.Tensor
|
||||
:param init_group_idx: The group idx to start the search with
|
||||
:type init_group_idx: Int32
|
||||
:param init_tile_count_searched: The number of tiles we have searched
|
||||
:type init_tile_count_searched: Int32
|
||||
:return: The updated search state
|
||||
:rtype: GroupedGemmGroupSearchState
|
||||
"""
|
||||
c_0 = Int32(0).ir_value()
|
||||
last_lane_idx = cute.arch.WARP_SIZE - 1
|
||||
|
||||
tile_count_searched = init_tile_count_searched
|
||||
start_group_idx = init_group_idx
|
||||
not_found = linear_idx >= tile_count_searched
|
||||
tile_count_prev_group = self.search_state.tile_count_prev_group
|
||||
while not_found:
|
||||
# get group to search for current lane
|
||||
cur_group_idx = start_group_idx + self.lane_idx
|
||||
# check if the group to be checked is out of range
|
||||
inside_group_bound = cur_group_idx < self.group_count
|
||||
cur_ntile_mn = c_0
|
||||
if inside_group_bound:
|
||||
# get problem size of current group
|
||||
cur_problem_mnkl = self._get_problem_for_group(
|
||||
problem_shape_mnkl, cur_group_idx
|
||||
)
|
||||
cur_ntile_mn = self._get_cluster_tile_count_mn(cur_problem_mnkl)
|
||||
# compute tile count from beginning to current group(included)
|
||||
total_cluster_tile_count_ps_per_thread = self._prefix_sum(cur_ntile_mn)
|
||||
cluster_tile_count_end_per_thread = (
|
||||
total_cluster_tile_count_ps_per_thread + tile_count_searched
|
||||
)
|
||||
|
||||
group_not_in_window = linear_idx >= cluster_tile_count_end_per_thread
|
||||
hitted_group_idx_in_search_window = cute.arch.popc(
|
||||
cute.arch.vote_ballot_sync(group_not_in_window)
|
||||
)
|
||||
not_found = hitted_group_idx_in_search_window == cute.arch.WARP_SIZE
|
||||
start_group_idx = hitted_group_idx_in_search_window + start_group_idx
|
||||
hit_the_1st_problem_in_search_window = (
|
||||
hitted_group_idx_in_search_window == c_0
|
||||
)
|
||||
tile_count_prev_group = tile_count_searched
|
||||
if hit_the_1st_problem_in_search_window == False:
|
||||
tile_count_prev_group = cute.arch.shuffle_sync(
|
||||
cluster_tile_count_end_per_thread,
|
||||
hitted_group_idx_in_search_window - 1,
|
||||
)
|
||||
|
||||
# If no matched group, then get new_cluster_tile_count_end from last lane
|
||||
# Otherwise, get new_cluster_tile_count_end from the hitted group
|
||||
lane_idx_for_cluster_tile_count_end = hitted_group_idx_in_search_window
|
||||
if not_found:
|
||||
lane_idx_for_cluster_tile_count_end = last_lane_idx
|
||||
tile_count_searched = cute.arch.shuffle_sync(
|
||||
cluster_tile_count_end_per_thread,
|
||||
lane_idx_for_cluster_tile_count_end,
|
||||
)
|
||||
|
||||
return GroupedGemmGroupSearchState(
|
||||
start_group_idx,
|
||||
tile_count_prev_group,
|
||||
tile_count_searched,
|
||||
)
|
||||
|
||||
def _group_search_and_load_problem_shape(
|
||||
self,
|
||||
linear_idx: Int32,
|
||||
problem_shape_mnkl: cute.Tensor,
|
||||
start_group_idx: Int32,
|
||||
tile_count_searched: Int32,
|
||||
) -> Tuple[Int32, cute.Tensor]:
|
||||
"""
|
||||
Perform group search and load problem shape for the matched group.
|
||||
|
||||
:param linear_idx: The linear index to be decomposed
|
||||
:type linear_idx: Int32
|
||||
:param problem_shape_mnkl: Tensor containing gemm problem size (M, N, K, L) for all groups
|
||||
:type problem_shape_mnkl: cute.Tensor
|
||||
:param start_group_idx: The group idx to start the search with
|
||||
:type start_group_idx: Int32
|
||||
:param tile_count_searched: The number of tiles we have searched
|
||||
:type tile_count_searched: Int32
|
||||
:return: A tuple containing the final group index and the problem shape tensor
|
||||
:rtype: Tuple[Int32, cute.Tensor]
|
||||
"""
|
||||
self.search_state = self._group_search(
|
||||
linear_idx,
|
||||
problem_shape_mnkl,
|
||||
start_group_idx,
|
||||
tile_count_searched,
|
||||
)
|
||||
# get final group search state
|
||||
final_group_idx = self.search_state.start_group_idx
|
||||
# let's revisit if it's better to broadcast problem_shape_mnk in group_search
|
||||
problem_mnkl = self._get_problem_for_group(problem_shape_mnkl, final_group_idx)
|
||||
return final_group_idx, problem_mnkl
|
||||
@@ -0,0 +1,174 @@
|
||||
# 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 cuda.bindings import driver, nvrtc
|
||||
|
||||
import cutlass.cute as cute
|
||||
|
||||
"""
|
||||
This class is used to get the hardware info of given GPU device.
|
||||
It provides methods to get the max active clusters for given cluster size.
|
||||
|
||||
Prerequisite:
|
||||
- CUDA driver is initialized via `driver.cuInit` or other CUDA APIs.
|
||||
- CUDA context is created via `driver.cuCtxCreate` or other CUDA APIs.
|
||||
|
||||
"""
|
||||
|
||||
|
||||
class HardwareInfo:
|
||||
"""
|
||||
device_id: CUDA device ID to get the hardware info.
|
||||
"""
|
||||
|
||||
def __init__(self, device_id: int = 0):
|
||||
count = self._checkCudaErrors(driver.cuDeviceGetCount())
|
||||
if device_id >= count:
|
||||
raise ValueError(
|
||||
f"Device ID {device_id} is out of range for device count {count}"
|
||||
)
|
||||
self.device_id = device_id
|
||||
self.device = self._checkCudaErrors(driver.cuDeviceGet(device_id))
|
||||
self.context = self._checkCudaErrors(driver.cuCtxGetCurrent())
|
||||
self.driver_version = self._checkCudaErrors(driver.cuDriverGetVersion())
|
||||
|
||||
# Getting the max active clusters for a given cluster size
|
||||
def get_max_active_clusters(self, cluster_size: int) -> int:
|
||||
self._get_device_function()
|
||||
if self._cuda_driver_version_lt(11, 8):
|
||||
raise RuntimeError(
|
||||
"CUDA Driver version < 11.8, cannot get _max_active_clusters"
|
||||
)
|
||||
if cluster_size <= 0 or cluster_size > 32:
|
||||
raise ValueError(
|
||||
f"Cluster size must be between 1 and 32, {cluster_size} is not supported"
|
||||
)
|
||||
|
||||
max_shared_memory_per_block = self._checkCudaErrors(
|
||||
driver.cuDeviceGetAttribute(
|
||||
driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MAX_SHARED_MEMORY_PER_BLOCK_OPTIN,
|
||||
self.device,
|
||||
)
|
||||
)
|
||||
self._checkCudaErrors(
|
||||
driver.cuFuncSetAttribute(
|
||||
self.kernel,
|
||||
driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_MAX_DYNAMIC_SHARED_SIZE_BYTES,
|
||||
max_shared_memory_per_block,
|
||||
)
|
||||
)
|
||||
max_dynamic_shared_memory = self._checkCudaErrors(
|
||||
driver.cuOccupancyAvailableDynamicSMemPerBlock(
|
||||
self.kernel, 1, 1 # numBlocks # blockSize
|
||||
)
|
||||
)
|
||||
max_active_blocks = self._checkCudaErrors(
|
||||
driver.cuOccupancyMaxActiveBlocksPerMultiprocessor(
|
||||
self.kernel, 1, max_dynamic_shared_memory # blockSize,
|
||||
)
|
||||
)
|
||||
# allow non-portable cluster size to support detection of non-portable cluster size
|
||||
self._checkCudaErrors(
|
||||
driver.cuFuncSetAttribute(
|
||||
self.kernel,
|
||||
driver.CUfunction_attribute.CU_FUNC_ATTRIBUTE_NON_PORTABLE_CLUSTER_SIZE_ALLOWED,
|
||||
1,
|
||||
)
|
||||
)
|
||||
# prepare launch configuration
|
||||
launch_config = driver.CUlaunchConfig()
|
||||
launch_config.blockDimX = 128
|
||||
launch_config.blockDimY = 1
|
||||
launch_config.blockDimZ = 1
|
||||
launch_config.sharedMemBytes = max_dynamic_shared_memory
|
||||
launch_config.numAttrs = 1
|
||||
# max possible cluster size is 32
|
||||
cluster_dims_attr = driver.CUlaunchAttribute()
|
||||
cluster_dims_attr.id = (
|
||||
driver.CUlaunchAttributeID.CU_LAUNCH_ATTRIBUTE_CLUSTER_DIMENSION
|
||||
)
|
||||
value = driver.CUlaunchAttributeValue()
|
||||
value.clusterDim.x = cluster_size
|
||||
value.clusterDim.y = 1
|
||||
value.clusterDim.z = 1
|
||||
cluster_dims_attr.value = value
|
||||
launch_config.attrs = [cluster_dims_attr]
|
||||
launch_config.gridDimX = cluster_size
|
||||
launch_config.gridDimY = max_active_blocks
|
||||
launch_config.gridDimZ = 1
|
||||
|
||||
num_clusters = self._checkCudaErrors(
|
||||
driver.cuOccupancyMaxActiveClusters(self.kernel, launch_config)
|
||||
)
|
||||
return num_clusters
|
||||
|
||||
def get_l2_cache_size_in_bytes(self) -> int:
|
||||
return self._checkCudaErrors(
|
||||
driver.cuDeviceGetAttribute(
|
||||
driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_L2_CACHE_SIZE,
|
||||
self.device,
|
||||
)
|
||||
)
|
||||
|
||||
def get_device_multiprocessor_count(self) -> int:
|
||||
return self._checkCudaErrors(
|
||||
driver.cuDeviceGetAttribute(
|
||||
driver.CUdevice_attribute.CU_DEVICE_ATTRIBUTE_MULTIPROCESSOR_COUNT,
|
||||
self.device,
|
||||
)
|
||||
)
|
||||
|
||||
def _checkCudaErrors(self, result) -> None:
|
||||
if result[0].value:
|
||||
raise RuntimeError(
|
||||
"CUDA error code={}({})".format(
|
||||
result[0].value, self._cudaGetErrorEnum(result[0])
|
||||
)
|
||||
)
|
||||
# CUDA APIs always return the status as the first element of the result tuple
|
||||
if len(result) == 1:
|
||||
return None
|
||||
elif len(result) == 2:
|
||||
return result[1]
|
||||
else:
|
||||
return result[1:]
|
||||
|
||||
def _cudaGetErrorEnum(self, error) -> str:
|
||||
if isinstance(error, driver.CUresult):
|
||||
err, name = driver.cuGetErrorName(error)
|
||||
return name if err == driver.CUresult.CUDA_SUCCESS else "<unknown>"
|
||||
elif isinstance(error, nvrtc.nvrtcResult):
|
||||
return nvrtc.nvrtcGetErrorString(error)[1]
|
||||
else:
|
||||
raise RuntimeError("Unknown error type: {}".format(error))
|
||||
|
||||
def _cuda_driver_version_ge(self, major: int, minor: int) -> bool:
|
||||
return self.driver_version >= (major * 1000 + 10 * minor)
|
||||
|
||||
def _cuda_driver_version_lt(self, major: int, minor: int) -> bool:
|
||||
return not self._cuda_driver_version_ge(major, minor)
|
||||
|
||||
@cute.kernel
|
||||
def _empty_kernel(self):
|
||||
return
|
||||
|
||||
@cute.jit
|
||||
def _host_function(self):
|
||||
self._empty_kernel().launch(
|
||||
grid=[1, 1, 1],
|
||||
block=[1, 1, 1],
|
||||
)
|
||||
|
||||
# get a empty kernel to compute occupancy
|
||||
def _get_device_function(self) -> None:
|
||||
self.compiled_kernel = cute.compile(self._host_function)
|
||||
self.module = next(iter(self.compiled_kernel.cuda_modules.modules)).cuda_module
|
||||
self.kernel = next(iter(self.compiled_kernel.cuda_modules.modules)).kernel_ptr
|
||||
@@ -0,0 +1,195 @@
|
||||
# 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, Tuple
|
||||
from enum import Enum
|
||||
|
||||
from cutlass.utils.layout import LayoutEnum
|
||||
from cutlass.cutlass_dsl import (
|
||||
Float16,
|
||||
BFloat16,
|
||||
Float8E5M2,
|
||||
Float8E4M3FN,
|
||||
Numeric,
|
||||
NumericMeta,
|
||||
dsl_user_op,
|
||||
)
|
||||
|
||||
import cutlass
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.nvgpu.common import CopyUniversalOp
|
||||
from cutlass.cute.nvgpu.warp import StMatrix8x8x16bOp
|
||||
from cutlass.cute.nvgpu.warpgroup import (
|
||||
MmaF16BF16Op,
|
||||
MmaF8Op,
|
||||
OperandMajorMode,
|
||||
OperandSource,
|
||||
)
|
||||
|
||||
@dsl_user_op
|
||||
def sm90_get_smem_store_op(
|
||||
layout_d: LayoutEnum,
|
||||
elem_ty_d: Type[Numeric],
|
||||
elem_ty_acc: Type[Numeric],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> cute.CopyAtom:
|
||||
"""
|
||||
Selects the largest vectorized smem store atom available subject to constraint of gmem layout.
|
||||
|
||||
Parameters:
|
||||
-----------
|
||||
layout_d : LayoutEnum
|
||||
The layout enum of the output tensor D.
|
||||
|
||||
elem_ty_d : Type[Numeric]
|
||||
The element type for output tensor D.
|
||||
|
||||
elem_ty_acc : Type[Numeric]
|
||||
The element type for accumulator.
|
||||
|
||||
Returns:
|
||||
--------
|
||||
Either SmemStoreMatrix or SimtSyncCopy, based on the input parameters.
|
||||
"""
|
||||
|
||||
def validate_type(ty, ty_name):
|
||||
if not isinstance(ty, NumericMeta):
|
||||
raise TypeError(f"{ty_name} must be a Numeric, but got {ty}")
|
||||
|
||||
validate_type(elem_ty_d, "elem_ty_d")
|
||||
validate_type(elem_ty_acc, "elem_ty_acc")
|
||||
|
||||
is_m_major = layout_d.is_m_major_c()
|
||||
|
||||
if elem_ty_d.width == 16:
|
||||
return cute.make_copy_atom(
|
||||
StMatrix8x8x16bOp(is_m_major, 4), elem_ty_d, loc=loc, ip=ip
|
||||
)
|
||||
else:
|
||||
return cute.make_copy_atom(CopyUniversalOp(), elem_ty_d, loc=loc, ip=ip)
|
||||
|
||||
|
||||
class SmemCapacity(Enum):
|
||||
SM90_SMEM_CAPACITY_BYTES = (228 - 1) * 1024
|
||||
|
||||
|
||||
# Dictionary to map compute capability to SMEM capacity
|
||||
SMEM_CAPACITY = {
|
||||
"sm90": SmemCapacity.SM90_SMEM_CAPACITY_BYTES.value,
|
||||
}
|
||||
|
||||
def make_trivial_tiled_mma(
|
||||
a_dtype: Type[Numeric],
|
||||
b_dtype: Type[Numeric],
|
||||
a_leading_mode: OperandMajorMode,
|
||||
b_leading_mode: OperandMajorMode,
|
||||
acc_dtype: Type[Numeric],
|
||||
atom_layout_mnk: Tuple[int, int, int],
|
||||
tiler_mn: Tuple[int, int],
|
||||
) -> cute.TiledMma:
|
||||
"""Make a tiled MMA atom with given data type, leading dimension, cta group and mma tile shape.
|
||||
By default, the MMA atom is created with SMEM operand source for A.
|
||||
|
||||
:param a_dtype: Data type of operand A.
|
||||
:type a_dtype: type[Numeric]
|
||||
:param b_dtype: Data type of operand B.
|
||||
:type b_dtype: type[Numeric]
|
||||
:param a_leading_mode: Leading dimension of operand A (1 for K, 0 for M/N).
|
||||
:type a_leading_mode: warpgroup.OperandMajorMode
|
||||
:param b_leading_mode: Leading dimension of operand B (1 for K, 0 for M/N).
|
||||
:type b_leading_mode: warpgroup.OperandMajorMode
|
||||
:param acc_dtype: Data type of the accumulator.
|
||||
:type acc_dtype: type[Numeric]
|
||||
:param atom_layout_mnk: A integer tuple describing the tiling of Atom across threads.
|
||||
:type atom_layout_mnk: Tuple[int, int, int]
|
||||
:param tiler_mn: The shape (M, N) of the cta tiler.
|
||||
:type tiler_mn: Tuple[int, int]
|
||||
|
||||
:return: A tiled MMA atom.
|
||||
:rtype: cute.TiledMma
|
||||
|
||||
:raises TypeError: If the data type is not supported.
|
||||
"""
|
||||
|
||||
if a_dtype in {Float16, BFloat16}:
|
||||
if cutlass.const_expr(a_dtype != b_dtype):
|
||||
raise TypeError(f"Type mismatch: {a_dtype} != {b_dtype}")
|
||||
if cutlass.const_expr(a_dtype.width != b_dtype.width):
|
||||
raise TypeError(f"Type width mismatch: {a_dtype.width} != {b_dtype.width}")
|
||||
|
||||
mma_op = MmaF16BF16Op(
|
||||
a_dtype,
|
||||
acc_dtype,
|
||||
(*tiler_mn, 16),
|
||||
OperandSource.SMEM,
|
||||
a_leading_mode,
|
||||
b_leading_mode,
|
||||
)
|
||||
elif a_dtype in {Float8E4M3FN, Float8E5M2} and b_dtype in {
|
||||
Float8E4M3FN,
|
||||
Float8E5M2,
|
||||
}:
|
||||
mma_op = MmaF8Op(
|
||||
a_dtype,
|
||||
b_dtype,
|
||||
acc_dtype,
|
||||
(*tiler_mn, 32),
|
||||
OperandSource.SMEM,
|
||||
a_leading_mode,
|
||||
b_leading_mode,
|
||||
)
|
||||
else:
|
||||
raise TypeError(f"unsupported a_dtype and b_dtype, got {a_dtype} and {b_dtype}")
|
||||
|
||||
return cute.make_tiled_mma(cute.make_mma_atom(mma_op), atom_layout_mnk)
|
||||
|
||||
def get_smem_layout_atom(
|
||||
layout: LayoutEnum,
|
||||
element_type: Type[Numeric],
|
||||
major_mode_size: int,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
"""Select the optimal shared memory layout atom based on parameters.
|
||||
|
||||
:param layout: Layout enum of the tensor
|
||||
:type layout: LayoutEnum
|
||||
:param element_type: Data type of the elements
|
||||
:type element_type: type[cutlass.Numeric]
|
||||
:param major_mode_size: Size of the major mode dimension
|
||||
:type major_mode_size: int
|
||||
|
||||
:return: Selected shared memory layout atom kind
|
||||
:rtype: cute.nvgpu.warpgroup.SmemLayoutAtomKind
|
||||
"""
|
||||
assert major_mode_size % 8 == 0
|
||||
sw128_num_contiguous_bits = 1024
|
||||
sw64_num_contiguous_bits = 512
|
||||
sw32_num_contiguous_bits = 256
|
||||
major_mode_size_bits = major_mode_size * element_type.width
|
||||
if layout.sm90_mma_major_mode() == OperandMajorMode.MN:
|
||||
if major_mode_size_bits % sw128_num_contiguous_bits == 0:
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.MN_SW128
|
||||
if major_mode_size_bits % sw64_num_contiguous_bits == 0:
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.MN_SW64
|
||||
if major_mode_size_bits % sw32_num_contiguous_bits == 0:
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.MN_SW32
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.MN_INTER
|
||||
if major_mode_size_bits % sw128_num_contiguous_bits == 0:
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW128
|
||||
if major_mode_size_bits % sw64_num_contiguous_bits == 0:
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW64
|
||||
if major_mode_size_bits % sw32_num_contiguous_bits == 0:
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_SW32
|
||||
return cute.nvgpu.warpgroup.SmemLayoutAtomKind.K_INTER
|
||||
@@ -0,0 +1,68 @@
|
||||
# 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 enum import Enum
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.nvgpu import warpgroup
|
||||
from cutlass.cute.nvgpu import tcgen05
|
||||
|
||||
|
||||
class LayoutEnum(Enum):
|
||||
ROW_MAJOR = "row_major"
|
||||
COL_MAJOR = "col_major"
|
||||
|
||||
def mma_major_mode(self):
|
||||
return (
|
||||
tcgen05.OperandMajorMode.K
|
||||
if self == LayoutEnum.ROW_MAJOR
|
||||
else tcgen05.OperandMajorMode.MN
|
||||
)
|
||||
|
||||
def sm90_mma_major_mode(self):
|
||||
return (
|
||||
warpgroup.OperandMajorMode.K
|
||||
if self == LayoutEnum.ROW_MAJOR
|
||||
else warpgroup.OperandMajorMode.MN
|
||||
)
|
||||
|
||||
def is_k_major_a(self):
|
||||
return self == LayoutEnum.ROW_MAJOR
|
||||
|
||||
def is_m_major_a(self):
|
||||
return self == LayoutEnum.COL_MAJOR
|
||||
|
||||
def is_k_major_b(self):
|
||||
return self == LayoutEnum.COL_MAJOR
|
||||
|
||||
def is_n_major_b(self):
|
||||
return self == LayoutEnum.ROW_MAJOR
|
||||
|
||||
def is_n_major_c(self):
|
||||
return self == LayoutEnum.ROW_MAJOR
|
||||
|
||||
def is_m_major_c(self):
|
||||
return self == LayoutEnum.COL_MAJOR
|
||||
|
||||
@staticmethod
|
||||
def from_tensor(tensor: cute.Tensor) -> "LayoutEnum":
|
||||
ret = None
|
||||
if tensor.leading_dim == 1:
|
||||
ret = LayoutEnum.ROW_MAJOR
|
||||
elif tensor.leading_dim == 0:
|
||||
ret = LayoutEnum.COL_MAJOR
|
||||
else:
|
||||
raise ValueError(f"Invalid leading dimension: {tensor.leading_dim}")
|
||||
|
||||
return ret
|
||||
|
||||
|
||||
__all__ = ["LayoutEnum"]
|
||||
@@ -0,0 +1,984 @@
|
||||
# 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 abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass
|
||||
from typing import Optional
|
||||
|
||||
from cutlass.cutlass_dsl import Boolean, Int32, Int64, T, if_generate, and_, or_
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
|
||||
import cutlass.cute as cute
|
||||
|
||||
|
||||
##############################################################################
|
||||
# Agent class
|
||||
##############################################################################
|
||||
|
||||
|
||||
class Agent(enum.Enum):
|
||||
"""
|
||||
Agent indicates what is participating in the pipeline synchronization.
|
||||
"""
|
||||
# Arbitrary grouping of N threads
|
||||
Thread = enum.auto()
|
||||
# Same as AsyncThread, but includes all threads in the block
|
||||
ThreadBlock = enum.auto()
|
||||
# Same as AsyncThread, but includes all threads in the cluster
|
||||
ThreadBlockCluster = enum.auto()
|
||||
|
||||
|
||||
class CooperativeGroup:
|
||||
"""
|
||||
CooperativeGroup contains size and alignment restrictions for an Agent.
|
||||
"""
|
||||
def __init__(self, agent: Agent, size: int = 1, alignment: int = 1):
|
||||
if agent is Agent.Thread:
|
||||
assert size > 0
|
||||
if size == 32:
|
||||
assert (
|
||||
size == alignment
|
||||
), "Error: Alignment does not match number of threads in a warp."
|
||||
elif size == 128:
|
||||
assert (
|
||||
size == alignment
|
||||
), "Error: Alignment does not match number of threads in a warpgroup."
|
||||
elif agent is Agent.ThreadBlock:
|
||||
assert False, "Error: Not yet supported."
|
||||
elif agent is Agent.ThreadBlockCluster:
|
||||
assert False, "Error: Not yet supported."
|
||||
else:
|
||||
# Should never reach this state
|
||||
size = 0
|
||||
|
||||
if size <= 0:
|
||||
raise ValueError(
|
||||
"Error: The number of threads in a CooperativeGroup must be more than 0."
|
||||
)
|
||||
|
||||
# Size indicates how many threads are participating in this CooperativeGroup
|
||||
self.size = size
|
||||
# Agent indicates the type of thread group
|
||||
self.agent = agent
|
||||
|
||||
|
||||
class _PipelineOp(enum.Enum):
|
||||
"""
|
||||
PipelineOp assigns an operation to an agent corresponding to a specific hardware feature.
|
||||
"""
|
||||
# async-threads
|
||||
AsyncThread = enum.auto()
|
||||
# Blackwell (SM100a) MMA instruction
|
||||
TCGen05Mma = enum.auto()
|
||||
# Tensor Memory Accelerator load
|
||||
TmaLoad = enum.auto()
|
||||
# TMA Store consuming smem produced by AsyncThread
|
||||
TmaStore = enum.auto()
|
||||
|
||||
|
||||
def _get_pipeline_op(type_str):
|
||||
return _PipelineOp(type_str)
|
||||
|
||||
|
||||
##############################################################################
|
||||
# SyncObjectArray class
|
||||
##############################################################################
|
||||
|
||||
|
||||
class SyncObjectArray(ABC):
|
||||
"""
|
||||
SyncObjectArray is an abstract base class for different types of hardware synchronizations (e.g. smem barriers, named barriers, fences)
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def wait(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def arrive(self):
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_barrier(self):
|
||||
pass
|
||||
|
||||
|
||||
class MbarrierArray(SyncObjectArray):
|
||||
"""
|
||||
MbarrierArray implements an abstraction for an array of smem barriers.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
barrier_storage: cute.Pointer,
|
||||
num_stages: int,
|
||||
agent: tuple[_PipelineOp, CooperativeGroup],
|
||||
tx_count: int = 0,
|
||||
):
|
||||
self.barrier_storage = barrier_storage
|
||||
self.tx_count = tx_count
|
||||
self.num_stages = num_stages
|
||||
self.op_type, self.cg = agent
|
||||
self.arrive_count = self.cg.size
|
||||
|
||||
if self.num_stages <= 0:
|
||||
raise ValueError("Error: Mbarrier stage count must be greater than 0.")
|
||||
if self.arrive_count <= 0:
|
||||
raise ValueError("Error: Mbarrier arrive count must be greater than 0.")
|
||||
if self.op_type is _PipelineOp.TmaLoad and self.tx_count <= 0:
|
||||
raise ValueError(
|
||||
"Error: Mbarrier tx count must be greater than 0 for TMA ops."
|
||||
)
|
||||
|
||||
# Using a tensor to store mbarrier i64 ptrs
|
||||
self.mbarrier_array = cute.make_fragment(cute.make_layout(num_stages), Int64)
|
||||
for i in range(num_stages):
|
||||
self.mbarrier_array[i] = _cute_ir.ptrtoint(
|
||||
T.i64(), (self.barrier_storage + i).value
|
||||
)
|
||||
|
||||
# Mbarrier initialization in constructor
|
||||
self.mbarrier_init()
|
||||
|
||||
# Mbarrier initialization
|
||||
def mbarrier_init(self):
|
||||
"""
|
||||
Initializes an array of mbarriers using warp 0.
|
||||
"""
|
||||
def then_body():
|
||||
for index in range(self.num_stages):
|
||||
cute.arch.mbarrier_init_arrive_cnt(
|
||||
_mbarrier_i64_to_ptr(self.mbarrier_array[index]), self.arrive_count
|
||||
)
|
||||
|
||||
warp_idx = cute.arch.warp_idx()
|
||||
warp_idx = cute.arch.make_warp_uniform(warp_idx)
|
||||
|
||||
if_generate(warp_idx == 0, then_body)
|
||||
|
||||
def arrive(self, index: int, dst: int):
|
||||
"""
|
||||
Select the arrive corresponding to this MbarrierArray's PipelineOp
|
||||
:param index: Index of the mbarrier in the array to arrive on
|
||||
:type index: int
|
||||
:param dst: Destination parameter for selective arrival, which can be either a mask or destination cta rank. When None, both TCGen05Mma and AsyncThread will arrive on their local mbarrier.
|
||||
- For TCGen05Mma, dst serves as a multicast mask (e.g., 0b1011 allows arrive signal to be multicast to CTAs in the cluster with rank = 0, 1, and 3).
|
||||
- For AsyncThread, dst serves as a destination cta rank (e.g., 3 means threads will arrive on the mbarrier with rank = 3 in the cluster).
|
||||
:type dst: int | None
|
||||
"""
|
||||
if self.op_type is _PipelineOp.AsyncThread:
|
||||
self.arrive_mbarrier(index, dst)
|
||||
elif self.op_type is _PipelineOp.TCGen05Mma:
|
||||
self.arrive_tcgen05mma(index, dst)
|
||||
elif self.op_type in [_PipelineOp.TmaLoad]:
|
||||
self.arrive_and_expect_tx(index, self.tx_count)
|
||||
else:
|
||||
print(_get_pipeline_op(self.op_type))
|
||||
assert False, "Error: MbarrierArray is not supported for this PipelineOp."
|
||||
|
||||
def arrive_mbarrier(self, index: int, dst_rank: int):
|
||||
if dst_rank is None:
|
||||
cute.arch.mbarrier_arrive(_mbarrier_i64_to_ptr(self.mbarrier_array[index]))
|
||||
else:
|
||||
cute.arch.mbarrier_arrive(
|
||||
_mbarrier_i64_to_ptr(self.mbarrier_array[index]), dst_rank
|
||||
)
|
||||
|
||||
def arrive_tcgen05mma(self, index: int, mask: int):
|
||||
if mask is None:
|
||||
with cute.arch.elect_one():
|
||||
cute.nvgpu.tcgen05.commit(
|
||||
_mbarrier_i64_to_ptr(self.mbarrier_array[index])
|
||||
)
|
||||
else:
|
||||
with cute.arch.elect_one():
|
||||
cute.nvgpu.tcgen05.commit(
|
||||
_mbarrier_i64_to_ptr(self.mbarrier_array[index]),
|
||||
mask,
|
||||
cute.nvgpu.tcgen05.CtaGroup.TWO,
|
||||
)
|
||||
|
||||
def arrive_and_expect_tx(self, index: int, tx_count: int):
|
||||
with cute.arch.elect_one():
|
||||
cute.arch.mbarrier_init_tx_bytes(
|
||||
_mbarrier_i64_to_ptr(self.mbarrier_array[index]), tx_count
|
||||
)
|
||||
|
||||
def try_wait(self, index: int, phase: int):
|
||||
return cute.arch.mbarrier_try_wait(
|
||||
_mbarrier_i64_to_ptr(self.mbarrier_array[index]), phase
|
||||
)
|
||||
|
||||
def wait(self, index: int, phase: int):
|
||||
cute.arch.mbarrier_wait(_mbarrier_i64_to_ptr(self.mbarrier_array[index]), phase)
|
||||
|
||||
def get_barrier(self, index: int) -> cute.Pointer:
|
||||
return _mbarrier_i64_to_ptr(self.mbarrier_array[index])
|
||||
|
||||
|
||||
class TmaStoreFence(SyncObjectArray):
|
||||
"""
|
||||
TmaStoreFence is used for a multi-stage epilogue buffer.
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
num_stages: int = 0,
|
||||
):
|
||||
if num_stages <= 0:
|
||||
raise ValueError("Mbarrier stage count must be greater than 0.")
|
||||
|
||||
self.num_stages = num_stages
|
||||
|
||||
def arrive(self):
|
||||
cute.arch.cp_async_bulk_commit_group()
|
||||
|
||||
def wait(self):
|
||||
cute.arch.cp_async_bulk_wait_group(self.num_stages - 1, read=True)
|
||||
|
||||
# TmaStoreFence doesn't have mbarriers
|
||||
def get_barrier(self):
|
||||
assert (
|
||||
False
|
||||
), "Error: TmaStoreFence doesn't use mbarriers and cannot return a barrier."
|
||||
|
||||
def tail(self):
|
||||
cute.arch.cp_async_bulk_wait_group(0, read=True)
|
||||
|
||||
|
||||
##############################################################################
|
||||
# PipelineState class
|
||||
##############################################################################
|
||||
|
||||
|
||||
class PipelineUserType(enum.Enum):
|
||||
Producer = enum.auto()
|
||||
Consumer = enum.auto()
|
||||
|
||||
|
||||
class PipelineState:
|
||||
"""
|
||||
Pipeline state contains an index and phase bit corresponding to the current position in the circular buffer.
|
||||
"""
|
||||
|
||||
def __init__(self, stages: int, count, index, phase):
|
||||
self._stages = stages
|
||||
self._count = count
|
||||
self._index = index
|
||||
self._phase = phase
|
||||
|
||||
def clone(self) -> "PipelineState":
|
||||
return PipelineState(self.stages, self._count, self.index, self.phase)
|
||||
|
||||
@property
|
||||
def index(self) -> Int32:
|
||||
return self._index
|
||||
|
||||
@property
|
||||
def count(self) -> Int32:
|
||||
return self._count
|
||||
|
||||
@property
|
||||
def stages(self) -> int:
|
||||
return self._stages
|
||||
|
||||
@property
|
||||
def phase(self) -> Int32:
|
||||
return self._phase
|
||||
|
||||
def reset_count(self):
|
||||
self._count = Int32(0)
|
||||
|
||||
def advance(self):
|
||||
self._index += 1
|
||||
self._count += 1
|
||||
|
||||
def then_body(index, phase):
|
||||
new_index = Int32(0)
|
||||
new_phase = phase ^ 1
|
||||
return new_index, new_phase
|
||||
|
||||
def else_body(index, phase):
|
||||
return index, phase
|
||||
|
||||
self._index, self._phase = if_generate(
|
||||
self._index == self.stages,
|
||||
then_body,
|
||||
else_body,
|
||||
[self.index, self.phase],
|
||||
[Int32, Int32],
|
||||
)
|
||||
|
||||
def reverse(self):
|
||||
self._index -= 1
|
||||
self._count -= 1
|
||||
|
||||
def then_body(index, phase):
|
||||
new_index = Int32(self.stages - 1)
|
||||
new_phase = phase ^ 1
|
||||
return new_index, new_phase
|
||||
|
||||
def else_body(index, phase):
|
||||
return index, phase
|
||||
|
||||
self._index, self._phase = if_generate(
|
||||
self._index == -1,
|
||||
then_body,
|
||||
else_body,
|
||||
[self.index, self.phase],
|
||||
[Int32, Int32],
|
||||
)
|
||||
|
||||
def __get_mlir_types__(self):
|
||||
return [self._count.type, self._index.type, self._phase.type]
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
count = self._count
|
||||
index = self._index
|
||||
phase = self._phase
|
||||
return [count.ir_value(), index.ir_value(), phase.ir_value()]
|
||||
|
||||
# This can be overridden by derived classes
|
||||
def __new_from_mlir_values__(self, values):
|
||||
return PipelineState(
|
||||
self.stages, Int32(values[0]), Int32(values[1]), Int32(values[2])
|
||||
)
|
||||
|
||||
|
||||
def make_pipeline_state(type: PipelineUserType, stages: int):
|
||||
"""
|
||||
Creates a pipeline state. Producers are assumed to start with an empty buffer and have a flipped phase bit of 1.
|
||||
"""
|
||||
if type is PipelineUserType.Producer:
|
||||
return PipelineState(
|
||||
stages,
|
||||
Int32(0),
|
||||
Int32(0),
|
||||
Int32(1),
|
||||
)
|
||||
elif type is PipelineUserType.Consumer:
|
||||
return PipelineState(
|
||||
stages,
|
||||
Int32(0),
|
||||
Int32(0),
|
||||
Int32(0),
|
||||
)
|
||||
else:
|
||||
assert (
|
||||
False
|
||||
), "Error: invalid PipelineUserType specified for make_pipeline_state."
|
||||
|
||||
|
||||
##############################################################################
|
||||
# Pipeline classes
|
||||
##############################################################################
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineAsync:
|
||||
"""
|
||||
PipelineAsync is a generic pipeline class where both the producer and consumer are
|
||||
AsyncThreads. It also serves as a base class for specialized pipeline classes.
|
||||
"""
|
||||
sync_object_array_full: SyncObjectArray
|
||||
sync_object_array_empty: SyncObjectArray
|
||||
num_stages: Int32
|
||||
producer_mask: Int32
|
||||
consumer_mask: Int32
|
||||
|
||||
@staticmethod
|
||||
def _make_sync_object_array(
|
||||
barrier_storage: cute.Pointer,
|
||||
num_stages: Int32,
|
||||
agent: tuple[_PipelineOp, CooperativeGroup],
|
||||
tx_count: int = 0,
|
||||
) -> SyncObjectArray:
|
||||
"""
|
||||
Returns a SyncObjectArray corresponding to an agent's PipelineOp.
|
||||
"""
|
||||
if agent[0] in [
|
||||
_PipelineOp.AsyncThread,
|
||||
_PipelineOp.TmaLoad,
|
||||
_PipelineOp.TCGen05Mma,
|
||||
]:
|
||||
return MbarrierArray(
|
||||
barrier_storage=barrier_storage,
|
||||
num_stages=num_stages,
|
||||
agent=agent,
|
||||
tx_count=tx_count,
|
||||
)
|
||||
elif agent[0] is _PipelineOp.TmaStore:
|
||||
# Path taken for AsyncTmaStore
|
||||
return TmaStoreFence(num_stages=num_stages)
|
||||
else:
|
||||
assert False, "Error: Invalid PipelineOp specified."
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
barrier_storage: cute.Pointer,
|
||||
num_stages: Int32,
|
||||
producer_group: CooperativeGroup,
|
||||
consumer_group: CooperativeGroup,
|
||||
producer_mask: Int32 = None,
|
||||
consumer_mask: Int32 = None,
|
||||
):
|
||||
"""
|
||||
This helper function computes any necessary attributes and returns an instance of PipelineAsync.
|
||||
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param producer_mask: Mask for signaling arrives for the producer agent
|
||||
:type producer_mask: Int32 | None
|
||||
:param consumer_mask: Mask for signaling arrives for the consumer agent
|
||||
:type consumer_mask: Int32 | None
|
||||
"""
|
||||
producer_type = _PipelineOp.AsyncThread
|
||||
consumer_type = _PipelineOp.AsyncThread
|
||||
|
||||
producer = (producer_type, producer_group)
|
||||
consumer = (consumer_type, consumer_group)
|
||||
|
||||
sync_object_array_full = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8), num_stages, producer
|
||||
)
|
||||
sync_object_array_empty = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
|
||||
)
|
||||
|
||||
pipeline_init_wait()
|
||||
|
||||
return PipelineAsync(
|
||||
sync_object_array_full,
|
||||
sync_object_array_empty,
|
||||
num_stages,
|
||||
producer_mask,
|
||||
consumer_mask,
|
||||
)
|
||||
|
||||
def producer_acquire(
|
||||
self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
|
||||
):
|
||||
if_generate(
|
||||
try_acquire_token is None or try_acquire_token == 0,
|
||||
lambda: self.sync_object_array_empty.wait(state.index, state.phase),
|
||||
)
|
||||
|
||||
def producer_try_acquire(self, state: PipelineState):
|
||||
return self.sync_object_array_empty.try_wait(state.index, state.phase)
|
||||
|
||||
def producer_commit(self, state: PipelineState):
|
||||
self.sync_object_array_full.arrive(state.index, self.producer_mask)
|
||||
|
||||
def consumer_wait(
|
||||
self, state: PipelineState, try_wait_token: Optional[Boolean] = None
|
||||
):
|
||||
if_generate(
|
||||
try_wait_token is None or try_wait_token == 0,
|
||||
lambda: self.sync_object_array_full.wait(state.index, state.phase),
|
||||
)
|
||||
|
||||
def consumer_try_wait(self, state: PipelineState):
|
||||
return self.sync_object_array_full.try_wait(state.index, state.phase)
|
||||
|
||||
def consumer_release(self, state: PipelineState):
|
||||
self.sync_object_array_empty.arrive(state.index, self.consumer_mask)
|
||||
|
||||
def producer_get_barrier(self, state: PipelineState) -> cute.Pointer:
|
||||
return self.sync_object_array_full.get_barrier(state.index)
|
||||
|
||||
def producer_tail(self, state: PipelineState):
|
||||
"""
|
||||
Make sure the last used buffer empty signal is visible to producer.
|
||||
Producer tail is usually executed by producer before exit, to avoid dangling
|
||||
mbarrier arrive signals after kernel exit.
|
||||
|
||||
:param state: The pipeline state that points to next useful buffer
|
||||
:type state: PipelineState
|
||||
"""
|
||||
# Assume state contains that next useful buffer
|
||||
# So we only need to advance to num_stages - 1 times to last used buffer
|
||||
for i in range(self.num_stages - 1):
|
||||
state.advance()
|
||||
self.producer_acquire(state)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineTmaAsync(PipelineAsync):
|
||||
"""
|
||||
PipelineTmaAsync is used for TMA producers and AsyncThread consumers (e.g. Hopper mainloops).
|
||||
"""
|
||||
is_signalling_thread: bool
|
||||
|
||||
@staticmethod
|
||||
def init_empty_barrier_arrive_signal(cta_layout_vmnk: cute.Layout):
|
||||
"""
|
||||
Initialize the empty barrier arrive signal
|
||||
This function returns the destination cta rank and a boolean indicating if the signalling thread is the same as the current thread
|
||||
"""
|
||||
# Logic to optimally schedule Empty Arrives
|
||||
cluster_shape_mnk = cta_layout_vmnk.shape
|
||||
tidx, _, _ = cute.arch.thread_idx()
|
||||
cta_rank_in_cluster = cute.arch.make_warp_uniform(
|
||||
cute.arch.block_idx_in_cluster()
|
||||
)
|
||||
|
||||
is_signalling_thread = tidx < cute.size(cluster_shape_mnk)
|
||||
dst_rank = tidx % cute.size(cluster_shape_mnk)
|
||||
m = cluster_shape_mnk[0]
|
||||
|
||||
# Check if same row
|
||||
is_same_row_l = dst_rank % m
|
||||
is_same_row_r = cta_rank_in_cluster % m
|
||||
is_same_row = is_same_row_l == is_same_row_r
|
||||
|
||||
# Check if same column
|
||||
is_same_col_l = dst_rank // m
|
||||
is_same_col_r = cta_rank_in_cluster // m
|
||||
|
||||
is_same_col = is_same_col_l == is_same_col_r
|
||||
|
||||
is_same_row_or_col = or_(is_same_row, is_same_col)
|
||||
is_signalling_thread_final = and_(is_signalling_thread, is_same_row_or_col)
|
||||
|
||||
return dst_rank, is_signalling_thread_final
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
barrier_storage: cute.Pointer,
|
||||
num_stages: Int32,
|
||||
producer_group: CooperativeGroup,
|
||||
consumer_group: CooperativeGroup,
|
||||
tx_count: int,
|
||||
cta_layout_vmnk: Optional[cute.Layout] = None,
|
||||
):
|
||||
"""
|
||||
This helper function computes any necessary attributes and returns an instance of PipelineTmaAsync.
|
||||
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
|
||||
:type tx_count: int
|
||||
:param cta_layout_vmnk: Layout of the cluster shape
|
||||
:type cta_layout_vmnk: cute.Layout | None
|
||||
"""
|
||||
producer_type = _PipelineOp.TmaLoad
|
||||
consumer_type = _PipelineOp.AsyncThread
|
||||
|
||||
producer = (producer_type, producer_group)
|
||||
consumer = (consumer_type, consumer_group)
|
||||
|
||||
sync_object_array_full = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8), num_stages, producer, tx_count
|
||||
)
|
||||
sync_object_array_empty = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
|
||||
)
|
||||
|
||||
dst_rank, is_signalling_thread = (
|
||||
PipelineTmaAsync.init_empty_barrier_arrive_signal(cta_layout_vmnk)
|
||||
)
|
||||
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
|
||||
dst_rank = None
|
||||
else:
|
||||
dst_rank = dst_rank
|
||||
|
||||
is_signalling_thread = is_signalling_thread
|
||||
producer_mask = None
|
||||
|
||||
pipeline_init_wait(cta_layout_vmnk)
|
||||
|
||||
return PipelineTmaAsync(
|
||||
sync_object_array_full,
|
||||
sync_object_array_empty,
|
||||
num_stages,
|
||||
producer_mask,
|
||||
dst_rank,
|
||||
is_signalling_thread,
|
||||
)
|
||||
|
||||
def producer_acquire(
|
||||
self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
|
||||
):
|
||||
"""
|
||||
TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
|
||||
"""
|
||||
if_generate(
|
||||
try_acquire_token is None or try_acquire_token == 0,
|
||||
lambda: self.sync_object_array_empty.wait(state.index, state.phase),
|
||||
)
|
||||
self.sync_object_array_full.arrive(state.index, self.producer_mask)
|
||||
|
||||
|
||||
def producer_commit(self, state: PipelineState):
|
||||
"""
|
||||
TMA producer commit is a NOP. The transaction barrier signals the commit upon completion of the TMA.
|
||||
"""
|
||||
pass
|
||||
|
||||
def consumer_release(self, state: PipelineState):
|
||||
"""
|
||||
TMA consumer release conditionally signals the empty buffer to the producer.
|
||||
"""
|
||||
if_generate(
|
||||
self.is_signalling_thread,
|
||||
lambda: self.sync_object_array_empty.arrive(
|
||||
state.index, self.consumer_mask
|
||||
),
|
||||
)
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineTmaUmma(PipelineAsync):
|
||||
"""
|
||||
PipelineTmaUmma is used for TMA producers and UMMA consumers (e.g. Blackwell mainloops).
|
||||
"""
|
||||
is_leader_cta: bool
|
||||
|
||||
@staticmethod
|
||||
def _compute_mcast_arrival_mask(cta_layout_vmnk: cute.Layout):
|
||||
"""
|
||||
Computes a mask for signaling arrivals to multicasting threadblocks.
|
||||
"""
|
||||
cta_rank_in_cluster = cute.arch.make_warp_uniform(
|
||||
cute.arch.block_idx_in_cluster()
|
||||
)
|
||||
cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
|
||||
|
||||
tma_mcast_mask_a = cute.nvgpu.cpasync.create_tma_multicast_mask(
|
||||
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=2
|
||||
)
|
||||
tma_mcast_mask_b = cute.nvgpu.cpasync.create_tma_multicast_mask(
|
||||
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mcast_mode=1
|
||||
)
|
||||
|
||||
block_in_cluster_coord_vmnk_peer = (
|
||||
cta_in_cluster_coord_vmnk[0] ^ 1,
|
||||
*cta_in_cluster_coord_vmnk[1:],
|
||||
)
|
||||
tma_mcast_mask_a_peer = cute.nvgpu.cpasync.create_tma_multicast_mask(
|
||||
cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=2
|
||||
)
|
||||
tma_mcast_mask_b_peer = cute.nvgpu.cpasync.create_tma_multicast_mask(
|
||||
cta_layout_vmnk, block_in_cluster_coord_vmnk_peer, mcast_mode=1
|
||||
)
|
||||
return (
|
||||
tma_mcast_mask_a
|
||||
| tma_mcast_mask_b
|
||||
| tma_mcast_mask_a_peer
|
||||
| tma_mcast_mask_b_peer
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _compute_is_leader_cta(cta_layout_vmnk: cute.Layout):
|
||||
"""
|
||||
Computes leader threadblocks for 2CTA kernels. For 1CTA, all threadblocks are leaders.
|
||||
"""
|
||||
bidx, bidy, _ = cute.arch.block_idx()
|
||||
|
||||
mma_coord_vmnk = (
|
||||
bidx % cute.size(cta_layout_vmnk, mode=[0]),
|
||||
bidx // cute.size(cta_layout_vmnk, mode=[0]),
|
||||
bidy,
|
||||
None,
|
||||
)
|
||||
return mma_coord_vmnk[0] == 0
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
barrier_storage: cute.Pointer,
|
||||
num_stages: Int32,
|
||||
producer_group: CooperativeGroup,
|
||||
consumer_group: CooperativeGroup,
|
||||
tx_count: int,
|
||||
cta_layout_vmnk: Optional[cute.Layout] = None,
|
||||
):
|
||||
"""
|
||||
This helper function computes any necessary attributes and returns an instance of PipelineTmaUmma.
|
||||
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param tx_count: Number of bytes expected to be written to the transaction barrier for one stage
|
||||
:type tx_count: int
|
||||
:param cta_layout_vmnk: Layout of the cluster shape
|
||||
:type cta_layout_vmnk: cute.Layout | None
|
||||
"""
|
||||
producer_type = _PipelineOp.TmaLoad
|
||||
consumer_type = _PipelineOp.TCGen05Mma
|
||||
|
||||
producer = (producer_type, producer_group)
|
||||
consumer = (consumer_type, consumer_group)
|
||||
|
||||
sync_object_array_full = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8), num_stages, producer, tx_count
|
||||
)
|
||||
sync_object_array_empty = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
|
||||
)
|
||||
|
||||
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
|
||||
# No mcast mask if not using clusters
|
||||
producer_mask = None
|
||||
# All threadblocks are leaders if not using clusters
|
||||
is_leader_cta = True
|
||||
else:
|
||||
producer_mask = PipelineTmaUmma._compute_mcast_arrival_mask(cta_layout_vmnk)
|
||||
is_leader_cta = PipelineTmaUmma._compute_is_leader_cta(cta_layout_vmnk)
|
||||
|
||||
consumer_mask = producer_mask
|
||||
|
||||
pipeline_init_wait(cta_layout_vmnk)
|
||||
|
||||
return PipelineTmaUmma(
|
||||
sync_object_array_full,
|
||||
sync_object_array_empty,
|
||||
num_stages,
|
||||
producer_mask,
|
||||
consumer_mask,
|
||||
is_leader_cta,
|
||||
)
|
||||
|
||||
def producer_acquire(
|
||||
self, state: PipelineState, try_acquire_token: Optional[Boolean] = None
|
||||
):
|
||||
"""
|
||||
TMA producer commit conditionally waits on buffer empty and sets the transaction barrier for leader threadblocks.
|
||||
"""
|
||||
if_generate(
|
||||
try_acquire_token is None or try_acquire_token == 0,
|
||||
lambda: self.sync_object_array_empty.wait(state.index, state.phase),
|
||||
)
|
||||
if_generate(
|
||||
self.is_leader_cta,
|
||||
lambda: self.sync_object_array_full.arrive(state.index, self.producer_mask),
|
||||
)
|
||||
|
||||
def producer_commit(self, state: PipelineState):
|
||||
"""
|
||||
TMA producer commit is a NOP. The transaction barrier signals the commit upon completion of the TMA.
|
||||
"""
|
||||
pass
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineUmmaAsync(PipelineAsync):
|
||||
"""
|
||||
PipelineTmaUmma is used for UMMA producers and AsyncThread consumers (e.g. Blackwell accumulator pipelines).
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def _compute_tmem_sync_mask(cta_layout_vmnk: cute.Layout):
|
||||
"""
|
||||
Computes a mask to signal completion of tmem buffers for 2CTA kernels.
|
||||
"""
|
||||
cta_rank_in_cluster = cute.arch.make_warp_uniform(
|
||||
cute.arch.block_idx_in_cluster()
|
||||
)
|
||||
cta_in_cluster_coord_vmnk = cta_layout_vmnk.get_flat_coord(cta_rank_in_cluster)
|
||||
return cute.make_layout_image_mask(
|
||||
cta_layout_vmnk, cta_in_cluster_coord_vmnk, mode=0
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def _compute_peer_cta_rank():
|
||||
"""
|
||||
Computes a mask to signal release of tmem buffers for 2CTA kernels.
|
||||
"""
|
||||
cta_rank_in_cluster = cute.arch.make_warp_uniform(
|
||||
cute.arch.block_idx_in_cluster()
|
||||
)
|
||||
return cta_rank_in_cluster // 2 * 2
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
barrier_storage: cute.Pointer,
|
||||
num_stages: Int32,
|
||||
producer_group: CooperativeGroup,
|
||||
consumer_group: CooperativeGroup,
|
||||
cta_layout_vmnk: Optional[cute.Layout] = None,
|
||||
):
|
||||
"""
|
||||
This helper function computes any necessary attributes and returns an instance of PipelineUmmaAsync.
|
||||
:param barrier_storage: Pointer to the smem address for this pipeline's mbarriers
|
||||
:type barrier_storage: cute.Pointer
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
:param consumer_group: CooperativeGroup for the consumer agent
|
||||
:type consumer_group: CooperativeGroup
|
||||
:param cta_layout_vmnk: Layout of the cluster shape
|
||||
:type cta_layout_vmnk: cute.Layout | None
|
||||
"""
|
||||
producer_type = _PipelineOp.TCGen05Mma
|
||||
consumer_type = _PipelineOp.AsyncThread
|
||||
|
||||
producer = (producer_type, producer_group)
|
||||
consumer = (consumer_type, consumer_group)
|
||||
|
||||
sync_object_array_full = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8), num_stages, producer
|
||||
)
|
||||
sync_object_array_empty = PipelineAsync._make_sync_object_array(
|
||||
barrier_storage.align(min_align=8) + num_stages, num_stages, consumer
|
||||
)
|
||||
|
||||
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
|
||||
# Set mask to None if not using clusters (i.e. 1CTA kernels)
|
||||
producer_mask = None
|
||||
else:
|
||||
producer_mask = PipelineUmmaAsync._compute_tmem_sync_mask(cta_layout_vmnk)
|
||||
|
||||
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk, mode=[0]) == 1:
|
||||
# Set mask to None if not using 2CTA intructions
|
||||
consumer_mask = None
|
||||
else:
|
||||
consumer_mask = PipelineUmmaAsync._compute_peer_cta_rank()
|
||||
|
||||
pipeline_init_wait(cta_layout_vmnk)
|
||||
|
||||
return PipelineUmmaAsync(
|
||||
sync_object_array_full,
|
||||
sync_object_array_empty,
|
||||
num_stages,
|
||||
producer_mask,
|
||||
consumer_mask,
|
||||
)
|
||||
|
||||
def producer_tail(self, state: PipelineState):
|
||||
"""
|
||||
Make sure the last used buffer empty signal is visible to producer.
|
||||
Producer tail is usually executed by producer before exit, to avoid dangling
|
||||
mbarrier arrive signals after kernel exit.
|
||||
|
||||
:param state: The pipeline state that points to next useful buffer
|
||||
:type state: PipelineState
|
||||
"""
|
||||
cta_rank_in_cluster = cute.arch.make_warp_uniform(
|
||||
cute.arch.block_idx_in_cluster()
|
||||
)
|
||||
is_leader_cta = cta_rank_in_cluster % 2 == 0
|
||||
|
||||
def then_body():
|
||||
# Assume state contains that next useful buffer
|
||||
# So we only need to advance to num_stages - 1 times to last used buffer
|
||||
for i in range(self.num_stages - 1):
|
||||
state.advance()
|
||||
self.producer_acquire(state)
|
||||
|
||||
if_generate(is_leader_cta, then_body)
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class PipelineTmaStore(PipelineAsync):
|
||||
"""
|
||||
PipelineTmaStore is used for synchronizing TMA stores in the epilogue. It does not use mbarriers.
|
||||
"""
|
||||
|
||||
@staticmethod
|
||||
def create(
|
||||
num_stages: Int32,
|
||||
producer_group: CooperativeGroup,
|
||||
):
|
||||
"""
|
||||
This helper function computes any necessary attributes and returns an instance of PipelineTmaStore.
|
||||
:param num_stages: Number of buffer stages for this pipeline
|
||||
:type num_stages: Int32
|
||||
:param producer_group: CooperativeGroup for the producer agent
|
||||
:type producer_group: CooperativeGroup
|
||||
"""
|
||||
producer_type = _PipelineOp.TmaStore
|
||||
|
||||
producer = (producer_type, producer_group)
|
||||
|
||||
sync_object_array_full = PipelineAsync._make_sync_object_array(
|
||||
None, num_stages, producer
|
||||
)
|
||||
|
||||
return PipelineTmaStore(sync_object_array_full, None, num_stages, None, None)
|
||||
|
||||
def producer_acquire(self):
|
||||
self.sync_object_array_full.wait()
|
||||
|
||||
def producer_commit(self):
|
||||
self.sync_object_array_full.arrive()
|
||||
|
||||
def consumer_wait(self):
|
||||
assert False, "Error: PipelineTmaStore does not have a consumer agent."
|
||||
|
||||
def consumer_release(self):
|
||||
assert False, "Error: PipelineTmaStore does not have a consumer agent."
|
||||
|
||||
def producer_tail(self):
|
||||
self.sync_object_array_full.tail()
|
||||
|
||||
|
||||
##############################################################################
|
||||
# Helper functions
|
||||
##############################################################################
|
||||
|
||||
|
||||
def pipeline_init_wait(cta_layout_vmnk: Optional[cute.Layout] = None):
|
||||
"""
|
||||
Fences the mbarrier init and syncs the threadblock or cluster
|
||||
"""
|
||||
cute.arch.mbarrier_init_fence()
|
||||
|
||||
if cta_layout_vmnk is None or cute.size(cta_layout_vmnk) == 1:
|
||||
# If not using clusters, sync the threadblock
|
||||
_sync(Agent.ThreadBlock)
|
||||
else:
|
||||
# If using clusters, sync the cluster
|
||||
_sync(Agent.ThreadBlockCluster)
|
||||
|
||||
|
||||
def _sync(group: Agent):
|
||||
"""
|
||||
Syncs all threads within an agent.
|
||||
"""
|
||||
if group is Agent.Thread:
|
||||
assert False, "Error: Not supported."
|
||||
elif group is Agent.ThreadBlock:
|
||||
cute.arch.sync_threads()
|
||||
elif group is Agent.ThreadBlockCluster:
|
||||
cute.arch.cluster_arrive()
|
||||
cute.arch.cluster_wait()
|
||||
else:
|
||||
assert (
|
||||
False
|
||||
), "Error: No explicit sync instruction exists. Please use barriers (named / mbarrier) instead."
|
||||
|
||||
|
||||
def _mbarrier_i64_to_ptr(val: Int64) -> cute.Pointer:
|
||||
"""
|
||||
Converts a smem pointer of type Int64 to cute.Pointer with 8B alignment
|
||||
"""
|
||||
return cute.make_ptr(
|
||||
Int64,
|
||||
val.ir_value(),
|
||||
mem_space=_cute_ir.AddressSpace.smem,
|
||||
assumed_align=8,
|
||||
)
|
||||
@@ -0,0 +1,217 @@
|
||||
# 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, Union, overload
|
||||
|
||||
from cutlass.cutlass_dsl import Int8, Numeric, NumericMeta
|
||||
|
||||
import cutlass.cute as cute
|
||||
from cutlass.cute.arch import get_dyn_smem
|
||||
|
||||
|
||||
class SmemAllocator:
|
||||
"""
|
||||
A class for managing shared memory allocation on GPU.
|
||||
|
||||
This class manages a chunk of shared memory and provide APIs for sub-allocation
|
||||
inside the chunk.
|
||||
|
||||
Attributes
|
||||
----------
|
||||
_base : cute.Pointer as i8 typed dynamic value
|
||||
The current base address of the shared memory.
|
||||
|
||||
_allocated_bytes:
|
||||
The bytes allocated in shared memory.
|
||||
|
||||
Methods
|
||||
-------
|
||||
allocate(num_bytes, alignment)
|
||||
Allocates num_bytes in the shared memory with the given byte alignment.
|
||||
|
||||
allocate_value(value_ty, num_elems)
|
||||
Allocates num_elems of value_ty values in the shared memory.
|
||||
|
||||
allocate_tensor(value_ty, layout, alignment)
|
||||
Allocates a tensor in the shared memory with given layout and byte alignment.
|
||||
|
||||
Notes
|
||||
-----
|
||||
This class is responsible for managing the allocation of tensors in shared memory.
|
||||
"""
|
||||
|
||||
def __init__(self):
|
||||
"""
|
||||
Initializes the SmemAllocator instance with dynamic smem base ptr,
|
||||
which is i8 type and aligned to 1024.
|
||||
|
||||
"""
|
||||
self._base = get_dyn_smem(Int8, alignment=1024)
|
||||
self._allocated_bytes = 0
|
||||
|
||||
@overload
|
||||
def allocate(self, size_or_type: int, byte_alignment: int): ...
|
||||
|
||||
@overload
|
||||
def allocate(self, size_or_type: cute.struct, byte_alignment: int): ...
|
||||
|
||||
def allocate(self, size_or_type, byte_alignment: int = 1) -> int:
|
||||
"""
|
||||
Allocates a block of memory with the specified size and byte alignment.
|
||||
|
||||
This method adjusts the base cute.Pointer to ensure that the allocated memory
|
||||
is aligned according to the specified byte alignment. It updates the internal
|
||||
state to reflect the new base cute.Pointer and the total allocated bytes.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
size_or_type : int or struct
|
||||
The number of bytes to allocate or struct class.
|
||||
byte_alignment : int
|
||||
The byte alignment requirement for the allocation. Defaults to 1 (no alignment).
|
||||
|
||||
Returns
|
||||
----------
|
||||
A cute.Pointer to the start of the allocated memory block or struct instance.
|
||||
|
||||
Raises
|
||||
----------
|
||||
ValueError
|
||||
If num_bytes is negative or if byte_alignmemt is less than 1.
|
||||
"""
|
||||
|
||||
if isinstance(size_or_type, cute.struct):
|
||||
alignment = max(byte_alignment, size_or_type.__alignof__())
|
||||
base_ptr = self.allocate(size_or_type.__sizeof__(), alignment)
|
||||
return size_or_type(base_ptr)
|
||||
|
||||
num_bytes = size_or_type
|
||||
if num_bytes < 0:
|
||||
raise ValueError("num_bytes must be non-negative")
|
||||
if byte_alignment < 1:
|
||||
raise ValueError("byte_alignment must be at least 1")
|
||||
|
||||
self._base = self._base.align(byte_alignment)
|
||||
ptr = self._base
|
||||
self._base += num_bytes
|
||||
if self._allocated_bytes % byte_alignment != 0:
|
||||
self._allocated_bytes += (
|
||||
byte_alignment - self._allocated_bytes % byte_alignment
|
||||
)
|
||||
self._allocated_bytes += num_bytes
|
||||
return ptr
|
||||
|
||||
def allocate_array(self, element_type: Type[Numeric], num_elems: int = 1):
|
||||
"""
|
||||
Allocates num_elems values of element_type in shared memory.
|
||||
|
||||
This method calls allocate() to return a byte ptr, pointing to start of shared
|
||||
memory. Then calls cute.recast_ptr() to recast this byte cute.Pointer to element_type.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
element_type : Type[Numeric]
|
||||
The type of the values in the tensor.
|
||||
num_elems : int, optional
|
||||
The number of elements for each allocation. Defaults to 1.
|
||||
|
||||
Returns
|
||||
----------
|
||||
A value_type cute.Pointer to the start of the allocated memory block.
|
||||
|
||||
Raises
|
||||
----------
|
||||
ValueError
|
||||
If num_elems is less than 1.
|
||||
"""
|
||||
if num_elems < 1:
|
||||
raise ValueError("num_elems must be at least 1")
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(
|
||||
f"value_ty must be a type of Numeric, but got {element_type}"
|
||||
)
|
||||
|
||||
ptr = self.allocate(
|
||||
element_type.width // 8 * num_elems, element_type.width // 8
|
||||
)
|
||||
|
||||
return cute.recast_ptr(ptr, dtype=element_type)
|
||||
|
||||
def allocate_tensor(
|
||||
self,
|
||||
element_type: Type[Numeric],
|
||||
layout: Union[int, cute.Layout, cute.ComposedLayout],
|
||||
byte_alignment: int = 1,
|
||||
swizzle: cute.Swizzle = None,
|
||||
):
|
||||
"""
|
||||
Allocates a tensor in the shared memory with value type, layout and byte alignment.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
element_type : Type[Numeric]
|
||||
The type of the values in the tensor.
|
||||
layout : int | DynamicInt | cute.Layout | cute.ComposedLayout
|
||||
The layout of the tensor.
|
||||
byte_alignment : int, optional
|
||||
The byte alignment requirement for the allocation. Defaults to 1 (no alignment).
|
||||
swizzle : cute.Swizzle
|
||||
A swizzle for the iterator (for position-dependent swizzling).
|
||||
|
||||
Returns
|
||||
-------
|
||||
tensor : cute.Tensor
|
||||
The allocated tensor with specified value type, layout and byte alignment.
|
||||
|
||||
Notes
|
||||
-----
|
||||
The base address is updated to point to the next available memory location.
|
||||
"""
|
||||
if not isinstance(element_type, NumericMeta):
|
||||
raise TypeError(
|
||||
f"value_ty must be a type of Numeric, but got {element_type}"
|
||||
)
|
||||
|
||||
if (
|
||||
isinstance(layout, cute.ComposedLayout)
|
||||
and isinstance(layout.inner, cute.Swizzle)
|
||||
) and (swizzle is not None):
|
||||
raise TypeError(
|
||||
f"iterator swizzle with swizzle layout is currently not supported"
|
||||
)
|
||||
|
||||
if isinstance(layout, int):
|
||||
layout = cute.make_layout(layout)
|
||||
|
||||
profile = layout(0)
|
||||
if isinstance(profile, tuple):
|
||||
raise TypeError(
|
||||
f"cannot allocate a shared memory tensor with a non-integer iterator"
|
||||
)
|
||||
|
||||
if not cute.is_static(layout.type):
|
||||
raise NotImplementedError(f"dynamic layout is not supported: {layout.type}")
|
||||
|
||||
# At least align the allocation to the natural alignment given by the element type
|
||||
if element_type.width // 8 > byte_alignment:
|
||||
byte_alignment = element_type.width // 8
|
||||
|
||||
# Relevant only for sub-byte data types: verify that the entire allocation is byte-aligned
|
||||
cosize_in_bits = cute.cosize(layout) * element_type.width
|
||||
assert isinstance(cosize_in_bits, int)
|
||||
if cosize_in_bits % 8 != 0:
|
||||
raise ValueError("invalid allocation that is not byte-aligned")
|
||||
|
||||
num_bytes = cosize_in_bits // 8
|
||||
ptr = self.allocate(num_bytes, byte_alignment)
|
||||
ptr = cute.recast_ptr(ptr, swizzle, dtype=element_type)
|
||||
res = cute.make_tensor(ptr, layout)
|
||||
return res
|
||||
@@ -0,0 +1,384 @@
|
||||
# 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 Tuple
|
||||
|
||||
from cutlass.cutlass_dsl import (
|
||||
Boolean,
|
||||
Integer,
|
||||
Int32,
|
||||
min,
|
||||
extract_mlir_values,
|
||||
new_from_mlir_values,
|
||||
dsl_user_op,
|
||||
)
|
||||
from cutlass._mlir import ir
|
||||
import cutlass.cute as cute
|
||||
|
||||
##############################################################################
|
||||
# Static persistent tile scheduler
|
||||
##############################################################################
|
||||
|
||||
|
||||
class WorkTileInfo:
|
||||
"""A class to represent information about a work tile.
|
||||
|
||||
:ivar tile_idx: The index of the tile.
|
||||
:type tile_idx: cute.Coord
|
||||
:ivar is_valid_tile: Whether the tile is valid.
|
||||
:type is_valid_tile: Boolean
|
||||
"""
|
||||
|
||||
def __init__(self, tile_idx: cute.Coord, is_valid_tile: Boolean):
|
||||
self._tile_idx = tile_idx
|
||||
self._is_valid_tile = Boolean(is_valid_tile)
|
||||
|
||||
def __extract_mlir_values__(self) -> list[ir.Value]:
|
||||
values = extract_mlir_values(self.tile_idx)
|
||||
values.extend(extract_mlir_values(self.is_valid_tile))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(self, values: list[ir.Value]) -> "WorkTileInfo":
|
||||
assert len(values) == 4
|
||||
new_tile_idx = new_from_mlir_values(self._tile_idx, values[:-1])
|
||||
new_is_valid_tile = new_from_mlir_values(self._is_valid_tile, [values[-1]])
|
||||
return WorkTileInfo(new_tile_idx, new_is_valid_tile)
|
||||
|
||||
@property
|
||||
def is_valid_tile(self) -> Boolean:
|
||||
"""Check latest tile returned by the scheduler is valid or not. Any scheduling
|
||||
requests after all tasks completed will return an invalid tile.
|
||||
|
||||
:return: The validity of the tile.
|
||||
:rtype: Boolean
|
||||
"""
|
||||
return self._is_valid_tile
|
||||
|
||||
@property
|
||||
def tile_idx(self) -> cute.Coord:
|
||||
"""
|
||||
Get the index of the tile.
|
||||
|
||||
:return: The index of the tile.
|
||||
:rtype: cute.Coord
|
||||
"""
|
||||
return self._tile_idx
|
||||
|
||||
|
||||
class PersistentTileSchedulerParams:
|
||||
"""A class to represent parameters for a persistent tile scheduler.
|
||||
|
||||
This class is designed to manage and compute the layout of clusters and tiles
|
||||
in a batched gemm problem.
|
||||
|
||||
:ivar cluster_shape_mn: Shape of the cluster in (m, n) dimensions (K dimension cta count must be 1).
|
||||
:type cluster_shape_mn: tuple
|
||||
:ivar problem_layout_ncluster_mnl: Layout of the problem in terms of
|
||||
number of clusters in (m, n, l) dimensions.
|
||||
:type problem_layout_ncluster_mnl: cute.Layout
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
problem_shape_ntile_mnl: cute.Shape,
|
||||
cluster_shape_mnk: cute.Shape,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
"""
|
||||
Initializes the PersistentTileSchedulerParams with the given parameters.
|
||||
|
||||
:param problem_shape_ntile_mnl: The shape of the problem in terms of
|
||||
number of CTA (Cooperative Thread Array) in (m, n, l) dimensions.
|
||||
:type problem_shape_ntile_mnl: cute.Shape
|
||||
:param cluster_shape_mnk: The shape of the cluster in (m, n) dimensions.
|
||||
:type cluster_shape_mnk: cute.Shape
|
||||
|
||||
:raises ValueError: If cluster_shape_k is not 1.
|
||||
"""
|
||||
|
||||
if cluster_shape_mnk[2] != 1:
|
||||
raise ValueError(f"unsupported cluster_shape_k {cluster_shape_mnk[2]}")
|
||||
|
||||
self.problem_shape_ntile_mnl = problem_shape_ntile_mnl
|
||||
# cluster_shape_mnk is kept for reconstruction
|
||||
self._cluster_shape_mnk = cluster_shape_mnk
|
||||
self.cluster_shape_mn = cluster_shape_mnk[:2]
|
||||
self._loc = loc
|
||||
|
||||
# By default, we follow m major (col-major) raster order, so make a col-major layout
|
||||
self.problem_layout_ncluster_mnl = cute.make_layout(
|
||||
cute.ceil_div(
|
||||
self.problem_shape_ntile_mnl, cluster_shape_mnk[:2], loc=loc, ip=ip
|
||||
),
|
||||
loc=loc,
|
||||
ip=ip,
|
||||
)
|
||||
|
||||
def __extract_mlir_values__(self):
|
||||
values, self._values_pos = [], []
|
||||
for obj in [self.problem_shape_ntile_mnl, self._cluster_shape_mnk]:
|
||||
obj_values = extract_mlir_values(obj)
|
||||
values += obj_values
|
||||
self._values_pos.append(len(obj_values))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(self, values):
|
||||
obj_list = []
|
||||
for obj, n_items in zip(
|
||||
[self.problem_shape_ntile_mnl, self._cluster_shape_mnk], self._values_pos
|
||||
):
|
||||
obj_list.append(new_from_mlir_values(obj, values[:n_items]))
|
||||
values = values[n_items:]
|
||||
return PersistentTileSchedulerParams(*(tuple(obj_list)), loc=self._loc)
|
||||
|
||||
@dsl_user_op
|
||||
def get_grid_shape(
|
||||
self, max_active_clusters: Int32, *, loc=None, ip=None
|
||||
) -> Tuple[Integer, Integer, Integer]:
|
||||
"""
|
||||
Computes the grid shape based on the maximum active clusters allowed.
|
||||
|
||||
:param max_active_clusters: The maximum number of active clusters that
|
||||
can run in one wave.
|
||||
:type max_active_clusters: Int32
|
||||
|
||||
:return: A tuple containing the grid shape in (m, n, persistent_clusters).
|
||||
- m: self.cluster_shape_m.
|
||||
- n: self.cluster_shape_n.
|
||||
- persistent_clusters: Number of persistent clusters that can run.
|
||||
"""
|
||||
|
||||
# Total ctas in problem size
|
||||
num_ctas_mnl = tuple(
|
||||
x * y
|
||||
for x, y in zip(
|
||||
self.problem_layout_ncluster_mnl.shape, self.cluster_shape_mn
|
||||
)
|
||||
) + (self.problem_layout_ncluster_mnl.shape[2],)
|
||||
|
||||
num_ctas_in_problem = cute.size(num_ctas_mnl, loc=loc, ip=ip)
|
||||
|
||||
num_ctas_per_cluster = cute.size(self.cluster_shape_mn, loc=loc, ip=ip)
|
||||
# Total ctas that can run in one wave
|
||||
num_ctas_per_wave = max_active_clusters * num_ctas_per_cluster
|
||||
|
||||
num_persistent_ctas = min(num_ctas_in_problem, num_ctas_per_wave)
|
||||
num_persistent_clusters = num_persistent_ctas // num_ctas_per_cluster
|
||||
|
||||
return (*self.cluster_shape_mn, num_persistent_clusters)
|
||||
|
||||
|
||||
class StaticPersistentTileScheduler:
|
||||
"""A scheduler for static persistent tile execution in CUTLASS/CuTe kernels.
|
||||
|
||||
:ivar params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl
|
||||
:type params: PersistentTileSchedulerParams
|
||||
:ivar num_persistent_clusters: Number of persistent clusters that can be launched
|
||||
:type num_persistent_clusters: Int32
|
||||
:ivar cta_id_in_cluster: ID of the CTA within its cluster
|
||||
:type cta_id_in_cluster: cute.Coord
|
||||
:ivar _num_tiles_executed: Counter for executed tiles
|
||||
:type _num_tiles_executed: Int32
|
||||
:ivar _current_work_linear_idx: Current cluster index
|
||||
:type _current_work_linear_idx: Int32
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
params: PersistentTileSchedulerParams,
|
||||
num_persistent_clusters: Int32,
|
||||
current_work_linear_idx: Int32,
|
||||
cta_id_in_cluster: cute.Coord,
|
||||
num_tiles_executed: Int32,
|
||||
):
|
||||
"""
|
||||
Initializes the StaticPersistentTileScheduler with the given parameters.
|
||||
|
||||
:param params: Tile schedule related params, including cluster shape and problem_layout_ncluster_mnl.
|
||||
:type params: PersistentTileSchedulerParams
|
||||
:param num_persistent_clusters: Number of persistent clusters that can be launched.
|
||||
:type num_persistent_clusters: Int32
|
||||
:param current_work_linear_idx: Current cluster index.
|
||||
:type current_work_linear_idx: Int32
|
||||
:param cta_id_in_cluster: ID of the CTA within its cluster.
|
||||
:type cta_id_in_cluster: cute.Coord
|
||||
:param num_tiles_executed: Counter for executed tiles.
|
||||
:type num_tiles_executed: Int32
|
||||
"""
|
||||
self.params = params
|
||||
self.num_persistent_clusters = num_persistent_clusters
|
||||
self._current_work_linear_idx = current_work_linear_idx
|
||||
self.cta_id_in_cluster = cta_id_in_cluster
|
||||
self._num_tiles_executed = num_tiles_executed
|
||||
|
||||
def __extract_mlir_values__(self) -> list[ir.Value]:
|
||||
values = extract_mlir_values(self.num_persistent_clusters)
|
||||
values.extend(extract_mlir_values(self._current_work_linear_idx))
|
||||
values.extend(extract_mlir_values(self.cta_id_in_cluster))
|
||||
values.extend(extract_mlir_values(self._num_tiles_executed))
|
||||
return values
|
||||
|
||||
def __new_from_mlir_values__(
|
||||
self, values: list[ir.Value]
|
||||
) -> "StaticPersistentTileScheduler":
|
||||
assert len(values) == 6
|
||||
new_num_persistent_clusters = new_from_mlir_values(
|
||||
self.num_persistent_clusters, [values[0]]
|
||||
)
|
||||
new_current_work_linear_idx = new_from_mlir_values(
|
||||
self._current_work_linear_idx, [values[1]]
|
||||
)
|
||||
new_cta_id_in_cluster = new_from_mlir_values(
|
||||
self.cta_id_in_cluster, values[2:5]
|
||||
)
|
||||
new_num_tiles_executed = new_from_mlir_values(
|
||||
self._num_tiles_executed, [values[5]]
|
||||
)
|
||||
return StaticPersistentTileScheduler(
|
||||
self.params,
|
||||
new_num_persistent_clusters,
|
||||
new_current_work_linear_idx,
|
||||
new_cta_id_in_cluster,
|
||||
new_num_tiles_executed,
|
||||
)
|
||||
|
||||
# called by host
|
||||
@dsl_user_op
|
||||
@staticmethod
|
||||
def create(
|
||||
params: PersistentTileSchedulerParams,
|
||||
block_idx: Tuple[Integer, Integer, Integer],
|
||||
grid_dim: Tuple[Integer, Integer, Integer],
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
):
|
||||
"""Initialize the static persistent tile scheduler.
|
||||
|
||||
:param params: Parameters for the persistent
|
||||
tile scheduler.
|
||||
:type params: PersistentTileSchedulerParams
|
||||
:param block_idx: The 3d block index in the format (bidx, bidy, bidz).
|
||||
:type block_idx: Tuple[Integer, Integer, Integer]
|
||||
:param grid_dim: The 3d grid dimensions for kernel launch.
|
||||
:type grid_dim: Tuple[Integer, Integer, Integer]
|
||||
|
||||
:return: A StaticPersistentTileScheduler object.
|
||||
:rtype: StaticPersistentTileScheduler
|
||||
"""
|
||||
params = params
|
||||
|
||||
# Calculate the number of persistent clusters by dividing the total grid size
|
||||
# by the number of CTAs per cluster
|
||||
num_persistent_clusters = cute.size(grid_dim, loc=loc, ip=ip) // cute.size(
|
||||
params.cluster_shape_mn, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
bidx, bidy, bidz = block_idx
|
||||
|
||||
# Initialize workload index equals to the cluster index in the grid
|
||||
current_work_linear_idx = Int32(bidz)
|
||||
|
||||
# CTA id in the cluster
|
||||
cta_id_in_cluster = (
|
||||
Int32(bidx % params.cluster_shape_mn[0]),
|
||||
Int32(bidy % params.cluster_shape_mn[1]),
|
||||
Int32(0),
|
||||
)
|
||||
# Initialize number of tiles executed to zero
|
||||
num_tiles_executed = Int32(0)
|
||||
return StaticPersistentTileScheduler(
|
||||
params,
|
||||
num_persistent_clusters,
|
||||
current_work_linear_idx,
|
||||
cta_id_in_cluster,
|
||||
num_tiles_executed,
|
||||
)
|
||||
|
||||
# called by host
|
||||
@staticmethod
|
||||
def get_grid_shape(
|
||||
params: PersistentTileSchedulerParams,
|
||||
max_active_clusters: Int32,
|
||||
*,
|
||||
loc=None,
|
||||
ip=None,
|
||||
) -> Tuple[Integer, Integer, Integer]:
|
||||
"""Calculates the grid shape to be launched on GPU using problem shape,
|
||||
threadblock shape, and active cluster size.
|
||||
|
||||
:param params: Parameters for grid shape calculation.
|
||||
:type params: PersistentTileSchedulerParams
|
||||
:param max_active_clusters: Maximum active clusters allowed.
|
||||
:type max_active_clusters: Int32
|
||||
|
||||
:return: The calculated 3d grid shape.
|
||||
:rtype: Tuple[Integer, Integer, Integer]
|
||||
"""
|
||||
|
||||
return params.get_grid_shape(max_active_clusters, loc=loc, ip=ip)
|
||||
|
||||
# private method
|
||||
def _get_current_work_for_linear_idx(
|
||||
self, current_work_linear_idx: Int32, *, loc=None, ip=None
|
||||
) -> WorkTileInfo:
|
||||
"""Compute current tile coord given current_work_linear_idx and cta_id_in_cluster.
|
||||
|
||||
:param current_work_linear_idx: The linear index of the current work.
|
||||
:type current_work_linear_idx: Int32
|
||||
|
||||
:return: An object containing information about the current tile coordinates
|
||||
and validity status.
|
||||
:rtype: WorkTileInfo
|
||||
"""
|
||||
|
||||
is_valid = current_work_linear_idx < cute.size(
|
||||
self.params.problem_layout_ncluster_mnl, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
cur_cluster_coord = self.params.problem_layout_ncluster_mnl.get_hier_coord(
|
||||
current_work_linear_idx, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
# cur_tile_coord is a tuple of i32 values
|
||||
cur_tile_coord = tuple(
|
||||
Int32(x) * Int32(z) + Int32(y)
|
||||
for x, y, z in zip(
|
||||
cur_cluster_coord,
|
||||
self.cta_id_in_cluster,
|
||||
(*self.params.cluster_shape_mn, Int32(1)),
|
||||
)
|
||||
)
|
||||
|
||||
return WorkTileInfo(cur_tile_coord, is_valid)
|
||||
|
||||
@dsl_user_op
|
||||
def get_current_work(self, *, loc=None, ip=None) -> WorkTileInfo:
|
||||
return self._get_current_work_for_linear_idx(
|
||||
self._current_work_linear_idx, loc=loc, ip=ip
|
||||
)
|
||||
|
||||
@dsl_user_op
|
||||
def initial_work_tile_info(self, *, loc=None, ip=None) -> WorkTileInfo:
|
||||
return self.get_current_work(loc=loc, ip=ip)
|
||||
|
||||
@dsl_user_op
|
||||
def advance_to_next_work(self, *, advance_count: int = 1, loc=None, ip=None):
|
||||
self._current_work_linear_idx += Int32(advance_count) * Int32(
|
||||
self.num_persistent_clusters
|
||||
)
|
||||
self._num_tiles_executed += Int32(1)
|
||||
|
||||
@property
|
||||
def num_tiles_executed(self) -> Int32:
|
||||
return self._num_tiles_executed
|
||||
@@ -0,0 +1,140 @@
|
||||
# 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 enum import Enum, auto
|
||||
from typing import Tuple
|
||||
|
||||
from cutlass.cutlass_dsl import const_expr
|
||||
|
||||
import cutlass._mlir.dialects.cute as _cute_ir
|
||||
import cutlass._mlir.dialects.cute_nvgpu as _cute_nvgpu_ir
|
||||
|
||||
import cutlass.cute as cute
|
||||
|
||||
|
||||
class TensorMapUpdateMode(Enum):
|
||||
"""
|
||||
Enum class defining tensor map update modes.
|
||||
|
||||
Modes:
|
||||
GMEM: Update tensormap in global memory
|
||||
SMEM: Load tensormap from global memory to shared memory,
|
||||
update it in shared memory, then store back to global memory
|
||||
"""
|
||||
|
||||
GMEM = auto() # Update tensormap in global memory
|
||||
SMEM = auto() # Update tensormap in shared memory
|
||||
|
||||
|
||||
@dataclass(frozen=True)
|
||||
class TensorMapManager:
|
||||
"""
|
||||
Manages TensorMap operations including initialization and updates.
|
||||
Provides utilities to convert tensormap pointer to across different memory spaces.
|
||||
"""
|
||||
|
||||
tensormap_update_mode: TensorMapUpdateMode
|
||||
bytes_per_tensormap: int
|
||||
|
||||
# convert given cute.Pointer or cutlass.Int64 to a cute.Pointer to tensormap.
|
||||
# address_space: the address space of the resulting tensormap pointer. It could be generic or gmem
|
||||
def get_tensormap_ptr(
|
||||
self,
|
||||
ptr: cute.Pointer,
|
||||
address_space=_cute_ir.AddressSpace.gmem,
|
||||
) -> cute.Pointer:
|
||||
if address_space not in [
|
||||
_cute_ir.AddressSpace.gmem,
|
||||
_cute_ir.AddressSpace.generic,
|
||||
]:
|
||||
raise ValueError(f"Invalid address space: {address_space} for tensormap")
|
||||
|
||||
gmem_ptr_i64 = ptr.toint().ir_value()
|
||||
gmem_ptr_i64_align_ty = _cute_ir.ConstrainedIntType.get(
|
||||
self.bytes_per_tensormap, gmem_ptr_i64.type.width
|
||||
)
|
||||
gmem_ptr_i64_align = _cute_ir.assume(gmem_ptr_i64_align_ty, gmem_ptr_i64)
|
||||
gmem_ptr_ty = _cute_ir.PtrType.get(
|
||||
_cute_nvgpu_ir.TmaDescriptorTiledType.get(),
|
||||
address_space,
|
||||
self.bytes_per_tensormap,
|
||||
)
|
||||
return _cute_ir.inttoptr(gmem_ptr_ty, gmem_ptr_i64_align)
|
||||
|
||||
# init tensormap pointed by dst_ptr with the one inside copy_atom.
|
||||
# dst_ptr should be pointing to a global memory location or a smem location
|
||||
# warp_id specifies which warp to perform the initialization
|
||||
@cute.jit
|
||||
def init_tensormap_from_atom(
|
||||
self, copy_atom: cute.CopyAtom, dst_ptr: cute.Pointer, warp_id: int
|
||||
) -> None:
|
||||
warp_idx = cute.arch.warp_idx()
|
||||
warp_idx = cute.arch.make_warp_uniform(warp_idx)
|
||||
if warp_idx == warp_id:
|
||||
with cute.arch.elect_one():
|
||||
cute.nvgpu.cpasync.copy_tensormap(copy_atom, dst_ptr)
|
||||
cute.arch.sync_warp()
|
||||
return
|
||||
|
||||
# Perform a fence operation to ensure previous `init_tensormap_from_atom` calls have been completed
|
||||
def fence_tensormap_initialization(
|
||||
self,
|
||||
) -> None:
|
||||
if self.tensormap_update_mode == TensorMapUpdateMode.GMEM:
|
||||
cute.arch.fence_acq_rel_cta()
|
||||
return
|
||||
|
||||
# Perform a fence operation to ensure previous `update_tensormap` calls have been completed
|
||||
def fence_tensormap_update(
|
||||
self,
|
||||
tensormap_ptr: cute.Pointer,
|
||||
) -> None:
|
||||
cute.nvgpu.cpasync.fence_tma_desc_acquire(tensormap_ptr)
|
||||
return
|
||||
|
||||
@cute.jit
|
||||
def update_tensormap(
|
||||
self,
|
||||
tensor_gmem: Tuple[cute.Tensor, ...],
|
||||
tma_copy_atom: Tuple[cute.CopyAtom, ...],
|
||||
tensormap_gmem_ptr: Tuple[cute.Pointer, ...],
|
||||
warp_id: int,
|
||||
tensormap_smem_ptr: Tuple[cute.Pointer, ...],
|
||||
) -> None:
|
||||
warp_idx = cute.arch.make_warp_uniform(cute.arch.warp_idx())
|
||||
# updates before touching tensormap in global memory
|
||||
if warp_idx == warp_id:
|
||||
if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM):
|
||||
for copy_atom, tensor, smem_ptr in zip(
|
||||
tma_copy_atom, tensor_gmem, tensormap_smem_ptr
|
||||
):
|
||||
cute.nvgpu.cpasync.update_tma_descriptor(
|
||||
copy_atom, tensor, smem_ptr
|
||||
)
|
||||
# wait until it's safe to update tensormap in global memory
|
||||
with cute.arch.elect_one():
|
||||
cute.arch.cp_async_bulk_commit_group()
|
||||
cute.arch.cp_async_bulk_wait_group(0, read=True)
|
||||
cute.arch.sync_warp()
|
||||
# updates to tensormap in global memory
|
||||
if const_expr(self.tensormap_update_mode == TensorMapUpdateMode.SMEM):
|
||||
for gmem_ptr, smem_ptr in zip(tensormap_gmem_ptr, tensormap_smem_ptr):
|
||||
cute.nvgpu.cpasync.cp_fence_tma_desc_release(gmem_ptr, smem_ptr)
|
||||
else:
|
||||
for copy_atom, tensor, gmem_ptr in zip(
|
||||
tma_copy_atom, tensor_gmem, tensormap_gmem_ptr
|
||||
):
|
||||
cute.nvgpu.cpasync.update_tma_descriptor(
|
||||
copy_atom, tensor, gmem_ptr
|
||||
)
|
||||
cute.arch.sync_warp()
|
||||
cute.nvgpu.cpasync.fence_tma_desc_release()
|
||||
Reference in New Issue
Block a user