// SPDX-License-Identifier: Apache-2.0
// K11 depthwise reader (sceneseg-p150 OPT_PLAN K11, tt/dw.py): the float32 TILE input map [1, 1, Hin*Win, C] (DRAM
// interleaved) of one depthwise kxk conv, gathered into one TILE page per (output pixel tile, channel tile, tap).
//
// Work: the NW = CT * PT items (channel tile ct, output pixel tile pt; ct-major) split over the first NC cores in
// row-major order, core i owning [i * B + min(i, E), + B + (i < E)) (no per-core RT args: probe P3). A core walks
// its items in chunks of at most CH pixel tiles of one ct. Per chunk:
//   1. the input rows the chunk needs are loaded into an L1 band: two planes per row (channel half h = 0 / 1, 64 B
//      per pixel = 16 fp32 channels), P zero columns on each side; for stride 2 each plane is stored column-parity
//      major (even padded columns, then odd), so every tap reads consecutive output pixels from consecutive band
//      pixels. A 16-pixel face half of an input tile is one contiguous DRAM run: stride 1 reads it straight into
//      the band; stride 2 reads the row into a staging row (CB 5, two rows ping-pong) and scatters its pixels to
//      the parity planes with 64 B local copies (64 B DRAM reads were latency-bound: R2 profile).
//   2. per output tile, the KK TILE pages (one per tap) are assembled from the band: face (rg, h) = 16 output
//      pixels of channel half h, copied in runs (split at output-row ends); out-of-image rows and pixels past
//      Hout*Wout read a zero block. Every copy is a local NoC read; ONE barrier per output tile (KK pages pushed
//      together, CB 0 holds 2 * KK pages), not per page.
// When ct changes the NWT weight tiles (tile ct * NWT + j; NWT = KK: every row of tile j = w[ct*32 .. +32, tap j];
// R7.1 DW_WPACK: NWT = ceil(KK / 8), rows 4l .. 4l+3 of tile j = tap 8j + l) and the bias tile ct are read into
// CBs 1 / 2 (resident until the next ct).
// CT args: [0] HIN, [1] WIN, [2] HOUT, [3] WOUT, [4] K, [5] S, [6] P, [7] CT, [8] PT, [9] NC, [10] GX, [11] CH,
//          [12] NRMAX (band rows), [13] STAGE (stride 2: 1 = staging rows, 0 = 64 B DRAM reads), [14] NWT (weight
//          tiles per channel tile), then the TensorAccessorArgs of x, w, b.
// Common RT args: [x_addr, w_addr, b_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t HIN = get_compile_time_arg_val(0);
constexpr uint32_t WIN = get_compile_time_arg_val(1);
constexpr uint32_t HOUT = get_compile_time_arg_val(2);
constexpr uint32_t WOUT = get_compile_time_arg_val(3);
constexpr uint32_t K = get_compile_time_arg_val(4);
constexpr uint32_t S = get_compile_time_arg_val(5);
constexpr uint32_t P = get_compile_time_arg_val(6);
constexpr uint32_t CT = get_compile_time_arg_val(7);
constexpr uint32_t PT = get_compile_time_arg_val(8);
constexpr uint32_t NC = get_compile_time_arg_val(9);
constexpr uint32_t GX = get_compile_time_arg_val(10);
constexpr uint32_t CH = get_compile_time_arg_val(11);
constexpr uint32_t NRMAX = get_compile_time_arg_val(12);
constexpr uint32_t STAGE = get_compile_time_arg_val(13);
constexpr uint32_t NWT = get_compile_time_arg_val(14);
constexpr auto x_args = TensorAccessorArgs<15>();
constexpr auto w_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();

constexpr uint32_t KK = K * K;
constexpr uint32_t HWO = HOUT * WOUT;
constexpr uint32_t NW = CT * PT;
constexpr uint32_t TB = 4096;
constexpr uint32_t WP = WIN + 2 * P;                      // padded input width (pixels)
constexpr uint32_t PH = (WP + 1) / 2;                     // stride 2: pixels of the even-column half of a plane row
constexpr uint32_t RS = (S == 1) ? WP : 2 * PH;           // pixels per plane row
constexpr uint32_t cb_x = 0, cb_w = 1, cb_b = 2, cb_band = 3, cb_zero = 4, cb_stage = 5;
constexpr uint32_t STAGE_ROW = WIN * 128;                 // one staging row: two planes of WIN pixels x 64 B

// L1 byte address of padded column j of band row r, channel half h
FORCE_INLINE uint32_t band_addr(uint32_t base, uint32_t r, uint32_t h, uint32_t j) {
    uint32_t pos;
    if constexpr (S == 1) {
        pos = j;
    } else {
        pos = (j & 1) * PH + (j >> 1);
    }
    return base + (((r * 2 + h) * RS) + pos) * 64;
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NW / NC, E = NW % NC;
    const uint32_t w0 = core * B + (core < E ? core : E);
    const uint32_t w1 = w0 + B + (core < E ? 1 : 0);
    const auto xa = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB);
    const auto wa = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), TB);
    const auto ba = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), TB);

    // the 1 KB zero block (16 pixels x 64 B)
    const uint32_t zero = get_write_ptr(cb_zero);
    {
        volatile tt_l1_ptr uint32_t* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero);
        for (uint32_t i = 0; i < 256; ++i) {
            z[i] = 0;
        }
    }
    const uint64_t zero_noc = get_noc_addr(zero);
    const uint32_t band = get_write_ptr(cb_band);
    const uint32_t stage = (S == 2 && STAGE) ? get_write_ptr(cb_stage) : 0;

    uint32_t cur_ct = 0xFFFFFFFF;
    uint32_t w = w0;
    while (w < w1) {
        const uint32_t ct = w / PT;
        const uint32_t pt0 = w - ct * PT;
        uint32_t pt1 = pt0 + CH;
        if (pt1 > PT) {
            pt1 = PT;
        }
        if (pt1 - pt0 > w1 - w) {
            pt1 = pt0 + (w1 - w);
        }
        if (ct != cur_ct) {
            cb_reserve_back(cb_w, NWT);
            uint32_t l1 = get_write_ptr(cb_w);
            for (uint32_t t = 0; t < NWT; ++t) {
                noc_async_read(wa.get_noc_addr(ct * NWT + t), l1, TB);
                l1 += TB;
            }
            cb_reserve_back(cb_b, 1);
            noc_async_read(ba.get_noc_addr(ct), get_write_ptr(cb_b), TB);
            noc_async_read_barrier();
            cb_push_back(cb_w, NWT);
            cb_push_back(cb_b, 1);
            cur_ct = ct;
        }
        // ---- 1. the band: padded input rows R0 .. R0 + NR - 1 (input row iy = R0 + r - P)
        const uint32_t p_first = pt0 * 32;
        uint32_t p_last = pt1 * 32 - 1;
        if (p_last >= HWO) {
            p_last = HWO - 1;
        }
        const uint32_t ya = p_first / WOUT, yb = p_last / WOUT;
        const uint32_t R0 = ya * S;
        const uint32_t NR = (yb - ya) * S + K;   // <= NRMAX by construction (host)
        for (uint32_t r = 0; r < NR; ++r) {
            const int32_t iy = (int32_t)(R0 + r) - (int32_t)P;
            if (iy < 0 || iy >= (int32_t)HIN) {
                continue;
            }
            for (uint32_t h = 0; h < 2; ++h) {
                for (uint32_t j = 0; j < P; ++j) {
                    noc_async_read(zero_noc, band_addr(band, r, h, j), 64);
                    noc_async_read(zero_noc, band_addr(band, r, h, P + WIN + j), 64);
                }
            }
            const uint32_t qa = (uint32_t)iy * WIN;
            const uint32_t qb = qa + WIN;
            if constexpr (S == 1) {
                uint32_t q = qa;
                while (q < qb) {
                    uint32_t ge = (q | 15) + 1;
                    if (ge > qb) {
                        ge = qb;
                    }
                    const uint32_t n = ge - q;
                    const uint32_t tile = (q >> 5) * CT + ct;
                    const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
                    for (uint32_t h = 0; h < 2; ++h) {
                        noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), band_addr(band, r, h, P + q - qa),
                                       n * 64);
                    }
                    q = ge;
                }
            } else if constexpr (STAGE == 0) {
                // stride 2 without staging (R2.1 form, knob DW_STAGE=0): one 64 B DRAM read per pixel and half
                for (uint32_t q = qa; q < qb; ++q) {
                    const uint32_t tile = (q >> 5) * CT + ct;
                    const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
                    for (uint32_t h = 0; h < 2; ++h) {
                        noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), band_addr(band, r, h, P + q - qa), 64);
                    }
                }
            } else {
                const uint32_t st = stage + (r & 1) * STAGE_ROW;
                uint32_t q = qa;
                while (q < qb) {
                    uint32_t ge = (q | 15) + 1;
                    if (ge > qb) {
                        ge = qb;
                    }
                    const uint32_t n = ge - q;
                    const uint32_t tile = (q >> 5) * CT + ct;
                    const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
                    for (uint32_t h = 0; h < 2; ++h) {
                        noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), st + h * (WIN * 64) + (q - qa) * 64,
                                       n * 64);
                    }
                    q = ge;
                }
                // the staging row is complete (and the copies out of the other staging row, issued one row ago,
                // are done before that row is refilled: this barrier covers them too)
                noc_async_read_barrier();
                noc_async_read_set_state(zero_noc);
                for (uint32_t h = 0; h < 2; ++h) {
                    for (uint32_t px = 0; px < WIN; ++px) {
                        noc_async_read_with_state(st + h * (WIN * 64) + px * 64, band_addr(band, r, h, P + px), 64);
                    }
                }
            }
        }
        noc_async_read_barrier();
        // ---- 2. KK TILE pages per output tile. The output pixels of the tile split into at most 8 segments (per
        // 16-pixel row group: runs inside one output row, or past Hout*Wout); a tap only shifts each segment's band
        // source, so the segments are computed once per tile. All copies are local: NoC state set once.
        noc_async_read_set_state(zero_noc);
        for (uint32_t pt = pt0; pt < pt1; ++pt) {
            uint32_t seg_d[8], seg_n[8], seg_y[8], seg_x[8], ns = 0;
            for (uint32_t rg = 0; rg < 2; ++rg) {
                const uint32_t p = pt * 32 + rg * 16;
                uint32_t y = p / WOUT;
                uint32_t x = p - y * WOUT;
                uint32_t i = 0;
                while (i < 16) {
                    seg_d[ns] = rg * 2048 + i * 64;
                    if (p + i >= HWO) {
                        seg_n[ns] = (16 - i) * 64;
                        seg_y[ns] = 0xFFFF;                               // zero segment
                        seg_x[ns] = 0;
                        ++ns;
                        break;
                    }
                    uint32_t n = WOUT - x;
                    if (n > 16 - i) {
                        n = 16 - i;
                    }
                    seg_n[ns] = n * 64;
                    seg_y[ns] = y;
                    seg_x[ns] = x;
                    ++ns;
                    i += n;
                    x += n;
                    if (x == WOUT) {
                        x = 0;
                        ++y;
                    }
                }
            }
            cb_reserve_back(cb_x, KK);
            uint32_t dst = get_write_ptr(cb_x);
            for (uint32_t dy = 0; dy < K; ++dy) {
                for (uint32_t dx = 0; dx < K; ++dx, dst += TB) {
                    for (uint32_t g = 0; g < ns; ++g) {
                        const uint32_t d0 = dst + seg_d[g];
                        const uint32_t rr = seg_y[g] * S + dy;                // padded row (huge for a zero segment)
                        if (rr >= P && rr < P + HIN) {
                            const uint32_t src = band_addr(band, rr - R0, 0, seg_x[g] * S + dx);
                            noc_async_read_with_state(src, d0, seg_n[g]);
                            noc_async_read_with_state(src + RS * 64, d0 + 1024, seg_n[g]);
                        } else {
                            noc_async_read_with_state(zero, d0, seg_n[g]);
                            noc_async_read_with_state(zero, d0 + 1024, seg_n[g]);
                        }
                    }
                }
            }
            noc_async_read_barrier();
            cb_push_back(cb_x, KK);
        }
        w += pt1 - pt0;
    }
}