// SPDX-License-Identifier: Apache-2.0
// MK-T reader (sceneseg-p150 OPT_PLAN 3.4, the D19 seg-tail module megakernel, tt/mkt.py), NCRISC.
//
// Core c (row-major logical index, c < NC) owns the n = RB + (c < RE) image rows ya .. ya + n - 1; band row r
// (0 .. n + 1) is image row ya - 1 + r (rows outside the image are zeros: the convs' zero padding).
//   0. the LOAD-time tensors once: decode_layer_9 weights (72 bf16 tiles, K tile (dy * 3 + dx) * 4 + k, N tiles 0 / 1)
//      and bias K tile (rows 0 / 1 = bf16 hi / lo, 2 N tiles), margin weights (18 tiles) and bias K tile (1), and
//      the 3 identical ones tiles (columns 0 / 1 = 1.0) the bias products use as activations.
//   1. phase 1, per tile column x (32 pixels): the column band, per band row and channel group g (16 channels):
//      34 pixel sticks of 32 B (pixel x0 - 1 .. x0 + 32; a TILE face of the bf16 map is 16 consecutive sticks of one
//      group, so the band is 2 face reads + 2 halo-stick reads per (row, group) from DRAM; matched 64-byte
//      alignment of source and destination), double buffered (the next column's DRAM reads are issued before this
//      column's tiles are built); then the column's activation tiles A[(r * 3 + dx) * 4 + k] (pixels x0 + dx - 1 ..
//      + 31, channels 32 k .. 32 k + 31) by four local 512-byte face copies each.
//   2. phase 2 (after the writer's token: the D store holds the 64-channel map of rows ya - 1 .. ya + n), per column:
//      the margin conv's activation tiles A[(r * 3 + dx) * 2 + k] from the D store, four local face copies each.
// Every NoC read is waited on before the push that hands it to compute. No semaphores here.
// CT args: [0] H, [1] W, [2] NC, [3] GX, [4] RB, [5] RE, then the TensorAccessorArgs of the input map, w9, b9, wm, bm, ones.
// Common RT args: [map, w9, b9, wm, bm, ones].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t H = get_compile_time_arg_val(0);
constexpr uint32_t W = 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 RB = get_compile_time_arg_val(4);
constexpr uint32_t RE = get_compile_time_arg_val(5);
constexpr auto in_args = TensorAccessorArgs<6>();
constexpr auto w9_args = TensorAccessorArgs<in_args.next_compile_time_args_offset()>();
constexpr auto b9_args = TensorAccessorArgs<w9_args.next_compile_time_args_offset()>();
constexpr auto wm_args = TensorAccessorArgs<b9_args.next_compile_time_args_offset()>();
constexpr auto bm_args = TensorAccessorArgs<wm_args.next_compile_time_args_offset()>();
constexpr auto one_args = TensorAccessorArgs<bm_args.next_compile_time_args_offset()>();

constexpr uint32_t XT = W / 32;                          // tile columns
constexpr uint32_t NMAX = RB + (RE ? 1 : 0);
constexpr uint32_t NRMAX = NMAX + 2;
constexpr uint32_t TB16 = 2048, TB32 = 4096, FACE = 512, STICK = 32;
constexpr uint32_t BS = 1152;                            // band region per (row, group)
constexpr uint32_t BUFB = NRMAX * 8 * BS;                // one column band
constexpr uint32_t RSB = (W + 2) * STICK;                // D store: bytes per (row slot, group)
constexpr uint32_t SLOTB = 4 * RSB;                      // D store: one image row (4 groups of 16 channels)
constexpr uint32_t cb_a = 0, cb_w9 = 1, cb_b9 = 2, cb_wm = 3, cb_bm = 4, cb_band = 5, cb_d = 6, cb_sync = 7,
                   cb_zero = 8, cb_one = 10;
static_assert(W % 32 == 0, "the map width must be whole tiles");
static_assert(NMAX >= 1 && NMAX <= 3, "1 to 3 rows per core (DEST: 2 n fp32 tiles + 2 bias tiles <= 8)");

#ifndef MKT_DBG
#define MKT_DBG 0
#endif

template <typename Acc>
FORCE_INLINE void load_const(const Acc& acc, uint32_t cb, uint32_t ntiles, uint32_t tb) {
    cb_reserve_back(cb, ntiles);
    const uint32_t l1 = get_write_ptr(cb);
    for (uint32_t i = 0; i < ntiles; ++i) {
        noc_async_read(acc.get_noc_addr(i), l1 + i * tb, tb);
    }
    noc_async_read_barrier();
    cb_push_back(cb, ntiles);
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    const uint32_t n = RB + (core < RE ? 1 : 0);
    const uint32_t ya = core * RB + (core < RE ? core : RE);
    const uint32_t nr = n + 2;
    const auto ain = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), TB16);

    // zero face (64-aligned) for padding sticks / out-of-image rows
    const uint32_t zero = (get_write_ptr(cb_zero) + 63u) & ~63u;
    {
        volatile tt_l1_ptr uint32_t* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero);
        for (uint32_t i = 0; i < FACE / 4; ++i) {
            z[i] = 0;
        }
        asm volatile("" ::: "memory");
        invalidate_l1_cache();
    }
    const uint64_t zero_noc = get_noc_addr(zero);
    const uint32_t band = (get_write_ptr(cb_band) + 63u) & ~63u;
    // band rows outside the image: zero once in both buffers (never written by the column loads)
    for (uint32_t b = 0; b < 2; ++b) {
        for (uint32_t r = 0; r < nr; ++r) {
            const int32_t yy = static_cast<int32_t>(ya + r) - 1;
            if (yy >= 0 && yy < static_cast<int32_t>(H)) {
                continue;
            }
            for (uint32_t g = 0; g < 8; ++g) {
                const uint32_t reg = band + b * BUFB + (r * 8 + g) * BS;
                noc_async_read(zero_noc, reg, FACE);
                noc_async_read(zero_noc, reg + FACE, FACE);
                noc_async_read(zero_noc, reg + 2 * FACE, BS - 2 * FACE);
            }
        }
    }
    load_const(TensorAccessor(w9_args, get_common_arg_val<uint32_t>(1), TB16), cb_w9, 72, TB16);
    load_const(TensorAccessor(b9_args, get_common_arg_val<uint32_t>(2), TB16), cb_b9, 2, TB16);
    load_const(TensorAccessor(wm_args, get_common_arg_val<uint32_t>(3), TB16), cb_wm, 18, TB16);
    load_const(TensorAccessor(bm_args, get_common_arg_val<uint32_t>(4), TB16), cb_bm, 1, TB16);
    load_const(TensorAccessor(one_args, get_common_arg_val<uint32_t>(5), TB16), cb_one, 3, TB16);

    // column band loads: region (r, g) of buffer b at band + b * BUFB + (r * 8 + g) * BS; stick s (pixel x0 + s - 1)
    // at region + 32 + 32 s: the left halo stick at +32, pixels x0 .. at +64 (64-aligned), the right halo at +1088
    auto load_band = [&](uint32_t x, uint32_t b) {
        const uint32_t buf = band + b * BUFB;
        for (uint32_t r = 0; r < nr; ++r) {
            const int32_t yy = static_cast<int32_t>(ya + r) - 1;
            if (yy < 0 || yy >= static_cast<int32_t>(H)) {
                continue;
            }
            const uint32_t trow = static_cast<uint32_t>(yy) * XT;
            for (uint32_t g = 0; g < 8; ++g) {
                const uint32_t reg = buf + (r * 8 + g) * BS;
                const uint32_t k = g >> 1, par = g & 1;
                const uint32_t tile = (trow + x) * 4 + k;
                noc_async_read(ain.get_noc_addr(tile, par * FACE), reg + 64, FACE);
                noc_async_read(ain.get_noc_addr(tile, (2 + par) * FACE), reg + 64 + FACE, FACE);
                if (x > 0) {
                    noc_async_read(ain.get_noc_addr(tile - 4, (2 + par) * FACE + 15 * STICK), reg + 32, STICK);
                } else {
                    noc_async_read(zero_noc, reg + 32, STICK);
                }
                if (x + 1 < XT) {
                    noc_async_read(ain.get_noc_addr(tile + 4, par * FACE), reg + 64 + 2 * FACE, STICK);
                } else {
                    noc_async_read(zero_noc, reg + 64 + 2 * FACE, STICK);
                }
            }
        }
    };

    // ---------------- phase 1: decode_layer_9 activation tiles
    constexpr uint32_t NA1 = 12;                         // tiles per band row
    load_band(0, 0);
    noc_async_read_barrier();
    for (uint32_t x = 0; x < XT; ++x) {
        if (x + 1 < XT) {
            load_band(x + 1, (x + 1) & 1);
        }
        const uint32_t buf = band + (x & 1) * BUFB;
        cb_reserve_back(cb_a, NRMAX * NA1);       // fixed push size: a CB push never wraps
        const uint32_t base = get_write_ptr(cb_a);
        for (uint32_t r = 0; r < nr; ++r) {
            for (uint32_t dx = 0; dx < 3; ++dx) {
                for (uint32_t k = 0; k < 4; ++k) {
                    const uint32_t dst = base + ((r * 3 + dx) * 4 + k) * TB16;
                    const uint32_t s0 = buf + (r * 8 + 2 * k) * BS + 32 + dx * STICK;
                    const uint32_t s1 = s0 + BS;
                    noc_async_read(get_noc_addr(s0), dst, FACE);
                    noc_async_read(get_noc_addr(s1), dst + FACE, FACE);
                    noc_async_read(get_noc_addr(s0 + FACE), dst + 2 * FACE, FACE);
                    noc_async_read(get_noc_addr(s1 + FACE), dst + 3 * FACE, FACE);
                }
            }
        }
        noc_async_read_barrier();                        // this column's tiles and the next column's band
        cb_push_back(cb_a, NRMAX * NA1);
    }

#if MKT_DBG != 1
    // ---------------- phase 2: margin conv activation tiles from the D store
    cb_wait_front(cb_sync, 1);
    cb_pop_front(cb_sync, 1);
    const uint32_t dst0 = get_write_ptr(cb_d);
    constexpr uint32_t NA2 = 6;
    for (uint32_t x = 0; x < XT; ++x) {
        cb_reserve_back(cb_a, NRMAX * NA2);
        const uint32_t base = get_write_ptr(cb_a);
        for (uint32_t r = 0; r < nr; ++r) {
            for (uint32_t dx = 0; dx < 3; ++dx) {
                for (uint32_t k = 0; k < 2; ++k) {
                    const uint32_t dst = base + ((r * 3 + dx) * 2 + k) * TB16;
                    const uint32_t s0 = dst0 + r * SLOTB + (2 * k) * RSB + (x * 32 + dx) * STICK;
                    const uint32_t s1 = s0 + RSB;
                    noc_async_read(get_noc_addr(s0), dst, FACE);
                    noc_async_read(get_noc_addr(s1), dst + FACE, FACE);
                    noc_async_read(get_noc_addr(s0 + FACE), dst + 2 * FACE, FACE);
                    noc_async_read(get_noc_addr(s1 + FACE), dst + 3 * FACE, FACE);
                }
            }
        }
        noc_async_read_barrier();
        cb_push_back(cb_a, NRMAX * NA2);
    }
#endif
}