// SPDX-License-Identifier: Apache-2.0
// kpre (tt/kpre.py), writer (BRISC). Per output row y of the core's pixels: the computed bf16 row (4 pseudo-tiles,
// 1280 x 3 values in pixel-major order, from the compute kernel) or zeros (the pad rows y >= RH) -> the core's own
// shard of the stem conv input [H W, 8] (3 channels + 5 zeros per 16-byte pixel; local L1 stores, as kin_rows).
#include <cstdint>

#include "api/dataflow/dataflow_api.h"
#include "kpre_common.h"

constexpr uint32_t N = get_compile_time_arg_val(0);
constexpr uint32_t S = get_compile_time_arg_val(1);
constexpr uint32_t W = get_compile_time_arg_val(2);
constexpr uint32_t RH = get_compile_time_arg_val(3);
constexpr uint32_t NC = get_compile_time_arg_val(4);
constexpr uint32_t GX = get_compile_time_arg_val(5);
constexpr uint32_t RA = get_compile_time_arg_val(6);
constexpr uint32_t RB = get_compile_time_arg_val(7);
constexpr uint32_t RD = get_compile_time_arg_val(8);
constexpr uint32_t SPLIT = get_compile_time_arg_val(9);   // 1: this RISC also builds the odd P rows (cb_p2)
constexpr uint32_t SWAP = get_compile_time_arg_val(10);
constexpr uint32_t ROW_BYTES = get_compile_time_arg_val(11);
constexpr uint32_t PERIODS = get_compile_time_arg_val(12);
constexpr uint32_t U16 = get_compile_time_arg_val(13);
constexpr auto img_args = TensorAccessorArgs<14>();
constexpr auto lut_args = TensorAccessorArgs<img_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_v = 16, cb_p2 = 5, cb_scratch2 = 6;
constexpr uint32_t TILE_WORDS = 1024;
constexpr uint32_t QT = 4;
constexpr uint32_t C_PAD = 8;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    const KpreGeom g = kpre_geom<N, S, W, RH, NC, RA, RB, RD>(core);
    if (!g.any) {
        return;
    }
    const uint32_t out_addr = get_common_arg_val<uint32_t>(4);
    const uint32_t shard0 = core * S;
    // SPLIT: the odd P rows (rlo + 1, rlo + 3, ...), each built before the output row that needs it is awaited
    uint32_t next = g.rlo + 1;
    uint32_t lut_l1 = 0, row_l1 = 0;
    const auto img = TensorAccessor(img_args, get_common_arg_val<uint32_t>(0), ROW_BYTES);
    if constexpr (SPLIT) {
        if (g.rows) {
            const auto lut_t = TensorAccessor(lut_args, get_common_arg_val<uint32_t>(1), 3 * 256 * 4);
            const uint32_t base = (get_write_ptr(cb_scratch2) + 63u) & ~63u;
            lut_l1 = base;
            row_l1 = base + 3072;
            noc_async_read(lut_t.get_noc_addr(0), lut_l1, 3 * 256 * 4);
            noc_async_read_barrier();
        }
    }
    for (uint32_t y = g.ya; y <= g.yb; ++y) {
        if constexpr (SPLIT) {
            if (y < RH) {
                const uint32_t need = kpre_r0<RA, RB, RD>(y) + 1;
                for (; next <= need && next <= g.rhi; next += 2) {
                    noc_async_read(img.get_noc_addr(next), row_l1, ROW_BYTES);
                    cb_reserve_back(cb_p2, 5);
                    noc_async_read_barrier();
                    uint32_t* d0 = reinterpret_cast<uint32_t*>(get_write_ptr(cb_p2));
                    asm volatile("" ::: "memory");
                    if constexpr (U16) {
                        kpre_build_row_u16<PERIODS, SWAP>(reinterpret_cast<const uint8_t*>(row_l1),
                                                          reinterpret_cast<uint16_t*>(d0));
                    } else {
                        kpre_build_row<PERIODS, SWAP>(reinterpret_cast<const uint8_t*>(row_l1),
                                                      reinterpret_cast<const uint32_t*>(lut_l1), d0);
                    }
                    asm volatile("" ::: "memory");
                    cb_push_back(cb_p2, 5);
                }
            }
        }
        const uint32_t a = y * W > g.p0 ? y * W : g.p0;
        const uint32_t e = (y + 1) * W < g.p1 ? (y + 1) * W : g.p1;
        uint32_t* dst = reinterpret_cast<uint32_t*>(out_addr + (a - shard0) * C_PAD * 2);
        uint32_t n = e - a;
        if (y >= RH) {
            for (uint32_t i = 0; i < n * 4; ++i) {
                dst[i] = 0;
            }
            continue;
        }
        cb_wait_front(cb_v, QT);
        // V row in phase-planar order: pixel x = 4k + p, channel c at [p * 1024 + 3k + c] (bf16)
        const uint16_t* v = reinterpret_cast<const uint16_t*>(get_read_ptr(cb_v));
        asm volatile("" ::: "memory");
#ifdef KPRE_NOWRITE
        n = 0;
#endif
        for (uint32_t x = a - y * W; n; ++x, --n) {
            const uint16_t* sp = v + (x & 3u) * 1024 + (x >> 2) * 3;
            dst[0] = static_cast<uint32_t>(sp[0]) | (static_cast<uint32_t>(sp[1]) << 16);
            dst[1] = sp[2];
            dst[2] = 0;
            dst[3] = 0;
            dst += 4;
        }
        asm volatile("" ::: "memory");
        cb_pop_front(cb_v, QT);
    }
}