// SPDX-License-Identifier: Apache-2.0
// kpre (tt/kpre.py), compute. The C14 x0.8 INTER_LINEAR resize of the normalised node image on the SFPU, in the
// exact float32 steps of the host twin (host.fast; each lerp is fma(x1 - x0, w, x0): the difference rounded to
// fp32, one fused multiply-add rounded once; probe logs/bevformer/opt_r4/probe_sfpu: SFPU sub / SFPMAD are the
// correctly rounded fp32 results on this domain, 0 of 262144 differ):
//   H row r (4 fp32 pseudo-tiles): H = (P1 - P0) * WX + P0          (P0 / P1 / WX from the reader)
//   output row y (4 bf16 pseudo-tiles): V = (H[r0 + 1] - H[r0]) * WY[y % 4] + H[r0], rounded to bf16 (RNE, as
//   torch's float32 -> bfloat16 cast of the host path) on the SFPU, then packed (exact)
// cb_h keeps all H rows of the core (no wrap of the 2-row window).
#include <cstdint>

#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#include "kpre_common.h"

#if defined(TRISC_MATH)
#include "llk_math_eltwise_unary_sfpu_macros.h"
#include "sfpi.h"

namespace ckernel::sfpu {
// DEST tiles 0 / 1 / 2 = x0, x1, w -> tile 0 = fma(x1 - x0, w, x0)
inline void kpre_lerp() {
#pragma GCC unroll 1
    for (int d = 0; d < 8; d++) {
        sfpi::vFloat x0 = sfpi::dst_reg[0];
        sfpi::vFloat x1 = sfpi::dst_reg[32];
        sfpi::vFloat w = sfpi::dst_reg[64];
        sfpi::vFloat dx = x1 - x0;
        sfpi::dst_reg[0] = dx * w + x0;
        sfpi::dst_reg++;
    }
}
// U16 rows: DEST tiles 0 / 1 = the uint16 bytes of planes p / p + 1 (low 16 bits of each 32-bit word), 2 = the
// per-lane -mean_c, 3 = w -> tile 0 = fma(x1 - x0, w, x0) with x = f32(u) + (-mean_c) (exact int -> fp32 cast, then
// the single fp32 add of the host's normalisation)
inline void kpre_lerp_u16() {
#pragma GCC unroll 1
    for (int d = 0; d < 8; d++) {
        sfpi::vInt u0 = sfpi::dst_reg[0];
        sfpi::vInt u1 = sfpi::dst_reg[32];
        sfpi::vFloat beta = sfpi::dst_reg[64];
        sfpi::vFloat w = sfpi::dst_reg[96];
        sfpi::vFloat x0 = sfpi::convert<sfpi::vFloat>(u0 & 0xFFFF) + beta;
        sfpi::vFloat x1 = sfpi::convert<sfpi::vFloat>(u1 & 0xFFFF) + beta;
        sfpi::vFloat dx = x1 - x0;
        sfpi::dst_reg[0] = dx * w + x0;
        sfpi::dst_reg++;
    }
}
// the same, then round to bf16 (RNE: + 0x7fff + lsb, low 16 bits cleared; the bf16 pack is then exact)
inline void kpre_lerp_bf16() {
#pragma GCC unroll 1
    for (int d = 0; d < 8; d++) {
        sfpi::vFloat x0 = sfpi::dst_reg[0];
        sfpi::vFloat x1 = sfpi::dst_reg[32];
        sfpi::vFloat w = sfpi::dst_reg[64];
        sfpi::vFloat dx = x1 - x0;
        sfpi::vFloat v = dx * w + x0;
        sfpi::vUInt bits = sfpi::as<sfpi::vUInt>(v);
        sfpi::vUInt lsb = (bits >> 16) & 1;
        bits = bits + 0x7fffU + lsb;
        bits = bits & 0xFFFF0000U;
        sfpi::dst_reg[0] = sfpi::as<sfpi::vFloat>(bits);
        sfpi::dst_reg++;
    }
}
}  // namespace ckernel::sfpu
#endif

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: odd rows (from rlo) come from the writer RISC in cb_p2
constexpr uint32_t U16 = get_compile_time_arg_val(10);    // 1: the P rows are raw uint16 planes (+ cb_beta)
constexpr uint32_t cb_p = 0, cb_wx = 1, cb_wy = 2, cb_h = 3, cb_p2 = 5, cb_beta = 7, cb_v = 16;
constexpr uint32_t QT = 4;   // pseudo-tiles per row

template <bool BF16>
inline void lerp_rows(uint32_t cb_x, uint32_t i0, uint32_t i1, uint32_t cb_w, uint32_t iw, uint32_t cb_o) {
    tile_regs_acquire();
#ifndef KPRE_NOCOPY
    copy_tile(cb_x, i0, 0);
    copy_tile(cb_x, i1, 1);
    copy_tile(cb_w, iw, 2);
#endif
#ifndef KPRE_NOSFPU
    if constexpr (BF16) {
        MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp_bf16, 0, VectorMode::RC)));
    } else {
        MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp, 0, VectorMode::RC)));
    }
#endif
    tile_regs_commit();
    tile_regs_wait();
    pack_tile(0, cb_o);
    tile_regs_release();
}

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.rows) {
        return;
    }
    compute_kernel_hw_startup(cb_p, cb_h);
    copy_init(cb_p);
    MATH((SFPU_UNARY_INIT(unused)));
    cb_wait_front(cb_wx, QT);
    cb_wait_front(cb_wy, QT);
    if constexpr (U16) {
        cb_wait_front(cb_beta, 1);
    }
    // cb_h holds every H row of the core (rlo .. rhi, sized on the host): rows are never popped before the end, so
    // the two rows an output row reads are always contiguous pages (a FIFO window could wrap the ring)
    uint32_t produced = g.rlo;   // next H row to compute
    for (uint32_t y = g.ya; y <= g.yv; ++y) {
        const uint32_t r0 = kpre_r0<RA, RB, RD>(y);
        while (produced <= r0 + 1) {
            // H row `produced`: (P1 - P0) * WX + P0, fp32
            // phase planes s = 0..4 of the node row: H plane p = lerp(plane p, plane p + 1, (2p + 1) / 8)
            const uint32_t cbp = (SPLIT && ((produced - g.rlo) & 1)) ? cb_p2 : cb_p;
            cb_wait_front(cbp, QT + 1);
            cb_reserve_back(cb_h, QT);
            pack_reconfig_data_format(cb_h);
            if constexpr (U16) {
                for (uint32_t t = 0; t < QT; ++t) {
                    tile_regs_acquire();
                    copy_tile_to_dst_init_short_with_dt(cb_wx, cbp);
                    copy_tile(cbp, t, 0);
                    copy_tile(cbp, t + 1, 1);
                    copy_tile_to_dst_init_short_with_dt(cbp, cb_wx);
                    copy_tile(cb_beta, 0, 2);
                    copy_tile(cb_wx, t, 3);
#ifndef KPRE_NOSFPU
                    MATH((SFPU_UNARY_INIT(unused)));
                    MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp_u16, 0,
                                                           VectorMode::RC)));
#endif
                    tile_regs_commit();
                    tile_regs_wait();
                    pack_tile(0, cb_h);
                    tile_regs_release();
                }
            } else {
                for (uint32_t t = 0; t < QT; ++t) {
                    lerp_rows<false>(cbp, t, t + 1, cb_wx, t, cb_h);
                }
            }
            cb_push_back(cb_h, QT);
            cb_pop_front(cbp, QT + 1);
            ++produced;
        }
        // output row y: (H[r0 + 1] - H[r0]) * WY + H[r0] -> bf16
        const uint32_t h0 = (r0 - g.rlo) * QT;
        cb_wait_front(cb_h, h0 + 2 * QT);
        cb_reserve_back(cb_v, QT);
        pack_reconfig_data_format(cb_v);
        for (uint32_t t = 0; t < QT; ++t) {
            lerp_rows<true>(cb_h, h0 + t, h0 + QT + t, cb_wy, y % 4, cb_v);
        }
        cb_push_back(cb_v, QT);
    }
    cb_pop_front(cb_h, (produced - g.rlo) * QT);
    if constexpr (U16) {
        cb_pop_front(cb_beta, 1);
    }
}