// SPDX-License-Identifier: Apache-2.0
// glue_k (tt/glue_k.py), compute: the stock LLK calls of the ops it replaces, in their order, on DEST fp32 (every
// fp32 input unpacked to DEST as fp32, as binary_ng / typecast / unary do for fp32 operands):
//   MODE 0: d0 = A, d1 = B, add_binary_tile (RNE) -> d0; HAS_C: d1 = C[col], add -> d0; typecast fp32 -> int32 ->
//           out0                      (= ttnn.add(A, B) [+ ttnn.add(., C row)] -> ttnn.typecast(int32))
//   MODE 1: d0 = A, d1 = S, mul_binary_tile -> d0; d1 = B, add -> d0; typecast -> out0; d0 = A, sigmoid<RC, 0> ->
//           out1                      (= ttnn.multiply(A, S), ttnn.add(., B), ttnn.typecast(int32); ttnn.sigmoid(A))
#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/eltwise_unary.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"

constexpr uint32_t MODE = get_compile_time_arg_val(0);
constexpr uint32_t NT = 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 NCOLS = get_compile_time_arg_val(4);
constexpr uint32_t HAS_C = get_compile_time_arg_val(5);
constexpr uint32_t cb0 = 0, cb1 = 1, cb2 = 2, cbo0 = 16, cbo1 = 17;
constexpr uint32_t F32 = static_cast<uint32_t>(DataFormat::Float32);
constexpr uint32_t I32 = static_cast<uint32_t>(DataFormat::Int32);

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t t0 = ci * NT / NC;
    const uint32_t t1 = (ci + 1) * NT / NC;
    if (t0 >= t1) {
        return;
    }
    compute_kernel_hw_startup(cb0, cbo0);
    if constexpr (MODE == 0 && HAS_C) {
        cb_wait_front(cb2, NCOLS);
    }
    for (uint32_t t = t0; t < t1; ++t) {
        cb_wait_front(cb0, 1);
        cb_wait_front(cb1, 1);
        if constexpr (MODE == 1) {
            cb_wait_front(cb2, 1);
        }
        tile_regs_acquire();
        reconfig_data_format_srca(cb0);
        copy_tile_to_dst_init_short(cb0);
        copy_tile(cb0, 0, 0);
        reconfig_data_format_srca(cb1);
        copy_tile_to_dst_init_short(cb1);
        copy_tile(cb1, 0, 1);
        if constexpr (MODE == 0) {
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
            if constexpr (HAS_C) {
                reconfig_data_format_srca(cb2);
                copy_tile_to_dst_init_short(cb2);
                copy_tile(cb2, t % NCOLS, 1);
                add_binary_tile_init();
                add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
            }
        } else {
            mul_binary_tile_init();
            mul_binary_tile(0, 1, 0);
            reconfig_data_format_srca(cb2);
            copy_tile_to_dst_init_short(cb2);
            copy_tile(cb2, 0, 1);
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
            reconfig_data_format_srca(cb0);
            copy_tile_to_dst_init_short(cb0);
            copy_tile(cb0, 0, 2);
            sigmoid_tile_init<false>();
            sigmoid_tile<VectorMode::RC, false>(2);
        }
        typecast_tile_init<F32, I32>();
        typecast_tile<F32, I32>(0);
        tile_regs_commit();
        cb_reserve_back(cbo0, 1);
        if constexpr (MODE == 1) {
            cb_reserve_back(cbo1, 1);
        }
        tile_regs_wait();
        pack_reconfig_data_format(cbo0);
        pack_tile(0, cbo0);
        if constexpr (MODE == 1) {
            pack_reconfig_data_format(cbo1);
            pack_tile(2, cbo1);
        }
        tile_regs_release();
        cb_push_back(cbo0, 1);
        if constexpr (MODE == 1) {
            cb_push_back(cbo1, 1);
        }
        cb_pop_front(cb0, 1);
        cb_pop_front(cb1, 1);
        if constexpr (MODE == 1) {
            cb_pop_front(cb2, 1);
        }
    }
}