// SPDX-License-Identifier: Apache-2.0
// STEM=fold kernel (meteor-p150 OPT2 round 3 R3.1, OPT_PLAN A6, tt/kin.py expand_fold): the camera input uploaded as
// whole uint8 rows (one [W * 3] page per image row) written as the stem's space-to-depth (fold 2) input: the image
// zero-padded by P on every side, cut into 2 x 2 pixel blocks, one block = one folded pixel of 16 bf16 channels
// (32 B page): channel (a * 2 + b) * 3 + c = pixel (2Y + a - P, 2X + b - P), colour c, as bf16 (0..255 exact);
// channels 12..15 = +0.0. The 7 x 7 / s2 / p3 stem conv on the image is then a 4 x 4 / s1 / p0 conv on this map
// (weights re-indexed on the host). Output [N * HF * WF, 16] bf16 ROW_MAJOR, HF = (H + 2P) / 2, WF = (W + 2P) / 2.
// One data-movement kernel on one RISC / NOC (reads and writes from the same NOC, SYS-1419).
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns the contiguous folded-row
// range [i * B + min(i, E), + B + (i < E)) of the RF = N * HF folded rows, B = RF / NC, E = RF % NC.
// CT args: [0] RF, [1] NC, [2] GX, [3] H, [4] W, [5] HF, [6] WF, [7] P, [8] T4, then the TensorAccessorArgs of the
// input and the output.
// T4 = 1 (OPT2 round 6 R6.1, METEOR_STEM_IN=t4): the output is the 4-D TILE tensor [N, HF, WF, 16] (padded to
// [N, HF, TPR * 32, 32], TPR = ceil(WF / 32) tiles per folded row; page = tile (N * HF + Y) * TPR + X / 32). A folded
// pixel's 16 channels are exactly one 32 B face row: pixel r = X % 32 of a tile sits at face (r / 16) * 2 (faces 0 /
// 2), row r % 16; faces 1 / 3 (channels 16..31) and the pad pixels X >= WF stay +0.0. The conv's DRAM-path slices
// then read whole tiles instead of 32 B pages (the values are the same: bit-identical stem output). Common RT args: [in_addr, out_addr]. CB 0: two slots x two input rows (64 B stride), CB 1: two output rows, CB 2: the
// 256-entry uint8 -> bf16 table.
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t RF = 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 H = get_compile_time_arg_val(3);
constexpr uint32_t W = get_compile_time_arg_val(4);
constexpr uint32_t HF = get_compile_time_arg_val(5);
constexpr uint32_t WF = get_compile_time_arg_val(6);
constexpr uint32_t P = get_compile_time_arg_val(7);
constexpr uint32_t T4 = get_compile_time_arg_val(8);
constexpr auto in_args = TensorAccessorArgs<9>();
constexpr auto out_args = TensorAccessorArgs<in_args.next_compile_time_args_offset()>();
constexpr uint32_t IN_ROW = W * 3;              // bytes of one uint8 image row
constexpr uint32_t IN_STRIDE = (IN_ROW + 63) / 64 * 64;   // L1 stride of a staged row (DRAM reads: 64 B aligned)
constexpr uint32_t PIX = 32;                    // bytes of one folded pixel (16 bf16 channels)
constexpr uint32_t TPR = (WF + 31) / 32;         // T4: tiles per folded row
constexpr uint32_t TILE_B = 2048;               // T4: one bf16 32 x 32 tile
constexpr uint32_t OUT_ROW = T4 ? TPR * TILE_B : WF * PIX;

// L1 byte offset of folded pixel x in a staged output row
inline uint32_t pix_off(uint32_t x) {
    if constexpr (T4) {
        const uint32_t r = x & 31;
        return (x >> 5) * TILE_B + (r >> 4) * 1024 + (r & 15) * PIX;
    } else {
        return x * PIX;
    }
}

// the image rows (2Y - P, 2Y + 1 - P) of folded row Y of camera n into slot `ib` (rows outside the image: not read)
template <typename A>
inline void read_rows(const A& in, uint32_t f, uint32_t ib) {
    const uint32_t n = f / HF, y = f % HF;
    for (uint32_t a = 0; a < 2; ++a) {
        const int32_t r = static_cast<int32_t>(2 * y + a) - static_cast<int32_t>(P);
        if (r >= 0 && r < static_cast<int32_t>(H)) {
            noc_async_read(in.get_noc_addr(n * H + static_cast<uint32_t>(r)), ib + a * IN_STRIDE, IN_ROW);
        }
    }
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = RF / NC, E = RF % NC;
    const uint32_t f0 = core * B + (core < E ? core : E);
    const uint32_t nf = B + (core < E ? 1 : 0);
    const auto in = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0));
    const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(1));
    // bf16 bits of the integers 0..255 (exact; the stock uint8 -> bf16 typecast)
    volatile tt_l1_ptr uint32_t* lut = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(2));  // CB 2
    lut[0] = 0;
    for (uint32_t v = 1; v < 256; ++v) {
        uint32_t e = 0;
        while ((v >> (e + 1)) != 0) {
            ++e;
        }
        lut[v] = ((127 + e) << 7) | ((v << (7 - e)) & 0x7F);
    }
    const uint32_t ibuf = get_write_ptr(0);
    const uint32_t obuf = get_write_ptr(1);
    // both output staging rows zeroed once: channels 12..15 (words 6, 7 of a pixel) are never written again (T4: nor
    // faces 1 / 3 and the pad pixels)
    for (uint32_t k = 0; k < 2; ++k) {
        volatile tt_l1_ptr uint32_t* o = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(obuf + k * OUT_ROW);
        for (uint32_t i = 0; i < OUT_ROW / 4; ++i) {
            o[i] = 0;
        }
    }
    if (nf == 0) {
        return;
    }
    read_rows(in, f0, ibuf);
    for (uint32_t j = 0; j < nf; ++j) {
        const uint32_t f = f0 + j;
        const uint32_t k = j & 1;
        const uint32_t ib = ibuf + k * 2 * IN_STRIDE;
        const uint32_t ob = obuf + k * OUT_ROW;
        noc_async_read_barrier();
        if (j + 1 < nf) {                        // prefetch the next folded row's image rows into the other slot
            read_rows(in, f + 1, ibuf + (k ^ 1) * 2 * IN_STRIDE);
        }
        if (j >= 2) {
            noc_async_writes_flushed();          // the writes of row j - 2 (same output slot) have left L1
        }
        const uint32_t y = f % HF;
        volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ob);
        for (uint32_t a = 0; a < 2; ++a) {
            const int32_t r = static_cast<int32_t>(2 * y + a) - static_cast<int32_t>(P);
            const bool row_ok = r >= 0 && r < static_cast<int32_t>(H);
            volatile tt_l1_ptr uint8_t* src = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(ib + a * IN_STRIDE);
            for (uint32_t x = 0; x < WF; ++x) {
                volatile tt_l1_ptr uint32_t* d = dst + pix_off(x) / 4 + a * 3;
                const int32_t c0 = static_cast<int32_t>(2 * x) - static_cast<int32_t>(P);
                if (!row_ok) {
                    d[0] = 0;
                    d[1] = 0;
                    d[2] = 0;
                } else if (c0 >= 0 && c0 + 1 < static_cast<int32_t>(W)) {   // both columns inside the image
                    volatile tt_l1_ptr uint8_t* s = src + 3 * c0;
                    d[0] = lut[s[0]] | (lut[s[1]] << 16);
                    d[1] = lut[s[2]] | (lut[s[3]] << 16);
                    d[2] = lut[s[4]] | (lut[s[5]] << 16);
                } else {                                                    // an edge block: per column
                    uint32_t v[6];
                    for (uint32_t i = 0; i < 6; ++i) {
                        const int32_t col = c0 + static_cast<int32_t>(i / 3);
                        v[i] = (col >= 0 && col < static_cast<int32_t>(W)) ? lut[src[3 * col + i % 3]] : 0;
                    }
                    d[0] = v[0] | (v[1] << 16);
                    d[1] = v[2] | (v[3] << 16);
                    d[2] = v[4] | (v[5] << 16);
                }
            }
        }
        if constexpr (T4) {
            const uint32_t t0 = f * TPR;
            for (uint32_t t = 0; t < TPR; ++t) {
                noc_async_write(ob + t * TILE_B, out.get_noc_addr(t0 + t), TILE_B);
            }
        } else {
            const uint32_t p0 = f * WF;
            for (uint32_t x = 0; x < WF; ++x) {
                noc_async_write(ob + x * PIX, out.get_noc_addr(p0 + x), PIX);
            }
        }
    }
    noc_async_write_barrier();
}