[TIRx] Flatten cuda/trn backend operator folder

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

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

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

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)))

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
 ---------------

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
 ---------------

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
 ---------------

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
 ---------------

python/tvm/backend/cuda/__init__.py

@@ -23,7 +23,16 @@
 
 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",
 ]

python/tvm/backend/cuda/operator/intrinsics/_schema.py → python/tvm/backend/cuda/intrinsics/_schema.py

@@ -43,8 +43,8 @@
 
 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

python/tvm/backend/cuda/lang/alloc_pool.py

@@ -69,15 +69,15 @@ 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
 
 
 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 @@ def alloc_sf(self, shape, dtype, *, sf_per_mma, sf_reuse=1):
         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 @@ def alloc(
 
     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,
         )

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
 
 

python/tvm/backend/cuda/operator/tile_primitive/elementwise/__init__.py → python/tvm/backend/cuda/tile_primitive/elementwise/__init__.py

@@ -30,7 +30,7 @@
 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

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",
 ]

python/tvm/backend/trn/operator/tile_primitive/private_alloc.py → 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 (

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.