// SPDX-License-Identifier: Apache-2.0
// IMG_IN kernel (meteor-p150 OPT round 3 R3.3, tt/kin.py): the camera input uploaded as whole uint8 rows (one
// [W * 3] page per image row, 2,304 B: H2D of 3,456 pages instead of 2.65 M pages of 3 B) expanded on the device
// into the stem conv's ROW_MAJOR input [N*H*W, CO] bf16 (one CO * 2 B page per pixel; channels 0..2 = the pixel's
// R, G, B as bf16 (0..255 are exact in bf16: the stock typecast), channels 3..CO-1 = +0.0, the zero channels the conv
// pads its 3-channel input with anyway).
// One data-movement kernel on one RISC / NOC (reads and writes from the same NOC: no two-NOC DRAM traffic, SYS-1419).
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns the contiguous image-row
// range [i * B + min(i, E), + B + (i < E)) of the R rows, B = R / NC, E = R % NC.
// CT args: [0] R, [1] NC, [2] GX, [3] W, [4] CO, then the TensorAccessorArgs of the input and the output.
// Common RT args: [in_addr, out_addr]. CB 0: two input rows (staging), CB 1: two output rows (staging), CB 2: the
// 256-entry uint8 -> bf16 table (1 KB, in L1 rather than on the small RISC stack).
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t R = 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 W = get_compile_time_arg_val(3);
constexpr uint32_t CO = get_compile_time_arg_val(4);
constexpr auto in_args = TensorAccessorArgs<5>();
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 PIX = CO * 2;                // bytes of one output pixel page
constexpr uint32_t OUT_ROW = W * PIX;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = R / NC, E = R % NC;
    const uint32_t r0 = core * B + (core < E ? core : E);
    const uint32_t nr = B + (core < E ? 1 : 0);
    // the accessors' page strides are the buffers' aligned page sizes (from the accessor CT args)
    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 (8 significant bits: 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);
    // the zero channels, written once into both output staging rows
    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 (nr == 0) {
        return;
    }
    noc_async_read(in.get_noc_addr(r0), ibuf, IN_ROW);
    for (uint32_t j = 0; j < nr; ++j) {
        const uint32_t k = j & 1;
        const uint32_t ib = ibuf + k * IN_ROW;
        const uint32_t ob = obuf + k * OUT_ROW;
        noc_async_read_barrier();
        if (j + 1 < nr) {                        // prefetch the next row into the other input slot
            noc_async_read(in.get_noc_addr(r0 + j + 1), ibuf + (k ^ 1) * IN_ROW, IN_ROW);
        }
        if (j >= 2) {
            noc_async_writes_flushed();          // the writes of row j - 2 (same output slot) have left L1
        }
        volatile tt_l1_ptr uint8_t* src = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(ib);
        volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ob);
        for (uint32_t w = 0; w < W; ++w) {
            dst[w * (PIX / 4)] = lut[src[3 * w]] | (lut[src[3 * w + 1]] << 16);
            dst[w * (PIX / 4) + 1] = lut[src[3 * w + 2]];   // channel 3 = +0.0; channels 4.. stay +0.0
        }
        const uint32_t p0 = (r0 + j) * W;
        for (uint32_t w = 0; w < W; ++w) {
            noc_async_write(ob + w * PIX, out.get_noc_addr(p0 + w), PIX);
        }
    }
    noc_async_write_barrier();
}