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

v4.4 tag release update. (#3032)

Browse Source
This commit is contained in:
Junkai-Wu
2026-02-14 12:27:58 +08:00
committed by GitHub
parent 01687cfba1
commit d4bbf728ca
140 changed files with 41624 additions and 3691 deletions
Download Patch File Download Diff File Expand all files Collapse all files

88
python/CuTeDSL/cutlass/cute/algorithm.py
View File

@@ -29,8 +29,14 @@ from .core import (
append_ones,
group_modes,
)
from .atom import MmaAtom, CopyAtom, make_atom
from .atom import (
MmaAtom,
CopyAtom,
make_atom,
_normalize_variadic_tensor_operand,
copy_atom_call,
)
from .nvgpu.common import CacheEvictionPriority
def _normalize_gemm_operand_list(
x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str
@@ -156,7 +162,7 @@ def basic_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
src.element_type.mlir_type, src.element_type.width
)
simt_copy = make_atom(simt_copy_ty, loc=loc, ip=ip)
return _cute_ir.copy(simt_copy, src.value, dst.value, loc=loc, ip=ip)
return _cute_ir.copy(simt_copy, [src.value], [dst.value], loc=loc, ip=ip)
s = size(dst, loc=loc, ip=ip)
# Always generate an scf.for Op when one of the tensors is dynamic
@@ -210,7 +216,14 @@ def _basic_copy_if_static(
@dsl_user_op
def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
def autovec_copy(
src: Tensor,
dst: Tensor,
*,
l1c_evict_priority: CacheEvictionPriority = CacheEvictionPriority.EVICT_NORMAL,
loc=None,
ip=None,
) -> None:
"""
Auto-vectorization SIMT copy policy.
@@ -263,11 +276,15 @@ def autovec_copy(src: Tensor, dst: Tensor, *, loc=None, ip=None) -> None:
# Dispatch to copy with atom
simt_type = _cute_nvgpu_ir.CopyAtomSIMTSyncCopyType.get(
src.element_type.mlir_type, num_bits_per_copy
src.element_type.mlir_type,
num_bits_per_copy,
0,
0,
l1c_evict_priority._to_ir(),
)
simt_copy = make_atom(simt_type, loc=loc, ip=ip)
return _cute_ir.copy(
simt_copy, tiled_src.value, tiled_dst.value, loc=loc, ip=ip
simt_copy, [tiled_src.value], [tiled_dst.value], loc=loc, ip=ip
)
# Failed to vectorize, use a basic copy
@@ -331,8 +348,8 @@ def _parse_auto_multicast_args(
@dsl_user_op
def copy(
atom: CopyAtom,
src: Tensor,
dst: Tensor,
src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
*,
pred: Optional[Tensor] = None,
loc=None,
@@ -343,10 +360,10 @@ def copy(
:param atom: Copy atom specifying the transfer operation
:type atom: CopyAtom
:param src: Source tensor with layout profile ``(V, Rest...)``
:type src: Tensor
:param dst: Destination tensor with layout profile ``(V, Rest...)``
:type dst: Tensor
:param src: Source tensor or list of tensors with layout profile ``(V, Rest...)``
:type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param dst: Destination tensor or list of tensors with layout profile ``(V, Rest...)``
:type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param pred: Optional predication tensor for conditional transfers, defaults to None
:type pred: Optional[Tensor], optional
:param loc: Source location information, defaults to None
@@ -374,6 +391,12 @@ def copy(
Source and destination tensors must be partitioned in accordance with the Copy Atom specifications.
Post-partitioning, both tensors will exhibit a ``(V, Rest...)`` layout profile.
The operands `src` and `dst` are variadic, each containing a variable number of tensors:
- For regular copy, `src` and `dst` contain single source and destination tensors respectively.
- For copy with auxiliary operands, `src` and `dst` contain the primary tensors followed by
their respective auxiliary tensors.
**Precondition:** The size of mode 1 must be equal for both source and destination tensors:
``size(src, mode=[1]) == size(dst, mode=[1])``
@@ -399,41 +422,54 @@ def copy(
for future releases.
"""
if isinstance(src.type, _cute_ir.MemRefType) and isinstance(
dst.type, _cute_ir.MemRefType
# Normalize src/dst to lists for variadic IR operands
src_list = _normalize_variadic_tensor_operand(src, "src")
dst_list = _normalize_variadic_tensor_operand(dst, "dst")
# Validate primary tensors (first element)
src_primary = src_list[0]
dst_primary = dst_list[0]
if isinstance(src_primary.type, _cute_ir.MemRefType) and isinstance(
dst_primary.type, _cute_ir.MemRefType
):
if src.element_type.width != dst.element_type.width:
if src_primary.element_type.width != dst_primary.element_type.width:
raise TypeError(
"`copy` currently only supports equal source and destination "
"element type bit width"
)
if rank(src) != rank(dst):
if rank(src_primary) != rank(dst_primary):
raise ValueError(
"Expected source and destination tensors to have the same rank, "
f"but got {rank(src)} and {rank(dst)}"
f"but got {rank(src_primary)} and {rank(dst_primary)}"
)
# Canonicalize to at least rank-2 tensors
src = group_modes(append_ones(src, up_to_rank=2), 1)
dst = group_modes(append_ones(dst, up_to_rank=2), 1)
# Canonicalize all tensors to at least rank-2
src_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in src_list]
dst_list = [group_modes(append_ones(t, up_to_rank=2), 1) for t in dst_list]
if pred is not None:
pred = group_modes(append_ones(pred, up_to_rank=2), 1)
if is_static(src.shape[1]) and is_static(dst.shape[1]):
if size(src, mode=[1]) != size(dst, mode=[1]):
# Recompute primary references after canonicalization
src_primary = src_list[0]
dst_primary = dst_list[0]
if is_static(src_primary.shape[1]) and is_static(dst_primary.shape[1]):
if size(src_primary, mode=[1]) != size(dst_primary, mode=[1]):
raise ValueError(
"Expected source and destination tensors to have the same size in mode-1, "
f"but got {size(src, mode=[1])} and {size(dst, mode=[1])}"
f"but got {size(src_primary, mode=[1])} and {size(dst_primary, mode=[1])}"
)
multicast_attr_pairs = _parse_auto_multicast_args(kwargs)
value = atom._unpack(loc=loc, ip=ip, **kwargs)
if isinstance(pred, Tensor):
pred = pred.value
pred_value = pred.value if isinstance(pred, Tensor) else pred
op = _cute_ir.copy(value, src.value, dst.value, pred=pred, loc=loc, ip=ip)
src_vals = [t.value for t in src_list]
dst_vals = [t.value for t in dst_list]
op = _cute_ir.copy(value, src_vals, dst_vals, pred=pred_value, loc=loc, ip=ip)
for name, attr in multicast_attr_pairs:
op.attributes[name] = attr

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

@@ -13,7 +13,7 @@ from functools import partial
from typing import Any, Optional, Tuple, Union, Callable, Literal
from typing_extensions import deprecated
from cutlass.cutlass_dsl import T, dsl_user_op
from cutlass.cutlass_dsl import T, dsl_user_op, target_version
import cutlass.cutlass_dsl as cutlass_dsl
@@ -2557,3 +2557,42 @@ def cvt_f4e2m1x8_to_f16x8(src_vec8, *, loc=None, ip=None):
vec_f16x8_type = ir.VectorType.get([8], Float16.mlir_type, loc=loc)
vec_f16x8 = llvm.bitcast(vec_f16x8_type, vec_f32x4, loc=loc, ip=ip)
return vec_f16x8
@dsl_user_op
def mapa(ptr, cta_rank_in_cluster=0, *, loc=None, ip=None):
"""
Map a pointer to distributed shared memory across cluster.
Portable wrapper that uses the appropriate NVVM API based on CUDA version:
- CUDA 13.1+: Uses nvvm.mapa with dsmem address space
- CUDA 12.9: Uses nvvm.mapa_shared_cluster
Args:
ptr: Pointer to shared memory (llvm_ptr attribute will be used)
cta_rank_in_cluster: CTA rank within the cluster (default 0)
Returns:
Mapped LLVM pointer to shared memory
"""
if target_version(min_version="13.1"):
dsmem_ptr_ty = llvm.PointerType.get(7) # dsmem
smem_ptr_ty = llvm.PointerType.get(3) # smem
llvm_ptr = nvvm.mapa(
dsmem_ptr_ty,
ptr.llvm_ptr,
Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)
return llvm.addrspacecast(smem_ptr_ty, llvm_ptr, loc=loc, ip=ip)
else:
llvm_ptr = ptr.llvm_ptr
return nvvm.mapa_shared_cluster(
llvm_ptr.type,
llvm_ptr,
Int32(cta_rank_in_cluster).ir_value(loc=loc, ip=ip),
loc=loc,
ip=ip,
)

96
python/CuTeDSL/cutlass/cute/atom.py
View File

@@ -10,7 +10,7 @@
# is strictly prohibited.
from abc import ABC, ABCMeta, abstractmethod
from typing import Type, Union, Optional, Any, overload
from typing import Type, Union, Optional, Any, List, Tuple, overload
from .typing import Shape, Layout, Tile, Tensor, Numeric, Int32
from .core import (
@@ -1113,25 +1113,66 @@ def make_tiled_copy_C_atom(atom: CopyAtom, mma: TiledMma, *, loc=None, ip=None):
return _make_tiled_copy(atom, layout_tv, tiler_mn, loc=loc, ip=ip)
def _normalize_variadic_tensor_operand(
x: Union["Tensor", List["Tensor"], Tuple["Tensor", ...]], name: str
) -> List["Tensor"]:
"""Normalize a Tensor or sequence of Tensors to a list of Tensors.
Helper function for operations with variadic operands.
"""
if isinstance(x, Tensor):
return [x]
if isinstance(x, (list, tuple)):
if len(x) == 0:
raise ValueError(f"`{name}` must contain at least one Tensor")
if not all(isinstance(t, Tensor) for t in x):
raise TypeError(f"All elements of `{name}` must be Tensor")
return list(x) # type: ignore
raise TypeError(f"`{name}` must be a Tensor or a sequence of Tensors")
@dsl_user_op
def copy_atom_call(
atom: CopyAtom,
src: Tensor,
dst: Tensor,
src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]],
*,
pred: Optional[Tensor] = None,
loc=None,
ip=None,
**kwargs,
) -> None:
"""Executes a single copy atom operation between two tensors.
"""
Execute a single copy atom operation.
The copy_atom_call operation executes a copy atom with the given operands.
Source and destination tensors have layout profile ``(V)``.
The ``V-mode`` represents either:
- A singular mode directly consumable by the provided Copy Atom
- A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``,
For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``.
- Certain Atoms may require additional operation-specific keyword arguments.
- Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned
for future releases.
Both ``src`` and ``dst`` operands are variadic, containing a variable number of tensors:
- For regular copy, ``src`` and ``dst`` each contain a single tensor.
- For copy with auxiliary operands, they contain the main tensor followed by
auxiliary tensors. For example:
:param atom: Copy atom specifying the transfer operation
:type atom: CopyAtom
:param src: Source tensor with layout profile ``(V)``
:type src: Tensor
:param dst: Destination tensor with layout profile ``(V)``
:type dst: Tensor
:param src: Source tensor(s) with layout profile ``(V)``. Can be a single Tensor
or a list/tuple of Tensors for operations with auxiliary source operands.
:type src: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param dst: Destination tensor(s) with layout profile ``(V)``. Can be a single Tensor
or a list/tuple of Tensors for operations with auxiliary destination operands.
:type dst: Union[Tensor, List[Tensor], Tuple[Tensor, ...]]
:param pred: Optional predication tensor for conditional transfers, defaults to None
:type pred: Optional[Tensor], optional
:param loc: Source location information, defaults to None
@@ -1144,54 +1185,43 @@ def copy_atom_call(
:return: None
:rtype: None
The copy_atom_call operation executes a single copy atom with the given operands.
Source and destination tensors with layout profile like ``(V)``.
The ``V-mode`` represents either:
- A singular mode directly consumable by the provided Copy Atom
- A composite mode requiring recursive decomposition, structured as ``(V, Rest...)``,
For src/dst layout like ``(V, Rest...)``, the layout profile of ``pred`` must match ``(Rest...)``.
**Examples**:
.. code-block:: python
# Basic copy atom operation
# Regular copy atom operation
cute.copy_atom_call(copy_atom, src, dst)
# Predicated copy atom operation
cute.copy_atom_call(copy_atom, src, dst, pred=pred)
.. note::
- Certain Atoms may require additional operation-specific keyword arguments.
- Current implementation limits ``V-mode`` rank to 2 or less. Support for higher ranks is planned
for future releases.
"""
if isinstance(src.type, _cute_ir.MemRefType) and isinstance(
dst.type, _cute_ir.MemRefType
# Normalize src/dst to lists for variadic IR operands, while keeping old API working.
src_list = _normalize_variadic_tensor_operand(src, "src")
dst_list = _normalize_variadic_tensor_operand(dst, "dst")
# Validate first src/dst for element type width check
if isinstance(src_list[0].type, _cute_ir.MemRefType) and isinstance(
dst_list[0].type, _cute_ir.MemRefType
):
if src.element_type.width != dst.element_type.width:
if src_list[0].element_type.width != dst_list[0].element_type.width:
raise TypeError(
"`copy_atom_call` currently only supports equal source and destination "
"element type bit width"
)
if rank(src, mode=[0]) > 2 or rank(dst, mode=[0]) > 2:
if rank(src_list[0], mode=[0]) > 2 or rank(dst_list[0], mode=[0]) > 2:
raise NotImplementedError(
"V-mode (mode-0) with rank > 2 is not supported yet, "
f"but got rank(src, mode=[0]) = {rank(src, mode=[0])} and rank(dst, mode=[0]) = {rank(dst, mode=[0])}"
f"but got rank(src, mode=[0]) = {rank(src_list[0], mode=[0])} and rank(dst, mode=[0]) = {rank(dst_list[0], mode=[0])}"
)
value = atom._unpack(loc=loc, ip=ip, **kwargs)
if isinstance(pred, Tensor):
pred = pred.value
return _cute_ir.copy_atom_call(
value, src.value, dst.value, pred=pred, loc=loc, ip=ip
)
src_vals = [t.value for t in src_list]
dst_vals = [t.value for t in dst_list]
return _cute_ir.copy_atom_call(value, src_vals, dst_vals, pred=pred, loc=loc, ip=ip)
@dsl_user_op

3
python/CuTeDSL/cutlass/cute/experimental/README.md
View File

@@ -51,4 +51,7 @@
- `get_cta_v_map_ab` — Compute CTA-V map for A/B operands
- `get_cta_v_map_c` — Compute CTA-V map for C operand
- `make_tmem_layout_acc` — Derive TMEM accumulator buffer layout from a tiled MMA
- `make_tmem_layout_a` — Derive TMEM A-operand buffer layout from a tiled MMA
- `make_t2r_rmem_layout` — Derive per-thread RMEM buffer layout for the T2R epilogue copy

85
python/CuTeDSL/cutlass/cute/experimental/utils.py
View File

@@ -70,10 +70,93 @@ def get_cta_v_map_c(
:param gmem_tensor: Global-memory tensor being stored/loaded by TMA.
:type gmem_tensor: cute.Tensor
:param epi_tile: Epilogue tile layout describing the CTA’s output tile shape.
:param epi_tile: Epilogue tile layout describing the CTA's output tile shape.
:type epi_tile: cute.Layout
:returns: A layout suitable to pass as `cta_v_map=...` to `tma_store` / `tma_load`.
:rtype: cute.Layout
"""
ident = cute.core.make_identity_layout(gmem_tensor.shape, loc=loc, ip=ip)
return cute.core.composition(ident, epi_tile, loc=loc, ip=ip)
def make_tmem_layout_acc(
tiled_mma,
mnk_tiler,
acc_stage,
*,
loc=None,
ip=None,
):
"""Return TMEM accumulator buffer layout for a tiled MMA.
This is a small helper around ``tiled_mma.make_fragment_C(...).layout`` to
keep example code fragment-free at the call site.
:param tiled_mma: The MMA tiler (``cute.TiledMma``).
:type tiled_mma: cute.TiledMma
:param mnk_tiler: Full MNK tiler; only the MN components are used for C.
:type mnk_tiler: tuple
:param acc_stage: Accumulator pipeline stages.
:param loc: Optional location for DSL ops.
:param ip: Optional insertion point for DSL ops.
:return: Layout for the accumulator TMEM buffer.
:rtype: cute.Layout
"""
acc_shape = tiled_mma.partition_shape_C(mnk_tiler[:2], loc=loc, ip=ip)
acc_shape_staged = cute.append(acc_shape, acc_stage, loc=loc, ip=ip)
return tiled_mma.make_fragment_C(acc_shape_staged, loc=loc, ip=ip).layout
def make_tmem_layout_a(
tiled_mma,
mk_tiler,
stage,
*,
loc=None,
ip=None,
):
"""Return TMEM A operand buffer layout for a tiled MMA.
:param tiled_mma: The MMA tiler (``cute.TiledMma``).
:type tiled_mma: cute.TiledMma
:param mk_tiler: MK tiler used to shape the A operand.
:type mk_tiler: tuple
:param stage: Pipeline stages for the A operand buffer.
:param loc: Optional location for DSL ops.
:param ip: Optional insertion point for DSL ops.
:return: Layout for the A operand TMEM buffer.
:rtype: cute.Layout
"""
a_shape = tiled_mma.partition_shape_A(mk_tiler, loc=loc, ip=ip)
a_shape_staged = cute.append(a_shape, stage, loc=loc, ip=ip)
return tiled_mma.make_fragment_A(a_shape_staged, loc=loc, ip=ip).layout
def make_t2r_rmem_layout(
tiled_copy_t2r,
gC_mnl_epi,
tidx,
*,
loc=None,
ip=None,
):
"""Return RMEM buffer layout for the T2R epilogue destination.
Computes the per-thread RMEM buffer layout produced by a TMEM->RMEM copy
for a single epilogue iteration.
:param tiled_copy_t2r: The TMEM->RMEM tiled copy op (``cute.TiledCopy``).
:type tiled_copy_t2r: cute.TiledCopy
:param gC_mnl_epi: Global C tensor partitioned by epilogue tile.
:type gC_mnl_epi: cute.Tensor
:param tidx: Thread index for the copy slice.
:param loc: Optional location for DSL ops.
:param ip: Optional insertion point for DSL ops.
:return: Layout for the RMEM buffer.
:rtype: cute.Layout
"""
thr_copy_t2r = tiled_copy_t2r.get_slice(tidx)
tTR_gC = thr_copy_t2r.partition_D(gC_mnl_epi, loc=loc, ip=ip)
return cute.make_fragment_like(
tTR_gC[(None, None, None, 0, 0)].layout, loc=loc, ip=ip
)

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

@@ -19,6 +19,7 @@ __all__ = [
# copy.py
#
"Repetition",
"TmemLoadRedOp",
"Pack",
"Unpack",
"Ld16x64bOp",

107
python/CuTeDSL/cutlass/cute/nvgpu/tcgen05/copy.py
View File

@@ -26,6 +26,22 @@ from ...typing import Numeric
from .mma import CtaGroup
class TmemLoadRedOp(enum.Enum):
"""
An enumeration for the possible reduce operations for TMEM load operations.
"""
MAX = _cute_nvgpu_ir.TmemLoadRedOp.max
MAXABS = _cute_nvgpu_ir.TmemLoadRedOp.maxabs
MIN = _cute_nvgpu_ir.TmemLoadRedOp.min
MINABS = _cute_nvgpu_ir.TmemLoadRedOp.minabs
def __str__(self) -> str:
return f"{self.__class__.__name__}.{self.name}"
def __repr__(self) -> str:
return f"<{self.__class__.__name__}.{self.name}>"
class Repetition(enum.Enum):
"""
An enumeration for the number of repetitions of a given TMEM copy within the instruction.
@@ -390,6 +406,97 @@ class Ld32x32bTrait(Trait):
pass
@dataclass(frozen=True)
class LdRed16x32bx2Op(_LdBase):
"""
16x32bx2 TMEM load Reduce Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
This Operation corresponds to the ``.red`` and ``.16x32bx2`` qualifiers.
"""
redOp: TmemLoadRedOp = TmemLoadRedOp.MAX
nan: bool = False
half_split_off: int = 0
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "LdRed16x32bx2Trait":
"""
Create a trait object for the 16x32bx2 TMEM load Reduce operation.
:param copy_internal_type: The data type for the copy operation
:type copy_internal_type: Type[Numeric]
:param loc: MLIR location information for debugging, defaults to None
:type loc: optional
:param ip: MLIR insertion point for code generation, defaults to None
:type ip: optional
:param kwargs: Additional keyword arguments
:type kwargs: dict
:return: A trait object for this load operation
:rtype: LdRed16x32bx2Trait
"""
ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get(
copy_internal_type.mlir_type,
16,
32,
self.repeat.value,
self.redOp.value,
ir.UnitAttr.get() if self.nan else None,
ir.IntegerAttr.get(ir.IntegerType.get_signless(32), self.half_split_off),
)
return LdRed16x32bx2Trait(make_atom(ty, loc=loc, ip=ip))
class LdRed16x32bx2Trait(Trait):
pass
@dataclass(frozen=True)
class LdRed32x32bOp(_LdBase):
"""
32x32b TMEM load Reduce Operation.
See the `PTX documentation <https://docs.nvidia.com/cuda/parallel-thread-execution/#tcgen05-instructions-tcgen05-ld>`__.
This Operation corresponds to the ``red`` and ``.32x32`` qualifiers.
"""
redOp: TmemLoadRedOp = TmemLoadRedOp.MAX
nan: bool = False
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "LdRed32x32bTrait":
"""
Create a trait object for the 32x32b TMEM load Reduce operation.
:param copy_internal_type: The data type for the copy operation
:type copy_internal_type: Type[Numeric]
:param loc: MLIR location information for debugging, defaults to None
:type loc: optional
:param ip: MLIR insertion point for code generation, defaults to None
:type ip: optional
:param kwargs: Additional keyword arguments
:type kwargs: dict
:return: A trait object for this load operation
:rtype: LdRed32x32bTrait
"""
ty = _cute_nvgpu_ir.CopyAtomSM10xTmemLoadRedType.get(
copy_internal_type.mlir_type,
32,
32,
self.repeat.value,
self.redOp.value,
ir.UnitAttr.get() if self.nan else None,
None,
)
return LdRed32x32bTrait(make_atom(ty, loc=loc, ip=ip))
class LdRed32x32bTrait(Trait):
pass
@dataclass(frozen=True)
class _StBase(CopyOp):
"""

15
python/CuTeDSL/cutlass/cute/nvgpu/warp/copy.py
View File

@@ -103,15 +103,20 @@ class LdMatrix8x16x8bOp(BaseOp):
self,
"expects the 'num_matrices' Op parameter to be one of [1,2,4]",
)
if self.unpack_bits not in [4, 6]:
raise OpError(self, "Op unpack bits must be 4 or 6")
if self.unpack_bits not in [None, 4, 6]:
raise OpError(self, "Op unpack bits must be 4 or 6 or None")
def _make_trait(
self, copy_internal_type: Type[Numeric], *, loc=None, ip=None, **kwargs
) -> "LdMatrix8x16x8bTrait":
mode = _pack_shape((8, 16), loc=loc, ip=ip)
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8
if self.unpack_bits == 6:
# LdMatrix8x16x8b without unpacking doesn't exist
# but is equivalent to LdMatrix8x8x16b
mode_n = 8 if self.unpack_bits is None else 16
mode = _pack_shape((8, mode_n), loc=loc, ip=ip)
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u16
if self.unpack_bits == 4:
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u4x16p64to8
elif self.unpack_bits == 6:
sz_pattern = _cute_nvgpu_ir.LdsmSzPattern.u6x16p32to8
ty = _cute_nvgpu_ir.CopyAtomLdsmType.get(
copy_internal_type.mlir_type,

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

@@ -1,4 +1,4 @@
# SPDX-FileCopyrightText: Copyright (c) 2025 - 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# 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
@@ -141,6 +141,11 @@ class _Tensor(Tensor):
# If tensor is already a DLPack object, use it directly
if hasattr(tensor, "__dlpack_device__") and not hasattr(tensor, "__dlpack__"):
self._dlpack_data = tensor.__dlpack_device__()
elif enable_tvm_ffi:
import tvm_ffi
self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor)
self._dlpack_data = self._tvm_ffi_tensor.__dlpack__()
else:
try:
# we expect no stream sync. Because torch has different default behavior
@@ -149,11 +154,6 @@ class _Tensor(Tensor):
self._dlpack_data = tensor.__dlpack__(stream=-1)
except Exception:
self._dlpack_data = tensor.__dlpack__()
if enable_tvm_ffi:
import tvm_ffi
self._tvm_ffi_tensor = tvm_ffi.from_dlpack(tensor)
self._dlpack_data = self._tvm_ffi_tensor.__dlpack__()
self._dltensor_wrapper = None
self._assumed_align = assumed_align