// SPDX-License-Identifier: Apache-2.0
// scatter_k (tt/scatter_k.py), compute. Per 32-row block (DT tiles), the steps of the stock ops in their order, each
// result packed to bf16 like the stock ops' bf16 outputs (16-bit DEST, fp32 dest accumulation off, as binary_ng runs
// bf16 operands):
//   s1 = g0 + g1, s2 = s1 + g2: the FPU add (add_tiles);
//   out = s2 * scale: the SFPU multiply (both operands copied to DEST, mul_binary_tile) with the scale column broadcast
//   over the row by the reader (binary_ng's column-broadcast fill).
// Which unit each stock op used was measured against the stock ops (logs/bevformer/opt_r8/scatp_*.log): the SFPU add
// rounds ties to even where the stock add does not; the FPU multiply rounds differently from the stock multiply.
#include <cstdint>

#include "api/compute/bcast.h"
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/eltwise_binary.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/pack.h"

constexpr uint32_t DT = get_compile_time_arg_val(0);
constexpr uint32_t NWORK = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t cb_g0 = 0, cb_g1 = 1, cb_g2 = 2, cb_sc = 3, cb_s1 = 4, cb_s2 = 5, cb_out = 16;
constexpr uint32_t SB = 4;  // tiles per DEST batch (operand pairs: 8 of the 16-bit DEST half)
static_assert(DT % SB == 0, "scatter_k: D a multiple of 128");

// s = a + b: binary_ng's FPU add of bf16 operands (add_tiles; measured: the SFPU add's RNE does not reproduce it)
static inline void add_step(uint32_t a, uint32_t b, uint32_t o) {
    cb_wait_front(a, DT);
    cb_wait_front(b, DT);
    cb_reserve_back(o, DT);
    reconfig_data_format(a, b);
    pack_reconfig_data_format(o);
    add_init(a, b);
    for (uint32_t t0 = 0; t0 < DT; t0 += SB) {
        tile_regs_acquire();
        for (uint32_t i = 0; i < SB; ++i) {
            add_tiles(a, b, t0 + i, t0 + i, i);
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t i = 0; i < SB; ++i) {
            pack_tile<true>(i, o, t0 + i);
        }
        tile_regs_release();
    }
    cb_push_back(o, DT);
    cb_pop_front(a, DT);
    cb_pop_front(b, DT);
}

// o = a * bcast: binary_ng's SFPU multiply (copy both to DEST, mul_binary_tile; the broadcast tile filled by the
// reader; measured: the FPU multiply does not reproduce it)
static inline void mul_step(uint32_t a, uint32_t b, uint32_t o) {
    cb_wait_front(a, DT);
    cb_wait_front(b, 1);
    cb_reserve_back(o, DT);
    pack_reconfig_data_format(o);
    for (uint32_t t0 = 0; t0 < DT; t0 += SB) {
        tile_regs_acquire();
        reconfig_data_format_srca(a);
        copy_init(a);
        for (uint32_t i = 0; i < SB; ++i) {
            copy_tile(a, t0 + i, 2 * i);
        }
        reconfig_data_format_srca(b);
        copy_init(b);
        for (uint32_t i = 0; i < SB; ++i) {
            copy_tile(b, 0, 2 * i + 1);
            mul_binary_tile_init();
            mul_binary_tile(2 * i, 2 * i + 1, 2 * i);
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t i = 0; i < SB; ++i) {
            pack_tile<true>(2 * i, o, t0 + i);
        }
        tile_regs_release();
    }
    cb_push_back(o, DT);
    cb_pop_front(a, DT);
    cb_pop_front(b, 1);
}

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t w0 = ci * NWORK / NC;
    const uint32_t w1 = (ci + 1) * NWORK / NC;
    if (w0 >= w1) {
        return;
    }
    compute_kernel_hw_startup(cb_g0, cb_g1, cb_out);
    for (uint32_t w = w0; w < w1; ++w) {
#if defined(SCATTER_K_STOP) && SCATTER_K_STOP == 1     // debug: out = g0 + g1 (g2 / scale consumed unused)
        add_step(cb_g0, cb_g1, cb_out);
        cb_wait_front(cb_g2, DT);
        cb_pop_front(cb_g2, DT);
        cb_wait_front(cb_sc, 1);
        cb_pop_front(cb_sc, 1);
#elif defined(SCATTER_K_STOP) && SCATTER_K_STOP == 2   // debug: out = (g0 + g1) + g2
        add_step(cb_g0, cb_g1, cb_s1);
        add_step(cb_s1, cb_g2, cb_out);
        cb_wait_front(cb_sc, 1);
        cb_pop_front(cb_sc, 1);
#else
        add_step(cb_g0, cb_g1, cb_s1);
        add_step(cb_s1, cb_g2, cb_s2);
        mul_step(cb_s2, cb_sc, cb_out);
#endif
    }
}