// SPDX-License-Identifier: Apache-2.0
// Stem reader (sceneseg-p150 OPT_PLAN K-IN kernel form / MK0, tt/stem.py): the stem conv's im2col tiles straight from
// the uploaded uint8 squash (ROW_MAJOR, one page of WIN * 3 bytes per image row: B, G, R per pixel), as bf16
// (0..255 is exact in bf16). Output tile t = 32 consecutive output pixels of one output row (WOUT % 32 == 0); A row
// i = pixel i, column k = (dy * 3 + dx) * 3 + c (27 used, 5 zero) of the 3x3 stride-2 padding-1 window; taps outside
// the image are 0 (the zero padding of the uint8 input; the normalisation's mean term is in the bias map). Also the
// tile's fp32 bias-map tile (the stem's exact border-aware bias, tt/layers.py stem_fold) into CB 2, and once the two
// bf16 weight tiles (hi, lo: [k, cout]) into CB 1.
// Work split (no per-core RT args, probe P3): core i of the first NC (row-major) owns tiles [i * B + min(i, E), + B +
// (i < E)), B = T / NC, E = T % NC. The three input rows of an output row are read once per output row.
// CT args: [0] T, [1] NC, [2] GX, [3] HIN, [4] WIN, [5] WOUT, then the TensorAccessorArgs of the image, the weight
// tiles and the bias map. Common RT args: [image, weights, bias].
#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 HIN = get_compile_time_arg_val(3);
constexpr uint32_t WIN = get_compile_time_arg_val(4);
constexpr uint32_t WOUT = get_compile_time_arg_val(5);
constexpr auto im_args = TensorAccessorArgs<6>();
constexpr auto w_args = TensorAccessorArgs<im_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr uint32_t ROWB = WIN * 3;
constexpr uint32_t TPR = WOUT / 32;                 // output tiles per output row
constexpr uint32_t cb_a = 0, cb_w = 1, cb_b = 2, cb_rows = 8, cb_lut = 9;

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 / NC, E = T % NC;
    const uint32_t t0 = core * B + (core < E ? core : E);
    const uint32_t n = B + (core < E ? 1 : 0);
    if (n == 0) {
        return;
    }
    const auto aim = TensorAccessor(im_args, get_common_arg_val<uint32_t>(0), ROWB);
    const auto aw = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), 2048);
    const auto ab = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), 4096);
    cb_reserve_back(cb_w, 2);
    noc_async_read(aw.get_noc_addr(0), get_write_ptr(cb_w), 2048);
    noc_async_read(aw.get_noc_addr(1), get_write_ptr(cb_w) + 2048, 2048);
    // bf16 bits of 0..255 (exact): v = 1.m * 2^e, 8 significant bits fit the 7-bit mantissa + the implicit 1
    volatile tt_l1_ptr uint16_t* lut = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(get_write_ptr(cb_lut));
    lut[0] = 0;
    for (uint32_t v = 1; v < 256; ++v) {
        uint32_t e = 7;
        while (((v >> e) & 1) == 0) {
            --e;
        }
        const uint32_t bits = ((e + 127) << 23) | ((v << (23 - e)) & 0x7FFFFF);
        lut[v] = (uint16_t)(bits >> 16);
    }
    const uint32_t rows = get_write_ptr(cb_rows);
    volatile tt_l1_ptr uint8_t* rowp = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(rows);
    noc_async_read_barrier();
    cb_push_back(cb_w, 2);

    uint32_t cur = 0xFFFFFFFFu;
    bool valid[3] = {false, false, false};
    for (uint32_t t = t0; t < t0 + n; ++t) {
        const uint32_t oy = t / TPR, ox0 = (t % TPR) * 32;
        if (oy != cur) {
            cur = oy;
            for (uint32_t dy = 0; dy < 3; ++dy) {
                const int32_t iy = (int32_t)(2 * oy + dy) - 1;
                valid[dy] = iy >= 0 && iy < (int32_t)HIN;
                if (valid[dy]) {
                    noc_async_read(aim.get_noc_addr((uint32_t)iy), rows + dy * ROWB, ROWB);
                }
            }
            noc_async_read_barrier();
        }
        cb_reserve_back(cb_b, 1);
        noc_async_read(ab.get_noc_addr(t), get_write_ptr(cb_b), 4096);
        cb_reserve_back(cb_a, 1);
        volatile tt_l1_ptr uint32_t* a = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_a));
        for (uint32_t i = 0; i < 32; ++i) {
            const int32_t ixc = (int32_t)(2 * (ox0 + i));          // input column of dx = 1
            uint16_t v[32];
            uint32_t k = 0;
            for (uint32_t dy = 0; dy < 3; ++dy) {
                for (int32_t dx = -1; dx <= 1; ++dx) {
                    const int32_t ix = ixc + dx;
                    if (valid[dy] && ix >= 0 && ix < (int32_t)WIN) {
                        const uint32_t o = dy * ROWB + (uint32_t)ix * 3;
                        v[k] = lut[rowp[o]];
                        v[k + 1] = lut[rowp[o + 1]];
                        v[k + 2] = lut[rowp[o + 2]];
                    } else {
                        v[k] = v[k + 1] = v[k + 2] = 0;
                    }
                    k += 3;
                }
            }
            for (; k < 32; ++k) {
                v[k] = 0;
            }
            // row i of the tile: face (i / 16) * 2 + 0 holds k 0..15, face + 1 holds k 16..31; 16 bf16 per face row
            const uint32_t base = ((i >> 4) * 2) * 128 + (i & 15) * 8;     // in uint32 (2 bf16) units
            for (uint32_t j = 0; j < 8; ++j) {
                a[base + j] = (uint32_t)v[2 * j] | ((uint32_t)v[2 * j + 1] << 16);
                a[base + 128 + j] = (uint32_t)v[16 + 2 * j] | ((uint32_t)v[16 + 2 * j + 1] << 16);
            }
        }
        cb_push_back(cb_a, 1);
        noc_async_read_barrier();
        cb_push_back(cb_b, 1);
    }
}