// The Swin-layer engine: descriptor-driven e4m3 GEMM phases for the eleven
// cuda_dldn_engine_swin_* denoiser layers.
//
// Why: the translated kernels are ~83,373 instructions for 536 FP8 WMMAs, dominated by
// operand reassembly (9,124 selects, 7,976 LDS permutes, ~23,000 shifts/ands/perms).  An
// engine that keeps its activations as an e4m3 image and reads operands directly from it
// pays none of that.
//
// STATUS: every function here is validated by kernels/tools/rrswin_engine_k16_test.cpp, which
// runs on the GPU with checked HIP calls and compares against host models written independently:
// step 1 one k32 step, step 2 a 64-wide k chain as two k32 steps with f16 rounding between them,
// step 3 staging an accumulator as the next phase's operands (staged image dumped and compared
// byte for byte), step 4a/4b the weight preparation (slots checked directly, then through a GEMM
// against the documented image layout), step 4c a 32-row phase with two m-tiles, which is the
// case swin_gemm_tile's slot_step exists for, step 5/5a the phase loop itself, two phases chained
// through a staged image (the staged image is checked against the model as well), and step 6
// swin_act_pair, the clamped cubic activation, against a scalar model of the specification's
// formula at a one-ulp bar (the device uses packed f16x2 ops, as the PTX does; the bit-exact
// check for it is the oracle comparison).  The layout was measured, not assumed:
// kernels/tools/rrswin_operand_fit.cpp reports D's row index as VGPR i + 8*(lane>>4), its
// column index as lane & 15, and both operands' byte order as the identity.
// Step 3 of that test stages an accumulator and consumes it as the next phase's operands, and
// also dumps the staged image: every staged byte equalled the model's own e4m3 requantisation
// of D[k][col], so the accumulator-to-operand mapping ("D VGPR i is row i + 8h" is the next
// GEMM's k) is measured, not assumed.  Its earlier mismatches were all bugs in the test: a
// k32 step mislabelled k16, a weight array too small for a second GEMM's slots, and a second
// k16 chunk left holding the first GEMM's operands.
//
// Layout facts this file uses, and where each comes from:
//  1. A k16 WMMA operand is 8 e4m3 bytes = this lane's half of a 16-byte row.  "half h the K
//     values 8h .. 8h+7 in order"        -- kernels/common/wmma_layout.h
//  2. The lane/half selection is `8 * hf` with hf = lane >> 4, and the caller reads the row's
//     16 bytes from shared memory itself
//                                        -- kernels/tex/enc0_tail.hip (working native)
//  3. A weight slot is the same 16 bytes per lane, `kslot_from16` is the identity under
//     D4R_TEX_FP8 on gfx12, and `kop_image(img, idx)` picks this half
//                                        -- kernels/tex/tex_common.h
//  4. `k32_e4m3(c, a0, b0, a1, b1)` is one m16n8k32 e4m3 step as two native gfx12 FP8
//     WMMAs and is bit-identical to ZLUDA's gfx12 lowering of NVIDIA's mma
//                                        -- kernels/tex/tex_common.h (with tex_common's own
//                                           comment, and enc0_tail/dec0_head using it in
//                                           shipped, verified kernels)
#pragma once
#include <stddef.h>
#include "../tex/tex_common.h"

namespace rrswin
{

// The engine's LDS-typed register vectors: the images, the prepared slots and the exchange
// live in the caller's shared arena, and a pointer here that is typed with no address space
// is a generic pointer whose accesses lower to flat -- in the cut context, where this body
// is linked, that faults on the raw shared-offset value (VA ~0).  Everything the engine
// reads or writes in shared memory is therefore typed with address_space(3) end to end
// (lds_u8 bytes, lds_v2 the 8-byte k16 operand read, lds_v4/u8 the widened variants of the
// non-native-FP8 layout); global pointers (the weight image, the _prep writes) stay generic.
typedef __attribute__((address_space(3))) uint2_t lds_v2;
typedef __attribute__((address_space(3))) u4v lds_v4;
typedef __attribute__((address_space(3))) u8v lds_v8;

// This lane's operand for k16 chunk `chunk` of image row `row_base + (lane & 15)`; rows are
// `stride` bytes apart.  The row comes from the lane, as in enc0_tail (`m = l & 15`), so a
// caller cannot pass a row that does not match its lane.  A row holds consecutive k, 16
// bytes per k16 chunk, so a k32 image needs stride >= 32; a smaller stride silently reads
// the next row.
__attribute__((device)) static inline kop swin_operand(const lds_u8* img, int stride, int row_base, int chunk)
{
    const uint32_t l = lane_id();
    const uint32_t half_bytes = (l >= 16) ? 8u : 0u;
    const uint32_t row = (uint32_t)row_base + (l & 15u);
    const lds_v2* p = (const lds_v2*)(img + (size_t)row * (size_t)stride + 16u * (uint32_t)chunk + half_bytes);
    return kop_from8(*p);
}

// One m16n16 output tile of one k32 step.  The operand roles follow the working kernels
// (enc0_tail): the prepared weight slot is the A operand and the staged activation is the B
// operand, so the WMMA computes D^T relative to the token-major view.  Both operands are
// lane-indexed: A's row and B's column are both `lane & 15` (wmma_layout.h: "A operand lane
// l = row r of A, B operand lane l = column r of B", r = l & 15), which is why the weight
// slot is `slot_base + (lane & 15)` and the two k16 halves of that slot are 16 slots apart
// -- exactly enc0_tail's `kop_image(w, slot)` / `kop_image(w, slot + 16)` with
// `slot = base + m`.
// `slot_step` is the distance between a row's two k16 slots.  It is a parameter, not a
// constant, because a phase with more than one m-tile puts the next tile's rows in between:
// with slot_step fixed at 16 a 32-row phase would read tile 1's rows as tile 0's second k16.
// `k32_index` selects which 32-wide slice of the activations this step reads: a phase with more
// than one k32 step must advance the image's chunks too, or every step re-reads k 0..31.
// The engine's own slot reader: the prepared slots live in the caller's shared arena, so the
// read must be AS3-typed (tex_common.h's generic kop_image is for global slot banks and stays
// for those callers).  Byte offsets are the same as the generic one's: 8 bytes per kslot under
// the native-FP8 layout, 16 under the widened one.
__attribute__((device)) static inline kop kop_image_lds(const lds_u8* img, int idx)
{
    const int i = 2 * idx + (int)(lane_id() >> 4);
#if defined(D4R_TEX_FP8)
    return *(const lds_v2*)(img + 8 * (unsigned)i);
#else
    return *(const lds_v4*)(img + 16 * (unsigned)i);
#endif
}

__attribute__((device)) static inline f8v swin_gemm_tile(f8v c, const lds_u8* w, int slot_base,
                                                         int slot_step, const lds_u8* img, int stride,
                                                         int row_base, int k32_index = 0)
{
    const int slot = slot_base + (int)(lane_id() & 15u);
    const kop a0 = kop_image_lds(w, slot);
    const kop b0 = swin_operand(img, stride, row_base, 2 * k32_index);
    const kop a1 = kop_image_lds(w, slot + slot_step);
    const kop b1 = swin_operand(img, stride, row_base, 2 * k32_index + 1);
    return k32_e4m3(c, a0, b0, a1, b1);
}

// Stage an accumulator as the operand image of the next phase.  On gfx12 the accumulator's
// layout *is* the operand layout (wmma_layout.h: "D VGPR i is row i + 8h, which is the gfx12
// B-operand layout of the next GEMM (no exchange)"): this lane's VGPR i holds the value that
// belongs at k = i + 8*(lane >> 4) of column (lane & 15), so each lane requantises its own
// eight values to e4m3 and writes them at its own row's k-half -- no cross-lane movement.
// `row_base` is the first token column this tile covers and `chunk` the k16 chunk the tile's
// rows land in (a 16-row m-tile is one k16 chunk), so a phase with several tiles stages each
// into its own place.
__attribute__((device)) static inline void swin_stage_operand(lds_u8* img, int stride, int row_base,
                                                              int chunk, const f8v acc)
{
    const uint32_t l = lane_id();
    const uint32_t h = l >> 4, col = (uint32_t)row_base + (l & 15u);
    lds_u8* p = img + (size_t)col * (size_t)stride + 16u * (uint32_t)chunk + 8u * h;
#pragma unroll
    for (int i = 0; i < 8; i += 2)
    {
        const hv2 pair = {(half_t)acc[i], (half_t)acc[i + 1]};
        const uint32_t code = e4m3x2(__builtin_bit_cast(uint32_t, pair));
        p[i] = (uint8_t)(code & 0xffu);
        p[i + 1] = (uint8_t)(code >> 8);
    }
}

// ---------------------------------------------------------------- weight preparation
// One phase's weights as WMMA slots, so the phase loop can read them with kop_image and never
// touch the image layout.  Slot `base + 2*n_rows*s + n_rows*h + n` holds the k32 step s, k16
// half h and output row n, i.e. 16 consecutive e4m3 bytes of k 32s + 16h .. +15; the caller
// passes that same `n_rows` as swin_gemm_tile's slot_step.  The 2*n_rows per step is what keeps
// a step's two k16 halves adjacent for any row count.
//
// `waddr(k, n)` returns the byte offset of the weight image's element (k, n).  Layer bodies
// document their own; enc0's is
//   base + 512*(n>>4) + 64*(n&7) + 16*((k>>1)&3) + 8*((n>>3)&1) + 4*((k>>3)&1) + 2*((k>>4)&1) + (k&1)
// and a layer whose image differs passes its own lookup.  Run it once per launch (a `_prep`
// kernel, as kernels/tex/enc0_tail.hip does for its weights).
template <typename Lookup>
__attribute__((device)) static inline void swin_prep_phase(uint8_t* slots, int base, int n_rows,
                                                           int k_steps, const uint8_t* image,
                                                           Lookup waddr)
{
    for (int n = 0; n < n_rows; ++n)
        for (int s = 0; s < k_steps; ++s)
            for (int h = 0; h < 2; ++h)
                for (int j = 0; j < 16; ++j)
                    slots[(size_t)(16 * (base + 2 * n_rows * s + n_rows * h + n) + j)] =
                        image[waddr(32 * s + 16 * h + j, n)];
}

// ---------------------------------------------------------------- f16 pair helpers
// The packed-f16 plumbing kernels/rr/rrswin_enc0.hip defines for its epilogues.
__attribute__((device)) static inline half_t hf16(uint16_t bits) { return __builtin_bit_cast(half_t, bits); }
__attribute__((device)) static inline float h2f(uint32_t bits)
{
    return (float)__builtin_bit_cast(half_t, (uint16_t)bits);
}
__attribute__((device)) static inline hv2 hv2_of(uint32_t p) { return __builtin_bit_cast(hv2, p); }
__attribute__((device)) static inline uint32_t u32_of(hv2 v) { return __builtin_bit_cast(uint32_t, v); }
__attribute__((device)) static inline uint32_t pair_of(half_t a, half_t b)
{
    return (uint32_t)__builtin_bit_cast(uint16_t, a) | ((uint32_t)__builtin_bit_cast(uint16_t, b) << 16);
}

// ---------------------------------------------------------------- epilogues
// The clamped cubic activation every layer's MLP applies before requantising:
//   y = clamp(x, -2, +2);  out = x * (0.5 + y * (0.412109375 - 0.0810546875 * |y|))
// in f16 throughout, on the two halves of a packed pair.  (Not SiLU/GELU: g(+2) = 1 and
// g(-2) = 0.)  The constants are the f16 values of the PTX's 0f3ED306EB/0f3DA60DD6/0f3F000000/
// 0f40000000 as kernels/rr/rrswin_enc0.hip documents them.
__attribute__((device)) static inline uint32_t swin_act_pair(uint32_t x)
{
    const hv2 two = hv2_of(pair_of(hf16(0x4000), hf16(0x4000)));
    const hv2 slope = hv2_of(pair_of(hf16(0x2d30), hf16(0x2d30)));
    const hv2 mid = hv2_of(pair_of(hf16(0x3698), hf16(0x3698)));
    const hv2 half0 = hv2_of(pair_of(hf16(0x3800), hf16(0x3800)));
    hv2 y = __builtin_elementwise_max(hv2_of(x), -two);
    y = __builtin_elementwise_min(y, two);
    y = __builtin_elementwise_max(y, -y);
    const hv2 g = half0 + y * (mid - slope * y);
    return u32_of(hv2_of(x) * g);
}

// Stage NVIDIA's m16n8k32 A fragments as an image.  The fragment layout is the one
// kernels/rr/rrswin_enc0.hip documents and verified against the oracle:
//   reg 0 byte j -> row g,     k = 2t + (j&1) + 16*((j>>1)&1)
//   reg 1 byte j -> row g+8,   same k
//   reg 2 byte j -> row g,     k = 8 + 2t + (j&1) + 16*((j>>1)&1)
//   reg 3 byte j -> row g+8,   same k
// with g = lane>>2 and t = lane&3.  One m-tile's fragment is four registers per lane, so a
// phase's 4*m_tiles registers cover 16*m_tiles rows.  Each lane scatters its own 16 bytes, which
// is a shared-memory store per byte -- the whole point of the engine is that this happens once
// per phase instead of once per mma.
__attribute__((device)) static inline void swin_stage_fragment(lds_u8* img, int stride,
                                                               const uint32_t* frag, int m_tiles)
{
    const uint32_t l = lane_id();
    const uint32_t g = l >> 2, t = l & 3u;
    for (int mi = 0; mi < m_tiles; ++mi)
    {
        const uint32_t* f = frag + 4 * mi;
        lds_u8* r0 = img + (size_t)(16 * mi + g) * (size_t)stride;
        lds_u8* r1 = img + (size_t)(16 * mi + g + 8) * (size_t)stride;
#pragma unroll
        for (int j = 0; j < 4; ++j)
        {
            const uint32_t k0 = 2 * t + (j & 1) + 16u * ((j >> 1) & 1);
            const uint32_t k2 = 8u + k0;
            r0[k0] = (uint8_t)(f[0] >> (8 * j));
            r1[k0] = (uint8_t)(f[1] >> (8 * j));
            r0[k2] = (uint8_t)(f[2] >> (8 * j));
            r1[k2] = (uint8_t)(f[3] >> (8 * j));
        }
    }
}

// ---------------------------------------------------------------- the phase loop
// One phase: m_tiles x n_tiles tiles of a k32-steps GEMM, each tile a 16x16 m16n16 WMMA block,
// with the epilogue staging each tile's accumulator as the next phase's activations.
struct Phase
{
    int m_tiles;      // 16-row tiles of the weight image
    int n_tiles;      // 16-token tiles
    int k32_steps;    // k = 32 * k32_steps
    int slot_base;    // first weight slot of this phase (see swin_prep_phase)
    int epilogue;     // 0 stage the accumulator, 1 keep it in acc
    int dst_image;    // image the epilogue stages into (epilogue 0)
};

// Run one phase from `src`, staging requantized accumulators into `dst` or retaining them.
// `weight_rows` is the output row count used to prepare each pair of k16 weight slot banks.
// When retaining results, `acc` must hold m_tiles * n_tiles entries in row-major tile order.
__attribute__((device)) static inline void swin_phase(const Phase& ph, const lds_u8* w, int weight_rows,
                                                      const lds_u8* src, int src_stride,
                                                      lds_u8* dst, int dst_stride, f8v* acc)
{
    for (int mt = 0; mt < ph.m_tiles; ++mt)
        for (int nt = 0; nt < ph.n_tiles; ++nt)
        {
            f8v c = {0, 0, 0, 0, 0, 0, 0, 0};
            for (int step = 0; step < ph.k32_steps; ++step)
                c = swin_gemm_tile(c, w, ph.slot_base + 16 * mt + 2 * weight_rows * step, weight_rows,
                                   src, src_stride, 16 * nt, step);
            if (ph.epilogue == 0)
                swin_stage_operand(dst, dst_stride, 16 * nt, mt, c);
            else if (acc != nullptr)
                acc[mt * ph.n_tiles + nt] = c;
        }
}

// enc0's weight image layout, as its body documents it.
__attribute__((device)) static inline int swin_waddr_enc0(int k, int n)
{
    return 512 * (n >> 4) + 64 * (n & 7) + 16 * ((k >> 1) & 3) + 8 * ((n >> 3) & 1)
         + 4 * ((k >> 3) & 1) + 2 * ((k >> 4) & 1) + (k & 1);
}

// ------------------------------------------------------------------ the enc0 port
// Address space 3 (the shared arena) explicitly: a generic access lowers to flat, which
// faults (rrswin_enc0.hip's typedef comment).  The cut body consumes rrswin::lds_u32 and
// must not define its own.
typedef unsigned lds_u32 __attribute__((address_space(3)));

// ZLUDA's relu_keep_nan_halves on a packed pair, byte for byte the body's relu_pair
// (rrswin_enc0.hip): a negative or -0 half becomes +0, NaN kept, positive unchanged.
__attribute__((device)) static inline uint32_t swin_relu_pair(uint32_t p)
{
    if ((p & 0x8000u) != 0 && (p & 0x7fffu) <= 0x7c00u)
    {
        p &= 0xffff0000u;
    }
    if ((p & 0x80000000u) != 0 && ((p >> 16) & 0x7fffu) <= 0x7c00u)
    {
        p &= 0x0000ffffu;
    }
    return p;
}

// The dcell exchange (contract §3), in u32 units: tile block (mt*n_nt + nt) is 128 words,
// the writer's own lane slot l*4 + p inside it, the reader walking the same words at
// (8rh + g + 16*(nt&1))*4 + t.  Writer and reader are separate calls; the wave barriers
// between them belong to the caller.
__attribute__((device)) static inline void swin_dcell_exchange_write(lds_u32* lds,
                                                                     const uint32_t* pairs,
                                                                     int n_mt, int n_nt)
{
    const uint32_t l = lane_id() & 31u;
    for (int mt = 0; mt < n_mt; ++mt)
        for (int nt = 0; nt < n_nt; ++nt)
#pragma unroll
            for (int p = 0; p < 4; ++p)
                lds[(mt * n_nt + nt) * 128 + l * 4 + p] = pairs[(mt * n_nt + nt) * 4 + p];
}

// The matching reader: the epilogue's d[mi][rh][nt] from the same lds words.
__attribute__((device)) static inline void swin_dcell_exchange_read(const lds_u32* lds, uint32_t* d,
                                                                    int n_mi, int n_nt4, int n_nt)
{
    const uint32_t l = lane_id();
    const uint32_t g = (l >> 2) & 7u, t = l & 3u;
    for (int mi = 0; mi < n_mi; ++mi)
        for (int rh = 0; rh < 2; ++rh)
#pragma unroll
            for (int nt = 0; nt < n_nt4; ++nt)
                d[(mi * 2 + rh) * n_nt4 + nt] =
                    lds[((nt >> 1) * n_nt + mi) * 128 + (8 * rh + g + 16 * (nt & 1)) * 4 + t];
}

// ------------------------------------------------------------------ slots from wimg
// One phase's B fragments (the WMMA slot order, rrswin_enc0.hip's documented layout)
// scattered into the slots swin_gemm_tile reads, without touching the image layout:
// one tile's 8 u32 per lane wait at w + base + 512*t + 16*lane + 4*word for tile t and
// fragment word, and their bytes are (n, k) at
//   (t, lane, word, j) -> n = 16*t + (lane>>2) + 8*(word>>1), k = 2*(lane&3) + (j&1)
//                          + 16*(j>>1) + 8*(word&1)
//                                  -- gemm_native_unbiased, enc0 phases 1-20 being a
//                                     k32-steps variant of the same D4R_ZLUDA_WMMA_FP8 + bfly
//                                     shape (D layout, wmma_layout.h)
// The destination layout is the slot one: slot (n_rows*(k>>4) + n)*16 + (k&15) of the
// phase's own bank, with the bytes in k order, exactly swin_gemm_tile's read with
// slot_base = 16*mt and slot_step = n_rows.  Like the body's loops that make the same
// scatter, it is wave-wide (no barrier inside); the caller orders it against its GEMM.
__attribute__((device)) static inline void swin_slots_from_wimg(lds_u8* slots, const uint8_t* w,
                                                                int base, int n_rows)
{
    const uint32_t l = lane_id();
    for (int t = 0; t < n_rows / 16; ++t)
        for (int word = 0; word < 4; ++word)
        {
            const uint32_t n = 16u * (uint32_t)t + (l >> 2) + 8u * ((uint32_t)word >> 1);
            const uint32_t src = *(const uint32_t*)(w + base + 512 * t + 16 * l + 4 * word);
#pragma unroll
            for (int j = 0; j < 4; ++j)
            {
                const uint32_t k = 2u * (l & 3u) + (j & 1) + 16u * (j >> 1) + 8u * (word & 1);
                slots[(size_t)(n_rows * (k >> 4) + (int)n) * 16 + (k & 15)] =
                    (uint8_t)(src >> (8 * j));
            }
        }
}

// ---------------------------------------------- stage a tile's pairs with a transform
// swin_stage_operand from a tile's four f16 pairs (p = f16(D[2p]) | f16(D[2p+1]) << 16)
// with the epilogue's per-pair transform before e4m3x2: mode 0 plain, 1 the relu
// (swin_relu_pair), 2 the clamped cubic (swin_act_pair).  Both transforms are
// per-element, so each is bit-identical to the body's packed-register path.  This
// mirrors swin_stage_operand byte for byte, with the caller's pair array in place of
// the accumulator.
__attribute__((device)) static inline void swin_stage_operand_p(lds_u8* img, int stride,
                                                                int row_base, int chunk,
                                                                const uint32_t pairs[4], int mode)
{
    const uint32_t l = lane_id();
    const uint32_t h = l >> 4, col = (uint32_t)row_base + (l & 15u);
    lds_u8* p = img + (size_t)col * (size_t)stride + 16u * (uint32_t)chunk + 8u * h;
#pragma unroll
    for (int p_ = 0; p_ < 4; ++p_)
    {
        uint32_t pair = pairs[p_];
        if (mode == 1)
            pair = swin_relu_pair(pair);
        else if (mode == 2)
            pair = swin_act_pair(pair);
        const uint32_t code = e4m3x2(pair);
        p[2 * p_] = (uint8_t)(code & 0xffu);
        p[2 * p_ + 1] = (uint8_t)(code >> 8);
    }
}

} // namespace rrswin