[FIX][TIRx] Use physical order for Buffer.local views

docs/tirx/native_basics/cuda/buffers.rst

@@ -422,7 +422,8 @@ or hand you a pointer — they emit no runtime op of their own. The common ones:
    * - ``B.view(*shape, layout=…)``
      - reinterpret the same storage under a new shape/layout (no copy)
    * - ``B.local(*shape, layout=…)``
-     - the calling thread's private register slice of a ``local`` buffer
+     - the calling thread's private register slice of a ``local`` buffer,
+       in physical storage order by default
    * - ``B.permute(*dims)``
      - a view with axes permuted (a transposed layout)
    * - ``B.access_ptr(mask, …)``
@@ -471,13 +472,23 @@ sees the 256-element buffer as ``64×4``; ``A.permute(1, 0)`` transposes the axe
     At_ptr[(j * 4) + i]                       // permute: swapped strides
 
 **Registers — ``local``.** Decomposes a thread-axis ``local`` layout into the
-calling thread's flat register bundle (used pervasively by the tile primitives):
+calling thread's register bundle (used pervasively by the tile primitives).
+Both forms expose the raw physical storage span by default, including layout
+gaps and offsets: ``R.local()`` infers a flat 1-D span, while
+``R.local(d0, d1, ...)`` is a row-major reshape whose product must equal that
+span.  Pass an explicit ``layout=`` only when storage-iterator coordinates are
+required.  This mediated form is an escape hatch: the supplied layout
+interprets the requested shape, so that shape is not constrained to the raw
+span.  Without ``layout=``, omitting the shape always infers a one-dimensional
+physical storage span.  Only the explicit-``layout=`` compatibility form with
+no shape infers a one-dimensional shape from the logical storage size:
 
 .. code-block:: python
 
     R  = T.alloc_buffer((32, 8), "float32", scope="local", layout=TileLayout(S[(32, 8) : (1 @ laneid, 1)]))
-    Rl = R.local(8)          # this lane's 8 registers
+    R_flat = R.local()       # this lane's 8 registers, physical order
+    R_2d = R.local(2, 4)     # the same registers, row-major 2x4 reshape
 
 .. code-block:: c++
 
-    alignas(64) float Rl_ptr[8];             // the lane's private registers
+    alignas(64) float R_flat_ptr[8];          // the lane's private registers

docs/tirx/tile_primitives/copy/reg.rst

@@ -121,7 +121,7 @@ loop (not ``T.unroll`` — same flooding rationale as :doc:`gmem_smem`):
 
 .. code-block:: python
 
-    r_local = r_buf.local(*per_thread_r_shape)   # flat per-thread registers
+    r_local = r_buf.local()                      # raw per-thread physical span
     for f in range(total_outer):
         ds, dr = _outer_const_offsets(outer, f)               # shared / reg deltas
         s_ptr = _ptr_off(s_buf.ptr_to(s_zero_indices), _s_iter_off(f, ds, s_off))
@@ -131,6 +131,11 @@ loop (not ``T.unroll`` — same flooding rationale as :doc:`gmem_smem`):
         else:
             copy_op(r_ptr, s_ptr)     # shared/global -> register
 
+``dr`` and ``r_off_base`` are physical storage offsets, so the register alias
+must use the raw-span ``local()`` view.  This remains correct when storage
+iterators are permuted or leave gaps; every physical slot through
+``layout.storage().span()`` is directly addressable.
+
 Generated TIRx IR
 -----------------
 

docs/tirx/tile_primitives/gemm.rst

@@ -88,7 +88,7 @@ accumulate) — one ``m16n8k16`` atom (from ``test_gemm_mma_m16n8k_.py``):
         D_f = T.alloc_buffer((16, 8),  "float32", scope="local", layout=D_FRAG)
         A_reg = A_f.local(8)                              # stage A into the lane's 8 regs
         for s in T.unroll(8):
-            kp, kHi, rM = s % 2, (s // 2) % 2, s // 4
+            kp, rM, kHi = s % 2, (s // 2) % 2, s // 4
             A_reg[s] = A_g[lane // 4 + 8 * rM, 2 * (lane % 4) + kp + 8 * kHi]
         B_reg = B_f.local(4)                              # stage B into the lane's 4 regs
         for s in T.unroll(4):
@@ -109,9 +109,13 @@ group the operand sub-layouts (``D_M, D_N, A_M, A_K, B_K, B_N, C_*``) into the f
 m16n8k frame, anchoring A/C on D's M, B/C on D's N, and B on A's K. The first
 instruction that fits, with matching warp-tiling, wins.
 
-**2. Derive register layouts.** Each operand gets a per-lane register view: D/C as
-``[Mo, No, rM, rN]`` (4 f32), A as ``[Mo, Ko, rM, kHi, k_pack]``, B as
-``[Ko, No, kHi, k_pack]`` — the exact register order ``mma.sync`` expects.
+**2. Derive register layouts.** Each operand gets a mediated per-lane view:
+D/C as ``[Mo, No, rM, rN]`` (4 f32), A as
+``[Mo, Ko, rM, kHi, k_pack]``, and B as
+``[Ko, No, kHi, k_pack]``.  The corresponding raw physical register order
+exposed by the default ``local`` view is D/C ``[Mo, No, rM, rN]``, A
+``[Mo, Ko, kHi, rM, k_pack]``, and B ``[Ko, No, kHi, k_pack]`` — the PTX
+operand enumeration that ``mma.sync`` expects.
 
 **3. Emit the unrolled nest** — initialize D (from C if ``beta==1``, else 0), then
 accumulate over K in place, one ``mma`` per (m, n) tile:

python/tvm/backend/cuda/operator/tile_primitive/copy/reg.py

@@ -488,10 +488,6 @@ def _emit_reg(op_call: TilePrimitiveCall, sctx: DispatchContext) -> PrimFunc:
     vec_len = _choose_vec_len(elem_bits, atoms, r_p, s_p)
     vec_bits = vec_len * elem_bits
     outer = _split_atoms_for_vec(atoms, vec_len)
-    per_thread_r_total = 1
-    for it in r_iters:
-        per_thread_r_total *= int(it.extent)
-    per_thread_r_shape = [per_thread_r_total or 1]
 
     # Build the per-thread S offset OUTSIDE the impl using placeholder Vars
     # (one per thread axis). Inside the impl we'll declare the real scope_ids
@@ -554,7 +550,9 @@ def _s_iter_off(f, ds, s_off):
     def impl():
         s_off = _substitute_axes(s_off_template, placeholders, sctx)
         _setup_swizzle(s_off)
-        r_local = r_buf.local(*per_thread_r_shape)
+        # ``dr`` and ``r_off_base`` below are physical storage offsets, so use
+        # the raw-span local view rather than storage-iterator coordinates.
+        r_local = r_buf.local()
         # Keep as a serial TIR loop and let ptxas unroll downstream. An
         # explicit ``T.unroll`` materializes the per-iter scratch
         # (ds/dr/s_ptr/r_ptr, swizzle ``v_<n>[]`` signed-strides) as N

python/tvm/backend/cuda/operator/tile_primitive/reduction/local.py

@@ -137,8 +137,8 @@ def _gen_warp_shuffle_reduce(src, dst, reduce_width, local_elems, accum, op_type
     # fmt: off
     @T.prim_func(check_well_formed=False)
     def impl():
-        src_local = src.local(local_elems)
-        dst_local = dst.local(local_elems)
+        src_local = src.local(local_elems, layout=src.layout.storage())
+        dst_local = dst.local(local_elems, layout=dst.layout.storage())
         for k in T.serial(local_elems):
             if not is_same_buffer:
                 dst_local[k] = src_local[k]
@@ -357,8 +357,10 @@ def inner_shuffle(v, shuffle_mask):
     if need_save_accum:
         @T.prim_func(check_well_formed=False)
         def impl():
-            src_local = src.local(*src_local_shape)
-            dst_local = dst.local(*dst_local_shape)
+            # These dimensions are storage-iterator coordinates, so request
+            # the mediated layout explicitly.  Bare local() is physical-order.
+            src_local = src.local(*src_local_shape, layout=src.layout.storage())
+            dst_local = dst.local(*dst_local_shape, layout=dst.layout.storage())
             old_val = T.alloc_buffer([1], dtype, scope="local")
 
             for spa in T.serial(dst_local_total):
@@ -376,8 +378,10 @@ def impl():
     else:
         @T.prim_func(check_well_formed=False)
         def impl():
-            src_local = src.local(*src_local_shape)
-            dst_local = dst.local(*dst_local_shape)
+            # These dimensions are storage-iterator coordinates, so request
+            # the mediated layout explicitly.  Bare local() is physical-order.
+            src_local = src.local(*src_local_shape, layout=src.layout.storage())
+            dst_local = dst.local(*dst_local_shape, layout=dst.layout.storage())
 
             for spa in T.serial(dst_local_total):
                 dst_idx = T.meta_var(get_indices(spa, dst_local_st, dst_local_ext))

python/tvm/tirx/buffer.py

@@ -387,30 +387,59 @@ def _infer_shape(shape):
     def local(self, *shape, layout=None) -> "Buffer":
         """Create a thread-local view of this buffer.
 
+        By default, both the inferred and explicit-shape forms address the
+        raw physical storage span.  ``local()[k]`` is the k-th physical
+        storage element, including any gaps or layout offset, while
+        ``local(d0, d1, ...)`` is a row-major reshape of that same span.
+        Pass ``layout=`` to request a mediated view explicitly.  This is an
+        escape hatch whose shape is interpreted by the supplied layout.  When
+        that shape is explicit, the parent buffer does not need a layout.
+
         When called with no shape arguments, auto-infers a 1D shape from
-        the layout's non-thread component (i.e. ``layout.storage().shard``).
+        the span of the parent layout's non-thread component (i.e.
+        ``self.layout.storage().span()``).  The explicit-``layout=`` form
+        instead infers the parent layout's ``storage().size()`` for
+        compatibility.  Either inference requires the parent buffer to have a
+        layout.
 
         Parameters
         ----------
         shape : tuple of Expr
-            The shape of the local view for indexing. If omitted, a 1D
-            shape is computed automatically.
+            The shape of the local view for indexing.  Without ``layout=``,
+            its product must equal the per-thread physical storage span.
+            With an explicit layout, the shape is not constrained by the raw
+            span.  If omitted, a matching 1D shape is computed automatically.
 
         layout : optional
-            Override layout. If None, uses the storage layout
-            (parent layout with thread axes removed).
+            Override layout. If None, the default (identity) layout is used.
 
         Returns
         -------
         local : DeclBufferFrame
             The corresponding local buffer.
         """
         if not shape:
-            local_layout = self.layout.storage()
-            total = functools.reduce(
-                lambda x, y: x * y, [it.extent for it in local_layout.shard], 1
-            )
-            shape = (total,)
+            if self.layout is None:
+                raise ValueError(
+                    "Buffer.local cannot infer a shape because the parent buffer has layout=None; "
+                    "pass an explicit shape together with layout=..."
+                )
+            storage_layout = self.layout.storage()
+            local_extent = storage_layout.span() if layout is None else storage_layout.size()
+            shape = (local_extent,)
+        elif layout is None:
+            if self.layout is None:
+                raise ValueError(
+                    "Buffer.local without layout= cannot validate the physical storage span "
+                    "because the parent buffer has layout=None; pass an explicit layout=..."
+                )
+            local_extent = self.layout.storage().span()
+            shape_total = functools.reduce(lambda x, y: x * y, shape, 1)
+            if not tvm.arith.Analyzer().can_prove_equal(shape_total, local_extent):
+                raise ValueError(
+                    f"Local view shape {shape} has {shape_total} elements, "
+                    f"but the buffer has physical storage span {local_extent} per thread"
+                )
         return tvm.tirx.script.builder.decl_buffer(
             shape,
             self.dtype,
@@ -421,7 +450,7 @@ def local(self, *shape, layout=None) -> "Buffer":
             self.scope(),
             self.data_alignment,
             self.offset_factor,
-            self.layout.storage() if layout is None else layout,
+            "default" if layout is None else layout,
         )
 
     def permute(self, *dims) -> "Buffer":