// SPDX-License-Identifier: Apache-2.0
// S1 compute (sceneseg-p150 OPT_PLAN S1, tt/subpixel.py): per output tile of the sub-pixel up-block, with the EPI
// LLKs (epi_compute.cpp):
//   y = Zhi + Zlo           add_binary_tile<NearestEven> (fp32 SFPU, operands unpacked to DEST as fp32)
//   y = (y + Shi) + Slo     (S1_SKIP: the composed skip conv's two products)
//   y = y + bias            the reader's bias tile (interior row or the border pattern of this row tile)
//   y = gelu(y)             gelu_tile<false> (S1_ACT 2: ttnn.gelu's exact erf form)
//   out = bf16(y)           typecast_tile<Float32, Float16_b> (S1_OUT_BF16)
// CT args: [0] R, [1] CT, [2] NC, [3] GX (work split as the reader: units u = c * R + r, contiguous per core).
// Defines: S1_ACT, S1_OUT_BF16, S1_SKIP, S1_ZLO (the Z lo product present: two weight terms), S1_SLO (likewise for
// the skip conv; round 4 P2: 0 = one weight term, no add, no CB).
#include <cstdint>

#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/gelu.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/tile_move_copy.h"

using namespace ckernel;

constexpr uint32_t R = get_compile_time_arg_val(0);
constexpr uint32_t CT = 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_zh = 0, cb_zl = 1, cb_sh = 2, cb_b = 3, cb_sl = 6, cb_out = 16;

FORCE_INLINE void add_from(uint32_t cb) {
    copy_init(cb);
    copy_tile(cb, 0, 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 U = R * CT;
    constexpr uint32_t B = U / NC, E = U % NC;
    const uint32_t n = B + (core < E ? 1 : 0);
    if (n == 0) {
        return;
    }
    compute_kernel_hw_startup(cb_zh, cb_out);
    for (uint32_t t = 0; t < n; ++t) {
        cb_wait_front(cb_zh, 1);
#if S1_ZLO
        cb_wait_front(cb_zl, 1);
#endif
        cb_wait_front(cb_b, 1);
#if S1_SKIP
        cb_wait_front(cb_sh, 1);
#if S1_SLO
        cb_wait_front(cb_sl, 1);
#endif
#endif
        tile_regs_acquire();
        copy_init(cb_zh);
        copy_tile(cb_zh, 0, 0);
#if S1_ZLO
        add_from(cb_zl);
#endif
#if S1_SKIP
        add_from(cb_sh);
#if S1_SLO
        add_from(cb_sl);
#endif
#endif
        add_from(cb_b);
#if S1_ACT == 2
        gelu_tile_init<false>();
        gelu_tile<false>(0);
#endif
#if S1_OUT_BF16
        typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
        typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
#endif
        tile_regs_commit();
        cb_reserve_back(cb_out, 1);
        tile_regs_wait();
        pack_tile(0, cb_out);
        tile_regs_release();
        cb_push_back(cb_out, 1);
        cb_pop_front(cb_zh, 1);
#if S1_ZLO
        cb_pop_front(cb_zl, 1);
#endif
        cb_pop_front(cb_b, 1);
#if S1_SKIP
        cb_pop_front(cb_sh, 1);
#if S1_SLO
        cb_pop_front(cb_sl, 1);
#endif
#endif
    }
}