// SPDX-License-Identifier: Apache-2.0
// K11 depthwise compute (sceneseg-p150 OPT_PLAN K11, tt/dw.py): per output tile (32 pixels x 32 channels, float32)
//   acc = x_0 * w_0;  acc = x_t * w_t + acc (t = 1 .. KK-1, SFPMAD: one fp32 rounding per tap)
//   acc = acc + bias  (add_binary_tile, fp32 SFPU add, as EPI)
//   acc = silu(acc)   (silu_tile, ttnn.silu's LLK, as EPI)        [DW_ACT 1; 0 = none]
// x_t: the gathered input page of tap t (CB 0, unpacked to DEST as fp32), w_t: weight tile t of the current channel
// tile (CB 1, every row = the 32 channel weights of tap t), bias: CB 2. DEST tiles: 0 = acc, 1 = x, 2 = w.
// Work split and order as dw_reader.cpp; the weight / bias CBs are popped when the channel tile changes.
// CT args: [0] KK, [1] CT, [2] PT, [3] NC, [4] GX, [5] CH, [6] NWT (weight tiles per channel tile). Defines: DW_ACT,
// DW_WPACK.
// R7.1 DW_WPACK=1: the NWT = ceil(KK / 8) weight tiles hold 8 taps each (rows 4l .. 4l+3 = tap 8j + l); per output
// tile they are unpacked to DEST tiles 2 / 3 once (5x5: again for taps 16 .. 24) and tap t's SFPMADs read their weight
// vectors from its 4-row group (k11_tap_tile). An SFPU vector is 4 DEST rows x 8 columns of one parity (sfpi unit
// 2k = rows 4k' .. +3 even columns, 2k + 1 the odd ones), so the x vector of unit d takes the group's unit of the same
// parity: the same operands per SFPMAD as the one-tap-per-tile form (bit-identical).
#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/tile_move_copy.h"

#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
// one face (8 row groups) of DEST tile 0 = tile 1 * tile 2 (+ tile 0)
inline void k11_mul_face(const std::uint32_t, const std::uint32_t, const std::uint32_t) {
#pragma GCC unroll 8
    for (int d = 0; d < 8; d++) {
        sfpi::vFloat x = sfpi::dst_reg[32];
        sfpi::vFloat w = sfpi::dst_reg[64];
        sfpi::dst_reg[0] = x * w;
        sfpi::dst_reg++;
    }
}
inline void k11_mad_face(const std::uint32_t, const std::uint32_t, const std::uint32_t) {
#pragma GCC unroll 8
    for (int d = 0; d < 8; d++) {
        sfpi::vFloat x = sfpi::dst_reg[32];
        sfpi::vFloat w = sfpi::dst_reg[64];
        sfpi::vFloat a = sfpi::dst_reg[0];
        sfpi::dst_reg[0] = x * w + a;
        sfpi::dst_reg++;
    }
}
// R7.1: the whole tile in one call (VectorMode::None, base = DEST tile 0): tile 0 = tile 1 * w (+ tile 0), w = the
// tap's weight group at sfpi unit WR relative to the face base for faces 0 / 1, WR - 16 for faces 2 / 3.
template <bool FIRST>
inline void k11_tap_tile(const std::uint32_t, const std::uint32_t, const std::uint32_t, const std::uint32_t wr) {
#pragma GCC unroll 4
    for (int f = 0; f < 4; f++) {
        const int wo = (int)wr - (f >= 2 ? 16 : 0);
        sfpi::vFloat we = sfpi::dst_reg[wo];
        sfpi::vFloat wd = sfpi::dst_reg[wo + 1];
#pragma GCC unroll 8
        for (int d = 0; d < 8; d++) {
            sfpi::vFloat x = sfpi::dst_reg[32];
            if constexpr (FIRST) {
                sfpi::dst_reg[0] = x * ((d & 1) ? wd : we);
            } else {
                sfpi::vFloat a = sfpi::dst_reg[0];
                sfpi::dst_reg[0] = x * ((d & 1) ? wd : we) + a;
            }
            sfpi::dst_reg++;
        }
    }
}
}  // namespace ckernel::sfpu
#endif

using namespace ckernel;

constexpr uint32_t KK = get_compile_time_arg_val(0);
constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t PT = 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 CH = get_compile_time_arg_val(5);
constexpr uint32_t NWT = get_compile_time_arg_val(6);
constexpr uint32_t NW = CT * PT;
constexpr uint32_t cb_x = 0, cb_w = 1, cb_b = 2, 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 = NW / NC, E = NW % NC;
    const uint32_t w0 = core * B + (core < E ? core : E);
    const uint32_t w1 = w0 + B + (core < E ? 1 : 0);
    compute_kernel_hw_startup(cb_x, cb_out);
    uint32_t cur_ct = 0xFFFFFFFF;
    uint32_t w = w0;
    while (w < w1) {
        const uint32_t ct = w / PT;
        const uint32_t pt0 = w - ct * PT;
        uint32_t pt1 = pt0 + CH;
        if (pt1 > PT) {
            pt1 = PT;
        }
        if (pt1 - pt0 > w1 - w) {
            pt1 = pt0 + (w1 - w);
        }
        if (ct != cur_ct) {
            if (cur_ct != 0xFFFFFFFF) {
                cb_pop_front(cb_w, NWT);
                cb_pop_front(cb_b, 1);
            }
            cb_wait_front(cb_w, NWT);
            cb_wait_front(cb_b, 1);
            cur_ct = ct;
        }
        for (uint32_t pt = pt0; pt < pt1; ++pt) {
            tile_regs_acquire();
#if DW_WPACK == 1
            copy_init(cb_w);
            copy_tile(cb_w, 0, 2);
            if constexpr (NWT > 1) {
                copy_tile(cb_w, 1, 3);
            }
            copy_init(cb_x);
            add_binary_tile_init();
            for (uint32_t t = 0; t < KK; ++t) {
                if (t == 16) {
                    copy_init(cb_w);
                    copy_tile(cb_w, 2, 2);
                    if constexpr (NWT > 3) {
                        copy_tile(cb_w, 3, 3);
                    }
                    copy_init(cb_x);
                    add_binary_tile_init();
                }
                cb_wait_front(cb_x, 1);
                copy_tile(cb_x, 0, 1);
                const uint32_t wr = 64 + 32 * ((t >> 3) & 1) + 16 * ((t >> 2) & 1) + 2 * (t & 3);
                if (t == 0) {
                    MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, k11_tap_tile, (true), 0, 1, 0,
                                           VectorMode::None, wr)));
                } else {
                    MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, k11_tap_tile, (false), 0, 1, 0,
                                           VectorMode::None, wr)));
                }
                cb_pop_front(cb_x, 1);
            }
#else
            for (uint32_t t = 0; t < KK; ++t) {
                cb_wait_front(cb_x, 1);
                copy_init(cb_x);
                copy_tile(cb_x, 0, 1);
                copy_init(cb_w);
                copy_tile(cb_w, t, 2);
                add_binary_tile_init();
                if (t == 0) {
                    MATH((SFPU_BINARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, k11_mul_face, 0, 1, 0,
                                                            VectorMode::RC)));
                } else {
                    MATH((SFPU_BINARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, k11_mad_face, 0, 1, 0,
                                                            VectorMode::RC)));
                }
                cb_pop_front(cb_x, 1);
            }
#endif
            copy_init(cb_b);
            copy_tile(cb_b, 0, 1);
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
#if DW_ACT == 1
            silu_tile_init();
            silu_tile(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);
        }
        w += pt1 - pt0;
    }
    if (cur_ct != 0xFFFFFFFF) {
        cb_pop_front(cb_w, NWT);
        cb_pop_front(cb_b, 1);
    }
}