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 4a09f4ef41 [TIRx] Flatten cuda/trn backend operator folder (#20081)
4a09f4ef41 is described below
commit 4a09f4ef410e96c9dea1de26c3183dacdda5d753
Author: Bohan Hou <[email protected]>
AuthorDate: Sat Aug 1 19:16:12 2026 -0400
[TIRx] Flatten cuda/trn backend operator folder (#20081)
## Summary
Flatten the CUDA/TRN backend Python layout by removing the pass-through
`operator` package layer:
- `tvm.backend.cuda.operator.intrinsics` → `tvm.backend.cuda.intrinsics`
- `tvm.backend.cuda.operator.tile_primitive` →
`tvm.backend.cuda.tile_primitive`
- `tvm.backend.trn.operator.tile_primitive` →
`tvm.backend.trn.tile_primitive`
The `operator` package was an empty shell: its `__init__.py` only
re-exported its subpackages, it added a meaningless path segment, and it
shadowed the stdlib `operator` name. No compatibility shim is kept; all
in-repo imports (source, tests, docs) are rewritten, and the
`tvm.tirx.<backend>.*` alias redirect picks the new paths up
automatically. Cross-package references inside the moved packages are
relative imports that resolve unchanged after the one-level move.
Companion kernel-side PR: mlc-ai/tirx-kernels#13.
## Test
Full TIRx suite on 8x B200 (16 xdist workers): `2620 passed, 41 skipped,
3 xpassed` in 386s; `bench_suite --check-imports` and `registry --cc 10
--strict` import gates green; `pre-commit` green.
---
docs/tirx/api/backend.rst | 13 ++++++++++---
docs/tirx/tile_primitives/copy/fallback.rst | 2 +-
docs/tirx/tile_primitives/copy/gmem_smem.rst | 2 +-
docs/tirx/tile_primitives/copy/ldstmatrix.rst | 2 +-
docs/tirx/tile_primitives/copy/reg.rst | 2 +-
docs/tirx/tile_primitives/copy_async/dsmem.rst | 2 +-
docs/tirx/tile_primitives/copy_async/ldgsts.rst | 2 +-
.../tirx/tile_primitives/copy_async/tcgen05_cp.rst | 2 +-
.../tile_primitives/copy_async/tcgen05_ldst.rst | 2 +-
docs/tirx/tile_primitives/elementwise/reg.rst | 2 +-
docs/tirx/tile_primitives/elementwise/smem.rst | 2 +-
docs/tirx/tile_primitives/gemm.rst | 2 +-
docs/tirx/tile_primitives/gemm_async.rst | 4 ++--
docs/tirx/tile_primitives/permute_layout.rst | 2 +-
docs/tirx/tile_primitives/reduction/local.rst | 2 +-
docs/tirx/tile_primitives/reduction/shared.rst | 2 +-
.../tile_primitives/reduction/sm100_packed.rst | 2 +-
python/tvm/backend/cuda/__init__.py | 18 ++++++++++++++----
.../cuda/{operator => }/intrinsics/__init__.py | 0
.../cuda/{operator => }/intrinsics/_schema.py | 2 +-
.../cuda/{operator => }/intrinsics/cp_async.py | 0
.../backend/cuda/{operator => }/intrinsics/cvt.py | 0
.../cuda/{operator => }/intrinsics/header.py | 0
.../backend/cuda/{operator => }/intrinsics/math.py | 0
.../cuda/{operator => }/intrinsics/memory.py | 0
.../backend/cuda/{operator => }/intrinsics/misc.py | 0
.../backend/cuda/{operator => }/intrinsics/mma.py | 0
.../cuda/{operator => }/intrinsics/nvshmem.py | 0
.../cuda/{operator => }/intrinsics/registry.py | 0
.../backend/cuda/{operator => }/intrinsics/sync.py | 0
.../cuda/{operator => }/intrinsics/tcgen05.py | 0
.../cuda/{operator => }/intrinsics/types.py | 0
.../cuda/{operator => }/intrinsics/utils.py | 0
.../cuda/{operator => }/intrinsics/wgmma.py | 0
python/tvm/backend/cuda/lang/alloc_pool.py | 10 +++++-----
python/tvm/backend/cuda/lang/smem_desc.py | 2 +-
python/tvm/backend/cuda/operator/__init__.py | 19 -------------------
.../cuda/{operator => }/tile_primitive/__init__.py | 0
.../cuda/{operator => }/tile_primitive/common.py | 0
.../{operator => }/tile_primitive/copy/__init__.py | 0
.../{operator => }/tile_primitive/copy/_common.py | 0
.../tile_primitive/copy/_swizzle_iter.py | 0
.../{operator => }/tile_primitive/copy/fallback.py | 0
.../tile_primitive/copy/gmem_smem.py | 0
.../tile_primitive/copy/ld_stmatrix.py | 0
.../cuda/{operator => }/tile_primitive/copy/reg.py | 0
.../{operator => }/tile_primitive/copy/utils.py | 0
.../{operator => }/tile_primitive/copy/vec_auto.py | 0
.../tile_primitive/copy/vec_auto_gmem_smem.py | 0
.../tile_primitive/copy/vec_auto_reg.py | 0
.../tile_primitive/copy/vec_forced.py | 0
.../tile_primitive/copy_async/__init__.py | 0
.../tile_primitive/copy_async/dsmem.py | 0
.../tile_primitive/copy_async/ldgsts.py | 0
.../tile_primitive/copy_async/tcgen05_cp.py | 0
.../tile_primitive/copy_async/tcgen05_ldst.py | 0
.../tile_primitive/copy_async/tma.py | 0
.../tile_primitive/copy_async/utils.py | 0
.../tile_primitive/elementwise/__init__.py | 2 +-
.../tile_primitive/elementwise/_common.py | 0
.../tile_primitive/elementwise/ops/__init__.py | 0
.../tile_primitive/elementwise/ops/binary.py | 0
.../tile_primitive/elementwise/ops/cast.py | 0
.../tile_primitive/elementwise/ops/fma.py | 0
.../tile_primitive/elementwise/ops/unary.py | 0
.../tile_primitive/elementwise/reg.py | 0
.../tile_primitive/elementwise/register.py | 0
.../tile_primitive/elementwise/smem.py | 0
.../elementwise/vec_emit/__init__.py | 0
.../elementwise/vec_emit/binary_f32x2.py | 0
.../elementwise/vec_emit/cast_vec2.py | 0
.../elementwise/vec_emit/fma_f32x2.py | 0
.../tile_primitive/exec_scope_utils.py | 0
.../{operator => }/tile_primitive/gemm/__init__.py | 0
.../tile_primitive/gemm/mma_m16n8k_.py | 0
.../tile_primitive/gemm_async/__init__.py | 0
.../tile_primitive/gemm_async/tcgen05.py | 0
.../{operator => }/tile_primitive/gemm_utils.py | 0
.../{operator => }/tile_primitive/layout_utils.py | 0
.../tile_primitive/permute_layout/__init__.py | 0
.../permute_layout/warp_xor_swizzle.py | 0
.../tile_primitive/reduction/__init__.py | 0
.../tile_primitive/reduction/local.py | 0
.../tile_primitive/reduction/shared.py | 0
.../tile_primitive/reduction/sm100_packed.py | 0
.../tile_primitive/reduction/utils.py | 0
.../{operator => }/tile_primitive/tma_utils.py | 0
python/tvm/backend/trn/__init__.py | 14 +++++++++++---
python/tvm/backend/trn/operator/__init__.py | 22 ----------------------
.../trn/{operator => }/tile_primitive/__init__.py | 0
.../tile_primitive/binary/__init__.py | 0
.../tile_primitive/binary/default.py | 0
.../{operator => }/tile_primitive/binary/utils.py | 0
.../trn/{operator => }/tile_primitive/common.py | 0
.../tile_primitive/compose_op/__init__.py | 0
.../tile_primitive/compose_op/binary_chain.py | 0
.../tile_primitive/compose_op/binary_reduce.py | 0
.../tile_primitive/compose_op/compose_op.py | 0
.../tile_primitive/compose_op/reduce_negate.py | 0
.../tile_primitive/compose_op/unary_reduce.py | 0
.../tile_primitive/compose_op/utils.py | 0
.../{operator => }/tile_primitive/copy/__init__.py | 0
.../{operator => }/tile_primitive/copy/default.py | 0
.../trn/{operator => }/tile_primitive/dim_utils.py | 0
.../{operator => }/tile_primitive/gemm/__init__.py | 0
.../{operator => }/tile_primitive/gemm/default.py | 0
.../tile_primitive/instruction_generator.py | 0
.../{operator => }/tile_primitive/private_alloc.py | 6 +++---
.../tile_primitive/reduction/__init__.py | 0
.../tile_primitive/reduction/default.py | 0
.../tile_primitive/reduction/utils.py | 0
.../tile_primitive/select/__init__.py | 0
.../tile_primitive/select/default.py | 0
.../tile_primitive/unary/__init__.py | 0
.../{operator => }/tile_primitive/unary/default.py | 0
.../{operator => }/tile_primitive/unary/utils.py | 0
.../tile_primitive/unary/with_bias_scale.py | 0
.../tile_primitive/workspace_utils.py | 0
python/tvm/tirx/operator/intrinsics/_common.py | 2 +-
tests/python/tirx/codegen/test_ptx_cvt.py | 4 ++--
tests/python/tirx/codegen/test_ptx_ld_st_ops.py | 2 +-
.../tile_primitive/cuda/copy/test_fallback.py | 2 +-
.../tile_primitive/cuda/copy/test_gmem_smem.py | 2 +-
.../tile_primitive/cuda/copy/test_ld_stmatrix.py | 2 +-
.../operator/tile_primitive/cuda/copy/test_reg.py | 8 ++++----
.../tile_primitive/cuda/copy/test_swizzle_iter.py | 6 +++---
.../tile_primitive/cuda/copy_async/test_dsmem.py | 2 +-
.../cuda/copy_async/test_tcgen05_cp.py | 2 +-
.../tile_primitive/cuda/copy_async/test_tma.py | 4 ++--
.../tile_primitive/cuda/elementwise/test_unary.py | 6 +++---
.../cuda/gemm_async/test_gemm_async.py | 10 +++++-----
.../cuda/permute_layout/test_permute_layout.py | 2 +-
tests/python/tirx/test_alloc_pool.py | 2 +-
tests/python/tirx/test_layout.py | 4 ++--
134 files changed, 93 insertions(+), 109 deletions(-)
diff --git a/docs/tirx/api/backend.rst b/docs/tirx/api/backend.rst
index 6fd4e037de..e97c65c74b 100644
--- a/docs/tirx/api/backend.rst
+++ b/docs/tirx/api/backend.rst
@@ -51,11 +51,18 @@ tvm.backend.cuda.script
:imported-members:
:exclude-members: PrimExpr, Op, Call
-tvm.backend.cuda.operator
-*************************
-.. automodule:: tvm.backend.cuda.operator
+tvm.backend.cuda.intrinsics
+***************************
+.. automodule:: tvm.backend.cuda.intrinsics
:members:
:imported-members:
+ :no-index:
+ :exclude-members: PrimExpr, Op, Call
+
+tvm.backend.cuda.tile_primitive
+*******************************
+.. automodule:: tvm.backend.cuda.tile_primitive
+ :members:
:exclude-members: PrimExpr, Op, Call
tvm.backend.cuda.target_tags
diff --git a/docs/tirx/tile_primitives/copy/fallback.rst
b/docs/tirx/tile_primitives/copy/fallback.rst
index 7e78c87bfd..60469557f3 100644
--- a/docs/tirx/tile_primitives/copy/fallback.rst
+++ b/docs/tirx/tile_primitives/copy/fallback.rst
@@ -22,7 +22,7 @@ The ``fallback`` variant is the **priority-0 catch-all**: a
scalar, single-threa
copy that runs only when every faster variant (:doc:`gmem_smem`, :doc:`reg`,
:doc:`ldstmatrix`) has declined. It always works — for any valid copy at any
scope
— and is intentionally slow, so it emits a ``UserWarning`` when chosen. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py``.
+``python/tvm/backend/cuda/tile_primitive/copy/fallback.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/copy/gmem_smem.rst
b/docs/tirx/tile_primitives/copy/gmem_smem.rst
index d34578a995..e3275eb74b 100644
--- a/docs/tirx/tile_primitives/copy/gmem_smem.rst
+++ b/docs/tirx/tile_primitives/copy/gmem_smem.rst
@@ -23,7 +23,7 @@ shared** memory (either direction) when **neither side is a
register**. Because
neither operand carries a thread partition, the dispatch *synthesizes* one
from the
execution scope: it splits the region into ``[outer, threads, vec]`` and emits
a
serial loop of vectorized loads/stores. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py``.
+``python/tvm/backend/cuda/tile_primitive/copy/gmem_smem.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/copy/ldstmatrix.rst
b/docs/tirx/tile_primitives/copy/ldstmatrix.rst
index 11c4b2cf75..9af2a80131 100644
--- a/docs/tirx/tile_primitives/copy/ldstmatrix.rst
+++ b/docs/tirx/tile_primitives/copy/ldstmatrix.rst
@@ -24,7 +24,7 @@ instruction moves ``num`` 8×8 16-bit matrix tiles between
shared memory and the
warp's registers, with the hardware performing the lane↔element shuffle that an
MMA fragment needs. It only applies when the register and shared **layouts
match
the m8n8 fragment geometry**; otherwise the copy falls back to :doc:`reg`.
Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py``.
+``python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/copy/reg.rst
b/docs/tirx/tile_primitives/copy/reg.rst
index d195ffb5a4..5db0edf86f 100644
--- a/docs/tirx/tile_primitives/copy/reg.rst
+++ b/docs/tirx/tile_primitives/copy/reg.rst
@@ -24,7 +24,7 @@ register** (``local``) buffer and the other is ``shared*`` or
``global``. Unlike
register operand's layout: that layout's thread-axis iters already say which
thread
owns which logical coordinate, so the dispatch drops those axes, leaves each
thread
its private bundle of elements, and copies them in a vectorized serial loop.
Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py``.
+``python/tvm/backend/cuda/tile_primitive/copy/reg.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/copy_async/dsmem.rst
b/docs/tirx/tile_primitives/copy_async/dsmem.rst
index bec98c5bcd..8690ebec62 100644
--- a/docs/tirx/tile_primitives/copy_async/dsmem.rst
+++ b/docs/tirx/tile_primitives/copy_async/dsmem.rst
@@ -24,7 +24,7 @@ memory). One elected thread on the source CTA maps the
destination CTA's shared
address into its own address space (PTX ``mapa``) and issues a bulk copy
(``cp.async.bulk.shared::cluster``); the hardware decrements the *destination*
CTA's
mbarrier when the bytes land. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py``.
+``python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/copy_async/ldgsts.rst
b/docs/tirx/tile_primitives/copy_async/ldgsts.rst
index a49b7d3c02..826bc238b8 100644
--- a/docs/tirx/tile_primitives/copy_async/ldgsts.rst
+++ b/docs/tirx/tile_primitives/copy_async/ldgsts.rst
@@ -24,7 +24,7 @@ vectorized copy that the hardware completes in the
background, so the warp can k
computing while the load is in flight. It reuses the exact ``[outer, threads,
vec]``
partition of the synchronous :doc:`../copy/gmem_smem` variant; the differences
are
all in *what* is emitted and *when* it completes. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py``.
+``python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
index 26d8e39e11..505e85b00a 100644
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
+++ b/docs/tirx/tile_primitives/copy_async/tcgen05_cp.rst
@@ -24,7 +24,7 @@ memory** (Blackwell ``tmem``) through a generic planner
covering every
descriptor fields (ldo/sdo/swizzle) and the cp issue sequence are derived from
the two buffer layouts. The dispatch issues only the copy; the caller signals
completion with ``tcgen05.commit``. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py``.
+``python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py``.
Shape selection
---------------
diff --git a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
b/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
index 106e084445..8acf01e176 100644
--- a/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
+++ b/docs/tirx/tile_primitives/copy_async/tcgen05_ldst.rst
@@ -25,7 +25,7 @@ registers. One registration handles both directions — ``tmem
→ local`` lower
``tcgen05.ld``, ``local → tmem`` to ``tcgen05.st`` — and the dispatch picks the
widest instruction shape the register layout matches. As with the other async
variants, completion (``tcgen05.wait.ld`` / ``wait.st``) is the caller's.
Source:
-``python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py``.
+``python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/elementwise/reg.rst
b/docs/tirx/tile_primitives/elementwise/reg.rst
index 2deceb03eb..435da44771 100644
--- a/docs/tirx/tile_primitives/elementwise/reg.rst
+++ b/docs/tirx/tile_primitives/elementwise/reg.rst
@@ -23,7 +23,7 @@ The ``reg`` variant lowers an elementwise op (``sqrt``,
``exp``, ``add``,
:doc:`../copy/reg` variant the partition is *induced* by the operands' register
layout — the thread axes are dropped, leaving each thread its private bundle —
and
the op is applied to every register in that bundle. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py``.
+``python/tvm/backend/cuda/tile_primitive/elementwise/reg.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/elementwise/smem.rst
b/docs/tirx/tile_primitives/elementwise/smem.rst
index 4bf5e8b5f0..dab384d495 100644
--- a/docs/tirx/tile_primitives/elementwise/smem.rst
+++ b/docs/tirx/tile_primitives/elementwise/smem.rst
@@ -23,7 +23,7 @@ The ``smem`` variant lowers an elementwise op (``sqrt``,
``exp``, ``add``,
:doc:`../copy/gmem_smem` variant it *synthesizes* a ``[outer, threads, vec]``
partition from the execution scope, then applies the op to each (vectorized)
element. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py``.
+``python/tvm/backend/cuda/tile_primitive/elementwise/smem.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/gemm.rst
b/docs/tirx/tile_primitives/gemm.rst
index e737621dcb..97d7e7ed6b 100644
--- a/docs/tirx/tile_primitives/gemm.rst
+++ b/docs/tirx/tile_primitives/gemm.rst
@@ -24,7 +24,7 @@ fragments and the C/D accumulators **all live in registers**
— the caller stag
and B into register fragments first (typically via :doc:`copy/ldstmatrix`). The
dispatch tiles M/N/K into ``m16n8k`` atoms and emits one ``mma`` per output
tile,
accumulating over K in place. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py``. (For
the
+``python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py``. (For the
Blackwell async tensor-core path see :doc:`gemm_async`.)
What it accepts
diff --git a/docs/tirx/tile_primitives/gemm_async.rst
b/docs/tirx/tile_primitives/gemm_async.rst
index ef2250379f..1b28279773 100644
--- a/docs/tirx/tile_primitives/gemm_async.rst
+++ b/docs/tirx/tile_primitives/gemm_async.rst
@@ -25,7 +25,7 @@ memory**, and one elected thread launches the MMA, which runs
asynchronously; th
caller signals completion with ``tcgen05.commit`` against an mbarrier. It also
supports **block-scaled** low precision (fp8 / fp4 with per-block scale factors
``SFA`` / ``SFB`` in tensor memory). Source:
-``python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py``.
(For the
+``python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py``. (For the
synchronous warp-register path see :doc:`gemm`.)
What it accepts
@@ -80,7 +80,7 @@ into a tmem accumulator, after TMA-loading A/B into shared
(from
.. code-block:: python
from tvm.tirx.layout import S, TCol, TLane, TileLayout, tid_in_wg as
axis_tid_in_wg
- from tvm.tirx.cuda.operator.tile_primitive.tma_utils import
mma_shared_layout
+ from tvm.tirx.cuda.tile_primitive.tma_utils import mma_shared_layout
A_smem = T.alloc_buffer((3,128,64), "float16", scope="shared",
layout=mma_shared_layout("float16", 3, (3,128,64)))
B_smem = T.alloc_buffer((3,128,64), "float16", scope="shared",
layout=mma_shared_layout("float16", 3, (3,128,64)))
diff --git a/docs/tirx/tile_primitives/permute_layout.rst
b/docs/tirx/tile_primitives/permute_layout.rst
index 05186d4be1..e9f9ce6618 100644
--- a/docs/tirx/tile_primitives/permute_layout.rst
+++ b/docs/tirx/tile_primitives/permute_layout.rst
@@ -25,7 +25,7 @@ back under the destination layout, with a per-lane **XOR
swizzle** on the iterat
index chosen so that *both* the read and the write phase are shared-memory
bank-conflict-free. A ``warp_sync`` separates the two phases so the op is safe
even
when source and destination alias. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py``.
+``python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/reduction/local.rst
b/docs/tirx/tile_primitives/reduction/local.rst
index 184d8a395d..9df8413fdb 100644
--- a/docs/tirx/tile_primitives/reduction/local.rst
+++ b/docs/tirx/tile_primitives/reduction/local.rst
@@ -23,7 +23,7 @@ source and destination are register** (``local``) buffers. At
thread scope it is
plain sequential reduction over each thread's own elements; at warp scope, if
the
destination layout carries a ``laneid`` replica, it also folds across lanes
with a
``__shfl_xor`` tree. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py``.
+``python/tvm/backend/cuda/tile_primitive/reduction/local.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/reduction/shared.rst
b/docs/tirx/tile_primitives/reduction/shared.rst
index 5d61dc6c87..aa328ca146 100644
--- a/docs/tirx/tile_primitives/reduction/shared.rst
+++ b/docs/tirx/tile_primitives/reduction/shared.rst
@@ -23,7 +23,7 @@ The ``shared`` variant lowers a reduction (``sum`` / ``max``
/ ``min``) when
it partitions the threads into groups — one group per output position — has
each
thread gather a chunk of the reduction axis, then folds the group with an
adaptive
``__shfl_xor`` tree. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py``.
+``python/tvm/backend/cuda/tile_primitive/reduction/shared.py``.
What it accepts
---------------
diff --git a/docs/tirx/tile_primitives/reduction/sm100_packed.rst
b/docs/tirx/tile_primitives/reduction/sm100_packed.rst
index b1c2e5de44..cdd4315055 100644
--- a/docs/tirx/tile_primitives/reduction/sm100_packed.rst
+++ b/docs/tirx/tile_primitives/reduction/sm100_packed.rst
@@ -23,7 +23,7 @@ it pre-empts :doc:`local`) for a thread-scope reduction of a
1-D ``float32`` vec
of at least 8 elements to a scalar. It uses the SM100 packed math instructions
—
``add.f32x2`` for ``sum``, ``max3.f32`` / ``min3.f32`` for ``max`` / ``min`` —
to fold two (or three) lanes of data per instruction. Source:
-``python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py``.
+``python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py``.
What it accepts
---------------
diff --git a/python/tvm/backend/cuda/__init__.py
b/python/tvm/backend/cuda/__init__.py
index 0db092201c..c063c77bcf 100644
--- a/python/tvm/backend/cuda/__init__.py
+++ b/python/tvm/backend/cuda/__init__.py
@@ -23,7 +23,16 @@ from tvm_ffi.libinfo import load_lib_ctypes
from tvm.base import _LOADED_LIBS
-_LAZY_SUBMODULES = {"iket", "lang", "op", "operator", "script", "target_tags",
"transforms"}
+_LAZY_SUBMODULES = {
+ "iket",
+ "intrinsics",
+ "lang",
+ "op",
+ "script",
+ "target_tags",
+ "tile_primitive",
+ "transforms",
+}
def _detect_target_from_device(dev):
@@ -60,8 +69,8 @@ def register_backend():
for name, namespace in script_namespaces().items():
builder_ir.register_script_namespace(name, namespace)
- import_module(f"{__name__}.operator.intrinsics")
- import_module(f"{__name__}.operator.tile_primitive")
+ import_module(f"{__name__}.intrinsics")
+ import_module(f"{__name__}.tile_primitive")
import_module(f"{__name__}.target_tags")
@@ -93,13 +102,14 @@ def __getattr__(name: str):
__all__ = [
"iket",
+ "intrinsics",
"lang",
"op",
- "operator",
"register_backend",
"script",
"script_namespace",
"script_namespaces",
"target_tags",
+ "tile_primitive",
"transforms",
]
diff --git a/python/tvm/backend/cuda/operator/intrinsics/__init__.py
b/python/tvm/backend/cuda/intrinsics/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/__init__.py
rename to python/tvm/backend/cuda/intrinsics/__init__.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/_schema.py
b/python/tvm/backend/cuda/intrinsics/_schema.py
similarity index 98%
rename from python/tvm/backend/cuda/operator/intrinsics/_schema.py
rename to python/tvm/backend/cuda/intrinsics/_schema.py
index e9da9c3e4f..b848962ae8 100644
--- a/python/tvm/backend/cuda/operator/intrinsics/_schema.py
+++ b/python/tvm/backend/cuda/intrinsics/_schema.py
@@ -43,8 +43,8 @@ from __future__ import annotations
from collections.abc import Callable
+from tvm.backend.cuda.intrinsics.registry import register_codegen
from tvm.backend.cuda.op import cuda_func_call
-from tvm.backend.cuda.operator.intrinsics.registry import register_codegen
# C primitive type → TVM dtype string. Used when the caller specifies a
# non-void ``return_type`` but no explicit ``tvm_return_type`` — the helper
diff --git a/python/tvm/backend/cuda/operator/intrinsics/cp_async.py
b/python/tvm/backend/cuda/intrinsics/cp_async.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/cp_async.py
rename to python/tvm/backend/cuda/intrinsics/cp_async.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/cvt.py
b/python/tvm/backend/cuda/intrinsics/cvt.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/cvt.py
rename to python/tvm/backend/cuda/intrinsics/cvt.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/header.py
b/python/tvm/backend/cuda/intrinsics/header.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/header.py
rename to python/tvm/backend/cuda/intrinsics/header.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/math.py
b/python/tvm/backend/cuda/intrinsics/math.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/math.py
rename to python/tvm/backend/cuda/intrinsics/math.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/memory.py
b/python/tvm/backend/cuda/intrinsics/memory.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/memory.py
rename to python/tvm/backend/cuda/intrinsics/memory.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/misc.py
b/python/tvm/backend/cuda/intrinsics/misc.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/misc.py
rename to python/tvm/backend/cuda/intrinsics/misc.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/mma.py
b/python/tvm/backend/cuda/intrinsics/mma.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/mma.py
rename to python/tvm/backend/cuda/intrinsics/mma.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/nvshmem.py
b/python/tvm/backend/cuda/intrinsics/nvshmem.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/nvshmem.py
rename to python/tvm/backend/cuda/intrinsics/nvshmem.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/registry.py
b/python/tvm/backend/cuda/intrinsics/registry.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/registry.py
rename to python/tvm/backend/cuda/intrinsics/registry.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/sync.py
b/python/tvm/backend/cuda/intrinsics/sync.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/sync.py
rename to python/tvm/backend/cuda/intrinsics/sync.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/tcgen05.py
b/python/tvm/backend/cuda/intrinsics/tcgen05.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/tcgen05.py
rename to python/tvm/backend/cuda/intrinsics/tcgen05.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/types.py
b/python/tvm/backend/cuda/intrinsics/types.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/types.py
rename to python/tvm/backend/cuda/intrinsics/types.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/utils.py
b/python/tvm/backend/cuda/intrinsics/utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/utils.py
rename to python/tvm/backend/cuda/intrinsics/utils.py
diff --git a/python/tvm/backend/cuda/operator/intrinsics/wgmma.py
b/python/tvm/backend/cuda/intrinsics/wgmma.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/intrinsics/wgmma.py
rename to python/tvm/backend/cuda/intrinsics/wgmma.py
diff --git a/python/tvm/backend/cuda/lang/alloc_pool.py
b/python/tvm/backend/cuda/lang/alloc_pool.py
index 766713f5a5..fbb1f7d1ff 100644
--- a/python/tvm/backend/cuda/lang/alloc_pool.py
+++ b/python/tvm/backend/cuda/lang/alloc_pool.py
@@ -69,7 +69,7 @@ def _shape_product(shape):
def _auto_swizzle_mode(dtype):
"""Select the default MMA swizzle mode for a shared-memory allocation."""
- from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+ from tvm.backend.cuda.tile_primitive.tma_utils import SwizzleMode
del dtype
return SwizzleMode.SWIZZLE_128B_ATOM
@@ -77,7 +77,7 @@ def _auto_swizzle_mode(dtype):
def _swizzle_atom_bytes(swizzle_mode):
"""Return the row width (in bytes) of one swizzle atom for
*swizzle_mode*."""
- from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+ from tvm.backend.cuda.tile_primitive.tma_utils import SwizzleMode
return {
SwizzleMode.SWIZZLE_NONE: 0,
@@ -112,7 +112,7 @@ def _validate_mma_alloc_shape(shape, dtype, swizzle_mode):
Validation skipped when *swizzle_mode* is ``SWIZZLE_NONE`` (no atom).
"""
- from tvm.backend.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+ from tvm.backend.cuda.tile_primitive.tma_utils import SwizzleMode
if swizzle_mode == SwizzleMode.SWIZZLE_NONE:
return
@@ -301,7 +301,7 @@ class TMEMPool:
what gemm dispatch iterates over). When ``shape`` has 3 dims, the first
is treated as a pipe-depth outer.
"""
- from tvm.backend.cuda.operator.tile_primitive.gemm_async.tcgen05
import sf_tmem_layout
+ from tvm.backend.cuda.tile_primitive.gemm_async.tcgen05 import
sf_tmem_layout
if len(shape) == 2:
pipe_depth, rows, last = None, shape[0], shape[1]
@@ -447,7 +447,7 @@ class SMEMPool:
def alloc_tcgen05_mma_AB(self, shape, dtype="float16",
swizzle_mode="auto", align=1024):
"""Allocate MMA-compatible shared memory with an inferred swizzle
layout."""
- from tvm.backend.cuda.operator.tile_primitive.tma_utils import (
+ from tvm.backend.cuda.tile_primitive.tma_utils import (
SwizzleMode,
mma_shared_layout,
)
diff --git a/python/tvm/backend/cuda/lang/smem_desc.py
b/python/tvm/backend/cuda/lang/smem_desc.py
index d2561d17dc..a156ff5608 100644
--- a/python/tvm/backend/cuda/lang/smem_desc.py
+++ b/python/tvm/backend/cuda/lang/smem_desc.py
@@ -17,7 +17,7 @@
"""SMEM matrix descriptor helper for tcgen05 / wgmma."""
-from tvm.backend.cuda.operator.tile_primitive.common import
smem_desc_add_16B_offset
+from tvm.backend.cuda.tile_primitive.common import smem_desc_add_16B_offset
from tvm.script import tirx as T
diff --git a/python/tvm/backend/cuda/operator/__init__.py
b/python/tvm/backend/cuda/operator/__init__.py
deleted file mode 100644
index 744c9bafc9..0000000000
--- a/python/tvm/backend/cuda/operator/__init__.py
+++ /dev/null
@@ -1,19 +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.
-"""CUDA backend operator registrations and helpers."""
-
-__all__ = ["intrinsics", "tile_primitive"]
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/__init__.py
b/python/tvm/backend/cuda/tile_primitive/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/__init__.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/common.py
b/python/tvm/backend/cuda/tile_primitive/common.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/common.py
rename to python/tvm/backend/cuda/tile_primitive/common.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py
b/python/tvm/backend/cuda/tile_primitive/copy/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/copy/__init__.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py
b/python/tvm/backend/cuda/tile_primitive/copy/_common.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/_common.py
rename to python/tvm/backend/cuda/tile_primitive/copy/_common.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py
b/python/tvm/backend/cuda/tile_primitive/copy/_swizzle_iter.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/copy/_swizzle_iter.py
rename to python/tvm/backend/cuda/tile_primitive/copy/_swizzle_iter.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py
b/python/tvm/backend/cuda/tile_primitive/copy/fallback.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/fallback.py
rename to python/tvm/backend/cuda/tile_primitive/copy/fallback.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py
b/python/tvm/backend/cuda/tile_primitive/copy/gmem_smem.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/gmem_smem.py
rename to python/tvm/backend/cuda/tile_primitive/copy/gmem_smem.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py
b/python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/ld_stmatrix.py
rename to python/tvm/backend/cuda/tile_primitive/copy/ld_stmatrix.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py
b/python/tvm/backend/cuda/tile_primitive/copy/reg.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py
rename to python/tvm/backend/cuda/tile_primitive/copy/reg.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py
b/python/tvm/backend/cuda/tile_primitive/copy/utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/utils.py
rename to python/tvm/backend/cuda/tile_primitive/copy/utils.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto.py
rename to python/tvm/backend/cuda/tile_primitive/copy/vec_auto.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_gmem_smem.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_gmem_smem.py
rename to python/tvm/backend/cuda/tile_primitive/copy/vec_auto_gmem_smem.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_reg.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/vec_auto_reg.py
rename to python/tvm/backend/cuda/tile_primitive/copy/vec_auto_reg.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy/vec_forced.py
b/python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy/vec_forced.py
rename to python/tvm/backend/cuda/tile_primitive/copy/vec_forced.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/__init__.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/__init__.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/copy_async/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/copy_async/__init__.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy_async/dsmem.py
rename to python/tvm/backend/cuda/tile_primitive/copy_async/dsmem.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy_async/ldgsts.py
rename to python/tvm/backend/cuda/tile_primitive/copy_async/ldgsts.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_cp.py
rename to python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_cp.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/copy_async/tcgen05_ldst.py
rename to python/tvm/backend/cuda/tile_primitive/copy_async/tcgen05_ldst.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy_async/tma.py
rename to python/tvm/backend/cuda/tile_primitive/copy_async/tma.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py
b/python/tvm/backend/cuda/tile_primitive/copy_async/utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/copy_async/utils.py
rename to python/tvm/backend/cuda/tile_primitive/copy_async/utils.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
similarity index 95%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
index 576dc437b9..c4dcfb63b8 100644
--- a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py
+++ b/python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py
@@ -30,7 +30,7 @@ PTX/CUDA intrinsics live in ``vec_emit/`` (``binary_f32x2`` /
``cast_vec2``
from .register import *
# Suppress submodule-attribute leakage. Without an explicit ``__all__`` here,
-# ``from tvm.backend.cuda.operator.tile_primitive.elementwise import *`` (run
by
+# ``from tvm.backend.cuda.tile_primitive.elementwise import *`` (run by
# tile_primitive/__init__.py) re-exports the implicit submodule attributes
# (``ops``, ``reg``, ``smem``, ``vec_emit``) — and ``ops`` in particular
# shadows the top-level ``tile_primitive/ops.py`` (BinaryReduce / UnaryReduce
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/_common.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/_common.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/__init__.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/ops/__init__.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/binary.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/ops/binary.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/cast.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/cast.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/ops/cast.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/fma.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/fma.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/ops/fma.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/ops/unary.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/ops/unary.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/elementwise/reg.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/reg.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/register.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/register.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/register.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/register.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/elementwise/smem.py
rename to python/tvm/backend/cuda/tile_primitive/elementwise/smem.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/__init__.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/__init__.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/__init__.py
rename to
python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/__init__.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/binary_f32x2.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/binary_f32x2.py
rename to
python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/binary_f32x2.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/cast_vec2.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/cast_vec2.py
rename to
python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/cast_vec2.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py
b/python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/fma_f32x2.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/elementwise/vec_emit/fma_f32x2.py
rename to
python/tvm/backend/cuda/tile_primitive/elementwise/vec_emit/fma_f32x2.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py
b/python/tvm/backend/cuda/tile_primitive/exec_scope_utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/exec_scope_utils.py
rename to python/tvm/backend/cuda/tile_primitive/exec_scope_utils.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/gemm/__init__.py
b/python/tvm/backend/cuda/tile_primitive/gemm/__init__.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/gemm/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/gemm/__init__.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py
b/python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/gemm/mma_m16n8k_.py
rename to python/tvm/backend/cuda/tile_primitive/gemm/mma_m16n8k_.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/__init__.py
b/python/tvm/backend/cuda/tile_primitive/gemm_async/__init__.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/gemm_async/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/gemm_async/__init__.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py
b/python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/gemm_async/tcgen05.py
rename to python/tvm/backend/cuda/tile_primitive/gemm_async/tcgen05.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py
b/python/tvm/backend/cuda/tile_primitive/gemm_utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/gemm_utils.py
rename to python/tvm/backend/cuda/tile_primitive/gemm_utils.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py
b/python/tvm/backend/cuda/tile_primitive/layout_utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/layout_utils.py
rename to python/tvm/backend/cuda/tile_primitive/layout_utils.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/__init__.py
b/python/tvm/backend/cuda/tile_primitive/permute_layout/__init__.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/permute_layout/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/permute_layout/__init__.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py
b/python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/permute_layout/warp_xor_swizzle.py
rename to
python/tvm/backend/cuda/tile_primitive/permute_layout/warp_xor_swizzle.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/reduction/__init__.py
b/python/tvm/backend/cuda/tile_primitive/reduction/__init__.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/reduction/__init__.py
rename to python/tvm/backend/cuda/tile_primitive/reduction/__init__.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py
b/python/tvm/backend/cuda/tile_primitive/reduction/local.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py
rename to python/tvm/backend/cuda/tile_primitive/reduction/local.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py
b/python/tvm/backend/cuda/tile_primitive/reduction/shared.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/reduction/shared.py
rename to python/tvm/backend/cuda/tile_primitive/reduction/shared.py
diff --git
a/python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py
b/python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py
similarity index 100%
rename from
python/tvm/backend/cuda/operator/tile_primitive/reduction/sm100_packed.py
rename to python/tvm/backend/cuda/tile_primitive/reduction/sm100_packed.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py
b/python/tvm/backend/cuda/tile_primitive/reduction/utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/reduction/utils.py
rename to python/tvm/backend/cuda/tile_primitive/reduction/utils.py
diff --git a/python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py
b/python/tvm/backend/cuda/tile_primitive/tma_utils.py
similarity index 100%
rename from python/tvm/backend/cuda/operator/tile_primitive/tma_utils.py
rename to python/tvm/backend/cuda/tile_primitive/tma_utils.py
diff --git a/python/tvm/backend/trn/__init__.py
b/python/tvm/backend/trn/__init__.py
index 7650ea87d5..bb75669376 100644
--- a/python/tvm/backend/trn/__init__.py
+++ b/python/tvm/backend/trn/__init__.py
@@ -18,7 +18,15 @@
from importlib import import_module
-_LAZY_SUBMODULES = {"layout", "op", "operator", "pipeline", "script",
"target_tags", "transform"}
+_LAZY_SUBMODULES = {
+ "layout",
+ "op",
+ "pipeline",
+ "script",
+ "target_tags",
+ "tile_primitive",
+ "transform",
+}
def register_backend():
@@ -29,7 +37,7 @@ def register_backend():
for name, namespace in script_namespaces().items():
builder_ir.register_script_namespace(name, namespace)
- import_module(f"{__name__}.operator.tile_primitive")
+ import_module(f"{__name__}.tile_primitive")
trn_pipeline = import_module(f"{__name__}.pipeline")
import_module(f"{__name__}.target_tags")
import_module(f"{__name__}.transform")
@@ -57,12 +65,12 @@ def __getattr__(name: str):
__all__ = [
"layout",
"op",
- "operator",
"pipeline",
"register_backend",
"script",
"script_namespace",
"script_namespaces",
"target_tags",
+ "tile_primitive",
"transform",
]
diff --git a/python/tvm/backend/trn/operator/__init__.py
b/python/tvm/backend/trn/operator/__init__.py
deleted file mode 100644
index 2d9b93d998..0000000000
--- a/python/tvm/backend/trn/operator/__init__.py
+++ /dev/null
@@ -1,22 +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.
-"""Trainium backend operator package.
-
-Loaded by the Trainium backend registration hook.
-"""
-
-__all__ = ["tile_primitive"]
diff --git a/python/tvm/backend/trn/operator/tile_primitive/__init__.py
b/python/tvm/backend/trn/tile_primitive/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/__init__.py
rename to python/tvm/backend/trn/tile_primitive/__init__.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/binary/__init__.py
b/python/tvm/backend/trn/tile_primitive/binary/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/binary/__init__.py
rename to python/tvm/backend/trn/tile_primitive/binary/__init__.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/binary/default.py
b/python/tvm/backend/trn/tile_primitive/binary/default.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/binary/default.py
rename to python/tvm/backend/trn/tile_primitive/binary/default.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/binary/utils.py
b/python/tvm/backend/trn/tile_primitive/binary/utils.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/binary/utils.py
rename to python/tvm/backend/trn/tile_primitive/binary/utils.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/common.py
b/python/tvm/backend/trn/tile_primitive/common.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/common.py
rename to python/tvm/backend/trn/tile_primitive/common.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/compose_op/__init__.py
b/python/tvm/backend/trn/tile_primitive/compose_op/__init__.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/compose_op/__init__.py
rename to python/tvm/backend/trn/tile_primitive/compose_op/__init__.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py
b/python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_chain.py
rename to python/tvm/backend/trn/tile_primitive/compose_op/binary_chain.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py
b/python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/compose_op/binary_reduce.py
rename to python/tvm/backend/trn/tile_primitive/compose_op/binary_reduce.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/compose_op/compose_op.py
b/python/tvm/backend/trn/tile_primitive/compose_op/compose_op.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/compose_op/compose_op.py
rename to python/tvm/backend/trn/tile_primitive/compose_op/compose_op.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/compose_op/reduce_negate.py
b/python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/compose_op/reduce_negate.py
rename to python/tvm/backend/trn/tile_primitive/compose_op/reduce_negate.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py
b/python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/compose_op/unary_reduce.py
rename to python/tvm/backend/trn/tile_primitive/compose_op/unary_reduce.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/compose_op/utils.py
b/python/tvm/backend/trn/tile_primitive/compose_op/utils.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/compose_op/utils.py
rename to python/tvm/backend/trn/tile_primitive/compose_op/utils.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/copy/__init__.py
b/python/tvm/backend/trn/tile_primitive/copy/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/copy/__init__.py
rename to python/tvm/backend/trn/tile_primitive/copy/__init__.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/copy/default.py
b/python/tvm/backend/trn/tile_primitive/copy/default.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/copy/default.py
rename to python/tvm/backend/trn/tile_primitive/copy/default.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/dim_utils.py
b/python/tvm/backend/trn/tile_primitive/dim_utils.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/dim_utils.py
rename to python/tvm/backend/trn/tile_primitive/dim_utils.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/gemm/__init__.py
b/python/tvm/backend/trn/tile_primitive/gemm/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/gemm/__init__.py
rename to python/tvm/backend/trn/tile_primitive/gemm/__init__.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/gemm/default.py
b/python/tvm/backend/trn/tile_primitive/gemm/default.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/gemm/default.py
rename to python/tvm/backend/trn/tile_primitive/gemm/default.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py
b/python/tvm/backend/trn/tile_primitive/instruction_generator.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/instruction_generator.py
rename to python/tvm/backend/trn/tile_primitive/instruction_generator.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py
b/python/tvm/backend/trn/tile_primitive/private_alloc.py
similarity index 97%
rename from python/tvm/backend/trn/operator/tile_primitive/private_alloc.py
rename to python/tvm/backend/trn/tile_primitive/private_alloc.py
index 1ce89ef4f0..a0759c6d9c 100644
--- a/python/tvm/backend/trn/operator/tile_primitive/private_alloc.py
+++ b/python/tvm/backend/trn/tile_primitive/private_alloc.py
@@ -17,9 +17,9 @@
from typing import Any
-from tvm.backend.trn.operator.tile_primitive.common import init_analyzer,
nki_dim
-from tvm.backend.trn.operator.tile_primitive.dim_utils import get_ewise_dim_map
-from tvm.backend.trn.operator.tile_primitive.instruction_generator import
InstructionGenerator
+from tvm.backend.trn.tile_primitive.common import init_analyzer, nki_dim
+from tvm.backend.trn.tile_primitive.dim_utils import get_ewise_dim_map
+from tvm.backend.trn.tile_primitive.instruction_generator import
InstructionGenerator
from tvm.script import tirx as T
from tvm.tirx import Buffer, FloatImm, Stmt
from tvm.tirx.operator.tile_primitive.ops import (
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/reduction/__init__.py
b/python/tvm/backend/trn/tile_primitive/reduction/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/reduction/__init__.py
rename to python/tvm/backend/trn/tile_primitive/reduction/__init__.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/reduction/default.py
b/python/tvm/backend/trn/tile_primitive/reduction/default.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/reduction/default.py
rename to python/tvm/backend/trn/tile_primitive/reduction/default.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py
b/python/tvm/backend/trn/tile_primitive/reduction/utils.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/reduction/utils.py
rename to python/tvm/backend/trn/tile_primitive/reduction/utils.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/select/__init__.py
b/python/tvm/backend/trn/tile_primitive/select/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/select/__init__.py
rename to python/tvm/backend/trn/tile_primitive/select/__init__.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/select/default.py
b/python/tvm/backend/trn/tile_primitive/select/default.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/select/default.py
rename to python/tvm/backend/trn/tile_primitive/select/default.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/unary/__init__.py
b/python/tvm/backend/trn/tile_primitive/unary/__init__.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/unary/__init__.py
rename to python/tvm/backend/trn/tile_primitive/unary/__init__.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/unary/default.py
b/python/tvm/backend/trn/tile_primitive/unary/default.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/unary/default.py
rename to python/tvm/backend/trn/tile_primitive/unary/default.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/unary/utils.py
b/python/tvm/backend/trn/tile_primitive/unary/utils.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/unary/utils.py
rename to python/tvm/backend/trn/tile_primitive/unary/utils.py
diff --git
a/python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py
b/python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
similarity index 100%
rename from
python/tvm/backend/trn/operator/tile_primitive/unary/with_bias_scale.py
rename to python/tvm/backend/trn/tile_primitive/unary/with_bias_scale.py
diff --git a/python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py
b/python/tvm/backend/trn/tile_primitive/workspace_utils.py
similarity index 100%
rename from python/tvm/backend/trn/operator/tile_primitive/workspace_utils.py
rename to python/tvm/backend/trn/tile_primitive/workspace_utils.py
diff --git a/python/tvm/tirx/operator/intrinsics/_common.py
b/python/tvm/tirx/operator/intrinsics/_common.py
index a96bc8caf3..75f3d62652 100644
--- a/python/tvm/tirx/operator/intrinsics/_common.py
+++ b/python/tvm/tirx/operator/intrinsics/_common.py
@@ -17,7 +17,7 @@
"""Shared enum / value tables for PTX intrinsic schemas and user wrappers.
Single source of truth. Both ``tvm.tirx.op`` (user wrappers that validate
-arguments via ``_choice``) and ``tvm.tirx.cuda.operator.intrinsics.*``
+arguments via ``_choice``) and ``tvm.tirx.cuda.intrinsics.*``
(schema declarations using ``Choice(choices=...)`` / ``IntAttr(choices=...)``)
import from here.
diff --git a/tests/python/tirx/codegen/test_ptx_cvt.py
b/tests/python/tirx/codegen/test_ptx_cvt.py
index e53c762cf8..d36e1c0a81 100644
--- a/tests/python/tirx/codegen/test_ptx_cvt.py
+++ b/tests/python/tirx/codegen/test_ptx_cvt.py
@@ -18,8 +18,8 @@
import pytest
-from tvm.backend.cuda.operator.intrinsics.cvt import _CVT_FORMS,
_cvt_form_parts
-from tvm.backend.cuda.operator.intrinsics.registry import CODEGEN_REGISTRY
+from tvm.backend.cuda.intrinsics.cvt import _CVT_FORMS, _cvt_form_parts
+from tvm.backend.cuda.intrinsics.registry import CODEGEN_REGISTRY
# One concrete expansion for every syntax form in the PTX ``cvt`` grammar.
# Args are register placeholders followed by the form-specific attrs.
diff --git a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
index aadef8a9a7..1927b094da 100644
--- a/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
+++ b/tests/python/tirx/codegen/test_ptx_ld_st_ops.py
@@ -24,7 +24,7 @@ from tvm.ir import Op
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
-from tvm.tirx.cuda.operator.tile_primitive.copy._common import (
+from tvm.tirx.cuda.tile_primitive.copy._common import (
copy_ptx_form,
copy_ptx_ld_return_type,
)
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py
index d1a227d0f8..5381a9a802 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_fallback.py
@@ -37,7 +37,7 @@ from tvm.testing import env
# Force the fallback dispatch to register before any test compiles a kernel.
# Without this import, in fresh pytest workers the `copy/fallback` variant
# isn't yet registered when the dispatcher snapshots its registry.
-from tvm.tirx.cuda.operator.tile_primitive.copy import fallback as
_fallback_module # noqa: F401
+from tvm.tirx.cuda.tile_primitive.copy import fallback as _fallback_module #
noqa: F401
from tvm.tirx.layout import S, TileLayout
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
index a80e34e8ac..edd9e16296 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_gmem_smem.py
@@ -261,7 +261,7 @@ def test_copy_g2s_s2g(task, dtype, scope):
def _align(
g_layout, g_shape, s_layout, s_shape, elem_bits, thread_cnt,
g_region=None, s_region=None
):
- from tvm.tirx.cuda.operator.tile_primitive.copy._common import
align_layouts_gs
+ from tvm.tirx.cuda.tile_primitive.copy._common import align_layouts_gs
target = tvm.target.Target("cuda")
if g_region is None:
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
index 314753d6c1..4d9755d862 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_ld_stmatrix.py
@@ -488,7 +488,7 @@ def test_ldstmatrix_tcgen05_warpgroup_atom_emits_ldmatrix():
``GetScope`` would reject (which silently fell back to a scalar reg path).
Compile-only (no GPU): asserts the instruction appears in generated source.
"""
- from tvm.tirx.cuda.operator.tile_primitive.tma_utils import
mma_shared_layout
+ from tvm.tirx.cuda.tile_primitive.tma_utils import mma_shared_layout
from tvm.tirx.layout import tcgen05_atom_layout
m, k = 64, 64
diff --git a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
index e07dab5b44..8cd88f5220 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_reg.py
@@ -732,7 +732,7 @@ _TCGEN05_D_SLICE = (slice(0, 64), slice(0, 64))
def _tcgen05_d_epilogue_layouts():
- from tvm.tirx.cuda.operator.tile_primitive.tma_utils import
mma_shared_layout
+ from tvm.tirx.cuda.tile_primitive.tma_utils import mma_shared_layout
from tvm.tirx.layout import tcgen05_atom_layout
m, n = _TCGEN05_D_SHAPE
@@ -743,7 +743,7 @@ def _tcgen05_d_epilogue_layouts():
def _build_tcgen05_d_epilogue_deposit():
"""``Tx.wg.copy(smem[slice], d_reg[slice])``: R (tcgen05 atom) → S (128B
swizzle)."""
- from tvm.tirx.cuda.operator.tile_primitive.tma_utils import
mma_shared_layout
+ from tvm.tirx.cuda.tile_primitive.tma_utils import mma_shared_layout
from tvm.tirx.layout import tcgen05_atom_layout
m, n = _TCGEN05_D_SHAPE
@@ -781,7 +781,7 @@ def
test_reg_copy_tcgen05_d_epilogue_deposit_layout_pairing():
``d_reg``: ``(64,64)`` fp32 ``tcgen05_atom_layout("16x256b", ...)``.
``smem_cd_mma``: ``(64,64)`` fp32 ``mma_shared_layout(..., swizzle=128B)``.
"""
- from tvm.backend.cuda.operator.tile_primitive.copy.vec_auto_reg import (
+ from tvm.backend.cuda.tile_primitive.copy.vec_auto_reg import (
_split_thread_loop,
align_layouts_raw,
)
@@ -853,7 +853,7 @@ def _tcgen05_16x256b_row_col(tid_wg: T.int32, lane:
T.int32, reg_idx: T.int32):
def _build_tcgen05_d_epilogue_deposit_roundtrip():
"""Fill ``d_reg``, R→S deposit, S→R reload, dump via ``.local()`` to
gmem."""
- from tvm.tirx.cuda.operator.tile_primitive.tma_utils import
mma_shared_layout
+ from tvm.tirx.cuda.tile_primitive.tma_utils import mma_shared_layout
from tvm.tirx.layout import tcgen05_atom_layout
m, n = _TCGEN05_D_SHAPE
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py
index 98bf68e61f..ff43e433b5 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy/test_swizzle_iter.py
@@ -37,7 +37,7 @@ import pytest
import tvm
from tvm.tirx import Var as _TirVar
-from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
+from tvm.tirx.cuda.tile_primitive.copy._swizzle_iter import (
get_swizzle,
try_recognize,
)
@@ -323,7 +323,7 @@ def test_recognize_linear_iter_pure_case_1d():
of the swizzle period 2^(p+at+sw) (pure Case 1.D, swizzle has no XOR
effect). The iter is stored as a LinearIter (no bit decomposition).
"""
- from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
+ from tvm.tirx.cuda.tile_primitive.copy._swizzle_iter import (
_BitIter,
_LinearIter,
)
@@ -363,7 +363,7 @@ def test_emit_mixed_linear_bit_correctness():
"""Brute-force: for a mixed (LinearIter outer, BitIter inner) pattern,
emit_iter_offset's prediction must equal the actual swizzle output for
every (tid, k) — including the non-pow2 outer extent's coord 2."""
- from tvm.tirx.cuda.operator.tile_primitive.copy._swizzle_iter import (
+ from tvm.tirx.cuda.tile_primitive.copy._swizzle_iter import (
_LinearIter,
)
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py
index ce468c41ad..a639e349ca 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_dsmem.py
@@ -32,7 +32,7 @@ from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
from tvm.tirx import IntImm, Var
-from tvm.tirx.cuda.operator.tile_primitive.copy_async.dsmem import
copy_dsmem_impl
+from tvm.tirx.cuda.tile_primitive.copy_async.dsmem import copy_dsmem_impl
from tvm.tirx.exec_scope import ExecScope
from tvm.tirx.layout import S, TileLayout
from tvm.tirx.operator.tile_primitive.dispatcher import DispatchFail
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
index 657f3a0e31..b807562ca7 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tcgen05_cp.py
@@ -62,7 +62,7 @@ import tvm.testing
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
-from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode,
mma_shared_layout
+from tvm.tirx.cuda.tile_primitive.tma_utils import SwizzleMode,
mma_shared_layout
from tvm.tirx.layout import ComposeLayout, R, S, TCol, TileLayout, TLane
# Multicast replica lane offsets: (extent, stride) pairs on TLane. These are
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
index 318ffba032..bac77e9e2d 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/copy_async/test_tma.py
@@ -31,7 +31,7 @@ from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
from tvm.tirx import IntImm, StringImm, Var
-from tvm.tirx.cuda.operator.tile_primitive.copy_async.tma import (
+from tvm.tirx.cuda.tile_primitive.copy_async.tma import (
AutoIssueAxis,
IssueCoord,
ProofStatus,
@@ -46,7 +46,7 @@ from tvm.tirx.cuda.operator.tile_primitive.copy_async.tma
import (
copy_tma_explicit_impl,
validate_tensor_map_spec,
)
-from tvm.tirx.cuda.operator.tile_primitive.tma_utils import (
+from tvm.tirx.cuda.tile_primitive.tma_utils import (
mma_atom_layout,
mma_atom_shape,
mma_shared_layout,
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
index 62fb0c6db8..89bc27a084 100644
--- a/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
+++ b/tests/python/tirx/operator/tile_primitive/cuda/elementwise/test_unary.py
@@ -24,7 +24,7 @@ import tvm.testing
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
-from tvm.tirx.cuda.operator.tile_primitive.layout_utils import (
+from tvm.tirx.cuda.tile_primitive.layout_utils import (
cast_layout_supported_for_local as _cast_layout_supported_for_local,
)
from tvm.tirx.layout import S, TileLayout, laneid, tid_in_wg, tx, warpid
@@ -1111,7 +1111,7 @@ def test_cast_local_view_sliced(A_dtype, B_dtype,
slice_start, slice_end):
def test_cast_layout_partition_and_validation():
"""Partition table (simplified): partition structure and
_cast_layout_supported_for_local."""
- from tvm.tirx.cuda.operator.tile_primitive.layout_utils import (
+ from tvm.tirx.cuda.tile_primitive.layout_utils import (
get_layout_thread_local_partition as
_get_layout_thread_local_partition,
)
from tvm.tirx.layout import Axis, Iter
@@ -1226,7 +1226,7 @@ def test_cast_mixed_axes_and_subregion(slice_start,
slice_end):
def test_cast_joint_decomposition_extents_order():
"""Test joint decomposition uses thread dims in layout order with correct
extents."""
- from tvm.tirx.cuda.operator.tile_primitive.layout_utils import (
+ from tvm.tirx.cuda.tile_primitive.layout_utils import (
get_layout_thread_local_partition as
_get_layout_thread_local_partition,
)
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
index 02f7fc212d..1e1a504f12 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/gemm_async/test_gemm_async.py
@@ -34,8 +34,8 @@ from tvm.ir.type import PointerType, PrimType
from tvm.script import tirx as T
from tvm.script.tirx import tile as Tx
from tvm.testing import env
-from tvm.tirx.cuda.operator.tile_primitive.gemm_async import sf_tmem_layout
-from tvm.tirx.cuda.operator.tile_primitive.tma_utils import (
+from tvm.tirx.cuda.tile_primitive.gemm_async import sf_tmem_layout
+from tvm.tirx.cuda.tile_primitive.tma_utils import (
SwizzleMode,
mma_atom_layout,
mma_atom_shape,
@@ -2445,7 +2445,7 @@ def test_gemm_tcgen05_contiguous_kslice_partial_k(k_lo,
k_hi):
so the MMA accumulates exactly k in [lo, hi) -- enabling fine K-major
split-K.
Any MMA_K(16)-aligned [lo:hi] is supported.
"""
- from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+ from tvm.tirx.cuda.tile_primitive.tma_utils import SwizzleMode
M, N, K_alloc = 128, 128, 64
dtype = "float16"
@@ -3477,7 +3477,7 @@ def _make_gemm_tcgen05_call(
(1@TLane, 1@TCol) layout.
"""
from tvm.ir import Range
- from tvm.tirx.cuda.operator.tile_primitive.gemm_async.tcgen05 import (
+ from tvm.tirx.cuda.tile_primitive.gemm_async.tcgen05 import (
gemm_async_tcgen05_impl,
)
from tvm.tirx.exec_scope import ExecScope
@@ -3556,7 +3556,7 @@ def
test_gemm_tcgen05_preserves_explicit_tmem_lane_bases():
def test_gemm_tcgen05_preserves_block_scale_tmem_lane_bases():
"""SFA/SFB row offsets must reach the encoded TMEM address operands."""
from tvm.ir import Range
- from tvm.tirx.cuda.operator.tile_primitive.gemm_async.tcgen05 import (
+ from tvm.tirx.cuda.tile_primitive.gemm_async.tcgen05 import (
gemm_async_tcgen05_impl,
)
from tvm.tirx.exec_scope import ExecScope
diff --git
a/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py
b/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py
index b1fc86bd62..1ae7f1706e 100644
---
a/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py
+++
b/tests/python/tirx/operator/tile_primitive/cuda/permute_layout/test_permute_layout.py
@@ -46,7 +46,7 @@ from tvm.script.tirx import tile as Tx
from tvm.testing import env
# Helpers exposed by the dispatcher module for direct algorithm tests.
-from tvm.tirx.cuda.operator.tile_primitive.permute_layout.warp_xor_swizzle
import (
+from tvm.tirx.cuda.tile_primitive.permute_layout.warp_xor_swizzle import (
_bank_free,
_check_bijection,
_choose_xor_k,
diff --git a/tests/python/tirx/test_alloc_pool.py
b/tests/python/tirx/test_alloc_pool.py
index ebc5f751a6..9614531ec0 100644
--- a/tests/python/tirx/test_alloc_pool.py
+++ b/tests/python/tirx/test_alloc_pool.py
@@ -19,7 +19,7 @@
import pytest
from tvm.tirx.cuda.lang.alloc_pool import _validate_mma_alloc_shape
-from tvm.tirx.cuda.operator.tile_primitive.tma_utils import SwizzleMode
+from tvm.tirx.cuda.tile_primitive.tma_utils import SwizzleMode
# ---------------------------------------------------------------------------
# alloc_tcgen05_mma_AB shape validation: bad inputs raise actionable
ValueError instead of
diff --git a/tests/python/tirx/test_layout.py b/tests/python/tirx/test_layout.py
index 1ad21c9b9d..d388d5c951 100644
--- a/tests/python/tirx/test_layout.py
+++ b/tests/python/tirx/test_layout.py
@@ -29,7 +29,7 @@ from tvm.script import tirx as T
from tvm.script.ir_builder import IRBuilder
from tvm.script.ir_builder import tirx as Tx_builder
from tvm.tirx import Var
-from tvm.tirx.cuda.operator.tile_primitive.tma_utils import (
+from tvm.tirx.cuda.tile_primitive.tma_utils import (
SwizzleMode,
mma_shared_layout,
)
@@ -1748,7 +1748,7 @@ def test_slice():
def test_cuda_copy_extract_swizzle_tile_simplifies_constant_region_extents():
- from tvm.tirx.cuda.operator.tile_primitive.copy._common import
_extract_tile
+ from tvm.tirx.cuda.tile_primitive.copy._common import _extract_tile
zero = tvm.tirx.Mul(tvm.tirx.IntImm("int32", 0), tvm.tirx.IntImm("int32",
64))
end = tvm.tirx.Add(zero, tvm.tirx.IntImm("int32", 64))
