This is an automated email from the ASF dual-hosted git repository.
spectrometerHBH pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new f50a32f232 refactor(op)!: consolidate TIRx helper ownership (#20668)
f50a32f232 is described below
commit f50a32f232ae7cff6ca8b2f6545150a0da3ae5e7
Author: Bohan Hou <[email protected]>
AuthorDate: Sun Oct 11 17:49:56 2026 -0400
refactor(op)!: consolidate TIRx helper ownership (#20668)
Keep CUDA language helpers in the backend namespace and remove unused
pools, roles, descriptors, and schedulers.
Move effective layout inference onto TensorType and migrate CUDA and
Trainium callers. Correct the Trainium sub-layout factory call.
Move CUDA modifier constants into their sole consumer and remove the
legacy operator package and unused tables.
BREAKING CHANGE: Remove tvm.tirx.lang, tvm.tirx.operator, the tile
layout_utils helper, and T.SMEMPool/T.TMEMPool. Import supported language
helpers from tvm.backend.cuda.lang and query tensor.ty.effective_layout().
---
docs/reference/api/python/script/script.rst | 2 +-
docs/tirx/api/cuda.rst | 33 +--
docs/tirx/api/layout.rst | 10 +-
docs/tirx/native_basics/cuda/tensors.rst | 34 +--
docs/tirx/tile_primitives/copy_async/dsmem.rst | 4 +-
python/tvm/backend/cuda/lang/__init__.py | 14 --
python/tvm/backend/cuda/lang/alloc_pool.py | 240 +--------------------
python/tvm/backend/cuda/lang/smem_desc.py | 56 -----
python/tvm/backend/cuda/lang/tile_scheduler.py | 191 +---------------
python/tvm/backend/cuda/lang/warp_role.py | 150 -------------
python/tvm/backend/cuda/layout.py | 5 +-
python/tvm/backend/cuda/op.py | 28 +--
python/tvm/backend/cuda/tile/common.py | 7 +-
python/tvm/backend/cuda/tile/copy/ld_stmatrix.py | 9 +-
python/tvm/backend/cuda/tile/copy/vec_auto_reg.py | 17 +-
python/tvm/backend/cuda/tile/copy_async/dsmem.py | 9 +-
python/tvm/backend/cuda/tile/copy_async/ldgsts.py | 7 +-
.../tvm/backend/cuda/tile/copy_async/tcgen05_cp.py | 9 +-
.../backend/cuda/tile/copy_async/tcgen05_ldst.py | 25 ++-
python/tvm/backend/cuda/tile/copy_async/tma.py | 25 +--
.../tvm/backend/cuda/tile/elementwise/_common.py | 7 +-
python/tvm/backend/cuda/tile/elementwise/reg.py | 9 +-
python/tvm/backend/cuda/tile/elementwise/smem.py | 3 +-
python/tvm/backend/cuda/tile/gemm/mma_m16n8k_.py | 7 +-
python/tvm/backend/cuda/tile/gemm_async/tcgen05.py | 23 +-
python/tvm/backend/cuda/tile/gemm_utils.py | 6 +-
python/tvm/backend/trn/tile/binary/default.py | 3 +-
python/tvm/backend/trn/tile/binary/utils.py | 8 +-
.../backend/trn/tile/compose_op/binary_chain.py | 3 +-
.../backend/trn/tile/compose_op/binary_reduce.py | 3 +-
.../backend/trn/tile/compose_op/unary_reduce.py | 3 +-
python/tvm/backend/trn/tile/copy/default.py | 15 +-
python/tvm/backend/trn/tile/gemm/default.py | 23 +-
.../tvm/backend/trn/tile/instruction_generator.py | 10 +-
python/tvm/backend/trn/tile/private_alloc.py | 5 +-
python/tvm/backend/trn/tile/reduction/utils.py | 13 +-
python/tvm/backend/trn/tile/select/default.py | 9 +-
python/tvm/backend/trn/tile/unary/utils.py | 9 +-
python/tvm/tirx/lang/__init__.py | 16 --
python/tvm/tirx/lang/alloc_pool.py | 8 -
python/tvm/tirx/lang/pipeline.py | 8 -
python/tvm/tirx/lang/smem_desc.py | 8 -
python/tvm/tirx/lang/tile_scheduler.py | 8 -
python/tvm/tirx/lang/warp_role.py | 8 -
python/tvm/tirx/operator/__init__.py | 18 --
python/tvm/tirx/operator/intrinsics/_common.py | 59 -----
python/tvm/tirx/script/ir_builder/__init__.py | 15 --
python/tvm/tirx/tensor.py | 27 +++
python/tvm/tirx/tile/layout_utils.py | 47 ----
.../python/tirx/codegen/test_codegen_blackwell.py | 25 +--
tests/python/tirx/script/test_tirx_script_gpu.py | 64 ------
tests/python/tirx/test_layout.py | 4 +-
tests/python/tirx/test_tile_namespace.py | 27 ++-
.../python/tirx/tile/cuda/copy_async/test_dsmem.py | 5 +-
.../tirx/tile/cuda/copy_async/test_tcgen05_ldst.py | 8 -
tests/python/tirx/tile/trn/test_copy_trn.py | 12 ++
56 files changed, 236 insertions(+), 1165 deletions(-)
diff --git a/docs/reference/api/python/script/script.rst
b/docs/reference/api/python/script/script.rst
index c5061f65a8..0241752e98 100644
--- a/docs/reference/api/python/script/script.rst
+++ b/docs/reference/api/python/script/script.rst
@@ -53,7 +53,7 @@ Operator attributes re-exported by this namespace are
documented in :mod:`tvm.ti
.. automodule:: tvm.script.tirx
:members:
- :exclude-members: Range, meta_var, Var, Call, CommReducer, Reduce,
SMEMPool, TMEMPool, FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv,
FloorMod, LShift, RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min,
Max, EQ, NE, LT, LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle,
CallEffectKind, IterVar, ComposeLayout, DtypeConstructor, ExecScope, Iter,
Layout, LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr,
TensorMapEncodeTiledAttr, function, [...]
+ :exclude-members: Range, meta_var, Var, Call, CommReducer, Reduce,
FloatImm, IntImm, Cast, Add, Sub, Mul, Div, Mod, FloorDiv, FloorMod, LShift,
RShift, BitwiseAnd, BitwiseOr, BitwiseXor, BitwiseNot, Min, Max, EQ, NE, LT,
LE, GT, GE, And, Or, Not, Select, Ramp, Broadcast, Shuffle, CallEffectKind,
IterVar, ComposeLayout, DtypeConstructor, ExecScope, Iter, Layout,
LetAnnotation, LocalVectorAnnotation, TileLayout, CallFFIKernelAttr,
TensorMapEncodeTiledAttr, function, jit, inline, macro
.. autofunction:: tvm.script.tirx.function
diff --git a/docs/tirx/api/cuda.rst b/docs/tirx/api/cuda.rst
index 6362af7397..1c89493d27 100644
--- a/docs/tirx/api/cuda.rst
+++ b/docs/tirx/api/cuda.rst
@@ -88,8 +88,7 @@ CUDA language utilities
-----------------------
``tvm.backend.cuda.lang`` provides reusable objects for hand-written CUDA TIRx
-kernels, including schedulers, pipelines, roles, barriers, descriptors, and
-memory pools.
+kernels, including schedulers, pipelines, barriers, and shared-memory pools.
.. autoclass:: tvm.backend.cuda.lang.BaseTileScheduler
:members:
@@ -110,12 +109,6 @@ memory pools.
``init`` starts from a CTA index, ``next_tile`` advances by ``num_ctas``,
and ``valid`` reports whether work remains.
-.. autoclass:: tvm.backend.cuda.lang.GroupMajor3D
- :members:
-
-.. autoclass:: tvm.backend.cuda.lang.IndexedTripleTileScheduler
- :members:
-
.. autoclass:: tvm.backend.cuda.lang.MBarrier
:members:
@@ -125,39 +118,15 @@ memory pools.
.. autoclass:: tvm.backend.cuda.lang.PipelineState
:members:
-.. autoclass:: tvm.backend.cuda.lang.RankAwareGroupMajorTileScheduler
- :members:
-
.. autoclass:: tvm.backend.cuda.lang.SMEMPool
:members:
-.. autoclass:: tvm.backend.cuda.lang.SmemDescriptor
- :members:
-
.. autoclass:: tvm.backend.cuda.lang.TCGen05Bar
:members:
.. autoclass:: tvm.backend.cuda.lang.TMABar
:members:
-.. autoclass:: tvm.backend.cuda.lang.TMEMPool
- :members:
- :exclude-members: alloc
-
-.. py:method:: TMEMPool.alloc(shape, dtype="float32", *, layout=None,
cols=None)
-
- Allocate a tensor-memory tensor from the pool. ``layout`` supplies an
- explicit ``TileLayout``; without it, two-dimensional shapes use the default
- dense layout. ``cols`` overrides the inferred tensor-memory column count.
- Datapath-specific accumulators can instead use
- ``alloc_tcgen05_mma_D(...)`` / ``alloc_tcgen05_mma_C(...)``.
-
-.. autoclass:: tvm.backend.cuda.lang.WarpRole
- :members:
-
-.. autoclass:: tvm.backend.cuda.lang.WarpgroupRole
- :members:
-
IKET profiling
--------------
diff --git a/docs/tirx/api/layout.rst b/docs/tirx/api/layout.rst
index 9b81ed363a..f078a0358d 100644
--- a/docs/tirx/api/layout.rst
+++ b/docs/tirx/api/layout.rst
@@ -59,10 +59,12 @@ accepted by generic tensor constructors.
Omitting ``layout`` or passing ``layout=None`` leaves every storage scope
without a layout. ``layout="default"`` explicitly requests a compact
``TileLayout``. Ordinary views preserve shape, strides and element offsets
-without creating a layout. A tile instruction can infer a temporary layout
-for a compact global tensor, using the full tensor shape before slicing its
-operand region. Noncompact global tensors and other storage scopes require
-an explicit layout when the instruction uses one.
+without creating a layout. ``TensorType.effective_layout()`` returns an
+explicit layout or infers a temporary compact layout for a global tensor.
+Tile instructions query it through ``tensor.ty.effective_layout()``, using
+the full tensor shape before slicing an operand region. The query does not
+change the tensor type or its ``layout=None`` field. Noncompact global tensors
+and other storage scopes require an explicit layout when the instruction uses
one.
.. automodule:: tvm.tirx.layout
:members:
diff --git a/docs/tirx/native_basics/cuda/tensors.rst
b/docs/tirx/native_basics/cuda/tensors.rst
index b980b4bd07..2de8e146c0 100644
--- a/docs/tirx/native_basics/cuda/tensors.rst
+++ b/docs/tirx/native_basics/cuda/tensors.rst
@@ -222,22 +222,23 @@ with views decl'd at offsets inside it.
Pool sugar
~~~~~~~~~~
-``Tx.SMEMPool`` automates that arena bookkeeping — it bump-allocates the
offsets so
-you don't ``decl`` views by hand. Beyond ``alloc`` / ``commit``, it offers
+``tvm.backend.cuda.lang.SMEMPool`` automates that arena bookkeeping — it
+bump-allocates the offsets so you don't ``decl`` views by hand. Beyond
+``alloc`` / ``commit``, it offers
per-tensor ``align=``, an ``alloc_tcgen05_mma_AB`` helper that builds an
MMA-compatible
swizzle layout for you, and ``move_base_to`` to rewind the cursor and reuse
space:
.. code-block:: python
- pool = Tx.SMEMPool() # bump allocator over
shared.dyn
+ from tvm.backend.cuda.lang import SMEMPool
+
+ pool = SMEMPool() # bump allocator over
shared.dyn
As = pool.alloc((BM, BK), "float16", align=128) # carve a tile
Bs = pool.alloc((BK, BN), "float16", align=128)
Cs = pool.alloc_tcgen05_mma_AB((BM, BN), "float16") # MMA-compatible,
swizzle inferred
pool.commit() # finalize the pool's size
# pool.move_base_to(offset) rewinds the cursor to reuse space
-The TMEM pool (`Tensor memory`_, below) is layered on top of an ``SMEMPool``.
-
Per-thread local storage
------------------------
@@ -387,28 +388,7 @@ tensor as a view at a column offset, and one warp frees it
at the end:
addr[0], Tx.uint32(512))
You add any column offsets to ``addr[0]`` and manage the ``tmem_layout`` (the
-matching datapath A–G layout) yourself. This is the sequence the pool below
emits.
-
-Pool
-~~~~
-
-``Tx.TMEMPool`` wraps the warp-uniform alloc/dealloc and column
bump-allocation.
-Pass an explicit layout to ``alloc``, or use its operand-specific allocation
-helpers when the MMA datapath determines the layout:
-
-.. code-block:: python
-
- tmem_addr = pool.alloc((1,), "uint32") # pool = the kernel's smem
pool
- tmem_pool = Tx.TMEMPool(pool, total_cols=512, cta_group=cta_group,
- tmem_addr=tmem_addr)
- # Choose the layout required by the instruction that consumes the tensor:
- acc = tmem_pool.alloc((CTA_M, 512), "float32") # Layout D when CTA_M=128
- # Layout B: PTX M=128, cta_group=2, 64 logical rows in this CTA.
- # acc = tmem_pool.alloc_tcgen05_mma_D(
- # (64, N), "float32", M=128, cta_group=2)
- tmem_pool.commit() # emits tcgen05.alloc
(one warp)
- # ... use acc ...
- tmem_pool.dealloc() # emits tcgen05.dealloc
(one warp)
+matching datapath A–G layout) yourself.
See the :doc:`../../tile_primitives` walkthroughs for full examples.
diff --git a/docs/tirx/tile_primitives/copy_async/dsmem.rst
b/docs/tirx/tile_primitives/copy_async/dsmem.rst
index e37cefedd8..73874f3179 100644
--- a/docs/tirx/tile_primitives/copy_async/dsmem.rst
+++ b/docs/tirx/tile_primitives/copy_async/dsmem.rst
@@ -62,7 +62,7 @@ mbarrier and writes the result out (from ``test_dsmem.py``):
.. code-block:: python
- from tvm.backend.cuda.lang import MBarrier
+ from tvm.backend.cuda.lang import MBarrier, SMEMPool
shape, dtype, CLUSTER_N = (128, 64), "float16", 2
src_layout = dst_layout = TileLayout(S[128, 64])
@@ -77,7 +77,7 @@ mbarrier and writes the result out (from ``test_dsmem.py``):
cbx = Tx.cuda.cluster_cta_id("x")
_index = Tx.cuda.block_idx("x")
tid = Tx.cuda.thread_idx("x")
- pool = Tx.SMEMPool()
+ pool = SMEMPool()
src_raw = pool.alloc([8192], dtype, align=128)
src_smem = Tx.decl_tensor(
src_raw.data,
diff --git a/python/tvm/backend/cuda/lang/__init__.py
b/python/tvm/backend/cuda/lang/__init__.py
index eee5c6ba98..dac3923fb9 100644
--- a/python/tvm/backend/cuda/lang/__init__.py
+++ b/python/tvm/backend/cuda/lang/__init__.py
@@ -23,19 +23,12 @@ __all__ = [
"ClusterPersistentScheduler2D",
"FlashAttentionLPTScheduler",
"FlashAttentionLinearScheduler",
- "GroupMajor3D",
- "IndexedTripleTileScheduler",
"MBarrier",
"Pipeline",
"PipelineState",
- "RankAwareGroupMajorTileScheduler",
"SMEMPool",
- "SmemDescriptor",
"TCGen05Bar",
"TMABar",
- "TMEMPool",
- "WarpRole",
- "WarpgroupRole",
]
_HELPER_MODULES = {
@@ -46,16 +39,9 @@ _HELPER_MODULES = {
"ClusterPersistentScheduler2D": ".tile_scheduler",
"FlashAttentionLPTScheduler": ".tile_scheduler",
"FlashAttentionLinearScheduler": ".tile_scheduler",
- "GroupMajor3D": ".tile_scheduler",
- "IndexedTripleTileScheduler": ".tile_scheduler",
- "RankAwareGroupMajorTileScheduler": ".tile_scheduler",
"SMEMPool": ".alloc_pool",
"TCGen05Bar": ".pipeline",
"TMABar": ".pipeline",
- "TMEMPool": ".alloc_pool",
- "WarpRole": ".warp_role",
- "WarpgroupRole": ".warp_role",
- "SmemDescriptor": ".smem_desc",
}
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py
b/python/tvm/backend/cuda/lang/alloc_pool.py
index d613f6d715..1e9b3c24c6 100644
--- a/python/tvm/backend/cuda/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""SMEM and TMEM bump-allocator pools for TIRX kernels."""
+"""Shared-memory bump allocator for CUDA TIRx kernels."""
from __future__ import annotations
@@ -22,8 +22,6 @@ import functools
import operator
from tvm import DataType
-from tvm.backend.cuda.axis import TCol, TLane
-from tvm.tirx.layout import S, TileLayout
# ---------------------------------------------------------------------------
# ir_builder helpers — imported lazily to avoid circular deps at module level
@@ -48,23 +46,6 @@ def _get_ir():
_POOL_UNSET = object()
-def _default_tmem_layout(rows, cols):
- # Keep the default M=128 layout structurally identical to Layout D. The
- # tcgen05 copy dispatch recognizes default-allocated accumulators by that
- # structure, so single-source the named datapath instead of duplicating it.
- if rows == 128:
- from tvm.backend.cuda.layout import tmem_datapath_layout
-
- return tmem_datapath_layout("D", rows, cols)
- return TileLayout(S[(rows, cols) : (1 @ TLane, 1 @ TCol)])
-
-
-def _emit_stmt(expr):
- from tvm.tirx.script.ir_builder.parser_protocol import evaluate
-
- evaluate(expr)
-
-
def _shape_product(shape):
return functools.reduce(operator.mul, shape, 1)
@@ -181,225 +162,6 @@ def _meta_class(cls):
return _get_ir().meta_class(cls)
-# ---------------------------------------------------------------------------
-# TMEMPool
-# ---------------------------------------------------------------------------
-
-
-@_meta_class
-class TMEMPool:
- """Bump allocator over TMEM columns."""
-
- def __init__(
- self,
- pool,
- total_cols=512,
- *,
- cta_group=1,
- alloc_warp=0,
- dealloc_warp=None,
- tmem_addr=None,
- sync_after_alloc=True,
- ):
- # tcgen05 alloc/dealloc are warp-uniform PTX instructions: every lane
- # in the chosen warp must participate, and exactly one warp in the
- # CTA must execute them. The pool emits its own
- # ``if warp_id() == target_warp: tcgen05.alloc(...)``
- # guard, using the cta->warp scope id ``T.cuda.warp_id()``.
- # NOTE: synccheck currently false-deadlocks on kernels that declare a
- # second warp-scope id (cpusim binds only one warp var); the generated
- # CUDA is equivalent to ``thread_rank() // 32 == target_warp``.
- self.pool = pool
- self.total_cols = total_cols
- self.cta_group = cta_group
- self.alloc_warp = alloc_warp
- self.dealloc_warp = alloc_warp if dealloc_warp is None else
dealloc_warp
- self.sync_after_alloc = sync_after_alloc
- self.offset = 0
- self.max_offset = 0
- self._committed = False
- self._deallocated = False
- self._addr_buf = pool.alloc([1], "uint32", align=4) if tmem_addr is
None else tmem_addr
-
- def _addr_slot(self):
- try:
- return self._addr_buf[0]
- except TypeError:
- return self._addr_buf
-
- @property
- def addr(self):
- return self._addr_slot()
-
- def _emit_warp_guard(self, target_warp, emit):
- from tvm.script import tirx as T
-
- warp_id = T.cuda.warp_id()
- with T.if_(warp_id == target_warp):
- with T.then_():
- emit()
-
- def _resolve_cols(self, shape, dtype, cols, layout=None):
- if cols is not None:
- return cols
- bits = DataType(dtype).bits
- if layout is not None:
- # span("TCol") is in *element* (buffer dtype) units; one TMEM cell
- # holds 32 bits regardless of the element type.
- tcol_elems = int(layout.span("TCol"))
- tcol_bits = tcol_elems * bits
- assert tcol_bits % 32 == 0, (
- f"layout TCol span={tcol_elems} elems x {bits}b is not 32-bit
aligned"
- )
- return tcol_bits // 32
- assert len(shape) == 2, "TMEMPool.alloc() requires cols= for non-2D
TMEM buffers"
- total_bits = _shape_product(shape) * bits
- rows = shape[0]
- assert total_bits % (32 * rows) == 0, (
- f"Cannot infer TMEM columns from shape={shape}, dtype={dtype!r}; "
- "please pass cols= explicitly"
- )
- return total_bits // (32 * rows)
-
- def alloc(self, shape, dtype="float32", *, layout=None, cols=None):
- """Allocate a TMEM buffer.
-
- Parameters
- ----------
- shape, dtype, cols
- Standard buffer shape / dtype / column count.
- layout
- Explicit ``TileLayout``. Mutually exclusive with ``datapath``.
- datapath : str | None
- Optional tcgen05 datapath letter (``"D"`` for M=128 full datapath,
- ``"F"`` for M=64 non-``.ws`` scattered, or ``"B"`` for per-CTA
- M=64 ``.cta_group::2`` "2x2"). When provided, the buffer's
- layout is derived from ``tmem_datapath_layout(datapath, *shape)``
- so the row index reflects the *physical* TMEM lane occupation
- (PTX ISA §9.7.16.10.5). The downstream ``.16x*b`` / ``.32x32b``
- dispatches structurally check this layout to catch mismatched
- atoms (e.g. a ``.16x*b`` M=128 read against a Layout F buffer).
- Defaults to ``None``, which means Layout D's identity row→lane
- mapping — keep this for shape ``(128, X)`` buffers that hold
- an M=128 MMA accumulator.
- """
- ir = _get_ir()
- cols = self._resolve_cols(shape, dtype, cols, layout)
- col_start = self.offset
- col_end = col_start + cols
- assert col_end <= self.total_cols, f"TMEM overflow: {col_end} >
{self.total_cols}"
- if layout is None:
- assert len(shape) == 2, "TMEMPool.alloc() requires layout= for
non-2D TMEM buffers"
- layout = _default_tmem_layout(shape[0], shape[1])
- import tvm # pylint: disable=import-outside-toplevel
- from tvm.tirx.script.ir_builder.parser_protocol import (
- bind_, # pylint: disable=import-outside-toplevel
- )
-
- res = bind_(
- tvm.ir.Call(
- "tirx.cuda.decl_tmem",
- [col_start],
- ty_args=[ir.Tensor(shape, dtype, scope="tmem", layout=layout)],
- )
- )
- self.offset = col_end
- self.max_offset = max(self.max_offset, self.offset)
- return res
-
- def alloc_sf(self, shape, dtype, *, sf_per_mma, sf_reuse=1):
- """Allocate a tcgen05 block-scaled SF TMEM buffer with an inferred
layout.
-
- ``shape`` last two dims are ``(rows, SF_K * sf_reuse)`` (the last dim
is
- what gemm dispatch iterates over). When ``shape`` has 3 dims, the first
- is treated as a pipe-depth outer.
- """
- from tvm.backend.cuda.tile.gemm_async.tcgen05 import sf_tmem_layout
-
- if len(shape) == 2:
- pipe_depth, rows, last = None, shape[0], shape[1]
- elif len(shape) == 3:
- pipe_depth, rows, last = shape[0], shape[1], shape[2]
- else:
- raise ValueError(
- f"alloc_sf expects 2D (rows, SF_K*sf_reuse) or 3D "
- f"(pipe_depth, rows, SF_K*sf_reuse); got shape={shape}"
- )
- assert last % sf_reuse == 0, (
- f"alloc_sf: shape last dim {last} must be divisible by
sf_reuse={sf_reuse}"
- )
- SF_K = last // sf_reuse
- layout = sf_tmem_layout(
- rows=rows, SF_K=SF_K, sf_per_mma=sf_per_mma, sf_reuse=sf_reuse,
pipe_depth=pipe_depth
- )
- return self.alloc(shape, dtype, layout=layout)
-
- def alloc_tcgen05_mma_A(
- self, shape, dtype="bfloat16", *, M, cta_group, ws=False,
sparse=False, cols=None
- ):
- """Allocate a TMEM A operand (A-in-TMEM); layout resolved from MMA
- instruction params. ``M`` is the PTX tcgen05.mma instruction M
(256/128/
- 64), NOT per-CTA rows. See ``localdoc/claude_plan.txt`` S3a."""
- from tvm.backend.cuda.layout import tmem_mma_operand_layout
-
- layout = tmem_mma_operand_layout(
- "A", shape, dtype, M=M, cta_group=cta_group, ws=ws, sparse=sparse
- )
- return self.alloc(shape, dtype, layout=layout, cols=cols)
-
- def alloc_tcgen05_mma_D(
- self, shape, dtype="float32", *, M, cta_group, ws=False, sparse=False,
group=None, cols=None
- ):
- """Allocate a TMEM D/C accumulator operand. ``M`` = PTX instruction M.
- ``group=(s,2,n)`` gives a flat (m,N) buffer a grouped column layout."""
- from tvm.backend.cuda.layout import tmem_mma_operand_layout
-
- layout = tmem_mma_operand_layout(
- "D", shape, dtype, M=M, cta_group=cta_group, ws=ws, sparse=sparse,
group=group
- )
- return self.alloc(shape, dtype, layout=layout, cols=cols)
-
- alloc_tcgen05_mma_C = alloc_tcgen05_mma_D
-
- def move_base_to(self, col):
- self.offset = col
- self.max_offset = max(self.max_offset, self.offset)
-
- def commit(self):
- assert not self._committed, "TMEMPool.commit() can only be called once"
- from tvm.script import tirx as T
-
- def emit_alloc():
- _emit_stmt(
-
T.ptx[f"tcgen05.alloc.cta_group::{self.cta_group}.sync.aligned.shared::cta.b32"](
- T.address_of(self.addr), T.uint32(self.total_cols)
- )
- )
- if self.sync_after_alloc:
- _emit_stmt(T.cuda.warp_sync())
-
- self._emit_warp_guard(self.alloc_warp, emit_alloc)
- self._committed = True
-
- def dealloc(self):
- assert self._committed, "TMEMPool.dealloc() called before commit()"
- assert not self._deallocated, "TMEMPool.dealloc() can only be called
once"
- self._deallocated = True
- from tvm.script import tirx as T
-
- def emit_dealloc():
- _emit_stmt(
-
T.ptx[f"tcgen05.relinquish_alloc_permit.cta_group::{self.cta_group}.sync.aligned"]()
- )
- _emit_stmt(
-
T.ptx[f"tcgen05.dealloc.cta_group::{self.cta_group}.sync.aligned.b32"](
- self.addr, T.uint32(self.total_cols)
- )
- )
-
- self._emit_warp_guard(self.dealloc_warp, emit_dealloc)
-
-
# ---------------------------------------------------------------------------
# SMEMPool
# ---------------------------------------------------------------------------
diff --git a/python/tvm/backend/cuda/lang/smem_desc.py
b/python/tvm/backend/cuda/lang/smem_desc.py
deleted file mode 100644
index fe1bea3464..0000000000
--- a/python/tvm/backend/cuda/lang/smem_desc.py
+++ /dev/null
@@ -1,56 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""SMEM matrix descriptor helper for tcgen05 / wgmma."""
-
-from tvm.backend.cuda.tile.common import smem_desc_add_16B_offset
-from tvm.script import tirx as T
-
-
[email protected]_class
-class SmemDescriptor:
- """Encoded once via :meth:`init`, reused via :meth:`add_16B_offset`."""
-
- def __init__(self):
- self._buf = T.alloc_local([1], "uint64")
-
- @property
- def desc(self):
- return self._buf[0]
-
- @T.inline
- def init(self, smem_ptr, ldo, sdo, swizzle):
- T.cuda.tcgen05.encode_matrix_descriptor(
- T.address_of(self._buf[0]), smem_ptr, ldo, sdo, swizzle
- )
-
- def add_16B_offset(self, offset):
- return smem_desc_add_16B_offset(self._buf[0], offset)
-
- def make_lo_uniform(self):
- """Broadcast the lower 32 bits to all warp lanes via
``__shfl_sync``."""
- desc_lo = T.alloc_local([1], "uint32")
- desc_hi = T.alloc_local([1], "uint32")
- T.ptx.mov.b64(desc_lo[0], desc_hi[0], self._buf[0])
- T.ptx.shfl_sync.idx.b32(
- desc_lo[0],
- desc_lo[0],
- T.uint32(0),
- T.uint32(0x1F),
- T.uint32(0xFFFFFFFF),
- )
- T.ptx.mov.b64(self._buf[0], desc_lo[0], desc_hi[0])
diff --git a/python/tvm/backend/cuda/lang/tile_scheduler.py
b/python/tvm/backend/cuda/lang/tile_scheduler.py
index 279b4524c2..1521ad790a 100644
--- a/python/tvm/backend/cuda/lang/tile_scheduler.py
+++ b/python/tvm/backend/cuda/lang/tile_scheduler.py
@@ -14,7 +14,7 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-"""Reusable tile scheduler helpers for TIR tests/kernels.
+"""Reusable tile scheduler helpers for CUDA TIRx kernels.
These classes emit TIR via @T.inline. Decorate with @T.meta_class so that
instances are automatically treated as meta values inside @T.function.
@@ -465,195 +465,6 @@ class ClusterPersistentScheduler2D(BaseTileScheduler):
return self.linear_idx < self._TOTAL_TILES
-class GroupMajor3D(BaseTileScheduler):
- """
- 3D grouped-row scheduler (M,N,K) with tail handling on M.
-
- Args
- ----
- prefix: str
- m_tiles: int | T Expr # tiles along M (static or runtime)
- n_tiles: int # tiles along N (static)
- k_tiles: int # tiles along K (static)
- group_rows: int # rows per group along M
- step: int = 1 # default stride for next_tile()
- """
-
- def __init__(
- self, prefix: str, m_tiles, n_tiles: int, k_tiles: int, group_rows:
int, step: int = 1
- ):
- super().__init__(prefix)
- self._step = step
- self.tile_idx = T.local_scalar("int32")
- self.k_idx = T.local_scalar("int32")
-
- # ---- constants / primexprs baked once ----
- self._G = group_rows
- self._N = n_tiles
- self._K = k_tiles
-
- if isinstance(m_tiles, int):
- self._GROUPS = m_tiles // group_rows
- self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
- self._SAFE_FINAL_ROWS = max(self._FINAL_ROWS, 1)
- self._GROUP_SIZE = group_rows * n_tiles * k_tiles
- self._TOTAL = m_tiles * n_tiles * k_tiles
- else:
- self._GROUPS = T.truncdiv(m_tiles, group_rows)
- self._FINAL_ROWS = m_tiles - self._GROUPS * group_rows
- self._SAFE_FINAL_ROWS = T.max(self._FINAL_ROWS, 1)
- self._GROUP_SIZE = self._G * self._N * self._K
- self._TOTAL = m_tiles * n_tiles * k_tiles
-
- # handy composites used in macro
- self._FULL_BOUND = self._GROUPS * self._GROUP_SIZE
- self._HAS_FULL = self._GROUPS > 0
- self._HAS_TAIL = self._FINAL_ROWS > 0
-
- # fmt: off
- @T.inline
- def update_current_m_n_idx(self, linear_idx):
- # full-group formulas
- full_m: T.let = T.floordiv(linear_idx, self._GROUP_SIZE) * self._G +
T.floormod(
- linear_idx, self._G
- )
- full_n: T.let = T.floormod(T.floordiv(linear_idx, self._G), self._N)
- full_k: T.let = T.floordiv(T.floormod(linear_idx, self._GROUP_SIZE),
self._G * self._N)
-
- # tail formulas (relative to FULL_BOUND)
- # Use _SAFE_FINAL_ROWS (max(FINAL_ROWS, 1)) to avoid divide-by-zero
when there is no tail
- rem: T.let = linear_idx - self._FULL_BOUND
- tail_m: T.let = self._GROUPS * self._G + T.floormod(rem,
self._SAFE_FINAL_ROWS)
- tail_n: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS) % self._N
- tail_k: T.let = T.floordiv(rem, self._SAFE_FINAL_ROWS * self._N)
-
- # choose phase
- if self._HAS_FULL & (linear_idx < self._FULL_BOUND):
- self.m_idx = full_m
- self.n_idx = full_n
- self.k_idx = full_k
- elif self._HAS_TAIL:
- self.m_idx = tail_m
- self.n_idx = tail_n
- self.k_idx = tail_k
- else:
- self.m_idx = 0
- self.n_idx = 0
- self.k_idx = 0
-
- @T.inline
- def init(self, linear_init):
- self.linear_idx = linear_init
- self.tile_idx = 0
- self.update_current_m_n_idx(linear_init)
-
- @T.inline
- def next_tile(self):
- self.linear_idx = self.linear_idx + self._step
- self.tile_idx = self.tile_idx + 1
- self.update_current_m_n_idx(self.linear_idx)
-
- @T.inline
- def next_tile_stride(self, stride: int):
- self.linear_idx = self.linear_idx + stride
- self.tile_idx = self.tile_idx + 1
- self.update_current_m_n_idx(self.linear_idx)
- # fmt: on
-
- def valid(self):
- return self.linear_idx < self._TOTAL
-
-
-class RankAwareGroupMajorTileScheduler(BaseTileScheduler):
- """
- Group-major scheduler that applies a rank-aware remapping (remote rows
first).
- Kept as a thin adapter because it depends on NVSHMEM rank at device-side.
- """
-
- def __init__(
- self, prefix: str, m_clusters: int, n_clusters: int, group_size: int,
world_size: int
- ):
- super().__init__(prefix)
- self._m_clusters = m_clusters
- self._n_clusters = n_clusters
- self._group_size = group_size
- self._world_size = world_size
-
- @T.inline
- def update_current_m_n_idx(self, linear_idx):
- my_rank: T.let = T.nvshmem.my_pe()
- remote_m_clusters: T.let = self._m_clusters - self._m_clusters //
self._world_size
- group_rows: T.let = (remote_m_clusters // self._group_size) *
self._group_size
- final_rows: T.let = remote_m_clusters - group_rows
- group_repeat: T.let = self._group_size * self._n_clusters
- if linear_idx < group_rows * self._n_clusters and group_rows > 0:
- self.m_idx = (
- (linear_idx // group_repeat) * self._group_size
- + (linear_idx % self._group_size)
- + (my_rank + 1) * self._m_clusters // self._world_size
- ) % self._m_clusters
- self.n_idx = (linear_idx % group_repeat) // self._group_size
- elif linear_idx < remote_m_clusters * self._n_clusters:
- remainder_idx: T.let = linear_idx - group_rows * self._n_clusters
- self.m_idx = (
- group_rows
- + remainder_idx % final_rows
- + (my_rank + 1) * self._m_clusters // self._world_size
- ) % self._m_clusters
- self.n_idx = remainder_idx // final_rows
- else:
- remainder_idx: T.let = linear_idx - remote_m_clusters *
self._n_clusters
- self.m_idx = (
- remote_m_clusters
- + remainder_idx % (self._m_clusters // self._world_size)
- + (my_rank + 1) * self._m_clusters // self._world_size
- ) % self._m_clusters
- self.n_idx = remainder_idx // (self._m_clusters //
self._world_size)
-
- @T.inline
- def next_tile(self, stride: int):
- self.linear_idx = self.linear_idx + stride
- self.update_current_m_n_idx(self.linear_idx)
-
- def valid(self):
- return self.linear_idx < self._m_clusters * self._n_clusters
-
-
-class IndexedTripleTileScheduler(BaseTileScheduler):
- """Scheduler that maps linear_idx to (b_idx, h_idx, q_idx) via index
lists."""
-
- def __init__(self, prefix: str, b_indices, h_indices, q_indices,
tiles_indptr):
- super().__init__(prefix)
- self.b_indices = b_indices
- self.h_indices = h_indices
- self.q_indices = q_indices
- self.tiles_indptr = tiles_indptr
- self.q_idx = T.local_scalar("int32")
- self.h_idx = T.local_scalar("int32")
- self.b_idx = T.local_scalar("int32")
- self.linear_lim = T.local_scalar("int32")
-
- @T.inline
- def _load(self):
- self.q_idx = self.q_indices[self.linear_idx]
- self.h_idx = self.h_indices[self.linear_idx]
- self.b_idx = self.b_indices[self.linear_idx]
-
- @T.inline
- def init(self, sm):
- self.linear_idx = self.tiles_indptr[sm]
- self.linear_lim = self.tiles_indptr[sm + 1]
- self._load()
-
- @T.inline
- def next_tile(self):
- self.linear_idx = self.linear_idx + 1
- self._load()
-
- def valid(self):
- return self.linear_idx < self.linear_lim
-
-
class FlashAttentionLinearScheduler(BaseTileScheduler):
"""Linear 3D scheduler for flash attention (batch, head, m_block).
diff --git a/python/tvm/backend/cuda/lang/warp_role.py
b/python/tvm/backend/cuda/lang/warp_role.py
deleted file mode 100644
index add1db63c1..0000000000
--- a/python/tvm/backend/cuda/lang/warp_role.py
+++ /dev/null
@@ -1,150 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Warp role helpers for SM100 kernels.
-
-Simplifies the common pattern of dispatching warps to named roles
-with register budgets.
-
-Example::
-
- # Declare roles
- tma_warp = WarpRole(warp_id, 1, regs=48)
- store_warp = WarpRole(warp_id, 2, regs=48)
- mma_warp = WarpRole(warp_id, 0, regs=232, increase=True)
-
- # Use with context manager
- with tma_warp:
- # TMA load code
- with store_warp:
- # TMA store code
- with mma_warp:
- # MMA compute code
-"""
-
-from tvm.script import tirx as T
-
-
-class WarpRole:
- """A warp-level role that guards a block of code by warp_id comparison
- with optional register budget.
-
- Generates::
-
- if <warp_id_var> == <warp_id_val>:
- # if regs specified:
- T.ptx[f"setmaxnreg.{'inc' if <increase> else
'dec'}.sync.aligned.u32"](<regs>)
- <user code>
-
- The ``if`` guard narrows the active set to the single warp; individual
- tile-primitive calls inside ``<user code>`` carry their own exec scope via
- a scope-namespace prefix (e.g. ``Tx.warp.copy(...)``).
-
- Parameters
- ----------
- warp_id_var : Var
- The warp_id variable (from ``T.cuda.warp_id()``).
- warp_id_val : int
- Which warp index this role corresponds to.
- regs : int, optional
- Register budget (passed to ``T.ptx.setmaxnreg``).
- If None, no setmaxnreg is emitted.
- increase : bool
- Direction for ``setmaxnreg`` (default False = decrease).
- """
-
- def __init__(self, warp_id_var, warp_id_val, regs=None, increase=False):
- self.warp_id_var = warp_id_var
- self.warp_id_val = warp_id_val
- self.regs = regs
- self.increase = increase
-
- def __enter__(self):
- self._if_frame = T.if_(self.warp_id_var == self.warp_id_val)
- self._if_frame.__enter__()
- self._then_frame = T.then_()
- self._then_frame.__enter__()
- if self.regs is not None:
- T.evaluate(
- T.ptx[f"setmaxnreg.{'inc' if self.increase else
'dec'}.sync.aligned.u32"](self.regs)
- )
- return self
-
- def __exit__(self, *exc):
- self._then_frame.__exit__(*exc)
- self._if_frame.__exit__(*exc)
- return False
-
-
-class WarpgroupRole:
- """A warpgroup-level role that guards by wg_id comparison,
- with optional register budget.
-
- Generates (single wg_id)::
-
- if <wg_id_var> == <wg_id_val>:
- # if regs specified:
- T.ptx[f"setmaxnreg.{'inc' if <increase> else
'dec'}.sync.aligned.u32"](<regs>)
- <user code>
-
- Generates (range of wg_ids, e.g. ``wg_id_val=(0, 2)``)::
-
- if 0 <= <wg_id_var> and <wg_id_var> < 2:
- T.ptx[f"setmaxnreg.{'inc' if <increase> else
'dec'}.sync.aligned.u32"](<regs>)
- <user code>
-
- The ``if`` guard narrows the active set to the target warpgroup(s);
- individual tile-primitive calls inside ``<user code>`` carry their own exec
- scope via a scope-namespace prefix (e.g. ``Tx.wg.copy(...)``).
-
- Parameters
- ----------
- wg_id_var : Var
- The warpgroup_id variable (from ``T.cuda.warpgroup_id()``).
- wg_id_val : int or tuple[int, int]
- Which warpgroup index (int) or range ``(start, stop)`` this role
- corresponds to.
- regs : int, optional
- Register budget.
- increase : bool
- Direction for ``setmaxnreg`` (default False = decrease).
- """
-
- def __init__(self, wg_id_var, wg_id_val, regs=None, increase=False):
- self.wg_id_var = wg_id_var
- self.wg_id_val = wg_id_val
- self.regs = regs
- self.increase = increase
-
- def __enter__(self):
- if isinstance(self.wg_id_val, tuple):
- start, stop = self.wg_id_val
- self._if_frame = T.if_(T.And(T.LE(start, self.wg_id_var),
self.wg_id_var < stop))
- else:
- self._if_frame = T.if_(self.wg_id_var == self.wg_id_val)
- self._if_frame.__enter__()
- self._then_frame = T.then_()
- self._then_frame.__enter__()
- if self.regs is not None:
- T.evaluate(
- T.ptx[f"setmaxnreg.{'inc' if self.increase else
'dec'}.sync.aligned.u32"](self.regs)
- )
- return self
-
- def __exit__(self, *exc):
- self._then_frame.__exit__(*exc)
- self._if_frame.__exit__(*exc)
- return False
diff --git a/python/tvm/backend/cuda/layout.py
b/python/tvm/backend/cuda/layout.py
index 71ba7e8c04..330afed506 100644
--- a/python/tvm/backend/cuda/layout.py
+++ b/python/tvm/backend/cuda/layout.py
@@ -27,9 +27,8 @@ from tvm.tirx.layout import Axis, Iter, R, S, TileLayout
# ``tcgen05.mma`` writes its output matrix C into TMEM using one of several
# **datapath layouts** depending on the MMA's M dimension and ``.ws`` mode.
# Each layout determines *which* physical TMEM lanes (rows) the matrix
-# occupies; the leak in the original ``_default_tmem_layout`` was that it
-# always used the identity ``(rows, cols) : (1@TLane, 1@TCol)`` mapping,
-# which is correct only for Layout D (M=128 full datapath). For Layout F
+# occupies. The identity ``(rows, cols) : (1@TLane, 1@TCol)`` mapping
+# is correct only for Layout D (M=128 full datapath). For Layout F
# (M=64 non-``.ws``) the MMA writes scattered lanes
# ``{0..15, 32..47, 64..79, 96..111}`` — half of each warp's 32-lane
# partition — and the readback path (``.16x*b`` M=64 atom) has the matching
diff --git a/python/tvm/backend/cuda/op.py b/python/tvm/backend/cuda/op.py
index 74a0528393..5c8ef73f70 100644
--- a/python/tvm/backend/cuda/op.py
+++ b/python/tvm/backend/cuda/op.py
@@ -29,15 +29,6 @@ from tvm.ir.location import UNKNOWN_LOC, Location
from tvm.ir.op import _init_op_api, _make_op_api
from tvm.ir.type import PrimType, PtrType
from tvm.tirx.op import bitwise_and, call_intrin, ptr_byte_offset
-from tvm.tirx.operator.intrinsics._common import (
- CP_ASYNC_BULK_CACHE_HINT as _CP_ASYNC_BULK_CACHE_HINT,
-)
-from tvm.tirx.operator.intrinsics._common import MBARRIER_ARRIVE_SCOPE as
_MBARRIER_ARRIVE_SCOPE
-from tvm.tirx.operator.intrinsics._common import MBARRIER_ARRIVE_SEM as
_MBARRIER_ARRIVE_SEM
-from tvm.tirx.operator.intrinsics._common import MBARRIER_ARRIVE_SPACE as
_MBARRIER_ARRIVE_SPACE
-from tvm.tirx.operator.intrinsics._common import NVSHMEM_CMP as _NVSHMEM_CMP
-from tvm.tirx.operator.intrinsics._common import NVSHMEM_SIG_OP as
_NVSHMEM_SIG_OP
-from tvm.tirx.operator.intrinsics._common import TCGEN05_CTA_GROUP as
_TCGEN05_CTA_GROUP
tir = tirx
@@ -344,6 +335,11 @@ def cuda_wait_until(
)
+_MBARRIER_ARRIVE_SEM = ("", "release", "relaxed")
+_MBARRIER_ARRIVE_SCOPE = ("", "cta", "cluster")
+_MBARRIER_ARRIVE_SPACE = ("shared", "shared::cta", "shared::cluster")
+
+
def _validate_mbarrier_arrive_attrs(sem, scope, space, remote):
if (sem == "") != (scope == ""):
raise ValueError("mbarrier.arrive sem and scope must be specified
together")
@@ -590,6 +586,9 @@ def ptx_legacy_ldmatrix(
)
+_TCGEN05_CTA_GROUP = (1, 2)
+
+
@tvm_ffi.register_object("tirx.cuda.TCGen05InstrDescriptorAttrs")
class TCGen05InstrDescriptorAttrs(Attrs):
"""Static options for the dense tcgen05 instruction descriptor."""
@@ -810,11 +809,6 @@ def _static_str(value):
return None
-# See top-of-file imports for `_FENCE_SEM` etc. (re-exported from _common).
-# Note: TCGEN05_LDST_SHAPES values must stay in sync with the tcgen05.ld/.st
-# shape tokens in backend/cuda/ptx/table.py.
-
-
def timer_init(
profiler_buffer,
profiler_tag,
@@ -1108,6 +1102,8 @@ _PTX_CVT_RETURN_TYPE = {
}
+_CP_ASYNC_BULK_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal",
"evict_last_use")
+
_PTX_CACHE_POLICY = {
"evict_normal": 0x1000000000000000,
"evict_first": 0x12F0000000000000,
@@ -1407,6 +1403,10 @@ def nvshmem_putmem_nbi_block(dst, src, nelems, pe, *,
ty=None, loc: Location = U
)
+_NVSHMEM_CMP = ("eq", "ne", "gt", "ge", "lt", "le")
+_NVSHMEM_SIG_OP = ("set", "add")
+
+
def nvshmem_signal_op(sig_addr, signal, sig_op, pe, *, ty=None, loc: Location
= UNKNOWN_LOC):
"""TVM intrinsic to call nvshmem_signal_op()
diff --git a/python/tvm/backend/cuda/tile/common.py
b/python/tvm/backend/cuda/tile/common.py
index 1cdf7826be..f81e44e657 100644
--- a/python/tvm/backend/cuda/tile/common.py
+++ b/python/tvm/backend/cuda/tile/common.py
@@ -29,7 +29,6 @@ from tvm.sym.analyzer import Analyzer
from tvm.tirx import Function, Var
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
def next_power_of_2(x: int) -> int:
@@ -90,7 +89,7 @@ def validate_copy_op(
dst_buffer_region, src_buffer_region = op_call.args[:2]
src: Var = src_buffer_region.source
dst: Var = dst_buffer_region.source
- if not (effective_layout(src) and effective_layout(dst) and src.dtype ==
dst.dtype):
+ if not (src.ty.effective_layout() and dst.ty.effective_layout() and
src.dtype == dst.dtype):
return False
# Extract regions and validate dimensions
analyzer = Analyzer()
@@ -117,8 +116,8 @@ def get_vec_len(
src: Var = src_buffer_region.source
# layout=None (flat local buffer) is treated as trivial for vectorization
purposes
if not (
- (effective_layout(dst) is None or effective_layout(dst).is_trivial())
- and (effective_layout(src) is None or
effective_layout(src).is_trivial())
+ (dst.ty.effective_layout() is None or
dst.ty.effective_layout().is_trivial())
+ and (src.ty.effective_layout() is None or
src.ty.effective_layout().is_trivial())
):
return None
diff --git a/python/tvm/backend/cuda/tile/copy/ld_stmatrix.py
b/python/tvm/backend/cuda/tile/copy/ld_stmatrix.py
index 08f3ab75e7..2f40f79a36 100644
--- a/python/tvm/backend/cuda/tile/copy/ld_stmatrix.py
+++ b/python/tvm/backend/cuda/tile/copy/ld_stmatrix.py
@@ -29,7 +29,6 @@ from tvm.tirx import Function
from tvm.tirx.layout import ComposeLayout, S, TileLayout
from tvm.tirx.tile.dispatch import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ._common import ( # noqa: F401 (_carve_tail reserved for future
variants)
_carve_tail,
@@ -87,9 +86,9 @@ def _emit(op_call: TensorCall, sctx: DispatchContext) ->
Function:
r_buf = r_br.source
s_buf = s_br.source
r_shape = list(r_buf.shape)
- r_layout = effective_layout(r_buf)
+ r_layout = r_buf.ty.effective_layout()
s_shape = list(s_buf.shape)
- s_layout = effective_layout(s_buf)
+ s_layout = s_buf.ty.effective_layout()
r_region = [(r.min, r.min + r.extent) for r in r_br.region]
s_region = [(r.min, r.min + r.extent) for r in s_br.region]
@@ -102,7 +101,9 @@ def _emit(op_call: TensorCall, sctx: DispatchContext) ->
Function:
# Peel the S-side wrapper before structural checks, retaining its
# parameters so the selected TileLayout can be wrapped again at emit.
s_swizzle = (
- effective_layout(s_buf) if isinstance(effective_layout(s_buf),
ComposeLayout) else None
+ s_buf.ty.effective_layout()
+ if isinstance(s_buf.ty.effective_layout(), ComposeLayout)
+ else None
)
if s_swizzle is not None and s_swizzle.per_element < 3:
# ldmatrix/stmatrix .b16 reads/writes 8 fp16 = 128b per lane in one
diff --git a/python/tvm/backend/cuda/tile/copy/vec_auto_reg.py
b/python/tvm/backend/cuda/tile/copy/vec_auto_reg.py
index 2281c67839..539c687549 100644
--- a/python/tvm/backend/cuda/tile/copy/vec_auto_reg.py
+++ b/python/tvm/backend/cuda/tile/copy/vec_auto_reg.py
@@ -39,7 +39,6 @@ from tvm.tirx.expr import IntImm as _IntImm
from tvm.tirx.layout import ComposeLayout, Iter, TileLayout
from tvm.tirx.tile.dispatch import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..layout_utils import recompose_swizzle, strip_swizzle_to_tile
from ._common import _alignment_ok, copy_ptx_form, copy_ptx_ld_chain
@@ -96,7 +95,7 @@ def _r_side_layout_valid(op_call: TensorCall, sctx:
DispatchContext) -> tuple[bo
src: Var = op_call.src.source
dst: Var = op_call.dst.source
r_buf = src if src.scope() == "local" else dst
- layout = effective_layout(r_buf)
+ layout = r_buf.ty.effective_layout()
if layout is None:
return False, "R has no layout"
if layout.is_swizzle():
@@ -139,7 +138,7 @@ def _s_side_slice_ok(op_call: TensorCall) -> tuple[bool,
str | None]:
dst_br = op_call.dst
s_br = dst_br if src_br.source.scope() == "local" else src_br
s_buf: Var = s_br.source
- layout = effective_layout(s_buf)
+ layout = s_buf.ty.effective_layout()
if layout is None:
return False, "S has no layout"
region = [(r.min, r.min + r.extent) for r in s_br.region]
@@ -278,8 +277,8 @@ def _align_layouts(op_call: TensorCall, sctx:
DispatchContext):
r_region = [(r.min, r.min + r.extent) for r in r_br.region]
s_region = [(r.min, r.min + r.extent) for r in s_br.region]
with sctx.target:
- r_sliced = effective_layout(r_buf).slice(list(r_buf.shape), r_region)
- s_sliced = effective_layout(s_buf).slice(list(s_buf.shape), s_region)
+ r_sliced = r_buf.ty.effective_layout().slice(list(r_buf.shape),
r_region)
+ s_sliced = s_buf.ty.effective_layout().slice(list(s_buf.shape),
s_region)
return align_layouts_raw(r_sliced, s_sliced, s_region)
@@ -526,17 +525,17 @@ def _emit_reg(op_call: TensorCall, sctx: DispatchContext)
-> Function:
elem_bits = DataType(src.dtype).bits
max_vec_len = None
if (
- isinstance(effective_layout(s_buf), ComposeLayout)
- and int(effective_layout(s_buf).swizzle_len) > 0
+ isinstance(s_buf.ty.effective_layout(), ComposeLayout)
+ and int(s_buf.ty.effective_layout().swizzle_len) > 0
):
- max_vec_len = 1 << int(effective_layout(s_buf).per_element)
+ max_vec_len = 1 << int(s_buf.ty.effective_layout().per_element)
vec_len = _choose_vec_len(elem_bits, atoms, r_p, s_p, max_vec_len)
vec_bits = vec_len * elem_bits
outer = _split_atoms_for_vec(atoms, vec_len)
placeholders = _make_thread_placeholders(r_p)
s_apply_layout, thread_coords, s_apply_shape = _build_s_apply_layout(
- effective_layout(s_buf), r_p, s_p, outer, placeholders, sctx
+ s_buf.ty.effective_layout(), r_p, s_p, outer, placeholders, sctx
)
# A plain TileLayout is additive, so keep its loop-invariant thread base
# outside the serial copy loop. A real swizzle must still apply to the
diff --git a/python/tvm/backend/cuda/tile/copy_async/dsmem.py
b/python/tvm/backend/cuda/tile/copy_async/dsmem.py
index d53610b458..6ff4c9c674 100644
--- a/python/tvm/backend/cuda/tile/copy_async/dsmem.py
+++ b/python/tvm/backend/cuda/tile/copy_async/dsmem.py
@@ -25,7 +25,6 @@ from tvm.script import tirx as T
from tvm.tirx import Function, Var
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from .utils import find_contiguous_region, to_tile_layout
@@ -73,8 +72,8 @@ def copy_dsmem_impl(op_call: TensorCall, sctx:
DispatchContext) -> Function:
dtype_bytes = tvm.DataType(src_buf.dtype).bits // 8
# Get tile layouts for both buffers
- src_tile_layout = to_tile_layout(effective_layout(src_buf), src_buf.shape)
- dst_tile_layout = to_tile_layout(effective_layout(dst_buf), dst_buf.shape)
+ src_tile_layout = to_tile_layout(src_buf.ty.effective_layout(),
src_buf.shape)
+ dst_tile_layout = to_tile_layout(dst_buf.ty.effective_layout(),
dst_buf.shape)
# Slice layouts to copy region
src_region_tuples = [(src_st[i], src_st[i] + src_ext[i]) for i in
range(len(src_st))]
@@ -137,8 +136,8 @@ def copy_dsmem_impl(op_call: TensorCall, sctx:
DispatchContext) -> Function:
dst_off = dst_off + v * outer_dst_strides[j]
return src_off, dst_off
- src_tile = to_tile_layout(effective_layout(src_buf), src_buf.shape)
- dst_tile = to_tile_layout(effective_layout(dst_buf), dst_buf.shape)
+ src_tile = to_tile_layout(src_buf.ty.effective_layout(), src_buf.shape)
+ dst_tile = to_tile_layout(dst_buf.ty.effective_layout(), dst_buf.shape)
# fmt: off
@T.function(check_well_formed=False)
diff --git a/python/tvm/backend/cuda/tile/copy_async/ldgsts.py
b/python/tvm/backend/cuda/tile/copy_async/ldgsts.py
index 8dad16f203..fbdfc6e649 100644
--- a/python/tvm/backend/cuda/tile/copy_async/ldgsts.py
+++ b/python/tvm/backend/cuda/tile/copy_async/ldgsts.py
@@ -47,7 +47,6 @@ from tvm.tirx import Function, Var
from tvm.tirx.expr import IntImm as _IntImm
from tvm.tirx.tile.dispatch import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..copy._common import (
_TID_AXIS_FOR_SCOPE,
@@ -190,17 +189,17 @@ def _emit_ldgsts(op_call: TensorCall, sctx:
DispatchContext) -> Function:
with sctx.target:
g_p, s_p, vec_len = align_layouts_gs(
- effective_layout(g_buf),
+ g_buf.ty.effective_layout(),
g_buf.shape,
g_region,
- effective_layout(s_buf),
+ s_buf.ty.effective_layout(),
s_buf.shape,
s_region,
elem_bits,
thread_cnt,
vec_bits_candidates=_LDGSTS_VEC_BITS,
)
- s_apply_layout = recompose_swizzle(effective_layout(s_buf), s_p)
+ s_apply_layout = recompose_swizzle(s_buf.ty.effective_layout(), s_p)
vec_bits = vec_len * elem_bits
cp_size = vec_bits // 8 # cp.async cp_size is in bytes
diff --git a/python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py
b/python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py
index d7cff5a60c..d2b36eb7c1 100644
--- a/python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py
+++ b/python/tvm/backend/cuda/tile/copy_async/tcgen05_cp.py
@@ -120,7 +120,6 @@ from tvm.tirx.layout import ComposeLayout, TileLayout
from tvm.tirx.layout import m as m_axis
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
# Allowed .multicast qualifiers per shape (PTX ISA 8.8 §9.7.16.9.2): "" = none,
# 64x128b needs an explicit warpx2 pick.
@@ -366,9 +365,9 @@ def _plan_for_shape(op_call: TensorCall, shape: str,
multicast: str):
# C: slice + canonicalize.
s_region = [(r.min, r.min + r.extent) for r in src_region.region]
t_region = [(r.min, r.min + r.extent) for r in dst_region.region]
- s_sliced = effective_layout(s_buf).slice(list(s_buf.shape), s_region)
+ s_sliced = s_buf.ty.effective_layout().slice(list(s_buf.shape), s_region)
s = s_sliced.canonicalize()
- t = effective_layout(t_buf).slice(list(t_buf.shape),
t_region).canonicalize()
+ t = t_buf.ty.effective_layout().slice(list(t_buf.shape),
t_region).canonicalize()
# If s is ComposeLayout (swizzle over a tile), peel off the swizzle for
# stride analysis; record the swizzle object for cross-checks.
@@ -725,8 +724,8 @@ def _validate_smem_tmem_copy(op_call: TensorCall, sctx:
DispatchContext):
return (
src.scope().startswith("shared")
and dst.scope() == "tmem"
- and effective_layout(src) is not None
- and effective_layout(dst) is not None
+ and src.ty.effective_layout() is not None
+ and dst.ty.effective_layout() is not None
# Byte reinterpret is legal (nvfp4 stages sf as uint8, reads fp8); a
# true width change needs decompress, rejected in copy_smem_tmem_impl.
and DataType(src.dtype).bits == DataType(dst.dtype).bits
diff --git a/python/tvm/backend/cuda/tile/copy_async/tcgen05_ldst.py
b/python/tvm/backend/cuda/tile/copy_async/tcgen05_ldst.py
index e92da80662..46d409dc57 100644
--- a/python/tvm/backend/cuda/tile/copy_async/tcgen05_ldst.py
+++ b/python/tvm/backend/cuda/tile/copy_async/tcgen05_ldst.py
@@ -33,7 +33,6 @@ from tvm.tirx import Function, Var
from tvm.tirx.layout import S, TileLayout
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import get_st_extent
@@ -58,7 +57,7 @@ def _match_tcgen05_atom_layout(buf):
if rows not in (64, 128):
return None
dtype = buf.dtype
- layout_c = effective_layout(buf).canonicalize()
+ layout_c = buf.ty.effective_layout().canonicalize()
for shape in _TCGEN05_COL_FACTOR_FP32:
try:
cand = tcgen05_atom_layout(shape, (rows, cols),
dtype).canonicalize()
@@ -116,20 +115,18 @@ def _classify_tmem_datapath(tmem_buf):
"""Return ``(datapath, sub_slab)`` if ``tmem_buf.layout`` matches a known
tcgen05 datapath (PTX ISA §9.7.16.10.5), else ``None``.
- Layout D (M=128, identity row→lane) is the default returned by
- ``_default_tmem_layout``. Layout F (M=64 non-``.ws``, scattered) is the
- explicit opt-in produced by ``tmem_pool.alloc(..., datapath="F")``.
- Layout B is the per-CTA M=64, ``.cta_group::2`` "2x2" placement produced
- by ``tmem_pool.alloc(..., datapath="B")``.
+ ``tmem_datapath_layout`` constructs the supported layouts: Layout D
+ (M=128, identity row→lane), Layout F (M=64 non-``.ws``, scattered),
+ and Layout B (per-CTA M=64, ``.cta_group::2`` "2x2" placement).
The dispatch uses this to pair each ``.16x*b`` / ``.32x32b`` atom with a
compatible layout — see ``_check_tmem_layout_for_atom``.
``sub_slab`` is always 0 for Layout D and B. For Layout F it selects the
lower (0) or upper (1) 16-lane half of each warp's 32-lane partition.
"""
- if effective_layout(tmem_buf) is None:
+ if tmem_buf.ty.effective_layout() is None:
return None
- buf_layout = effective_layout(tmem_buf).canonicalize()
+ buf_layout = tmem_buf.ty.effective_layout().canonicalize()
rows = int(tmem_buf.shape[0])
if rows == 128:
cand = tmem_datapath_layout("D", 128, tmem_buf.shape[1]).canonicalize()
@@ -196,7 +193,7 @@ def _tmem_window(tmem_buf, tmem_region, atom_kind,
frag_rows, analyzer):
f"incompatible with atom={atom_kind!r} (frag_rows={frag_rows}). "
f"See PTX ISA §9.7.16.10.5 for datapath/atom pairings."
)
- window = effective_layout(tmem_buf).slice(tmem_buf.shape,
tmem_region.region).canonicalize()
+ window = tmem_buf.ty.effective_layout().slice(tmem_buf.shape,
tmem_region.region).canonicalize()
if datapath == "D":
base = TileLayout(S[(frag_rows, width) : (1 @ TLane, 1 @ TCol)])
else:
@@ -234,8 +231,8 @@ def copy_tmem_local_impl(op_call: TensorCall, sctx:
DispatchContext) -> Function
tmem_buf, local_buf = tmem_region.source, local_region.source
- assert effective_layout(tmem_buf) is not None
- assert effective_layout(local_buf) is not None
+ assert tmem_buf.ty.effective_layout() is not None
+ assert local_buf.ty.effective_layout() is not None
assert tmem_buf.dtype == local_buf.dtype
analyzer = Analyzer()
@@ -521,7 +518,9 @@ def _emit_datapath_b_path(
n_half = n_cols // 2
expected_local = tcgen05_atom_layout("32x32b", (64, n_cols),
local_buf.dtype).canonicalize()
try:
-
tvm.ir.assert_structural_equal(effective_layout(local_buf).canonicalize(),
expected_local)
+ tvm.ir.assert_structural_equal(
+ local_buf.ty.effective_layout().canonicalize(), expected_local
+ )
except (AssertionError, ValueError) as err:
raise ValueError(
"datapath B (.cta_group::2 M=64) requires a matching Layout B "
diff --git a/python/tvm/backend/cuda/tile/copy_async/tma.py
b/python/tvm/backend/cuda/tile/copy_async/tma.py
index 7b6b525703..65e891efd4 100644
--- a/python/tvm/backend/cuda/tile/copy_async/tma.py
+++ b/python/tvm/backend/cuda/tile/copy_async/tma.py
@@ -39,7 +39,6 @@ from tvm.tirx import Function, IntImm, Var, is_tensor_var
from tvm.tirx.layout import Layout, TileLayout
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ...op import _is_static_unicast_cta_mask, _resolve_cache_policy
from ..layout_utils import strip_swizzle_to_tile
@@ -274,7 +273,7 @@ def _layout_offset(layout: TileLayout):
def _slice_layout(buffer: Var, starts, extents, label: str) ->
tuple[TileLayout, TileLayout]:
- tile = _to_tile_layout(effective_layout(buffer), buffer.shape)
+ tile = _to_tile_layout(buffer.ty.effective_layout(), buffer.shape)
_assert_plain_memory_layout(tile, label)
region = [(start, start + extent) for start, extent in zip(starts,
extents)]
sliced = tile.slice(list(buffer.shape), region)
@@ -300,7 +299,7 @@ def _slice_global_layout(buffer: Var, starts, extents) ->
tuple[TileLayout, Tile
"""
analyzer = Analyzer()
- tile = _to_tile_layout(effective_layout(buffer), buffer.shape)
+ tile = _to_tile_layout(buffer.ty.effective_layout(), buffer.shape)
_assert_plain_memory_layout(tile, "global")
try:
grouped, separators = tile.group(list(buffer.shape))
@@ -433,7 +432,7 @@ def _elements_to_bytes(
def _buffer_base(buffer: Var, *, auto: bool, stage: str):
analyzer = Analyzer()
- layout = _to_tile_layout(effective_layout(buffer), buffer.shape)
+ layout = _to_tile_layout(buffer.ty.effective_layout(), buffer.shape)
layout_offset = _layout_offset(layout)
element_offset = analyzer.simplify(buffer.elem_offset + layout_offset)
byte_offset = _elements_to_bytes(
@@ -1180,7 +1179,7 @@ def _auto_spec_for_prefix(
return TMAPlan(
spec=spec,
issue_axes=tuple(issue_axes),
- shared_layout=_to_tile_layout(effective_layout(s_buf), s_buf.shape),
+ shared_layout=_to_tile_layout(s_buf.ty.effective_layout(),
s_buf.shape),
)
@@ -1664,11 +1663,13 @@ def _build_auto_plan(op_call: TensorCall, sctx:
DispatchContext) -> TMAPlan:
s_starts, s_extents = _copy_region_parts(shared_region.region)
g_starts, g_extents = _copy_region_parts(global_region.region)
try:
- swizzle = get_swizzle_mode_from_layout(effective_layout(s_buf))
+ swizzle = get_swizzle_mode_from_layout(s_buf.ty.effective_layout())
except ValueError as error:
_auto_fail("shared-swizzle", str(error))
if swizzle is None:
- _auto_fail("shared-layout", f"cannot recognize shared swizzle in
{effective_layout(s_buf)}")
+ _auto_fail(
+ "shared-layout", f"cannot recognize shared swizzle in
{s_buf.ty.effective_layout()}"
+ )
try:
_, sliced_smem_with_offset = _slice_layout(s_buf, s_starts, s_extents,
"shared")
except ValueError as error:
@@ -1763,13 +1764,13 @@ def _explicit_smem_layout(s_buf: Var, starts, extents,
swizzle: SwizzleMode):
f"{byte_offset}B violates {atom_alignment}B swizzle/base alignment"
)
return (
- _to_tile_layout(effective_layout(s_buf), s_buf.shape),
+ _to_tile_layout(s_buf.ty.effective_layout(), s_buf.shape),
analyzer.simplify(s_buf.elem_offset + offset),
)
def _direct_global_layout(g_buf: Var):
- layout = effective_layout(g_buf)
+ layout = g_buf.ty.effective_layout()
if not isinstance(layout, TileLayout):
fail(
"tma_explicit stage=global-layout: global Var/view layout must be "
@@ -1995,11 +1996,11 @@ def _build_explicit_plan(op_call: TensorCall, sctx:
DispatchContext):
fail("tma_explicit gather4 shared public axis 0 must describe four
rows")
try:
- swizzle = get_swizzle_mode_from_layout(effective_layout(s_buf))
+ swizzle = get_swizzle_mode_from_layout(s_buf.ty.effective_layout())
except ValueError as error:
fail(f"tma_explicit stage=shared-swizzle: {error}")
if swizzle is None:
- fail(f"tma_explicit cannot recognize shared swizzle in
{effective_layout(s_buf)}")
+ fail(f"tma_explicit cannot recognize shared swizzle in
{s_buf.ty.effective_layout()}")
shared_layout, smem_base_offset = _explicit_smem_layout(s_buf, s_starts,
s_extents, swizzle)
main_spec = _explicit_spec_for_gmem(
g_buf=g_buf,
@@ -2270,7 +2271,7 @@ def _validate_tma_copy_op(op_call: TensorCall, _sctx:
DispatchContext) -> bool:
dst_region, src_region = op_call.args[:2]
src = src_region.source
dst = dst_region.source
- if effective_layout(src) is None or effective_layout(dst) is None:
+ if src.ty.effective_layout() is None or dst.ty.effective_layout() is None:
return False
src_scope, dst_scope = src.scope(), dst.scope()
return (src_scope == "global" and dst_scope.startswith("shared")) or (
diff --git a/python/tvm/backend/cuda/tile/elementwise/_common.py
b/python/tvm/backend/cuda/tile/elementwise/_common.py
index fa9c710c87..eeceb1e396 100644
--- a/python/tvm/backend/cuda/tile/elementwise/_common.py
+++ b/python/tvm/backend/cuda/tile/elementwise/_common.py
@@ -35,7 +35,6 @@ from tvm.runtime import DataType
from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
from tvm.tirx.layout import Axis, Iter, TileLayout
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import get_indices, get_st_extent
@@ -206,7 +205,7 @@ def preprocess_operand(op_br, anchor_tshape):
Raises ``ValueError`` if the lift is not broadcast-compatible (caller
should have verified via ``shape_broadcast_compat`` in the predicate).
"""
- op_layout = effective_layout(op_br.source)
+ op_layout = op_br.source.ty.effective_layout()
op_shape = op_br.source.shape
op_region = [(r.min, r.min + r.extent) for r in op_br.region]
sliced = op_layout.slice(list(op_shape), op_region).canonicalize()
@@ -287,7 +286,7 @@ def align_operands_to_anchor(anchor_br, layout_others_br):
Uses ``_align_layouts_no_post_canon`` (not copy's ``align_layouts_raw``
directly) so ``anchor_p.shard`` length matches ``op_seps`` groupings.
"""
- anchor_layout = effective_layout(anchor_br.source)
+ anchor_layout = anchor_br.source.ty.effective_layout()
anchor_shape = anchor_br.source.shape
anchor_region = [(r.min, r.min + r.extent) for r in anchor_br.region]
@@ -302,7 +301,7 @@ def align_operands_to_anchor(anchor_br, layout_others_br):
return anchor_p, per_op_aligned
for op_br in layout_others_br:
- op_layout = effective_layout(op_br.source)
+ op_layout = op_br.source.ty.effective_layout()
op_shape = op_br.source.shape
op_region = [(r.min, r.min + r.extent) for r in op_br.region]
r_p, op_p, op_seps = _align_layouts_no_post_canon(
diff --git a/python/tvm/backend/cuda/tile/elementwise/reg.py
b/python/tvm/backend/cuda/tile/elementwise/reg.py
index ff958ec9d2..e17105d4e2 100644
--- a/python/tvm/backend/cuda/tile/elementwise/reg.py
+++ b/python/tvm/backend/cuda/tile/elementwise/reg.py
@@ -41,7 +41,6 @@ from tvm.tirx.layout import TileLayout
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.dispatch import fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import get_st_extent
from ..copy._common import _carve_tail, _verify_s_tail_contig
@@ -64,7 +63,7 @@ from .vec_emit import _emit_vec
# Predicate
# -----------------------------------------------------------------------------
def _validate_anchor_layout(anchor_br) -> tuple[bool, str | None]:
- layout = effective_layout(anchor_br.source)
+ layout = anchor_br.source.ty.effective_layout()
if layout.is_swizzle():
return False, "anchor layout is swizzle"
if not isinstance(layout, TileLayout):
@@ -90,7 +89,7 @@ def _validate_scope_level_anchor(anchor_br, sctx:
DispatchContext) -> tuple[bool
# Canonicalize the sliced anchor with the target so warp/lane axes fuse.
st, ext = get_st_extent(anchor_br)
sliced = get_sublayout_from_region(
- effective_layout(anchor_br.source), anchor_br.source.shape, st, ext
+ anchor_br.source.ty.effective_layout(), anchor_br.source.shape, st, ext
)
with sctx.target:
canon = sliced.canonicalize()
@@ -156,7 +155,7 @@ def _check_layout_operands_agree(plan, sctx) -> tuple[bool,
str | None]:
st, ext = get_st_extent(br)
with sctx.target:
sliced = get_sublayout_from_region(
- effective_layout(br.source), br.source.shape, st, ext
+ br.source.ty.effective_layout(), br.source.shape, st, ext
)
canon = sliced.canonicalize()
sig = layout_signature(canon)
@@ -195,7 +194,7 @@ def is_reg_ewise(spec):
for br in buffer_regions(plan):
if br.source.scope() != "local":
return False, f"operand scope {br.source.scope()} != local"
- if effective_layout(br.source) is None:
+ if br.source.ty.effective_layout() is None:
return False, f"operand {br} has no layout"
if spec.check_extras is not None:
ok2, reason2 = spec.check_extras(plan.extras,
compute_dtype_of(plan))
diff --git a/python/tvm/backend/cuda/tile/elementwise/smem.py
b/python/tvm/backend/cuda/tile/elementwise/smem.py
index 24dc497865..623c6f85b6 100644
--- a/python/tvm/backend/cuda/tile/elementwise/smem.py
+++ b/python/tvm/backend/cuda/tile/elementwise/smem.py
@@ -35,7 +35,6 @@ from tvm.tirx import Function
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.dispatch import fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import get_indices, get_st_extent, get_thread_cnt
from ._common import (
@@ -76,7 +75,7 @@ def is_smem_ewise(spec):
for br in buffer_regions(plan):
if not br.source.scope().startswith("shared"):
return False, f"operand scope {br.source.scope()} != shared*"
- if effective_layout(br.source) is None:
+ if br.source.ty.effective_layout() is None:
return False, "shared operand has no layout"
if spec.check_extras is not None:
ok2, reason2 = spec.check_extras(plan.extras,
compute_dtype_of(plan))
diff --git a/python/tvm/backend/cuda/tile/gemm/mma_m16n8k_.py
b/python/tvm/backend/cuda/tile/gemm/mma_m16n8k_.py
index 46441e10c4..1d4b492897 100644
--- a/python/tvm/backend/cuda/tile/gemm/mma_m16n8k_.py
+++ b/python/tvm/backend/cuda/tile/gemm/mma_m16n8k_.py
@@ -25,7 +25,6 @@ from tvm.tirx import Function
from tvm.tirx.layout import TileLayout
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
@dataclass(frozen=True)
@@ -248,8 +247,8 @@ def _full_active_lanes(op: TensorCall, sctx:
DispatchContext):
def _no_replica(op: TensorCall, sctx: DispatchContext):
"""All operand layouts must have no replica (no broadcast/duplicated
axes)."""
for region, name in zip(op.args[:4], ("D", "A", "B", "C")):
- if effective_layout(region.source).replica:
- return False, f"{name} layout has replica
{effective_layout(region.source).replica}"
+ if region.source.ty.effective_layout().replica:
+ return False, f"{name} layout has replica
{region.source.ty.effective_layout().replica}"
return True
@@ -335,7 +334,7 @@ def gemm_cuda_mma_dispatch(op: TensorCall, sctx:
DispatchContext) -> Function:
# ICHECK when the layout can't be tiled as shape2d -- guard both.
canon = None
try:
- sliced = effective_layout(buf).slice(buf.shape, region.region)
+ sliced = buf.ty.effective_layout().slice(buf.shape, region.region)
if sliced is not None:
canon = sliced.canonicalize()
except Exception as e: # ICHECK failure -> layout not tileable as
shape2d
diff --git a/python/tvm/backend/cuda/tile/gemm_async/tcgen05.py
b/python/tvm/backend/cuda/tile/gemm_async/tcgen05.py
index 5e5282cf01..e54a64f68b 100644
--- a/python/tvm/backend/cuda/tile/gemm_async/tcgen05.py
+++ b/python/tvm/backend/cuda/tile/gemm_async/tcgen05.py
@@ -40,7 +40,6 @@ from tvm.tirx import op as tirx_op
from tvm.tirx.layout import ComposeLayout, Iter, R, S, TileLayout
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ...cpp.descriptors import (
_check_tcgen05_mma_matrix_shape,
@@ -333,14 +332,14 @@ def _layout_matches_datapath_f(tmem_buf) -> bool:
F C buffers (the canonical pairing for M=64 outputs that are read back
via ``.16x*b`` M=64; see PTX ISA §9.7.16.10.5).
"""
- if effective_layout(tmem_buf) is None or int(tmem_buf.shape[0]) != 64:
+ if tmem_buf.ty.effective_layout() is None or int(tmem_buf.shape[0]) != 64:
return False
try:
base = tmem_datapath_layout("F", 64, tmem_buf.shape[1])
expected = TileLayout.from_iters(
- base.shard, base.replica, effective_layout(tmem_buf).offset
+ base.shard, base.replica, tmem_buf.ty.effective_layout().offset
).canonicalize()
-
tvm.ir.assert_structural_equal(effective_layout(tmem_buf).canonicalize(),
expected)
+
tvm.ir.assert_structural_equal(tmem_buf.ty.effective_layout().canonicalize(),
expected)
return True
except (AssertionError, ValueError):
return False
@@ -471,10 +470,10 @@ def gemm_async_tcgen05_impl(op_call: TensorCall, sctx:
DispatchContext) -> Funct
f"got SFA_scope={SFA_scope}, SFB_scope={SFB_scope}"
)
SFA_type, SFB_type = _dtype_name(SFA_buffer.dtype),
_dtype_name(SFB_buffer.dtype)
- SFA_slice_layout = effective_layout(SFA_buffer).slice(
+ SFA_slice_layout = SFA_buffer.ty.effective_layout().slice(
SFA_buffer.shape, SFA_buffer_region.region
)
- SFB_slice_layout = effective_layout(SFB_buffer).slice(
+ SFB_slice_layout = SFB_buffer.ty.effective_layout().slice(
SFB_buffer.shape, SFB_buffer_region.region
)
SFA_elem_per_col = 32 // DataType(SFA_type).bits
@@ -519,9 +518,9 @@ def gemm_async_tcgen05_impl(op_call: TensorCall, sctx:
DispatchContext) -> Funct
C_st, C_extent = get_st_extent(C_buffer_region)
_, A_extent = get_st_extent(A_buffer_region)
_, B_extent = get_st_extent(B_buffer_region)
- A_slice_layout = effective_layout(A_buffer).slice(A_buffer.shape,
A_buffer_region.region)
- B_slice_layout = effective_layout(B_buffer).slice(B_buffer.shape,
B_buffer_region.region)
- C_slice_layout = effective_layout(C_buffer).slice(C_buffer.shape,
C_buffer_region.region)
+ A_slice_layout = A_buffer.ty.effective_layout().slice(A_buffer.shape,
A_buffer_region.region)
+ B_slice_layout = B_buffer.ty.effective_layout().slice(B_buffer.shape,
B_buffer_region.region)
+ C_slice_layout = C_buffer.ty.effective_layout().slice(C_buffer.shape,
C_buffer_region.region)
# Extract pre-swizzle tile layout for descriptor offset computation.
# strip_swizzle_to_tile rebuilds an identity tile over the slice extents
for
# a bare swizzle (ComposeLayout with a trivial tile), so .apply() maps over
@@ -723,7 +722,7 @@ def gemm_async_tcgen05_impl(op_call: TensorCall, sctx:
DispatchContext) -> Funct
# contiguous slice -- this is what enables fine K-major split-K.
cax = len(region) - 1 # innermost (contiguous) axis
elem_per_16b = 128 // DataType(dtype).bits
- phys_mode = get_swizzle_mode_from_layout(effective_layout(buf))
+ phys_mode = get_swizzle_mode_from_layout(buf.ty.effective_layout())
desc_region = list(region)
if phys_mode in (
SwizzleMode.SWIZZLE_128B_ATOM,
@@ -750,7 +749,7 @@ def gemm_async_tcgen05_impl(op_call: TensorCall, sctx:
DispatchContext) -> Funct
)
desc_region[cax] = tvm.ir.Range.from_min_extent(0, rounded)
- slice_layout = effective_layout(buf).slice(buf.shape, desc_region)
+ slice_layout = buf.ty.effective_layout().slice(buf.shape, desc_region)
# Strip unit dims to get the 2D matrix shape.
shape_2d = [int(r.extent) for r in desc_region if int(r.extent) != 1]
assert len(shape_2d) == 2, (
@@ -1015,7 +1014,7 @@ def gemm_async_tcgen05_impl(op_call: TensorCall, sctx:
DispatchContext) -> Funct
elif (
M == 64
and int(C_buffer.shape[0]) == 64
- and effective_layout(C_buffer) is not None
+ and C_buffer.ty.effective_layout() is not None
and _layout_matches_datapath_f(C_buffer)
):
base = tmem_datapath_layout("F", 64, N)
diff --git a/python/tvm/backend/cuda/tile/gemm_utils.py
b/python/tvm/backend/cuda/tile/gemm_utils.py
index 383328a698..0428fbaf55 100644
--- a/python/tvm/backend/cuda/tile/gemm_utils.py
+++ b/python/tvm/backend/cuda/tile/gemm_utils.py
@@ -21,7 +21,6 @@ from tvm.sym.analyzer import Analyzer
from tvm.tirx import Var
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
def validate_gemm_op(op_call: TensorCall, sctx: DispatchContext) -> bool:
@@ -31,7 +30,10 @@ def validate_gemm_op(op_call: TensorCall, sctx:
DispatchContext) -> bool:
A: Var = A_buffer_region.source
B: Var = B_buffer_region.source
if not (
- effective_layout(C) and effective_layout(A) and effective_layout(B)
and A.dtype == B.dtype
+ C.ty.effective_layout()
+ and A.ty.effective_layout()
+ and B.ty.effective_layout()
+ and A.dtype == B.dtype
):
return False
# Extract regions and validate dimensions
diff --git a/python/tvm/backend/trn/tile/binary/default.py
b/python/tvm/backend/trn/tile/binary/default.py
index a5270bc91c..6d3ee356ce 100644
--- a/python/tvm/backend/trn/tile/binary/default.py
+++ b/python/tvm/backend/trn/tile/binary/default.py
@@ -22,7 +22,6 @@ from tvm.tirx import FloatImm, Function
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.common import MapOpType
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import init_analyzer, nki_dim
from ..instruction_generator import InstructionGenerator
@@ -55,7 +54,7 @@ def binary_trn(op: TensorCall, binary_op: MapOpType, sctx:
DispatchContext) -> F
p_var = T.Var("P", "int32")
b_var = T.Var("B", "int32")
f_var = T.Var("F", "int32")
- p_size = effective_layout(dst).size("P")
+ p_size = dst.ty.effective_layout().size("P")
inst_size_limit = op.options.get("max_inst_size", 512)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
inst_gen.bind_inst_iter(_dst, p_var, p_size, 1, False)
diff --git a/python/tvm/backend/trn/tile/binary/utils.py
b/python/tvm/backend/trn/tile/binary/utils.py
index 8b9f932cd9..da57fa7f86 100644
--- a/python/tvm/backend/trn/tile/binary/utils.py
+++ b/python/tvm/backend/trn/tile/binary/utils.py
@@ -24,7 +24,6 @@ from tvm.ir import TensorRegion
from tvm.sym.analyzer import Analyzer
from tvm.tirx import FloatImm
from tvm.tirx.tile.common import MapOpType
-from tvm.tirx.tile.layout_utils import effective_layout
from ..dim_utils import get_ewise_dim_map
from ..instruction_generator import InstructionGenerator
@@ -74,9 +73,10 @@ def try_find_inst_nary(
valid_tensors = all(
[
- effective_layout(dst) and all(effective_layout(src) for src in
srcs if src is not None),
- is_trainium_layout(effective_layout(dst)),
- all(is_trainium_layout(effective_layout(src)) for src in srcs if
src is not None),
+ dst.ty.effective_layout()
+ and all(src.ty.effective_layout() for src in srcs if src is not
None),
+ is_trainium_layout(dst.ty.effective_layout()),
+ all(is_trainium_layout(src.ty.effective_layout()) for src in srcs
if src is not None),
dst.scope() == "trn.sbuf",
all(src.scope() in ["trn.sbuf", "trn.psum"] for src in srcs if src
is not None),
]
diff --git a/python/tvm/backend/trn/tile/compose_op/binary_chain.py
b/python/tvm/backend/trn/tile/compose_op/binary_chain.py
index e36cbd7976..9daa384c92 100644
--- a/python/tvm/backend/trn/tile/compose_op/binary_chain.py
+++ b/python/tvm/backend/trn/tile/compose_op/binary_chain.py
@@ -22,7 +22,6 @@ from tvm.script import tirx as T
from tvm.tirx import Function
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from tvm.tirx.tile.views import BinaryChain
from ..binary.utils import InstType, try_find_inst_nary
@@ -62,7 +61,7 @@ def binary_chain_trn(op: TensorCall, sctx: DispatchContext)
-> Function | None:
p_var = T.Var("P", "int32")
b_var = T.Var("B", "int32")
f_var = T.Var("F", "int32")
- p_size = effective_layout(output.source).size("P")
+ p_size = output.source.ty.effective_layout().size("P")
inst_size_limit = op.options.get("max_inst_size", 512)
inst_repr.bound_inst_size(inst_size_limit, analyzer)
inst_gen.bind_inst_iter(output, p_var, p_size, 1, False)
diff --git a/python/tvm/backend/trn/tile/compose_op/binary_reduce.py
b/python/tvm/backend/trn/tile/compose_op/binary_reduce.py
index 434e630a60..cbf0e68942 100644
--- a/python/tvm/backend/trn/tile/compose_op/binary_reduce.py
+++ b/python/tvm/backend/trn/tile/compose_op/binary_reduce.py
@@ -22,7 +22,6 @@ from tvm.script import tirx as T
from tvm.tirx import Function
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from tvm.tirx.tile.views import BinaryReduce
from ..binary.utils import InstType, try_find_inst_nary
@@ -82,7 +81,7 @@ def binary_reduce_trn(op: TensorCall, sctx: DispatchContext)
-> Function | None:
f_var = T.Var("F", "int32")
reduction_b_var = T.Var("rB", "int32")
spatial_b_var = T.Var("sB", "int32")
- p_size = effective_layout(binary_output.source).size("P")
+ p_size = binary_output.source.ty.effective_layout().size("P")
inst_gen.bind_inst_iter(binary_output, p_var, p_size, 1, False)
inst_gen.bind_inst_iter(binary_output, f_var, inst_repr.size,
inst_repr.stride, True)
reduction_b_extent = inst_gen.fill_in_block_dim(binary_output,
reduction_b_var, reduce_axes)
diff --git a/python/tvm/backend/trn/tile/compose_op/unary_reduce.py
b/python/tvm/backend/trn/tile/compose_op/unary_reduce.py
index 9e5eb3df9a..8b14a6c60c 100644
--- a/python/tvm/backend/trn/tile/compose_op/unary_reduce.py
+++ b/python/tvm/backend/trn/tile/compose_op/unary_reduce.py
@@ -22,7 +22,6 @@ from tvm.script import tirx as T
from tvm.tirx import Function
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from tvm.tirx.tile.views import UnaryReduce
from ..binary.utils import try_find_inst_nary
@@ -80,7 +79,7 @@ def unary_reduce_trn(op: TensorCall, sctx: DispatchContext)
-> Function | None:
f_var = T.Var("F", "int32")
reduction_b_var = T.Var("rB", "int32")
spatial_b_var = T.Var("sB", "int32")
- p_size = effective_layout(unary_output.source).size("P")
+ p_size = unary_output.source.ty.effective_layout().size("P")
inst_gen.bind_inst_iter(unary_output, p_var, p_size, 1, False)
inst_gen.bind_inst_iter(unary_output, f_var, inst_repr.size,
inst_repr.stride, True)
reduction_b_extent = inst_gen.fill_in_block_dim(unary_output,
reduction_b_var, reduce_axes)
diff --git a/python/tvm/backend/trn/tile/copy/default.py
b/python/tvm/backend/trn/tile/copy/default.py
index b5b51ef754..264ee406ae 100644
--- a/python/tvm/backend/trn/tile/copy/default.py
+++ b/python/tvm/backend/trn/tile/copy/default.py
@@ -22,7 +22,6 @@ from tvm.script import tirx as T
from tvm.tirx import Function
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import init_analyzer, nki_dim
from ..dim_utils import get_ewise_dim_map
@@ -41,19 +40,19 @@ def copy_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
# Check for valid tensor configurations
valid_config = all(
[
- effective_layout(src) and effective_layout(dst),
+ src.ty.effective_layout() and dst.ty.effective_layout(),
src.scope() in ["global", "trn.sbuf", "trn.psum"],
dst.scope() in ["global", "trn.sbuf", "trn.psum"],
src.scope() != "global" or dst.scope() != "global",
- (src.scope() == "global" and isinstance(effective_layout(src),
T.TileLayout))
+ (src.scope() == "global" and isinstance(src.ty.effective_layout(),
T.TileLayout))
or (
src.scope() in ["trn.sbuf", "trn.psum"]
- and is_trainium_layout(effective_layout(src))
+ and is_trainium_layout(src.ty.effective_layout())
),
- (dst.scope() == "global" and isinstance(effective_layout(dst),
T.TileLayout))
+ (dst.scope() == "global" and isinstance(dst.ty.effective_layout(),
T.TileLayout))
or (
dst.scope() in ["trn.sbuf", "trn.psum"]
- and is_trainium_layout(effective_layout(dst))
+ and is_trainium_layout(dst.ty.effective_layout())
),
]
)
@@ -85,7 +84,7 @@ def copy_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
"prepare identity and use matmul explicitly"
)
- if is_trainium_layout(effective_layout(src)):
+ if is_trainium_layout(src.ty.effective_layout()):
inst = inst_gen.find_max_inst_size_from_one_region(src_region)
inst = inst_gen.fit_inst_tile_to_region(inst, dst_region)
src_to_dst = True
@@ -114,7 +113,7 @@ def copy_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
from_region, _to_region = src_region, dst_region
else:
from_region, _to_region = dst_region, src_region
- p_size = effective_layout(from_region.source).size("P")
+ p_size = from_region.source.ty.effective_layout().size("P")
inst_gen.bind_inst_iter(from_region, p_var, p_size, 1, is_free_dim=False)
inst_gen.bind_inst_iter(from_region, f_var, inst.size, inst.stride,
is_free_dim=True)
b_extent = inst_gen.fill_in_block_dim(from_region, b_var)
diff --git a/python/tvm/backend/trn/tile/gemm/default.py
b/python/tvm/backend/trn/tile/gemm/default.py
index dbc9dbd888..8ee8672b18 100644
--- a/python/tvm/backend/trn/tile/gemm/default.py
+++ b/python/tvm/backend/trn/tile/gemm/default.py
@@ -27,7 +27,6 @@ from tvm.sym.analyzer import Analyzer
from tvm.tirx import Function
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import init_analyzer
from ..dim_utils import normalize_and_group
@@ -57,7 +56,7 @@ def get_pf_dim_from_tensor_region(
assert len(non_unit_dims) == 2, "Only 2D matrix is supported for gemm"
layout, seps = normalize_and_group(
- effective_layout(tensor_region.source), tensor_region.source.ty.shape
+ tensor_region.source.ty.effective_layout(),
tensor_region.source.ty.shape
)
# Determine partition and free dimensions based on operator kind
if operator_kind == OperatorKind.A:
@@ -142,19 +141,19 @@ def matmul_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
# Validate tensor properties
assert all(
[
- effective_layout(A) and effective_layout(B) and
effective_layout(C),
+ A.ty.effective_layout() and B.ty.effective_layout() and
C.ty.effective_layout(),
A.ty.dtype == B.ty.dtype,
A.scope() == "trn.sbuf" and B.scope() == "trn.sbuf",
C.scope() == "trn.psum" or C.scope() == "trn.sbuf",
- is_trainium_layout(effective_layout(A)),
- is_trainium_layout(effective_layout(B)),
- is_trainium_layout(effective_layout(C)),
- effective_layout(A).size("P") == effective_layout(B).size("P"),
+ is_trainium_layout(A.ty.effective_layout()),
+ is_trainium_layout(B.ty.effective_layout()),
+ is_trainium_layout(C.ty.effective_layout()),
+ A.ty.effective_layout().size("P") ==
B.ty.effective_layout().size("P"),
]
), "Invalid tensor layout and scope"
- p_size = effective_layout(A).size("P")
- assert p_size == effective_layout(B).size("P"), "Partition size mismatch"
+ p_size = A.ty.effective_layout().size("P")
+ assert p_size == B.ty.effective_layout().size("P"), "Partition size
mismatch"
# Get partition and free dimensions
lhs_p_dim, lhs_f_dim = get_pf_dim_from_tensor_region(
@@ -204,12 +203,14 @@ def matmul_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
reduction_b = T.Var("reduction_b", "int32")
lhs_b = T.Var("lhs_b", "int32")
rhs_b = T.Var("rhs_b", "int32")
- lhs_f_size = effective_layout(C).size("P")
+ lhs_f_size = C.ty.effective_layout().size("P")
inst_gen.bind_inst_iter(
B_tensor_region, rhs_f, inst_repr.size, inst_repr.stride,
is_free_dim=True
)
inst_gen.bind_inst_iter(C_tensor_region, lhs_f, lhs_f_size, 1,
is_free_dim=False)
- inst_gen.bind_inst_iter(A_tensor_region, p, effective_layout(A).size("P"),
1, is_free_dim=False)
+ inst_gen.bind_inst_iter(
+ A_tensor_region, p, A.ty.effective_layout().size("P"), 1,
is_free_dim=False
+ )
reduction_b_extent = inst_gen.fill_in_block_dim(A_tensor_region,
reduction_b, [lhs_p_dim])
lhs_b_extent = inst_gen.fill_in_block_dim(A_tensor_region, lhs_b,
[lhs_f_dim])
rhs_b_extent = inst_gen.fill_in_block_dim(B_tensor_region, rhs_b,
[rhs_f_dim])
diff --git a/python/tvm/backend/trn/tile/instruction_generator.py
b/python/tvm/backend/trn/tile/instruction_generator.py
index 41bdef401a..495e82ebe6 100644
--- a/python/tvm/backend/trn/tile/instruction_generator.py
+++ b/python/tvm/backend/trn/tile/instruction_generator.py
@@ -25,14 +25,12 @@ from operator import mul
import tvm_ffi
-import tvm
from tvm.backend.trn.layout import is_trainium_layout
from tvm.ir import Range, TensorRegion, TensorRegionType
from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
from tvm.tirx import Expr, IntImm, Var, is_tensor_var
-from tvm.tirx.layout import Iter
-from tvm.tirx.tile.layout_utils import effective_layout
+from tvm.tirx.layout import Iter, TileLayout
from .dim_utils import DimensionMapper, RangeInfo, normalize_and_group
@@ -138,7 +136,7 @@ class InstructionGenerator:
return tensor_region
def _get_sub_layout(self, tensor_region: TensorRegion):
- layout = effective_layout(tensor_region.source)
+ layout = tensor_region.source.ty.effective_layout()
layout, seps = normalize_and_group(layout,
tensor_region.source.ty.shape)
tiled_range_infos_per_dim = []
new_shard = []
@@ -173,9 +171,7 @@ class InstructionGenerator:
assert False, f"Cannot analyze physical tensor region for:
{tensor_region}"
new_shard += reversed(reversed_shard)
new_seps.append(len(reversed_shard) + new_seps[-1])
- new_tile_layout = effective_layout(tvm.tirx).TileLayout.from_iters( #
pylint: disable=no-member
- new_shard, [], dict()
- )
+ new_tile_layout = TileLayout.from_iters(new_shard, [], dict())
return new_tile_layout, new_seps
def _init_bind_iters(self):
diff --git a/python/tvm/backend/trn/tile/private_alloc.py
b/python/tvm/backend/trn/tile/private_alloc.py
index 7176d870c3..3914d68337 100644
--- a/python/tvm/backend/trn/tile/private_alloc.py
+++ b/python/tvm/backend/trn/tile/private_alloc.py
@@ -24,7 +24,6 @@ from tvm.tirx import FloatImm, IntImm, Var
from tvm.tirx.tile import DispatchContext
from tvm.tirx.tile.dispatch import f_op_dispatcher
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from tvm.tirx.tile.views import (
BinaryReduce,
Exp,
@@ -72,7 +71,7 @@ def alloc_const_bias_trn(
if not isinstance(bias, (FloatImm)):
return {}
bias_key = ("const_bias", _scalar_dtype(bias), float(bias.value).hex())
- par_size = effective_layout(op.dsts[0].source).size("P")
+ par_size = op.dsts[0].source.ty.effective_layout().size("P")
max_inst_size = op.options.get("max_inst_size", 512)
if isinstance(max_inst_size, int | IntImm) and int(max_inst_size) == -1:
raise ValueError("Constant bias workspace allocation requires a finite
max_inst_size")
@@ -126,7 +125,7 @@ def alloc_acc_psum_trn(
) -> dict[str, Any]:
if "acc_psum" in op.workspaces or op.dsts[0].source.scope() == "trn.psum":
return {}
- par_size = effective_layout(op.dsts[0].source).size("P")
+ par_size = op.dsts[0].source.ty.effective_layout().size("P")
acc_psum = T.Var("acc_psum", T.Tensor((8, par_size, 512), "float32",
scope="trn.psum"))
# no reuse opportunity
tensor_dict[acc_psum] = (
diff --git a/python/tvm/backend/trn/tile/reduction/utils.py
b/python/tvm/backend/trn/tile/reduction/utils.py
index 8e04ee1ef4..7d7763571a 100644
--- a/python/tvm/backend/trn/tile/reduction/utils.py
+++ b/python/tvm/backend/trn/tile/reduction/utils.py
@@ -23,7 +23,6 @@ from tvm.tirx import Function
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.common import ReduceOpType
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import init_analyzer, nki_dim
from ..dim_utils import get_reduction_dim_map
@@ -41,7 +40,7 @@ def generate_intermediate_tensor(
Returns:
Tuple[Optional[tensor], int]: The intermediate tensor and reduction
factor size.
"""
- intermediate_shape =
[effective_layout(dst_tensor_region.source).size("P"), rfactor_size]
+ intermediate_shape =
[dst_tensor_region.source.ty.effective_layout().size("P"), rfactor_size]
if "partial_reduce" in workspace:
intermediate_tensor = workspace["partial_reduce"]
@@ -92,12 +91,12 @@ def reduction_trn(
# Layout validation
assert all(
[
- effective_layout(src) and effective_layout(dst),
+ src.ty.effective_layout() and dst.ty.effective_layout(),
src.scope() == "trn.sbuf" or src.scope() == "trn.psum",
dst.scope() == "trn.sbuf",
- is_trainium_layout(effective_layout(src)),
- is_trainium_layout(effective_layout(dst)),
- effective_layout(src).size("P") == effective_layout(dst).size("P"),
+ is_trainium_layout(src.ty.effective_layout()),
+ is_trainium_layout(dst.ty.effective_layout()),
+ src.ty.effective_layout().size("P") ==
dst.ty.effective_layout().size("P"),
]
), "Invalid layout"
@@ -110,7 +109,7 @@ def reduction_trn(
assert analyzer.can_prove(inst_repr.size > 1), "Instruction size must be
greater than 1"
# Get partition size and extents
- p_size = effective_layout(src).size("P")
+ p_size = src.ty.effective_layout().size("P")
f_var = T.Var("F", "int32")
p_var = T.Var("P", "int32")
spatial_b_var = T.Var("sB", "int32")
diff --git a/python/tvm/backend/trn/tile/select/default.py
b/python/tvm/backend/trn/tile/select/default.py
index b21db5cc97..ec0cb0500e 100644
--- a/python/tvm/backend/trn/tile/select/default.py
+++ b/python/tvm/backend/trn/tile/select/default.py
@@ -23,7 +23,6 @@ from tvm.script import tirx as T
from tvm.tirx import FloatImm, Function
from tvm.tirx.tile import DispatchContext, fail
from tvm.tirx.tile.instruction import TensorCall
-from tvm.tirx.tile.layout_utils import effective_layout
from tvm.tirx.tile.views import Select
from ..common import init_analyzer, nki_dim
@@ -60,10 +59,10 @@ def select_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
# Validate tensor layout and scope
tensor_conditions = [
- effective_layout(dst.source) and effective_layout(true_value.source),
+ dst.source.ty.effective_layout() and
true_value.source.ty.effective_layout(),
dst.source.scope() == "trn.sbuf" and true_value.source.scope() ==
"trn.sbuf",
- is_trainium_layout(effective_layout(true_value.source)),
- is_trainium_layout(effective_layout(dst.source)),
+ is_trainium_layout(true_value.source.ty.effective_layout()),
+ is_trainium_layout(dst.source.ty.effective_layout()),
]
if not all(tensor_conditions):
@@ -96,7 +95,7 @@ def select_trn(op: TensorCall, sctx: DispatchContext) ->
Function | None:
p_var = T.Var("p", "int32")
b_var = T.Var("b", "int32")
f_var = T.Var("f", "int32")
- p_size = effective_layout(dst.source).size("P")
+ p_size = dst.source.ty.effective_layout().size("P")
inst_gen.bind_inst_iter(dst, f_var, inst_repr.size, inst_repr.stride, True)
inst_gen.bind_inst_iter(dst, p_var, p_size, 1, False)
b_extent = inst_gen.fill_in_block_dim(dst, b_var)
diff --git a/python/tvm/backend/trn/tile/unary/utils.py
b/python/tvm/backend/trn/tile/unary/utils.py
index 029fe42885..7fad0514b7 100644
--- a/python/tvm/backend/trn/tile/unary/utils.py
+++ b/python/tvm/backend/trn/tile/unary/utils.py
@@ -23,7 +23,6 @@ from tvm.script import tirx as T
from tvm.sym.analyzer import Analyzer
from tvm.tirx import FloatImm
from tvm.tirx.tile.common import MapOpType
-from tvm.tirx.tile.layout_utils import effective_layout
from ..common import nki_dim
from ..dim_utils import get_ewise_dim_map
@@ -56,11 +55,11 @@ def try_find_inst_unary(
# Validate tensor layouts and scopes
valid_layout_scope = all(
[
- effective_layout(src) and effective_layout(dst),
+ src.ty.effective_layout() and dst.ty.effective_layout(),
src.scope() in ("trn.sbuf", "trn.psum"),
dst.scope() == "trn.sbuf",
- is_trainium_layout(effective_layout(src)),
- is_trainium_layout(effective_layout(dst)),
+ is_trainium_layout(src.ty.effective_layout()),
+ is_trainium_layout(dst.ty.effective_layout()),
]
)
@@ -139,7 +138,7 @@ def generate_unary_func(
):
"""Generate a function that implements a unary operation."""
# Prepare parameters
- p_size = effective_layout(dst_tensor_region.source).size("P")
+ p_size = dst_tensor_region.source.ty.effective_layout().size("P")
# Apply instruction size limits if specified
inst_size_limit = config.get("max_inst_size", 512)
diff --git a/python/tvm/tirx/lang/__init__.py b/python/tvm/tirx/lang/__init__.py
deleted file mode 100644
index 13a83393a9..0000000000
--- a/python/tvm/tirx/lang/__init__.py
+++ /dev/null
@@ -1,16 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
diff --git a/python/tvm/tirx/lang/alloc_pool.py
b/python/tvm/tirx/lang/alloc_pool.py
deleted file mode 100644
index b45bcef60c..0000000000
--- a/python/tvm/tirx/lang/alloc_pool.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0.
-"""Compatibility redirect for CUDA allocation pool helpers."""
-
-from tvm.backend.cuda.lang.alloc_pool import * # noqa: F403 # pylint:
disable=wildcard-import
diff --git a/python/tvm/tirx/lang/pipeline.py b/python/tvm/tirx/lang/pipeline.py
deleted file mode 100644
index ee6380f818..0000000000
--- a/python/tvm/tirx/lang/pipeline.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0.
-"""Compatibility redirect for CUDA pipeline helpers."""
-
-from tvm.backend.cuda.lang.pipeline import * # noqa: F403 # pylint:
disable=wildcard-import
diff --git a/python/tvm/tirx/lang/smem_desc.py
b/python/tvm/tirx/lang/smem_desc.py
deleted file mode 100644
index d4a5f9686d..0000000000
--- a/python/tvm/tirx/lang/smem_desc.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0.
-"""Compatibility redirect for CUDA shared-memory descriptors."""
-
-from tvm.backend.cuda.lang.smem_desc import * # noqa: F403 # pylint:
disable=wildcard-import
diff --git a/python/tvm/tirx/lang/tile_scheduler.py
b/python/tvm/tirx/lang/tile_scheduler.py
deleted file mode 100644
index 145b17581c..0000000000
--- a/python/tvm/tirx/lang/tile_scheduler.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0.
-"""Compatibility redirect for CUDA tile schedulers."""
-
-from tvm.backend.cuda.lang.tile_scheduler import * # noqa: F403 # pylint:
disable=wildcard-import
diff --git a/python/tvm/tirx/lang/warp_role.py
b/python/tvm/tirx/lang/warp_role.py
deleted file mode 100644
index bb5927fded..0000000000
--- a/python/tvm/tirx/lang/warp_role.py
+++ /dev/null
@@ -1,8 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0.
-"""Compatibility redirect for CUDA warp role helpers."""
-
-from tvm.backend.cuda.lang.warp_role import * # noqa: F403 # pylint:
disable=wildcard-import
diff --git a/python/tvm/tirx/operator/__init__.py
b/python/tvm/tirx/operator/__init__.py
deleted file mode 100644
index 3db58184cc..0000000000
--- a/python/tvm/tirx/operator/__init__.py
+++ /dev/null
@@ -1,18 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-
-"""Low-level operator helpers."""
diff --git a/python/tvm/tirx/operator/intrinsics/_common.py
b/python/tvm/tirx/operator/intrinsics/_common.py
deleted file mode 100644
index 9e2deea574..0000000000
--- a/python/tvm/tirx/operator/intrinsics/_common.py
+++ /dev/null
@@ -1,59 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Shared enum / value tables for the PTX intrinsic user wrappers.
-
-Single source of truth for the modifier values ``tvm.tirx.op`` accepts, which
-it validates via ``_choice``. The ``T.ptx`` dialect does not read these -- it
-derives its legal modifier tokens from its own instruction table.
-
-Adding a new modifier value requires changing exactly one place.
-"""
-
-# Memory ordering / scope -----------------------------------------------------
-CLUSTER_BARRIER_SEM = ("", "release", "relaxed")
-MBARRIER_COMPLETE_TX_SEM = ("relaxed",)
-MBARRIER_COMPLETE_TX_SCOPE = ("cta", "cluster")
-MBARRIER_COMPLETE_TX_SPACE = ("shared", "shared::cta", "shared::cluster")
-MBARRIER_ARRIVE_SEM = ("", "release", "relaxed")
-MBARRIER_ARRIVE_SCOPE = ("", "cta", "cluster")
-MBARRIER_ARRIVE_SPACE = ("shared", "shared::cta", "shared::cluster")
-
-# CTA group (used by tcgen05 and TMA) -----------------------------------------
-TCGEN05_CTA_GROUP = (1, 2)
-
-# NVSHMEM ---------------------------------------------------------------------
-NVSHMEM_CMP = ("eq", "ne", "gt", "ge", "lt", "le")
-NVSHMEM_SIG_OP = ("set", "add")
-
-# Floating-point rounding -----------------------------------------------------
-F32X2_ROUND = ("rz", "rn", "rm", "rp")
-
-# cp.async (non-bulk) ---------------------------------------------------------
-CP_ASYNC_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal")
-CP_ASYNC_PREFETCH_SIZE = (-1, 64, 128, 256)
-CP_ASYNC_FILL_MODE = ("", "zero")
-
-# cp.async.bulk (TMA) ---------------------------------------------------------
-CP_ASYNC_BULK_CACHE_HINT = ("", "evict_last", "evict_first", "evict_normal",
"evict_last_use")
-CP_ASYNC_BULK_RED_OP = ("add", "min", "max", "inc", "dec", "and", "or", "xor")
-
-# ldmatrix / stmatrix ---------------------------------------------------------
-LDMATRIX_DTYPE = (".b16", ".b8")
-LDMATRIX_NUM = (1, 2, 4)
-
-# tcgen05.ld / tcgen05.st -----------------------------------------------------
-TCGEN05_LDST_SHAPES = ("16x32bx2", "16x64b", "16x128b", "16x256b", "32x32b")
diff --git a/python/tvm/tirx/script/ir_builder/__init__.py
b/python/tvm/tirx/script/ir_builder/__init__.py
index 79007c81c8..e872ab7a85 100644
--- a/python/tvm/tirx/script/ir_builder/__init__.py
+++ b/python/tvm/tirx/script/ir_builder/__init__.py
@@ -36,19 +36,6 @@ from .utils import (
tensor_indices as tensor_indices,
)
-
-def SMEMPool(*args, **kwargs):
- from tvm.backend.cuda.lang.alloc_pool import SMEMPool as pool
-
- return pool(*args, **kwargs)
-
-
-def TMEMPool(*args, **kwargs):
- from tvm.backend.cuda.lang.alloc_pool import TMEMPool as pool
-
- return pool(*args, **kwargs)
-
-
supports_mutable_declarations = True
__all__ = [
@@ -58,8 +45,6 @@ __all__ = [
"dynamic",
"dtype",
"constexpr",
- "SMEMPool",
- "TMEMPool",
"bool",
"tensor_indices",
"frame_scope",
diff --git a/python/tvm/tirx/tensor.py b/python/tvm/tirx/tensor.py
index b193d5a963..8782348db5 100644
--- a/python/tvm/tirx/tensor.py
+++ b/python/tvm/tirx/tensor.py
@@ -83,6 +83,33 @@ class TensorType(Type):
"sub": lambda ty, expr: ty._sub(expr),
}
+ def effective_layout(self):
+ """Return the explicit layout or infer a compact global layout.
+
+ Inference uses the full tensor shape before a region is sliced and
+ leaves this type unchanged. Other storage scopes and noncompact global
+ tensors require an explicit layout.
+ """
+ if self.layout is not None:
+ return self.layout
+ if self.storage_scope != "global":
+ raise ValueError(f"An explicit layout is required for
{self.storage_scope} tensors")
+
+ from tvm.sym import Analyzer
+
+ from .layout import S, TileLayout
+
+ analyzer = Analyzer()
+ expected = 1
+ if self.strides:
+ for extent, stride in reversed(list(zip(self.shape,
self.strides))):
+ if not analyzer.can_prove_equal(extent, 1) and not
analyzer.can_prove_equal(
+ stride, expected
+ ):
+ raise ValueError("An explicit layout is required for
noncompact global tensors")
+ expected = analyzer.simplify(expected * extent)
+ return TileLayout(S[tuple(self.shape)])
+
def scope(self, expr):
"""Return the storage scope associated with this tensor.
Returns
diff --git a/python/tvm/tirx/tile/layout_utils.py
b/python/tvm/tirx/tile/layout_utils.py
deleted file mode 100644
index 1c1c838e77..0000000000
--- a/python/tvm/tirx/tile/layout_utils.py
+++ /dev/null
@@ -1,47 +0,0 @@
-# Licensed to the Apache Software Foundation (ASF) under one
-# or more contributor license agreements. See the NOTICE file
-# distributed with this work for additional information
-# regarding copyright ownership. The ASF licenses this file
-# to you under the Apache License, Version 2.0 (the
-# "License"); you may not use this file except in compliance
-# with the License. You may obtain a copy of the License at
-#
-# http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing,
-# software distributed under the License is distributed on an
-# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
-# KIND, either express or implied. See the License for the
-# specific language governing permissions and limitations
-# under the License.
-"""Layouts inferred only while lowering a tile instruction."""
-
-from tvm.sym import Analyzer
-from tvm.tirx.layout import S, TileLayout
-
-
-def effective_layout(tensor):
- """Return an explicit layout or the compact layout of a global tensor.
-
- Inference uses the complete source shape before a region is sliced. The
- source type and its pointer, strides and element offset remain unchanged.
- Other storage scopes and noncompact global tensors require an explicit
layout.
- """
- if tensor.layout is not None:
- return tensor.layout
- if tensor.scope() != "global":
- raise ValueError(
- f"tile instruction requires an explicit layout for
{tensor.scope()} tensor"
- )
- analyzer = Analyzer()
- expected = 1
- if tensor.strides:
- for extent, stride in reversed(list(zip(tensor.shape,
tensor.strides))):
- if not analyzer.can_prove_equal(extent, 1) and not
analyzer.can_prove_equal(
- stride, expected
- ):
- raise ValueError(
- "tile instruction requires an explicit layout for a
noncompact global tensor"
- )
- expected = analyzer.simplify(expected * extent)
- return TileLayout(S[tuple(tensor.shape)])
diff --git a/tests/python/tirx/codegen/test_codegen_blackwell.py
b/tests/python/tirx/codegen/test_codegen_blackwell.py
index 8ffbf0fdf7..9a3a80bef7 100644
--- a/tests/python/tirx/codegen/test_codegen_blackwell.py
+++ b/tests/python/tirx/codegen/test_codegen_blackwell.py
@@ -24,6 +24,7 @@ import tvm_ffi
import tvm
import tvm.testing
from tvm.backend.cuda.cpp.descriptors import _get_tcgen05_mma_kind
+from tvm.backend.cuda.lang import SMEMPool
from tvm.script import tirx as T
from tvm.testing import env
@@ -156,13 +157,13 @@ def test_mbarrier_try_wait_once_codegen():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
def test_mbarrier_remote_view_codegen():
- from tvm.tirx.lang.pipeline import MBarrier
+ from tvm.backend.cuda.lang.pipeline import MBarrier
# fmt: off
@T.function
def test_remote_view():
T.device_entry(launch=T.cuda.LaunchConfig(grid=1 * 2, block=(128,),
cluster=(2,)))
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = MBarrier(pool, 1)
pool.commit()
remote_bar = bar.remote_view(0)
@@ -188,14 +189,14 @@ def test_mbarrier_remote_view_codegen():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
def test_mbarrier_local_arrive_forwards_predicate_and_count():
- from tvm.tirx.lang.pipeline import MBarrier
+ from tvm.backend.cuda.lang.pipeline import MBarrier
# fmt: off
@T.function
def test_local_arrive():
T.device_entry(launch=T.cuda.LaunchConfig(grid=1, block=(32,)))
thread = T.cuda.thread_idx('x')
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = MBarrier(pool, 1)
pool.commit()
bar.arrive(0, pred=(thread == 0), count=2)
@@ -221,13 +222,13 @@ def
test_mbarrier_local_arrive_forwards_predicate_and_count():
@pytest.mark.gpu
@pytest.mark.skipif(not env.has_cuda_compute(10), reason="need cuda compute >=
10.0")
def test_tma_mbarrier_remote_view_codegen():
- from tvm.tirx.lang.pipeline import TMABar
+ from tvm.backend.cuda.lang.pipeline import TMABar
# fmt: off
@T.function
def test_remote_view():
T.device_entry(launch=T.cuda.LaunchConfig(grid=1 * 2, block=(128,),
cluster=(2,)))
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = TMABar(pool, 1)
pool.commit()
remote_bar = bar.remote_view(0)
@@ -245,14 +246,14 @@ def test_tma_mbarrier_remote_view_codegen():
def test_mbarrier_remote_view_rejects_invalid_operations():
- from tvm.tirx.lang.pipeline import MBarrier, TMABar
+ from tvm.backend.cuda.lang.pipeline import MBarrier, TMABar
with pytest.raises(ValueError, match=r"remote_view\(\) cannot be
initialized"):
# fmt: off
@T.function
def invalid_init():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(2,), block=(128,)))
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = MBarrier(pool, 1)
bar.remote_view(0).init(1)
# fmt: on
@@ -262,7 +263,7 @@ def test_mbarrier_remote_view_rejects_invalid_operations():
@T.function
def invalid_wait():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(2,), block=(128,)))
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = MBarrier(pool, 1)
bar.remote_view(0).wait(0, 0)
# fmt: on
@@ -272,7 +273,7 @@ def test_mbarrier_remote_view_rejects_invalid_operations():
@T.function
def ambiguous_mbarrier_arrive():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(2,), block=(128,)))
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = MBarrier(pool, 1)
bar.remote_view(0).arrive(0, remote=1)
# fmt: on
@@ -282,7 +283,7 @@ def test_mbarrier_remote_view_rejects_invalid_operations():
@T.function
def ambiguous_tma_arrive():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(2,), block=(128,)))
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = TMABar(pool, 1)
bar.remote_view(0).arrive(0, tx_count=128, remote=1)
# fmt: on
@@ -295,7 +296,7 @@ def test_mbarrier_remote_view_rejects_invalid_operations():
@T.function
def nested_remote_view():
T.device_entry(launch=T.cuda.LaunchConfig(grid=(2,), block=(128,)))
- pool = T.SMEMPool()
+ pool = SMEMPool()
bar = MBarrier(pool, 1)
bar.remote_view(0).remote_view(1)
# fmt: on
diff --git a/tests/python/tirx/script/test_tirx_script_gpu.py
b/tests/python/tirx/script/test_tirx_script_gpu.py
index 12735ecb04..c853523243 100644
--- a/tests/python/tirx/script/test_tirx_script_gpu.py
+++ b/tests/python/tirx/script/test_tirx_script_gpu.py
@@ -431,70 +431,6 @@ def test_roundtrip_persistent_not_present():
assert "persistent" not in code, f"persistent should NOT appear:\n{code}"
-def test_warp_role():
- """WarpRole should emit guarded warp scopes plus setmaxnreg."""
- from tvm.tirx.lang.warp_role import WarpRole
-
- # fmt: off
- @T.function
- def test(A: T.Tensor((128,), 'float32', scope='global')) -> None:
-
- T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=4 * 128))
- cta_id = T.cuda.block_idx('x')
- wg_id = T.cuda.warpgroup_id()
- warp_id = T.cuda.warp_in_warpgroup()
- lane_id = T.cuda.lane_id()
- with WarpRole(warp_id, 1, regs=48):
- T.cuda.tile.mov(A[0:32], T.float32(0), scope='cta')
- with WarpRole(warp_id, 0, regs=232, increase=True):
- T.cuda.tile.mov(A[32:64], T.float32(1), scope='cta')
- # fmt: on
-
- code = test.script()
- warp_name = next(
- line.partition(":")[0].strip()
- for line in code.splitlines()
- if " = T.cuda.warp_in_warpgroup()" in line
- )
- assert f"{warp_name} == 1" in code, f"should have warp_id==1
guard:\n{code}"
- assert f"{warp_name} == 0" in code, f"should have warp_id==0
guard:\n{code}"
- assert "setmaxnreg" in code, f"should have setmaxnreg:\n{code}"
- assert f"if {warp_name} == 1:" in code, f"should have warp_id==1
if-guard:\n{code}"
- assert f"if {warp_name} == 0:" in code, f"should have warp_id==0
if-guard:\n{code}"
- # The printed code is valid TIR — it should parse back
- assert from_source(code).script() == code
- assert_structural_equal(test, from_source(code))
-
-
-def test_warpgroup_role():
- """WarpgroupRole should emit guarded warpgroup scope plus setmaxnreg."""
- from tvm.tirx.lang.warp_role import WarpgroupRole
-
- # fmt: off
- @T.function
- def test(A: T.Tensor((128,), 'float32', scope='global')) -> None:
-
- T.device_entry(launch=T.cuda.LaunchConfig(grid=(1,), block=4 * 128))
- cta_id = T.cuda.block_idx('x')
- wg_id = T.cuda.warpgroup_id()
- warp_id_in_wg = T.cuda.warp_in_warpgroup()
- lane_id = T.cuda.lane_id()
- with WarpgroupRole(wg_id, 2, regs=200, increase=True):
- T.cuda.tile.mov(A[0:32], T.float32(0), scope='cta')
- # fmt: on
-
- code = test.script()
- group_name = next(
- line.partition(":")[0].strip()
- for line in code.splitlines()
- if " = T.cuda.warpgroup_id()" in line
- )
- assert f"{group_name} == 2" in code, f"should have wg_id==2 guard:\n{code}"
- assert "setmaxnreg" in code, f"should have setmaxnreg:\n{code}"
- assert from_source(code).script() == code
- assert_structural_equal(test, from_source(code))
-
-
def test_roundtrip_cuda_func_call_source_code():
"""cuda_func_call with a multiline trailing source operand must print with
inline string literal, not as a metadata reference."""
diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py
index 093efbf4b9..9d57bccc46 100644
--- a/tests/python/tirx/test_layout.py
+++ b/tests/python/tirx/test_layout.py
@@ -24,13 +24,13 @@ import pytest
import tvm
from tvm.backend.cuda.axis import Axis as CUDAAxis
from tvm.backend.cuda.axis import TCol, TLane, laneid, tid_in_wg, tx, warpid,
wgid, wid_in_wg
+from tvm.backend.cuda.lang import SMEMPool
from tvm.backend.cuda.layout import tmem_mma_operand_layout, wg_local_layout
from tvm.backend.trn.axis import Axis as TRNAxis
from tvm.backend.trn.axis import F, P
from tvm.backend.trn.layout import to_psum_layout, trainium_layout
from tvm.ir import assert_structural_equal
from tvm.ir.type import PrimType, PtrType
-from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.sym import Analyzer
from tvm.tirx import Var
@@ -1478,7 +1478,7 @@ def test_pool_allocator_alloc_mma():
def alloc_layout(shape, dtype, swizzle_mode="auto"):
with IRBuilder():
with Tx_builder.function():
- pool = T.SMEMPool(Var("smem_ptr", PtrType(PrimType("uint8"))))
+ pool = SMEMPool(Var("smem_ptr", PtrType(PrimType("uint8"))))
buf = pool.alloc_tcgen05_mma_AB(shape, dtype,
swizzle_mode=swizzle_mode)
return buf.layout
diff --git a/tests/python/tirx/test_tile_namespace.py
b/tests/python/tirx/test_tile_namespace.py
index 267c5a582f..29f20dacbf 100644
--- a/tests/python/tirx/test_tile_namespace.py
+++ b/tests/python/tirx/test_tile_namespace.py
@@ -31,7 +31,6 @@ from tvm.tirx import decl_tensor
from tvm.tirx.exec_scope import ExecScope
from tvm.tirx.layout import Axis, S, TileLayout
from tvm.tirx.script import ir_builder as I
-from tvm.tirx.tile.layout_utils import effective_layout
def test_generic_layout_without_backend_autoload():
@@ -49,7 +48,10 @@ assert Axis.m.backend is None
assert T.Axis.m.same_as(Axis.m)
assert TileLayout(S[8]).apply(3, shape=[8])["m"] == 3
assert decl_tensor((8,), "float32").layout is None
-assert T.Tensor((8,), "float32").layout is None
+tensor_type = T.Tensor((8,), "float32")
+assert tensor_type.layout is None
+tvm.ir.assert_structural_equal(tensor_type.effective_layout(),
TileLayout(S[8]))
+assert tensor_type.layout is None
assert "tvm.backend.cuda.axis" not in sys.modules
assert "tvm.backend.trn.axis" not in sys.modules
""",
@@ -107,25 +109,36 @@ def test_omitted_and_explicit_none_layout(scope):
assert allocated.layout is None and declared.layout is None
[email protected](
+ "shape,strides",
+ [((8, 16), None), ((8, 16), (16, 1)), ((8, 1, 16), (16, 99, 1))],
+)
+def test_effective_layout_on_tensor_type(shape, strides):
+ tensor_type = T.Tensor(shape, "float32", strides=strides)
+ before = save_json(tensor_type)
+ assert_structural_equal(tensor_type.effective_layout(),
TileLayout(S[shape]))
+ assert save_json(tensor_type) == before
+
+
def test_effective_global_layout_preserves_source_and_region():
tensor = decl_tensor((8, 16), "float32", elem_offset=13)
before = save_json(tensor)
region = tensor[2:4, 3:7]
- layout = effective_layout(region.source)
+ layout = region.source.ty.effective_layout()
assert int(layout.apply(2, 3, shape=tensor.shape)["m"]) == 35
sliced = layout.slice(tensor.shape, region.region)
assert int(sliced.apply(0, 0)["m"]) == 35
assert int(sliced.apply(1, 3)["m"]) == 54
assert save_json(tensor) == before
compact = decl_tensor((8, 16), "float32", strides=(16, 1))
- assert_structural_equal(effective_layout(compact), layout)
+ assert_structural_equal(compact.ty.effective_layout(), layout)
strided = decl_tensor((8, 16), "float32", strides=(32, 1))
with pytest.raises(ValueError, match="noncompact"):
- effective_layout(strided)
+ strided.ty.effective_layout()
with pytest.raises(ValueError, match="explicit layout"):
- effective_layout(decl_tensor((8, 16), "float32", scope="shared"))
+ decl_tensor((8, 16), "float32", scope="shared").ty.effective_layout()
explicit = decl_tensor((8, 16), "float32", scope="shared", layout=layout)
- assert effective_layout(explicit).same_as(layout)
+ assert explicit.ty.effective_layout().same_as(layout)
def test_none_layout_views_keep_offsets_and_strides():
diff --git a/tests/python/tirx/tile/cuda/copy_async/test_dsmem.py
b/tests/python/tirx/tile/cuda/copy_async/test_dsmem.py
index 32bd38feff..8cc372b8d6 100644
--- a/tests/python/tirx/tile/cuda/copy_async/test_dsmem.py
+++ b/tests/python/tirx/tile/cuda/copy_async/test_dsmem.py
@@ -30,6 +30,7 @@ import tvm_ffi
import tvm
import tvm.testing
+from tvm.backend.cuda.lang import SMEMPool
from tvm.ir import TensorRegion, TensorRegionType
from tvm.script import tirx as T
from tvm.testing import env
@@ -134,7 +135,7 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected):
For non-fail cases: also runs a 2-CTA cluster kernel via T.copy_async
dispatch (using src_spec as layout for both CTAs) and verifies correctness.
"""
- from tvm.tirx.lang.pipeline import MBarrier
+ from tvm.backend.cuda.lang.pipeline import MBarrier
src_layout = TileLayout(src_spec)
dst_layout = TileLayout(dst_spec)
@@ -169,7 +170,7 @@ def test_dsmem(shape, dtype, src_spec, dst_spec, expected):
)
cbx = T.cuda.cluster_cta_id('x')
tid = T.cuda.thread_idx('x')
- pool = T.SMEMPool()
+ pool = SMEMPool()
# src_smem: CTA 0 writes here, dispatch reads from here
src_raw = pool.alloc([src_phys], dtype, align=128)
src_smem = T.decl_tensor(
diff --git a/tests/python/tirx/tile/cuda/copy_async/test_tcgen05_ldst.py
b/tests/python/tirx/tile/cuda/copy_async/test_tcgen05_ldst.py
index 43469b9e60..f082b9fc01 100644
--- a/tests/python/tirx/tile/cuda/copy_async/test_tcgen05_ldst.py
+++ b/tests/python/tirx/tile/cuda/copy_async/test_tcgen05_ldst.py
@@ -45,7 +45,6 @@ import tvm
import tvm.testing
from tvm.backend.cuda.axis import TCol, TLane
from tvm.backend.cuda.axis import tid_in_wg as axis_tid_in_wg
-from tvm.backend.cuda.lang.alloc_pool import _default_tmem_layout
from tvm.backend.cuda.layout import tcgen05_atom_layout, tmem_datapath_layout
from tvm.script import tirx as T
from tvm.testing import env
@@ -391,13 +390,6 @@ def test_tmem_datapath_layout_rejects_invalid_sub_slab():
tmem_datapath_layout("D", 128, 8, sub_slab=1)
-def test_default_tmem_layout_M128_matches_datapath_D():
- """The default accumulator layout and named Layout D must not drift."""
- tvm.ir.assert_structural_equal(
- _default_tmem_layout(128, 32), tmem_datapath_layout("D", 128, 32)
- )
-
-
@pytest.mark.parametrize("shape", ["16x64b", "16x128b", "16x256b"])
@pytest.mark.parametrize("rep", [1, 2, 4])
def test_tcgen05_atom_layout_apply_matches_decompose_fp32(shape, rep):
diff --git a/tests/python/tirx/tile/trn/test_copy_trn.py
b/tests/python/tirx/tile/trn/test_copy_trn.py
index feafe84982..af50f0e27d 100644
--- a/tests/python/tirx/tile/trn/test_copy_trn.py
+++ b/tests/python/tirx/tile/trn/test_copy_trn.py
@@ -21,8 +21,11 @@ import tvm_ffi
import tvm
import tvm.testing
from tvm.backend.trn.axis import F, P
+from tvm.backend.trn.tile.instruction_generator import InstructionGenerator
from tvm.ir import assert_structural_equal as _assert_structural_equal
from tvm.script import tirx as T
+from tvm.sym import Analyzer
+from tvm.tirx import decl_tensor
from tvm.tirx.layout import S, TileLayout
target = tvm.target.Target("aws/trn1/trn1.2xlarge")
@@ -49,6 +52,15 @@ def assert_structural_equal(lhs, rhs, *args, **kwargs):
_assert_structural_equal(lhs, rhs, *args, **kwargs)
+def test_instruction_generator_compact_global_layout():
+ tensor = decl_tensor((8, 16), "float32")
+ region = tensor[0:8, 0:16]
+ generator = InstructionGenerator((region,), Analyzer())
+ layout = generator.split_layout_views[region]
+ assert int(layout.apply(3, 5, shape=tensor.shape)["m"]) == 53
+ assert tensor.layout is None
+
+
def test_simple_copy():
src_shape = [128, 512]
src_layout = T.TileLayout(T.S[(128, 512) : (512, 1)])
