// SPDX-License-Identifier: Apache-2.0
// EPI reader (sceneseg-p150 OPT_PLAN EPI, tt/epi.py): the NP fp32 term outputs of one conv / linear (DRAM or L1
// interleaved TILE, one shape) into CBs 0 .. NP-1, and once per program the CT fp32 bias tiles (the bias row
// replicated over the 32 rows of a tile on the host: what binary_ng's reader builds with fill_tile_with_first_row)
// into CB 3, kept resident for the whole program.
// EPI-split (OPT_REPORT R2.3) operands: HAS_RES = a residual map of the output's shape (tile t) into CB 4, read with
// the parts; HAS_ROW = the CT tiles of a [1, 1, 1, C] fp32 row tensor (the SE gate, computed on the device) into
// CB 5, row 0 replicated over the 32 rows here (binary_ng's fill_tile_with_first_row), resident.
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores in row-major order owns the
// contiguous tile range [i * B + min(i, E), + B + (i < E)) of the T output tiles, B = T / NC, E = T % NC.
// A part may be a column range of a wider tensor (ST2: the [W_hi | W_lo] conv's [R, 2C] output read as its two
// products): output tile t = (r, c) = (t / CT, t % CT) reads tile r * W_i + O_i + c of part i.
// Resident tiles (R4.4 EPI-B, knob SCENESEG_EPI_BIAS_LOCAL): with BLOCAL a core reads only the bias / row tiles of the
// output columns its tile range touches (min(n, CT) tiles from column t0 % CT, wrapping) into their CB slots; the other
// slots stay unread (never used by its compute). Without it every core reads all CT tiles (CT x 4 KB per core: 17.7 MB
// per program on the 1152-channel maps).
// CT args: [0] T, [1] NC, [2] GX, [3] CT (output column tiles), [4] NP (0..3), [5] HAS_BIAS, [6] NB (tiles per read
//          batch), [7 + 2i] / [8 + 2i] W_i / O_i (column tiles / column-tile offset of part i, i < 3), [13] HAS_RES,
//          [14] HAS_ROW, [15] UNIT (tiles per work unit: the split hands out whole units; 1, or the tiles of one image
//          row for the K8 class map), [16] BLOCAL, then the TensorAccessorArgs of part 0, part 1, part 2 (repeats of part 0 when
//          NP < 3), the bias, the residual and the row (part 0 when absent). Common RT args: [p0, p1, p2, bias, res,
//          row].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t T = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t CT = get_compile_time_arg_val(3);
constexpr uint32_t NP = get_compile_time_arg_val(4);
constexpr uint32_t HAS_BIAS = get_compile_time_arg_val(5);
constexpr uint32_t NB = get_compile_time_arg_val(6);
constexpr uint32_t W0 = get_compile_time_arg_val(7), O0 = get_compile_time_arg_val(8);
constexpr uint32_t W1 = get_compile_time_arg_val(9), O1 = get_compile_time_arg_val(10);
constexpr uint32_t W2 = get_compile_time_arg_val(11), O2 = get_compile_time_arg_val(12);
constexpr uint32_t HAS_RES = get_compile_time_arg_val(13);
constexpr uint32_t HAS_ROW = get_compile_time_arg_val(14);
constexpr uint32_t UNIT = get_compile_time_arg_val(15);
constexpr uint32_t BLOCAL = get_compile_time_arg_val(16);
constexpr auto p0_args = TensorAccessorArgs<17>();
constexpr auto p1_args = TensorAccessorArgs<p0_args.next_compile_time_args_offset()>();
constexpr auto p2_args = TensorAccessorArgs<p1_args.next_compile_time_args_offset()>();
constexpr auto bs_args = TensorAccessorArgs<p2_args.next_compile_time_args_offset()>();
constexpr auto rs_args = TensorAccessorArgs<bs_args.next_compile_time_args_offset()>();
constexpr auto rw_args = TensorAccessorArgs<rs_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t cb_bias = 3, cb_res = 4, cb_row = 5;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = T / UNIT / NC, E = T / UNIT % NC;
    const uint32_t t0 = (core * B + (core < E ? core : E)) * UNIT;
    const uint32_t n = (B + (core < E ? 1 : 0)) * UNIT;
    if (n == 0) {
        return;
    }
    const auto a0 = TensorAccessor(p0_args, get_common_arg_val<uint32_t>(0), TB);
    const auto a1 = TensorAccessor(p1_args, get_common_arg_val<uint32_t>(1), TB);
    const auto a2 = TensorAccessor(p2_args, get_common_arg_val<uint32_t>(2), TB);
    const auto ar = TensorAccessor(rs_args, get_common_arg_val<uint32_t>(4), TB);
    // the resident column tiles this core reads: NR from column C0 (wrapping), or all CT
    const uint32_t NR = (BLOCAL && n < CT) ? n : CT;
    const uint32_t C0 = (BLOCAL && n < CT) ? t0 % CT : 0;
    if constexpr (HAS_BIAS) {
        const auto ab = TensorAccessor(bs_args, get_common_arg_val<uint32_t>(3), TB);
        cb_reserve_back(cb_bias, CT);
        const uint32_t base = get_write_ptr(cb_bias);
        for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
            noc_async_read(ab.get_noc_addr(c), base + c * TB, TB);
        }
        noc_async_read_barrier();
        cb_push_back(cb_bias, CT);
    }
    if constexpr (HAS_ROW) {
        const auto aw = TensorAccessor(rw_args, get_common_arg_val<uint32_t>(5), TB);
        cb_reserve_back(cb_row, CT);
        const uint32_t base = get_write_ptr(cb_row);
        for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
            noc_async_read(aw.get_noc_addr(c), base + c * TB, TB);
        }
        noc_async_read_barrier();
        // row 0 of faces 0 / 1 down their 16 rows, then faces 2 / 3 = faces 0 / 1
        for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
            const uint32_t tl = base + c * TB;
            for (uint32_t f = 0; f < 2; ++f) {
                const uint64_t src = get_noc_addr(tl + f * 1024);
                for (uint32_t r = 1; r < 16; ++r) {
                    noc_async_read(src, tl + f * 1024 + r * 64, 64);
                }
            }
        }
        noc_async_read_barrier();
        for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
            const uint32_t tl = base + c * TB;
            noc_async_read(get_noc_addr(tl), tl + 2048, 2048);
        }
        noc_async_read_barrier();
        cb_push_back(cb_row, CT);
    }
    for (uint32_t t = t0; t < t0 + n; t += NB) {
        const uint32_t nb = (t + NB <= t0 + n) ? NB : t0 + n - t;
        uint32_t w0 = 0, w1 = 0, w2 = 0, wr = 0;
        if constexpr (NP > 0) {
            cb_reserve_back(0, nb);
            w0 = get_write_ptr(0);
        }
        if constexpr (NP > 1) {
            cb_reserve_back(1, nb);
            w1 = get_write_ptr(1);
        }
        if constexpr (NP > 2) {
            cb_reserve_back(2, nb);
            w2 = get_write_ptr(2);
        }
        if constexpr (HAS_RES) {
            cb_reserve_back(cb_res, nb);
            wr = get_write_ptr(cb_res);
        }
        for (uint32_t j = 0; j < nb; ++j) {
            const uint32_t r = (t + j) / CT, c = (t + j) % CT;
            if constexpr (NP > 0) {
                noc_async_read(a0.get_noc_addr(r * W0 + O0 + c), w0, TB);
                w0 += TB;
            }
            if constexpr (NP > 1) {
                noc_async_read(a1.get_noc_addr(r * W1 + O1 + c), w1, TB);
                w1 += TB;
            }
            if constexpr (NP > 2) {
                noc_async_read(a2.get_noc_addr(r * W2 + O2 + c), w2, TB);
                w2 += TB;
            }
            if constexpr (HAS_RES) {
                noc_async_read(ar.get_noc_addr(t + j), wr, TB);
                wr += TB;
            }
        }
        noc_async_read_barrier();
        if constexpr (NP > 0) {
            cb_push_back(0, nb);
        }
        if constexpr (NP > 1) {
            cb_push_back(1, nb);
        }
        if constexpr (NP > 2) {
            cb_push_back(2, nb);
        }
        if constexpr (HAS_RES) {
            cb_push_back(cb_res, nb);
        }
    }
}