// SPDX-License-Identifier: Apache-2.0
// dcn_k (tt/dcn_k.py), compute. Per work item (32 pixels x one tap; see dcn_k_reader.cpp):
//   1. DEST0..2 <- cb_geo (fx, fy, mask in column 0; fp32, UnpackToDestFp32); SFPU: the 4 corner weights
//      mask * wx * wy (fp32) -> cb_w (4 fp32 tiles, column 0 valid, the other columns 0).
//   2. for each corner c: tilize its 32 gathered rows (cb_rm0 for NW / NE, cb_rm1 for SW / SE) into CT tiles (cb_xt),
//      DEST = xt * bcast_col(W_c) (HiFi4; the fp32 weight as tf32), packer L1 accumulation into cb_acc (CT fp32).
//   3. cb_acc -> bf16 -> cb_out (CT tiles: the item's im2col tiles).
#include <cstdint>

#include "api/compute/bcast.h"
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.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"
#include "api/compute/tilize.h"
#ifdef DCN_K_GLUE
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/typecast.h"
#endif

#if defined(TRISC_MATH)
#include "llk_math_eltwise_unary_sfpu_macros.h"
#include "sfpi.h"

namespace ckernel::sfpu {
inline void dcn_k_w_face() {
#pragma GCC unroll 2
    for (int d = 0; d < 8; d++) {
        sfpi::vFloat fx = sfpi::dst_reg[0];
        sfpi::vFloat fy = sfpi::dst_reg[32];
        sfpi::vFloat a = sfpi::dst_reg[64];
        sfpi::vFloat gx = 1.0f - fx;
        sfpi::vFloat ay0 = a * (1.0f - fy);
        sfpi::vFloat ay1 = a * fy;
        sfpi::dst_reg[0] = gx * ay0;
        sfpi::dst_reg[32] = fx * ay0;
        sfpi::dst_reg[64] = gx * ay1;
        sfpi::dst_reg[96] = fx * ay1;
        sfpi::dst_reg++;
    }
}
}  // namespace ckernel::sfpu
#endif

constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t NWORK = 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_rm0 = 0, cb_rm1 = 1, cb_geo = 2, cb_xt = 3, cb_w = 4, cb_acc = 5, cb_out = 16;
constexpr uint32_t DB = 4;  // tiles per DEST section (fp32)
#ifdef DCN_K_GLUE
// BEVFORMER_DCN_GLUE=kernel (D20 round 2): per new block, glue_k's MODE-1 steps (loc = int32(off * S + B) with the
// SFPU mul / RNE add / typecast, sig = sigmoid(off)) -> loc to both readers (cb 6 / 7), sig to PART 0 (cb 8)
constexpr uint32_t cb_l0 = 6, cb_l1 = 7, cb_sg = 8, cb_goff = 9, cb_gb = 10, cb_gs = 11;
constexpr uint32_t F32 = static_cast<uint32_t>(DataFormat::Float32);
constexpr uint32_t I32 = static_cast<uint32_t>(DataFormat::Int32);

inline void glue_block() {
    cb_wait_front(cb_goff, 1);
    cb_wait_front(cb_gb, 1);
    tile_regs_acquire();
    reconfig_data_format_srca(cb_goff);
    copy_tile_to_dst_init_short(cb_goff);
    copy_tile(cb_goff, 0, 0);
    reconfig_data_format_srca(cb_gs);
    copy_tile_to_dst_init_short(cb_gs);
    copy_tile(cb_gs, 0, 1);
    mul_binary_tile_init();
    mul_binary_tile(0, 1, 0);
    reconfig_data_format_srca(cb_gb);
    copy_tile_to_dst_init_short(cb_gb);
    copy_tile(cb_gb, 0, 1);
    add_binary_tile_init();
    add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
    reconfig_data_format_srca(cb_goff);
    copy_tile_to_dst_init_short(cb_goff);
    copy_tile(cb_goff, 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(cb_l0, 1);
    cb_reserve_back(cb_l1, 1);
    cb_reserve_back(cb_sg, 1);
    tile_regs_wait();
    pack_reconfig_data_format(cb_l0);
    pack_tile(0, cb_l0);
    pack_tile(0, cb_l1);
    pack_reconfig_data_format(cb_sg);
    pack_tile(2, cb_sg);
    tile_regs_release();
    cb_push_back(cb_l0, 1);
    cb_push_back(cb_l1, 1);
    cb_push_back(cb_sg, 1);
    cb_pop_front(cb_goff, 1);
    cb_pop_front(cb_gb, 1);
}
#endif
static_assert(CT % DB == 0, "dcn_k: C a multiple of 128");

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t w0 = ci * NWORK / NC;
    const uint32_t w1 = (ci + 1) * NWORK / NC;
    if (w0 == w1) {
        return;
    }
    compute_kernel_hw_startup(cb_geo, cb_geo, cb_w);
#ifdef DCN_K_GLUE
    cb_wait_front(cb_gs, 1);
    uint32_t cur_b = 0xFFFFFFFF;
#endif
    for (uint32_t w = w0; w < w1; ++w) {
#ifdef DCN_K_GLUE
        if (w / 9 != cur_b) {
            cur_b = w / 9;
            glue_block();
        }
#endif
#ifdef DCN_K_NEAREST
        // BEVFORMER_DCN_SAMPLE=nearest (D20 round 3): one gathered corner per item (even items from cb_rm0, odd from
        // cb_rm1), tilized, times the mask (cb_geo tile, column 0; column-broadcast multiply), packed to bf16
        {
            const uint32_t src = ((w - w0) & 1) ? cb_rm1 : cb_rm0;
            cb_wait_front(src, CT);
            cb_reserve_back(cb_xt, CT);
            reconfig_data_format_srca(src);
            pack_reconfig_data_format(cb_xt);
#ifdef DCN_K_FAST_TILIZE
            fast_tilize_init(src, CT, cb_xt);
            fast_tilize_block(src, CT, cb_xt, 0, 0);
            fast_tilize_uninit(src, cb_xt, CT);
#else
            tilize_init(src, CT, cb_xt);
            tilize_block(src, CT, cb_xt, 0, 0);
            tilize_uninit(src, cb_xt);
#endif
            cb_push_back(cb_xt, CT);
            cb_pop_front(src, CT);
            cb_wait_front(cb_xt, CT);
            cb_wait_front(cb_geo, 1);
            cb_reserve_back(cb_out, CT);
            reconfig_data_format_srca(cb_xt);
            reconfig_data_format_srcb(cb_geo);
            pack_reconfig_data_format(cb_out);
            mul_bcast_cols_init(cb_xt, cb_geo);
            for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
                tile_regs_acquire();
                for (uint32_t i = 0; i < DB; ++i) {
                    mul_tiles_bcast_cols(cb_xt, cb_geo, t0 + i, 0, i);
                }
                tile_regs_commit();
                tile_regs_wait();
                for (uint32_t i = 0; i < DB; ++i) {
                    pack_tile<true>(i, cb_out, t0 + i);
                }
                tile_regs_release();
            }
            cb_push_back(cb_out, CT);
            cb_pop_front(cb_xt, CT);
            cb_pop_front(cb_geo, 1);
        }
        continue;
#endif
        // ---- 1. corner weights
        cb_wait_front(cb_geo, 3);
        cb_reserve_back(cb_w, 4);
        tile_regs_acquire();
        reconfig_data_format_srca(cb_geo);
        copy_init(cb_geo);
        copy_tile(cb_geo, 0, 0);
        copy_tile(cb_geo, 1, 1);
        copy_tile(cb_geo, 2, 2);
        MATH((SFPU_UNARY_INIT(unused)));
        MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, dcn_k_w_face, 0, VectorMode::RC)));
        tile_regs_commit();
        tile_regs_wait();
        pack_reconfig_data_format(cb_w);
        for (uint32_t c = 0; c < 4; ++c) {
            pack_tile(c, cb_w);
        }
        tile_regs_release();
        cb_push_back(cb_w, 4);
        cb_pop_front(cb_geo, 3);
        cb_wait_front(cb_w, 4);
        // ---- 2. the 4 corners
#ifdef DCN_K_BATCH
        // BEVFORMER_DCN_BATCH=1: the 4 corners tilized back to back (one tilize init / uninit per item), then the
        // 4 weighted multiplies with one init (the same products, the same accumulation order)
#ifndef DCN_K_MM
        cb_reserve_back(cb_acc, CT);
#endif
        cb_wait_front(cb_rm0, 2 * CT);
        cb_wait_front(cb_rm1, 2 * CT);
        cb_reserve_back(cb_xt, 4 * CT);
        reconfig_data_format_srca(cb_rm0);
        pack_reconfig_data_format(cb_xt);
#ifdef DCN_K_NOTILIZE
        // (anatomy probe only: no tilize, garbage tiles)
#elif defined(DCN_K_FAST_TILIZE)
        fast_tilize_init(cb_rm0, CT, cb_xt);
        for (uint32_t c = 0; c < 4; ++c) {
            fast_tilize_block(c < 2 ? cb_rm0 : cb_rm1, CT, cb_xt, (c & 1) * CT, c * CT);
        }
        fast_tilize_uninit(cb_rm0, cb_xt, CT);
#elif defined(DCN_K_PAIRS)
        // BEVFORMER_DCN_PAIRS: each reader's block is 32 rows x 2C ([NW | NE] / [SW | SE]): 2 CT tiles wide
        tilize_init(cb_rm0, 2 * CT, cb_xt);
        tilize_block(cb_rm0, 2 * CT, cb_xt, 0, 0);
        tilize_block(cb_rm1, 2 * CT, cb_xt, 0, 2 * CT);
        tilize_uninit(cb_rm0, cb_xt);
#else
        tilize_init(cb_rm0, CT, cb_xt);
        for (uint32_t c = 0; c < 4; ++c) {
            tilize_block(c < 2 ? cb_rm0 : cb_rm1, CT, cb_xt, (c & 1) * CT, c * CT);
        }
        tilize_uninit(cb_rm0, cb_xt);
#endif
        cb_push_back(cb_xt, 4 * CT);
        cb_pop_front(cb_rm0, 2 * CT);
        cb_pop_front(cb_rm1, 2 * CT);
        cb_wait_front(cb_xt, 4 * CT);
#ifdef DCN_K_MM
        // BEVFORMER_DCN_SUM=mm (D20 round 2): the reader put fx / fy / mask on the tile diagonal, so cb_w holds the 4
        // corner weights as diagonal matrices diag(w_c) (fp32); the weighted corner sum is one DEST accumulation of
        // 4 matmuls per output tile, out = sum_c diag(w_c) @ X_c, packed once to bf16 (no packer-L1 passes, no copy)
        cb_reserve_back(cb_out, CT);
        reconfig_data_format_srca(cb_xt);
        reconfig_data_format_srcb(cb_w);
        matmul_init(cb_w, cb_xt);
        pack_reconfig_data_format(cb_out);
        for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
            tile_regs_acquire();
#ifndef DCN_K_NOMUL
            for (uint32_t c = 0; c < 4; ++c) {
                for (uint32_t i = 0; i < DB; ++i) {
                    matmul_tiles(cb_w, cb_xt, c, c * CT + t0 + i, i);
                }
            }
#endif
            tile_regs_commit();
            tile_regs_wait();
            for (uint32_t i = 0; i < DB; ++i) {
                pack_tile<true>(i, cb_out, t0 + i);
            }
            tile_regs_release();
        }
        cb_push_back(cb_out, CT);
        cb_pop_front(cb_xt, 4 * CT);
        cb_pop_front(cb_w, 4);
        continue;
#endif
        reconfig_data_format_srca(cb_xt);
        reconfig_data_format_srcb(cb_w);
        pack_reconfig_data_format(cb_acc);
        mul_bcast_cols_init(cb_xt, cb_w);
        for (uint32_t c = 0; c < 4; ++c) {
            pack_reconfig_l1_acc(c == 0 ? 0 : 1);
            for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
                tile_regs_acquire();
                for (uint32_t i = 0; i < DB; ++i) {
                    mul_tiles_bcast_cols(cb_xt, cb_w, c * CT + t0 + i, c, i);
                }
                tile_regs_commit();
                tile_regs_wait();
                for (uint32_t i = 0; i < DB; ++i) {
                    pack_tile<true>(i, cb_acc, t0 + i);
                }
                tile_regs_release();
            }
        }
        pack_reconfig_l1_acc(0);
        cb_pop_front(cb_xt, 4 * CT);
#else
        cb_reserve_back(cb_acc, CT);
        for (uint32_t c = 0; c < 4; ++c) {
            const uint32_t src = c < 2 ? cb_rm0 : cb_rm1;
            if ((c & 1) == 0) {
                cb_wait_front(src, 2 * CT);
            }
            cb_reserve_back(cb_xt, CT);
            reconfig_data_format_srca(src);
            pack_reconfig_data_format(cb_xt);
#ifndef DCN_K_NOTILIZE
#ifdef DCN_K_FAST_TILIZE
            // BEVFORMER_DCN_TILIZE=fast: the BH fast tilize (unpack row chunks of 4 tiles; a lossless bf16 reorder)
            fast_tilize_init(src, CT, cb_xt);
            fast_tilize_block(src, CT, cb_xt, (c & 1) * CT, 0);
            fast_tilize_uninit(src, cb_xt, CT);
#else
            tilize_init(src, CT, cb_xt);
            tilize_block(src, CT, cb_xt, (c & 1) * CT, 0);
            tilize_uninit(src, cb_xt);
#endif
#endif
            cb_push_back(cb_xt, CT);
            cb_wait_front(cb_xt, CT);
            reconfig_data_format_srca(cb_xt);
            reconfig_data_format_srcb(cb_w);
            pack_reconfig_data_format(cb_acc);
            mul_bcast_cols_init(cb_xt, cb_w);
            pack_reconfig_l1_acc(c == 0 ? 0 : 1);
#ifdef DCN_K_NOMUL
            for (uint32_t t0 = 0; t0 < 0; t0 += DB) {
#else
            for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
#endif
                tile_regs_acquire();
                for (uint32_t i = 0; i < DB; ++i) {
                    mul_tiles_bcast_cols(cb_xt, cb_w, t0 + i, c, i);
                }
                tile_regs_commit();
                tile_regs_wait();
                for (uint32_t i = 0; i < DB; ++i) {
                    pack_tile<true>(i, cb_acc, t0 + i);
                }
                tile_regs_release();
            }
            pack_reconfig_l1_acc(0);
            cb_pop_front(cb_xt, CT);
            if (c & 1) {
                cb_pop_front(src, 2 * CT);
            }
        }
#endif
        cb_push_back(cb_acc, CT);
        cb_pop_front(cb_w, 4);
        // ---- 3. fp32 -> bf16 im2col tiles
        cb_wait_front(cb_acc, CT);
        cb_reserve_back(cb_out, CT);
        reconfig_data_format_srca(cb_acc);
        copy_init(cb_acc);
        pack_reconfig_data_format(cb_out);
        for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
            tile_regs_acquire();
            for (uint32_t i = 0; i < DB; ++i) {
                copy_tile(cb_acc, t0 + i, i);
            }
            tile_regs_commit();
            tile_regs_wait();
            for (uint32_t i = 0; i < DB; ++i) {
                pack_tile<true>(i, cb_out, t0 + i);
            }
            tile_regs_release();
        }
        cb_push_back(cb_out, CT);
        cb_pop_front(cb_acc, CT);
    }
}