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

permute_layout
==============

``permute_layout`` rearranges a warp's data from a source ``TileLayout`` to a
destination one — typically an in-place transpose. The single CUDA variant
(``warp_xor_swizzle``) stages each lane's elements through registers and writes them
back under the destination layout, with a per-lane **XOR swizzle** on the iteration
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/tile_primitive/permute_layout/warp_xor_swizzle.py``.

What it accepts
---------------

The predicate (``_why_reject``) gates the variant:

.. code-block:: python

    if sctx.scope_kind != "warp":                  return "scope is not 'warp'"
    if src_buf.dtype != dst_buf.dtype:             return "dtype mismatch"
    if src_ext_i != dst_ext_i:                     return "extent mismatch"
    if dtype_bytes not in (1, 2, 4, 8, 16):        return "unsupported dtype byte width"
    if not isinstance(src_buf.layout, TileLayout): return "src not a plain TileLayout"
    if not isinstance(dst_buf.layout, TileLayout): return "dst not a plain TileLayout"
    # + layouts must slice/canonicalize; _choose_xor_k must find a valid k (else fail)

.. list-table::
   :header-rows: 1
   :widths: 22 78

   * - Property
     - Requirement
   * - target / scope / priority
     - ``cuda``; **warp** scope only; priority ``20``
   * - operands
     - equal dtype, equal (compile-time) extents; both plain ``TileLayout`` (no
       swizzle wrapper); dtype byte width ∈ {1, 2, 4, 8, 16}
   * - bank-freedom
     - ``_choose_xor_k`` must find an XOR-bit count ``k ∈ [0, log2(P)]`` that makes
       **both** phases bank-conflict-free, else the dispatch declines (``fail``)

Demonstration program
----------------------

A warp transposes the inner ``4×32`` block of a scale-factor tile — source layout
strides ``(…, 32, 1)``, destination ``(…, 1, 4)`` — for two pipeline stages (the
canonical SF-transpose, from ``test_permute_layout.py``):

.. code-block:: python

    pipe, blk, dtype = 2, 128, "float32"; high = 1
    shape = (pipe, high, 4, 32)
    pre  = TileLayout(S[shape : (blk, 128, 32, 1)])   # source
    post = TileLayout(S[shape : (blk, 128, 1, 4)])    # destination (4↔32 transposed)

    @T.prim_func
    def f(A: T.handle, B: T.handle):
        A_buf = T.match_buffer(A, shape, dtype, layout=pre)
        B_buf = T.match_buffer(B, shape, dtype, layout=post)
        T.device_entry(); T.cta_id([1]); T.thread_id([32])
        for s in T.serial(0, pipe):
            Tx.warp.permute_layout(B_buf[s, 0:1, 0:4, 0:32], A_buf[s, 0:1, 0:4, 0:32])

Algorithm
---------

**1. Align the two layouts.** Both layouts are sliced to the region and
canonicalized; if their shards differ in structure (a linear layout collapses to 1-D
under canon, a transposed one keeps its multi-dim shape) the source is regrouped to
the destination's shape. From the destination shard come the iteration ``extent``
and the per-side strides ``src_str`` / ``dst_str``. ``P`` = elements per lane =
``prod(extent) / 32`` (here ``4``).

**2. Choose the XOR swizzle.** ``_choose_xor_k`` simulates the shared-memory bank
pattern at shard granularity for ``k = 0, 1, … log2(P)`` and picks the smallest
``k`` whose ``shift`` / ``mask`` make *both* phases conflict-free (here ``shift = 3``,
``mask = 3``).

**3. Emit two register-staged phases.** Each lane reads its ``P`` elements through
the source layout into registers (the swizzle permutes which register holds which
iteration), a ``warp_sync`` follows, then the registers are written back through the
destination layout:

.. code-block:: python

    regs = T.alloc_buffer((P,), dtype, scope="local")
    for r in T.unroll(0, P):                                   # read via src layout
        j   = r ^ ((lane_id >> shift) & mask)
        idx = decompose(lane_id + j * 32, extent)
        regs[r] = src_buf[project(idx, src_st)]
    T.cuda.warp_sync()
    for r in T.unroll(0, P):                                   # write via dst layout
        j   = r ^ ((lane_id >> shift) & mask)
        idx = decompose(lane_id + j * 32, extent)
        dst_buf[project(idx, dst_st)] = regs[r]
    T.cuda.warp_sync()

Generated TIRx IR
-----------------

.. code-block:: python

    regs[r] = A_buf[s*128 + (r ^ ((tx >> 3) & 3)) % 4 * 32 + tx]   # phase 1 (src order)
    T.cuda.warp_sync()
    B_buf[s*128 + tx * 4 + (r ^ ((tx >> 3) & 3)) % 4] = regs[r]    # phase 2 (dst order)
    T.cuda.warp_sync()

Generated CUDA
--------------

.. code-block:: c++

    alignas(64) float regs_ptr[4];
    regs_ptr[0] = A_buf_ptr[(s*128) + (((0 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
    regs_ptr[1] = A_buf_ptr[(s*128) + (((1 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
    regs_ptr[2] = A_buf_ptr[(s*128) + (((2 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
    regs_ptr[3] = A_buf_ptr[(s*128) + (((3 ^ ((threadIdx.x >> 3) & 3)) & 3) * 32) + threadIdx.x];
    __syncwarp();
    // ... 4 transposed writes into B_buf_ptr, then __syncwarp();

Each lane owns column ``threadIdx.x`` and stages its 4 rows through ``regs``; the
``(threadIdx.x >> 3)`` XOR rotates the register order per lane-group of 8 so the
write phase hits distinct banks. Verified on ``sm_100a`` — the ``4×32`` block is
transposed for every pipeline stage.

How inputs change the algorithm
-------------------------------

.. list-table::
   :header-rows: 1
   :widths: 28 72

   * - input
     - effect
   * - layout strides (the permutation)
     - define ``extent`` / ``src_str`` / ``dst_str`` and hence ``P`` and the
       per-element index math (the transpose pattern)
   * - dtype byte width
     - feeds the bank simulation in ``_choose_xor_k``; 4-byte dtypes always admit a
       valid ``k`` (one element per bank), while stride-1 **sub-4-byte** reads can
       pack several lanes into one bank and make the dispatch ``fail``
   * - chosen ``k``
     - sets ``shift`` / ``mask`` of the XOR swizzle (``k = 0`` ⇒ no swizzle)
   * - ``P`` (= elements/lane)
     - the number of staged registers and unrolled iterations per phase
