// SPDX-License-Identifier: Apache-2.0
// S1 reader (sceneseg-p150 OPT_PLAN S1, tt/subpixel.py): the operands of the sub-pixel up-block epilogue, per output
// tile of the 2H x 2W map (32 pixels x 32 channels):
//   CB 0 / CB 1  the Z conv's hi / lo products of the tile's 32 pixels, gathered by the shifted depth-to-space
//                Y[2h + a, 2w + b] = Z[(a, b), h + a, w + b]. Z: ZSPLIT 1 = one [(H+1)(W+1), 8 Cout] fp32 TILE
//                tensor, columns term-major then phase (2a + b) then Cout; ZSPLIT 2 = one [(H+1)(W+1), 4 Cout]
//                tensor per term, columns phase-major; ZSPLIT 4 = one [(H+1)(W+1), 2 Cout] tensor per phase,
//                columns term-major. Whole 4 KB Z tiles are read from DRAM into an L1 cache
//                (2 terms x 4 phases x S slots, slot = Z row tile % S), then each output row (two 64 B face rows) is
//                copied locally by the NoC; never 64 B DRAM reads (OPT_REPORT R2.2).
//   CB 2 / CB 6  the composed skip conv's hi / lo products (S: [4HW, 2 Cout] fp32 TILE), when HAS_SKIP
//   CB 3         the tile's bias: the interior bias tile (resident in CB 9, copied locally), or for a border row
//                tile its pattern tile (DRAM), from the per-row-tile index (CB 10: 0 = interior, k = pattern k - 1)
// Work split (no per-core RT args, probe P3): units u = c * R + r (c = channel tile, r = output row tile), core i of
// the first NC (row-major) owns the contiguous range [i * B + min(i, E), + B + (i < E)), B = U / NC, E = U % NC.
// Units go in batches of up to NB tiles (one channel tile, no cache-slot conflict): the batch's Z misses are read
// with one barrier, then the batch's pages are filled (CB pages addressed with wrap-around) with one barrier and
// pushed one by one (a producer push must not cross the CB end).
// CT args: [0] R, [1] CT, [2] NC, [3] GX, [4] H, [5] W, [6] HAS_SKIP, [7] IDXB (bytes of the index page), [8] NB,
//          [9] ZSPLIT (1 / 2 / 4), [10] ZSHIFT (1: Z on the (H+1) x (W+1) grid of the 2x2 conv, phase (a, b) of
//          output (2h + a, 2w + b) at Z pixel (h + a, w + b); 0: Z on the H x W grid of the 3x3 form, at (h, w)),
//          [11] ZTERMS (2: Z holds the hi and lo weight-term products; 1: one product, no CB 1, round 4 P2),
//          [12] STERMS (the same for the skip conv: 2 = [Shi | Slo], 1 = S only, no CB 6),
//          [13] BLOCAL (R4.4 EPI-B: read only the interior bias tiles of the channel tiles this core's units cover,
//          c = u0 / R .. (u1 - 1) / R; 0 = all CT),
//          then the TensorAccessorArgs of Z0..Z3 (Z0 repeated past ZSPLIT), S (Z0 when
//          absent), the interior bias, the patterns, the index.
// Common RT args: [z0, z1, z2, z3, s, bias, patterns, index].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t R = get_compile_time_arg_val(0);
constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t H = get_compile_time_arg_val(4);
constexpr uint32_t W = get_compile_time_arg_val(5);
constexpr uint32_t HAS_SKIP = get_compile_time_arg_val(6);
constexpr uint32_t IDXB = get_compile_time_arg_val(7);
constexpr uint32_t NB = get_compile_time_arg_val(8);
constexpr uint32_t ZSPLIT = get_compile_time_arg_val(9);
constexpr uint32_t ZSHIFT = get_compile_time_arg_val(10);
constexpr uint32_t ZTERMS = get_compile_time_arg_val(11);
constexpr uint32_t STERMS = get_compile_time_arg_val(12);
constexpr uint32_t BLOCAL = get_compile_time_arg_val(13);
constexpr auto z0_args = TensorAccessorArgs<14>();
constexpr auto z1_args = TensorAccessorArgs<z0_args.next_compile_time_args_offset()>();
constexpr auto z2_args = TensorAccessorArgs<z1_args.next_compile_time_args_offset()>();
constexpr auto z3_args = TensorAccessorArgs<z2_args.next_compile_time_args_offset()>();
constexpr auto s_args = TensorAccessorArgs<z3_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();
constexpr auto p_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr auto i_args = TensorAccessorArgs<p_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t W2 = 2 * W, WP = W + ZSHIFT, SW = STERMS * CT;
constexpr uint32_t ZW = ZTERMS * 4 * CT / ZSPLIT;           // column tiles of a Z tensor
constexpr uint32_t ZTERM = ZSPLIT == 1 ? 4 * CT : (ZSPLIT == 4 ? CT : 0);   // column-tile offset of the lo term
constexpr uint32_t ZPH = ZSPLIT == 4 ? 0 : CT;              // column-tile offset per phase
constexpr uint32_t S = 4;                                   // cache slots per (term, phase)
constexpr uint32_t NSLOT = 2 * 4 * S;
constexpr uint32_t cb_zh = 0, cb_zl = 1, cb_sh = 2, cb_b = 3, cb_sl = 6, cb_cache = 8, cb_bint = 9, cb_idx = 10;
constexpr uint32_t cb_src = 11;                             // L1 scratch: NB x 32 source-row offsets (uint32)

FORCE_INLINE uint32_t face_off(uint32_t row) { return (row >> 4) * 2048 + (row & 15) * 64; }

// address of page j after the write pointer of a CB, wrapped at its end
FORCE_INLINE uint32_t page_addr(uint32_t cb, uint32_t j) {
    const auto& ci = get_local_cb_interface(cb);
    uint32_t a = ci.fifo_wr_ptr + j * TB;
    if (a >= ci.fifo_limit) {
        a -= ci.fifo_size;
    }
    return a;
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t U = R * CT;
    constexpr uint32_t B = U / NC, E = U % NC;
    const uint32_t u0 = core * B + (core < E ? core : E);
    const uint32_t u1 = u0 + B + (core < E ? 1 : 0);
    if (u1 == u0) {
        return;
    }
    const auto az0 = TensorAccessor(z0_args, get_common_arg_val<uint32_t>(0), TB);
    const auto az1 = TensorAccessor(z1_args, get_common_arg_val<uint32_t>(1), TB);
    const auto az2 = TensorAccessor(z2_args, get_common_arg_val<uint32_t>(2), TB);
    const auto az3 = TensorAccessor(z3_args, get_common_arg_val<uint32_t>(3), TB);
    const auto as = TensorAccessor(s_args, get_common_arg_val<uint32_t>(4), TB);
    const auto ab = TensorAccessor(b_args, get_common_arg_val<uint32_t>(5), TB);
    const auto ap = TensorAccessor(p_args, get_common_arg_val<uint32_t>(6), TB);
    const auto ai = TensorAccessor(i_args, get_common_arg_val<uint32_t>(7), IDXB);
    // NoC address of a Z tile of term tm (0 hi, 1 lo) and phase ph (tile = row tile * ZW + column tile)
    auto zaddr = [&](uint32_t tm, uint32_t ph, uint32_t tile) -> uint64_t {
        if constexpr (ZSPLIT == 1) {
            return az0.get_noc_addr(tile);
        } else if constexpr (ZSPLIT == 2) {
            return tm == 0 ? az0.get_noc_addr(tile) : az1.get_noc_addr(tile);
        } else {
            return ph == 0 ? az0.get_noc_addr(tile)
                           : (ph == 1 ? az1.get_noc_addr(tile)
                                      : (ph == 2 ? az2.get_noc_addr(tile) : az3.get_noc_addr(tile)));
        }
    };

    const uint32_t idx_l1 = get_write_ptr(cb_idx);
    noc_async_read(ai.get_noc_addr(0), idx_l1, IDXB);
    const uint32_t bint = get_write_ptr(cb_bint);
    const uint32_t c_lo = BLOCAL ? u0 / R : 0, c_hi = BLOCAL ? (u1 - 1) / R : CT - 1;
    for (uint32_t c = c_lo; c <= c_hi; ++c) {
        noc_async_read(ab.get_noc_addr(c), bint + c * TB, TB);
    }
    noc_async_read_barrier();
    volatile tt_l1_ptr uint32_t* idx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(idx_l1);
    const uint32_t cache = get_write_ptr(cb_cache);
    const uint64_t local = get_noc_addr(cache);

    uint32_t tag[NSLOT / 2];                               // Z row tile held by slot (phase, k) (both terms alike)
    uint32_t used[NSLOT / 2];                              // batch that last used the slot
    for (uint32_t i = 0; i < NSLOT / 2; ++i) {
        tag[i] = 0xFFFFFFFFu;
        used[i] = 0xFFFFFFFFu;
    }
    // per batch tile and row: cache offset of the source row (L1 scratch, not the RISC stack)
    volatile tt_l1_ptr uint32_t* src = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_src));
    uint32_t miss[NSLOT / 2], miss_rt[NSLOT / 2], miss_ph[NSLOT / 2];
    uint32_t cur_c = 0xFFFFFFFFu;
    uint32_t batch = 0;
    uint32_t u = u0;
    while (u < u1) {
        const uint32_t c = u / R;
        if (c != cur_c) {
            cur_c = c;
            for (uint32_t i = 0; i < NSLOT / 2; ++i) {
                tag[i] = 0xFFFFFFFFu;
            }
        }
        // ---- plan the batch: up to NB tiles of channel tile c without a slot conflict
        uint32_t k = 0, nmiss = 0;
        while (k < NB && u + k < u1 && (u + k) / R == c) {
            const uint32_t r = (u + k) % R;
            uint32_t p = r * 32;
            uint32_t y = p / W2;
            uint32_t x = p - y * W2;
            volatile tt_l1_ptr uint32_t* rows = src + k * 32;
            bool ok = true;
            for (uint32_t i = 0; i < 32 && ok; ++i) {
                const uint32_t a = y & 1, b = x & 1;
                const uint32_t zp = ((y >> 1) + a * ZSHIFT) * WP + (x >> 1) + b * ZSHIFT;
                const uint32_t zrt = zp >> 5;
                const uint32_t ph = 2 * a + b;
                const uint32_t sl = ph * S + (zrt % S);
                if (tag[sl] != zrt) {
                    if (used[sl] == batch) {
                        ok = false;                        // the slot holds a row tile this batch still needs
                        break;
                    }
                    tag[sl] = zrt;
                    miss[nmiss] = sl;
                    miss_rt[nmiss] = zrt * ZW + ph * ZPH + c;
                    miss_ph[nmiss] = ph;
                    ++nmiss;
                }
                used[sl] = batch;
                rows[i] = sl * TB + face_off(zp & 31);
                if (++x == W2) {
                    x = 0;
                    ++y;
                }
            }
            if (!ok) {
                break;
            }
            ++k;
        }
        // (k >= 1 always: a fresh batch has every slot free for one tile's <= 8 (phase, row tile) pairs: a tile
        // spans at most two image rows, 2W >= 32 is checked on the host)
        for (uint32_t m = 0; m < nmiss; ++m) {
            const uint32_t dst = cache + miss[m] * TB;
            noc_async_read(zaddr(0, miss_ph[m], miss_rt[m]), dst, TB);                              // hi term
            if constexpr (ZTERMS == 2) {
                noc_async_read(zaddr(1, miss_ph[m], miss_rt[m] + ZTERM), dst + (NSLOT / 2) * TB, TB);  // lo term
            }
        }
        noc_async_read_barrier();
        // ---- fill the batch's pages
        cb_reserve_back(cb_zh, k);
        if constexpr (ZTERMS == 2) {
            cb_reserve_back(cb_zl, k);
        }
        cb_reserve_back(cb_b, k);
        if constexpr (HAS_SKIP) {
            cb_reserve_back(cb_sh, k);
            if constexpr (STERMS == 2) {
                cb_reserve_back(cb_sl, k);
            }
        }
        noc_async_read_set_state(local);
        for (uint32_t j = 0; j < k; ++j) {
            const uint32_t dh = page_addr(cb_zh, j);
            const uint32_t dl = ZTERMS == 2 ? page_addr(cb_zl, j) : 0;
            for (uint32_t i = 0; i < 32; ++i) {
                const uint32_t so = cache + src[j * 32 + i];
                const uint32_t fo = face_off(i);
                noc_async_read_with_state(so, dh + fo, 64);
                noc_async_read_with_state(so + 1024, dh + fo + 1024, 64);
                if constexpr (ZTERMS == 2) {
                    noc_async_read_with_state(so + (NSLOT / 2) * TB, dl + fo, 64);
                    noc_async_read_with_state(so + (NSLOT / 2) * TB + 1024, dl + fo + 1024, 64);
                }
            }
        }
        for (uint32_t j = 0; j < k; ++j) {
            const uint32_t r = (u + j) % R;
            const uint32_t pi = idx[r];
            const uint32_t db = page_addr(cb_b, j);
            if (pi == 0) {
                noc_async_read(get_noc_addr(bint + c * TB), db, TB);
            } else {
                noc_async_read(ap.get_noc_addr((pi - 1) * CT + c), db, TB);
            }
            if constexpr (HAS_SKIP) {
                noc_async_read(as.get_noc_addr(r * SW + c), page_addr(cb_sh, j), TB);
                if constexpr (STERMS == 2) {
                    noc_async_read(as.get_noc_addr(r * SW + CT + c), page_addr(cb_sl, j), TB);
                }
            }
        }
        noc_async_read_barrier();
        for (uint32_t j = 0; j < k; ++j) {
            cb_push_back(cb_zh, 1);
            if constexpr (ZTERMS == 2) {
                cb_push_back(cb_zl, 1);
            }
            cb_push_back(cb_b, 1);
            if constexpr (HAS_SKIP) {
                cb_push_back(cb_sh, 1);
                if constexpr (STERMS == 2) {
                    cb_push_back(cb_sl, 1);
                }
            }
        }
        u += k;
        ++batch;
    }
}