// SPDX-License-Identifier: Apache-2.0
// SE-K compute (sceneseg-p150 OPT_PLAN SE-K, tt/sek.py): one squeeze-excitation MLP on the fp32 row sum x
// (the 1 / (H W) of the mean is folded into W1), with the port's precision scheme (tt/layers.py Linear: bf16 terms
// of the activation and of the weights, the products hi x W_hi, hi x W_lo, lo x W_hi, fp32 bias, fp32 activation)
// but all products of a layer accumulated in one fp32 DEST tile (the stock path rounded each product's matmul to
// fp32 separately and summed them in the epilogue):
//   x -> (x_hi, x_lo)                  typecast_tile<Float32, Float16_b> + sub_binary_tile (layers.split_activation)
//   h = silu(sum_k x_hi W1hi + x_hi W1lo + x_lo W1hi + b1)        (every core, all NT1 tiles)
//   h -> (h_hi, h_lo)
//   g = sigmoid(sum_k h_hi W2hi + h_hi W2lo + h_lo W2hi + b2)     (this core's fc2 output tiles)
// CT args: [0] KT1, [1] NT1, [2] NT2, [3] NC, [4] GX.
#include <cstdint>

#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/copy_dest_values.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/matmul.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"

using namespace ckernel;

constexpr uint32_t KT1 = get_compile_time_arg_val(0);
constexpr uint32_t NT1 = get_compile_time_arg_val(1);
constexpr uint32_t NT2 = get_compile_time_arg_val(2);
constexpr uint32_t NC = get_compile_time_arg_val(3);
constexpr uint32_t GX = get_compile_time_arg_val(4);
constexpr uint32_t cb_x = 0, cb_w1 = 1, cb_b1 = 2, cb_w2 = 3, cb_b2 = 4, cb_out = 16;
constexpr uint32_t cb_xhi = 24, cb_xlo = 25, cb_hhi = 26, cb_hlo = 27;

// DEST 0 holds y (fp32): DEST 1 = bf16(y), DEST 2 = bf16(y - DEST 1), packed to the bf16 CBs hi / lo
FORCE_INLINE void split_pack(uint32_t cb_hi, uint32_t cb_lo) {
    copy_dest_values_init();
    copy_dest_values<DataFormat::Float32>(0, 1);
    copy_dest_values<DataFormat::Float32>(0, 2);
    typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
    typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(1);
    sub_binary_tile_init();
    sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(2, 1, 2);
    typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
    typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(2);
    tile_regs_commit();
    cb_reserve_back(cb_hi, 1);
    cb_reserve_back(cb_lo, 1);
    tile_regs_wait();
    pack_reconfig_data_format(cb_hi);
    pack_tile(1, cb_hi);
    pack_tile(2, cb_lo);
    tile_regs_release();
    cb_push_back(cb_hi, 1);
    cb_push_back(cb_lo, 1);
}

// srcA: the bf16 weights for matmul (srcA = in1), the fp32 CBs for the copies to DEST: reconfigured explicitly
// around each copy (as bmm_large_block_zm_fused_bias_activation.cpp does around its partials reload)
FORCE_INLINE void add_bias(uint32_t cb, uint32_t idx) {
    reconfig_data_format_srca(cb);
    copy_init(cb);
    copy_tile(cb, idx, 1);
    add_binary_tile_init();
    add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NT2 / NC, E = NT2 % NC;
    const uint32_t nn = B + (core < E ? 1 : 0);
    compute_kernel_hw_startup(cb_xhi, cb_w1, cb_out);
    // ---- x -> bf16 terms
    cb_wait_front(cb_x, KT1);
    for (uint32_t k = 0; k < KT1; ++k) {
        tile_regs_acquire();
        reconfig_data_format_srca(cb_x);
        copy_init(cb_x);
        copy_tile(cb_x, k, 0);
        split_pack(cb_xhi, cb_xlo);
    }
    cb_pop_front(cb_x, KT1);
    // ---- fc1 + SiLU -> h terms
    cb_wait_front(cb_xhi, KT1);
    cb_wait_front(cb_xlo, KT1);
    cb_wait_front(cb_w1, 2 * KT1 * NT1);
    cb_wait_front(cb_b1, NT1);
    for (uint32_t n = 0; n < NT1; ++n) {
        tile_regs_acquire();
        reconfig_data_format(cb_xhi, cb_w1);
        matmul_init(cb_xhi, cb_w1);
        for (uint32_t k = 0; k < KT1; ++k) {
            matmul_tiles(cb_xhi, cb_w1, k, k * NT1 + n, 0);                      // x_hi W1_hi
            matmul_tiles(cb_xhi, cb_w1, k, (KT1 + k) * NT1 + n, 0);              // x_hi W1_lo
        }
        matmul_init(cb_xlo, cb_w1);
        for (uint32_t k = 0; k < KT1; ++k) {
            matmul_tiles(cb_xlo, cb_w1, k, k * NT1 + n, 0);                      // x_lo W1_hi
        }
        add_bias(cb_b1, n);
        silu_tile_init();
        silu_tile(0);
        split_pack(cb_hhi, cb_hlo);
    }
    cb_pop_front(cb_xhi, KT1);
    cb_pop_front(cb_xlo, KT1);
    cb_pop_front(cb_w1, 2 * KT1 * NT1);
    cb_pop_front(cb_b1, NT1);
    // ---- fc2 + sigmoid on this core's output tiles
    cb_wait_front(cb_hhi, NT1);
    cb_wait_front(cb_hlo, NT1);
    cb_wait_front(cb_w2, 2 * NT1 * nn);
    cb_wait_front(cb_b2, nn);
    for (uint32_t j = 0; j < nn; ++j) {
        tile_regs_acquire();
        reconfig_data_format(cb_hhi, cb_w2);
        matmul_init(cb_hhi, cb_w2);
        for (uint32_t k = 0; k < NT1; ++k) {
            matmul_tiles(cb_hhi, cb_w2, k, k * nn + j, 0);                       // h_hi W2_hi
            matmul_tiles(cb_hhi, cb_w2, k, (NT1 + k) * nn + j, 0);               // h_hi W2_lo
        }
        matmul_init(cb_hlo, cb_w2);
        for (uint32_t k = 0; k < NT1; ++k) {
            matmul_tiles(cb_hlo, cb_w2, k, k * nn + j, 0);                       // h_lo W2_hi
        }
        add_bias(cb_b2, j);
        sigmoid_tile_init<false>();
        sigmoid_tile<VectorMode::RC, false>(0);
        tile_regs_commit();
        cb_reserve_back(cb_out, 1);
        tile_regs_wait();
        pack_reconfig_data_format(cb_out);
        pack_tile(0, cb_out);
        tile_regs_release();
        cb_push_back(cb_out, 1);
    }
    cb_pop_front(cb_hhi, NT1);
    cb_pop_front(cb_hlo, NT1);
    cb_pop_front(cb_w2, 2 * NT1 * nn);
    cb_pop_front(cb_b2, nn);
}