// SPDX-License-Identifier: Apache-2.0
// BRES reader v2 (meteor-p150 OPT2 round 5, tt/bresize.py, METEOR_BRES=cache): bres_reader.cpp with reader-private
// L1 caches of the input blocks and the weight tiles. v1 read every entry's KM x CT input tiles and its weight tile
// from DRAM: a x2 resize reads each input tile ~10 times (the output blocks of one core share their input rows), so
// the kernel was bound by its DRAM reads. v2 keeps the input blocks (all NX x CT tiles of block k) in a direct-mapped
// cache of S slots (tag k) and the weight tiles in one of SW slots (tag w): a block's misses are read
// from DRAM, then its KM x CT input tiles and KM weight tiles are copied from the caches into CB 0 / CB 1 by NOC
// reads of this core's own L1, in the same order as v1. Compute and writer are v1's: the same tiles in the same order,
// so the output is bit-identical to v1.
// The v2 table (bresize.cache_plan) gives per output block and entry e the words [k_e, w_e, sx_e, sw_e]: the input
// block, the weight tile and their cache slots (k % S / w % SW, probed on the host so that no two distinct input
// blocks (weight tiles) of one output block share a slot: the misses of a block never overwrite a tile the same block
// still copies). A slot whose tag is not the key is refilled from DRAM. One barrier after the misses, one after the
// copies (the copies of block b complete before block b + 1 refills a slot).
// CT args: [0] NB, [1] NC, [2] GX, [3] KM, [4] CT, [5] PW, [6] NX, [7] S, [8] SW, then the TensorAccessorArgs of x
// (x_hi), of the weight tiles, of the table and of x_lo (x again when NX = 1).
// Common RT args: [x_addr, w_addr, table_addr, lo_addr]. CBs: 0 x tiles, 1 weight tiles, 8 table page (scratch),
// 9 input cache (S * NX * CT tiles, scratch), 10 weight cache (SW tiles, scratch), 11 tags (S + SW words, scratch).
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NB = 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 KM = get_compile_time_arg_val(3);
constexpr uint32_t CT = get_compile_time_arg_val(4);
constexpr uint32_t PW = get_compile_time_arg_val(5);
constexpr uint32_t NX = get_compile_time_arg_val(6);
constexpr uint32_t S = get_compile_time_arg_val(7);
constexpr uint32_t SW = get_compile_time_arg_val(8);
constexpr auto x_args = TensorAccessorArgs<9>();
constexpr auto w_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto t_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr auto l_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048;
constexpr uint32_t SLOT = NX * CT * TB16;           // bytes of one cached input block
constexpr uint32_t cb_x = 0, cb_w = 1, cb_tab = 8, cb_xc = 9, cb_wc = 10, cb_tags = 11;
constexpr uint32_t NONE = 0xFFFFFFFFu;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NB / NC, E = NB % NC;
    const uint32_t b0 = core * B + (core < E ? core : E);
    const uint32_t n = B + (core < E ? 1 : 0);
    if (n == 0) {
        return;
    }
    const auto ax = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB16);
    const auto aw = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), TB16);
    const auto at = TensorAccessor(t_args, get_common_arg_val<uint32_t>(2), PW * 4);
    const auto al = TensorAccessor(l_args, get_common_arg_val<uint32_t>(3), TB16);
    const uint32_t tab = get_write_ptr(cb_tab);
    const uint32_t xc = get_write_ptr(cb_xc);
    const uint32_t wc = get_write_ptr(cb_wc);
    volatile tt_l1_ptr uint32_t* ent = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(tab);
    volatile tt_l1_ptr uint32_t* xtag = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_tags));
    volatile tt_l1_ptr uint32_t* wtag = xtag + S;
    for (uint32_t s = 0; s < S; ++s) {
        xtag[s] = NONE;
    }
    for (uint32_t s = 0; s < SW; ++s) {
        wtag[s] = NONE;
    }
    for (uint32_t b = b0; b < b0 + n; ++b) {
        noc_async_read(at.get_noc_addr(b), tab, PW * 4);
        noc_async_read_barrier();
        // misses: DRAM -> the caches
        for (uint32_t e = 0; e < KM; ++e) {
            const uint32_t k = ent[4 * e], w = ent[4 * e + 1], s = ent[4 * e + 2], sw = ent[4 * e + 3];
            if (xtag[s] != k) {
                xtag[s] = k;
                uint32_t d = xc + s * SLOT;
                for (uint32_t ct = 0; ct < CT; ++ct) {
                    noc_async_read(ax.get_noc_addr(k * CT + ct), d, TB16);
                    d += TB16;
                }
                if constexpr (NX > 1) {
                    for (uint32_t ct = 0; ct < CT; ++ct) {
                        noc_async_read(al.get_noc_addr(k * CT + ct), d, TB16);
                        d += TB16;
                    }
                }
            }
            if (wtag[sw] != w) {
                wtag[sw] = w;
                noc_async_read(aw.get_noc_addr(w), wc + sw * TB16, TB16);
            }
        }
        noc_async_read_barrier();
        // copies: the caches -> CB 1 (KM weight tiles) and CB 0 (entry-major x_hi tiles, then x_lo), v1's order
        cb_reserve_back(cb_w, KM);
        cb_reserve_back(cb_x, NX * KM * CT);
        uint32_t dw = get_write_ptr(cb_w), dx = get_write_ptr(cb_x);
        for (uint32_t e = 0; e < KM; ++e) {
            const uint32_t s = ent[4 * e + 2], sw = ent[4 * e + 3];
            noc_async_read(get_noc_addr(wc + sw * TB16), dw, TB16);
            dw += TB16;
            noc_async_read(get_noc_addr(xc + s * SLOT), dx, CT * TB16);
            dx += CT * TB16;
        }
        if constexpr (NX > 1) {
            for (uint32_t e = 0; e < KM; ++e) {
                const uint32_t s = ent[4 * e + 2];
                noc_async_read(get_noc_addr(xc + s * SLOT + CT * TB16), dx, CT * TB16);
                dx += CT * TB16;
            }
        }
        noc_async_read_barrier();
        cb_push_back(cb_w, KM);
        cb_push_back(cb_x, NX * KM * CT);
    }
}