// SPDX-License-Identifier: Apache-2.0
// TFUSE compute (meteor-p150 OPT2 MK-TF, tt/tgate.py fuse_mm): the temporal fuse's merged 1x1 conv and its gate glue
// in one pass per 32-row block (the TGATE kernel, tgate_compute.cpp, with the conv as its first phase):
//   phase 1, the conv:  g_j = sum_k x_k @ W[k, j]   (j = 0..3, k = 0..2; FPU matmuls, bf16 x (SrcB) times fp32 W
//                       (SrcA, TF32) into fp32 DEST, HiFi4: the stock fp32-weight matmul's operands), packed fp32:
//                       j = 0..2 -> CB 0 (t), j = 3 -> CB 1 (the logit differences without their bias)
//   phase 2, the glue:  E = exp(d + gb_hi + gb_lo)   (SFPU, fp32; the gate bias as two bf16 terms, CB 7)
//                       S = E @ ONES, G = SCALE / (1 + S), out_k = bf16_rne(relu(t_k * G + b_k))  (as TGATE)
// SrcA stays fp32 throughout (W, every copy source, the ones tile); SrcB is bf16 for phase 1 (x) and fp32 for the
// E @ ONES matmul: reconfigured twice per block. Every CB reaches DEST through SrcA / SrcB (TF32, as TGATE).
// CBs: 8 x (3 tiles bf16), 6 W (12 fp32), 0 t (3, loop-back), 1 d (loop-back), 2 ones, 3 bias (3), 7 gate bias (2),
// 4 E (loop-back), 5 G (loop-back), 16 out (bf16).
// CT args: [0] NB, [1] NC, [2] GX, [3] SCALE (fp32 bits).
#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/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/binop_with_scalar.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/exp.h"
#include "api/compute/eltwise_unary/recip.h"
#include "api/compute/eltwise_unary/relu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/matmul.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"

using namespace ckernel;

constexpr uint32_t NB = 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 SCALE = get_compile_time_arg_val(3);
constexpr uint32_t ONE = 0x3F800000u;
constexpr uint32_t cb_t = 0, cb_l = 1, cb_one = 2, cb_bias = 3, cb_e = 4, cb_g = 5, cb_w = 6, cb_gb = 7, cb_x = 8,
                   cb_out = 16;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NB / NC, E = NB % NC;
    const uint32_t b0 = core * B + (core < E ? core : E);
    const uint32_t n = B + (core < E ? 1 : 0);
    if (n == 0) {
        return;
    }
    // matmul operand order: in0 (x) -> SrcB, in1 (W) -> SrcA
    compute_kernel_hw_startup<SrcOrder::Reverse>(cb_x, cb_w, cb_t);
    cb_wait_front(cb_w, 12);
    cb_wait_front(cb_one, 1);
    cb_wait_front(cb_bias, 3);
    cb_wait_front(cb_gb, 2);
    for (uint32_t b = b0; b < b0 + n; ++b) {
        // phase 1: g_j = x @ W[:, j]
        reconfig_data_format_srcb(cb_x);
        cb_wait_front(cb_x, 3);
        for (uint32_t j = 0; j < 4; ++j) {
            const uint32_t dst = j < 3 ? cb_t : cb_l;
            tile_regs_acquire();
            matmul_init(cb_x, cb_w);
            for (uint32_t k = 0; k < 3; ++k) {
                matmul_tiles(cb_x, cb_w, k, k * 4 + j, 0);
            }
            tile_regs_commit();
            cb_reserve_back(dst, 1);
            tile_regs_wait();
            pack_reconfig_data_format(dst);
            pack_tile(0, dst);
            tile_regs_release();
            cb_push_back(dst, 1);
        }
        cb_pop_front(cb_x, 3);
        reconfig_data_format_srcb(cb_e);
        // E = exp(d + gb_hi + gb_lo)
        cb_wait_front(cb_l, 1);
        tile_regs_acquire();
        copy_init(cb_l);
        copy_tile(cb_l, 0, 0);
        copy_init(cb_gb);
        copy_tile(cb_gb, 0, 1);
        add_binary_tile_init();
        add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
        copy_init(cb_gb);
        copy_tile(cb_gb, 1, 1);
        add_binary_tile_init();
        add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
        exp_tile_init();
        exp_tile(0);
        tile_regs_commit();
        cb_reserve_back(cb_e, 1);
        tile_regs_wait();
        pack_reconfig_data_format(cb_e);
        pack_tile(0, cb_e);
        tile_regs_release();
        cb_push_back(cb_e, 1);
        cb_pop_front(cb_l, 1);
        // G = SCALE / (1 + E @ ONES)
        cb_wait_front(cb_e, 1);
        tile_regs_acquire();
        matmul_init(cb_e, cb_one);
        matmul_tiles(cb_e, cb_one, 0, 0, 0);
        binop_with_scalar_tile_init();
        add_unary_tile(0, ONE);
        recip_tile_init();
        recip_tile(0);
        binop_with_scalar_tile_init();
        mul_unary_tile(0, SCALE);
        tile_regs_commit();
        cb_reserve_back(cb_g, 1);
        tile_regs_wait();
        pack_reconfig_data_format(cb_g);
        pack_tile(0, cb_g);
        tile_regs_release();
        cb_push_back(cb_g, 1);
        cb_pop_front(cb_e, 1);
        // out_k = bf16(relu(t_k * G + b_k))
        cb_wait_front(cb_g, 1);
        cb_wait_front(cb_t, 3);
        cb_reserve_back(cb_out, 3);
        for (uint32_t k = 0; k < 3; ++k) {
            tile_regs_acquire();
            copy_init(cb_t);
            copy_tile(cb_t, k, 0);
            copy_init(cb_g);
            copy_tile(cb_g, 0, 1);
            mul_binary_tile_init();
            mul_binary_tile(0, 1, 0);
            copy_init(cb_bias);
            copy_tile(cb_bias, k, 1);
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
            relu_tile_init();
            relu_tile(0);
            typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
            typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
            tile_regs_commit();
            tile_regs_wait();
            pack_reconfig_data_format(cb_out);
            pack_tile(0, cb_out);
            tile_regs_release();
        }
        cb_push_back(cb_out, 3);
        cb_pop_front(cb_t, 3);
        cb_pop_front(cb_g, 1);
    }
}