// SPDX-License-Identifier: Apache-2.0
// msda_mk (tt/msda_mk.py), compute. Per 32-query block b of the core's range (see msda_mk_dm.cpp):
//   FRONT(b) (run one block ahead of the items):
//     1. [cols | logits] = x @ W: for each N subblock and each K tile k (in0_block_w 1), DEST = x[k] * W[k, n]
//        (matmul_block, HiFi4, fp32 DEST), packed to the fp32 partials with packer L1 accumulation from k = 1 on;
//        then partials + bias row (FPU row broadcast) packed fp32: the steps of ttnn.linear's fused-bias matmul
//        (bmm_large_block_zm_fused_bias_activation: FUSE_BIAS, PACKER_L1_ACC, FP32_DEST_ACC_EN, in0_block_w 1) for
//        one output tile, as the stock TUNED_BEV / pinned 1D configs of the compact sampling and attention
//        projections run them (both have in0_block_w 1, so their columns side by side are one matmul);
//     2. loc = int32((cols + base) [+ shift row]): glue_k MODE 0's LLK calls (copy to DEST fp32, add_binary_tile RNE,
//        typecast fp32 -> int32), packed into cb_loc0 and cb_loc1 (one copy per data-movement RISC); the logit
//        tiles go to cb_lgc (PART 0 rebuilds each head's softmax input from them).
//   ITEMS (b, j): msda_k's compute steps with the in-kernel softmax and the pixel-pair maps (msda_k_compute.cpp SM 1,
//     PAIRS 1), unchanged.
#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/eltwise_binary.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"
#include "api/compute/tilize.h"
#include "api/compute/reduce.h"
#include "api/compute/softmax.h"
#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_compute.hpp"
#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/chain.hpp"
#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/convenience.hpp"
#include "ttnn/cpp/ttnn/kernel_lib/eltwise/unary/math.hpp"

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

namespace ckernel::sfpu {
// DEST 0, 1, 2, 3 = fx, fy, softmax, half mask -> the NW, NE, SW, SE weights (msda_k_w_face_sm)
inline void msda_mk_w_face_sm() {
#pragma GCC unroll 2
    for (int d = 0; d < 8; d++) {
#ifdef MSDA_MK_FRAC16
        // BEVFORMER_MSDA_FRAC16: the reader wrote 1 + q / 2^16 (exact); 1 + f - 1 = f exactly (Sterbenz)
        sfpi::vFloat fx = sfpi::dst_reg[0] - 1.0f;
        sfpi::vFloat fy = sfpi::dst_reg[32] - 1.0f;
#else
        sfpi::vFloat fx = sfpi::dst_reg[0];
        sfpi::vFloat fy = sfpi::dst_reg[32];
#endif
        sfpi::vFloat a = sfpi::dst_reg[64] * sfpi::dst_reg[96];
        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

namespace ckl = compute_kernel_lib;

constexpr uint32_t P = get_compile_time_arg_val(0);
constexpr uint32_t NB = get_compile_time_arg_val(1);
constexpr uint32_t H = 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 KT = get_compile_time_arg_val(5);
constexpr uint32_t LOC_TW = get_compile_time_arg_val(6);
constexpr uint32_t AT_TW = get_compile_time_arg_val(7);
constexpr uint32_t HAS_SHIFT = get_compile_time_arg_val(8);
// ACT (BEVFORMER_SCA_ACTIVE): the work runs over the active blocks only; their count comes from the per-calibration
// nblk page that the PART-0 data-movement kernel pushes to cb_cnt (entries 0..SUB-1, clamped to NB / SUB)
constexpr uint32_t ACT = get_compile_time_arg_val(9);
constexpr uint32_t SUB = get_compile_time_arg_val(10);
constexpr uint32_t cb_cnt = 7;
// RAW (BEVFORMER_SCA_BASE=kernel): the base arrives as the raw grid base (one tile per block, its 2P point columns
// duplicated into both column halves by the reader) and the front computes base * S + O itself (cb_aff: the S and O
// rows replicated over the tile, both halves) with binary_ng's SFPU mul / add (RNE) before the cols + base add:
// the per-frame multiply / add / compact-column programs in front of the SCA disappear
constexpr uint32_t RAW = get_compile_time_arg_val(11);
constexpr uint32_t cb_aff = 31;
constexpr uint32_t BASE_TILES = RAW ? 1 : LOC_TW;
constexpr uint32_t NOUT = LOC_TW + AT_TW;
// N subblock (fp32 DEST half: 4 tiles): 2 when it splits the loc / logit tiles evenly, else 1 (the per-tile steps
// do not depend on it)
constexpr uint32_t SN = (NOUT % 2 == 0 && LOC_TW % 2 == 0) ? 2 : 1;

constexpr uint32_t cb_in0 = 0, cb_in1 = 1, cb_geo0 = 2, cb_geo1 = 3, cb_e = 4, cb_w = 5, cb_acc = 6, cb_out = 16;
constexpr uint32_t cb_lg = 9, cb_scmax = 10, cb_scsum = 11, cb_msk = 12, cb_sm = 13, cb_max = 14, cb_exps = 15,
                   cb_rsum = 17;
constexpr uint32_t cb_tl0 = 19, cb_tl1 = 20;
constexpr uint32_t cb_x = 21, cb_wsa = 22, cb_bias = 23, cb_part = 24, cb_base = 25, cb_shift = 26, cb_cols = 27,
                   cb_loc0 = 28, cb_loc1 = 29, cb_lgc = 30;
constexpr uint32_t F32 = static_cast<uint32_t>(DataFormat::Float32);
constexpr uint32_t I32 = static_cast<uint32_t>(DataFormat::Int32);

static inline void softmax_tile() {
    reconfig_data_format(cb_lg, cb_lg);
    pack_reconfig_data_format(cb_exps);
    copy_init(cb_lg);
    compute_kernel_lib::reduce<PoolType::MAX, ReduceDim::REDUCE_ROW, cb_lg, cb_scmax, cb_max,
                               compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop,
                               compute_kernel_lib::ReduceDataFormatReconfigMode::INPUT>(
        compute_kernel_lib::ReduceInputBlockShape::row(1));
    ckl::eltwise_chain(
        ckl::IterationShape::tiles(1).block_size(1),
        ckl::BinaryFpu<ckl::BinaryFpuOp::Sub,
                       ckl::input(cb_lg, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd, ckl::InputTileMapping::Block,
                                  ckl::DataFormatReconfig::Disabled),
                       ckl::input(cb_max, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd)>{},
        ckl::Exp<static_cast<ckl::Approx>(0), ckl::Dst::D0>{},
        ckl::PackTile<ckl::output(cb_exps, ckl::ReservePolicy::PerBlockSize, ckl::PushPolicy::PerBlockSize,
                                  ckl::DataFormatReconfig::Disabled)>{});
    cb_wait_front(cb_exps, 1);
    compute_kernel_lib::reduce<PoolType::SUM, ReduceDim::REDUCE_ROW, cb_exps, cb_scsum, cb_rsum,
                               compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop>(
        compute_kernel_lib::ReduceInputBlockShape::row(1), compute_kernel_lib::ReduceInputMemoryLayout::contiguous(),
        compute_kernel_lib::NoAccumulation{}, [](uint32_t) {
            recip_tile_init<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>();
            recip_tile<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>(0);
        });
    ckl::mul<ckl::input(cb_exps, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd, ckl::InputTileMapping::Block),
             ckl::input(cb_rsum, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd),
             ckl::output(cb_sm, ckl::ReservePolicy::PerBlockSize, ckl::PushPolicy::PerBlockSize)>(
        ckl::IterationShape::tiles(1).block_size(1));
}

constexpr uint32_t PH = P / 2;
constexpr uint32_t NS = PH * 4;

// FRONT of one block: projections + bias -> cb_cols (LOC_TW) / cb_lgc (AT_TW), then the loc glue -> cb_loc0 / 1
static inline void front() {
    // ---- 1. matmul, K blocks of one tile, packer L1 accumulation into the fp32 partials
    cb_wait_front(cb_x, KT);
    reconfig_data_format_srca(cb_wsa);
    reconfig_data_format_srcb(cb_x);
    matmul_block_init(cb_x, cb_wsa, 0, SN, 1, KT);
    pack_reconfig_data_format(cb_part);
    cb_reserve_back(cb_part, NOUT);
    for (uint32_t n0 = 0; n0 < NOUT; n0 += SN) {
#ifdef MSDA_MK_NOFRONTMM
        for (uint32_t k = 0; k < 1; ++k) {
#else
        for (uint32_t k = 0; k < KT; ++k) {
#endif
            tile_regs_acquire();
            matmul_block(cb_x, cb_wsa, k, k * NOUT + n0, 0, 0, SN, 1, KT);
            tile_regs_commit();
            tile_regs_wait();
            pack_reconfig_l1_acc(k == 0 ? 0 : 1);
            for (uint32_t i = 0; i < SN; ++i) {
                pack_tile<true>(i, cb_part, n0 + i);
            }
            tile_regs_release();
        }
    }
    pack_reconfig_l1_acc(0);
    cb_push_back(cb_part, NOUT);
    cb_pop_front(cb_x, KT);
    // ---- bias (row broadcast, FPU) -> fp32 cols / logits
    cb_wait_front(cb_part, NOUT);
    reconfig_data_format_srca(cb_part);
    reconfig_data_format_srcb(cb_bias);
    add_bcast_rows_init(cb_part, cb_bias);
    pack_reconfig_data_format(cb_cols);
    cb_reserve_back(cb_cols, LOC_TW);
    cb_reserve_back(cb_lgc, AT_TW);
    for (uint32_t n0 = 0; n0 < NOUT; n0 += SN) {
        tile_regs_acquire();
        for (uint32_t i = 0; i < SN; ++i) {
            add_tiles_bcast_rows(cb_part, cb_bias, n0 + i, n0 + i, i);
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t i = 0; i < SN; ++i) {
            const uint32_t n = n0 + i;
            if (n < LOC_TW) {
                pack_tile<true>(i, cb_cols, n);
            } else {
                pack_tile<true>(i, cb_lgc, n - LOC_TW);
            }
        }
        tile_regs_release();
    }
    cb_push_back(cb_cols, LOC_TW);
    cb_push_back(cb_lgc, AT_TW);
    cb_pop_front(cb_part, NOUT);
    // ---- 2. loc = int32((cols + base) [+ shift]) (glue_k MODE 0)
    cb_wait_front(cb_cols, LOC_TW);
    cb_wait_front(cb_base, BASE_TILES);
    cb_reserve_back(cb_loc0, LOC_TW);
    cb_reserve_back(cb_loc1, LOC_TW);
    pack_reconfig_data_format(cb_loc0);
    for (uint32_t t = 0; t < LOC_TW; ++t) {
        tile_regs_acquire();
        reconfig_data_format_srca(cb_cols);
        copy_init(cb_cols);
        copy_tile(cb_cols, t, 0);
        reconfig_data_format_srca(cb_base);
        copy_init(cb_base);
        copy_tile(cb_base, RAW ? 0 : t, 1);
        if constexpr (RAW) {
            reconfig_data_format_srca(cb_aff);
            copy_init(cb_aff);
            copy_tile(cb_aff, 0, 2);
            mul_binary_tile_init();
            mul_binary_tile(1, 2, 1);
            reconfig_data_format_srca(cb_aff);
            copy_init(cb_aff);
            copy_tile(cb_aff, 1, 2);
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(1, 2, 1);
        }
        add_binary_tile_init();
        add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
        if constexpr (HAS_SHIFT) {
            reconfig_data_format_srca(cb_shift);
            copy_init(cb_shift);
            copy_tile(cb_shift, t, 1);
            add_binary_tile_init();
            add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
        }
        typecast_tile_init<F32, I32>();
        typecast_tile<F32, I32>(0);
        tile_regs_commit();
        tile_regs_wait();
        pack_tile<true>(0, cb_loc0, t);
        pack_tile<true>(0, cb_loc1, t);
        tile_regs_release();
    }
    cb_push_back(cb_loc0, LOC_TW);
    cb_push_back(cb_loc1, LOC_TW);
    cb_pop_front(cb_cols, LOC_TW);
    cb_pop_front(cb_base, BASE_TILES);
}

// one MSDA item (msda_k_compute.cpp with SM 1, PAIRS 1)
static inline void item() {
#ifdef MSDA_MK_NOSOFTMAX
    // (anatomy probe only: no softmax, garbage weights)
    cb_wait_front(cb_lg, 1);
    cb_pop_front(cb_lg, 1);
    cb_reserve_back(cb_sm, 1);
    cb_push_back(cb_sm, 1);
#else
    softmax_tile();
#endif
    cb_wait_front(cb_sm, 1);
    // ---- 1. corner weights
    cb_reserve_back(cb_w, 8);
    for (uint32_t h = 0; h < 2; ++h) {
        const uint32_t cg = h ? cb_geo1 : cb_geo0;
        cb_wait_front(cg, 3);
#ifdef MSDA_MK_NOWEIGHTS
        cb_pop_front(cg, 3);      // (anatomy probe only: no corner weights)
        continue;
#endif
        tile_regs_acquire();
        reconfig_data_format_srca(cg);
        copy_init(cg);
        copy_tile(cg, 0, 0);
        copy_tile(cg, 1, 1);
        reconfig_data_format_srca(cb_sm);
        copy_init(cb_sm);
        copy_tile(cb_sm, 0, 2);
        reconfig_data_format_srca(cb_msk);
        copy_init(cb_msk);
        copy_tile(cb_msk, h, 3);
        MATH((SFPU_UNARY_INIT(unused)));
        MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, msda_mk_w_face_sm, 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<true>(c, cb_w, h * 4 + c);
        }
        tile_regs_release();
        cb_pop_front(cg, 3);
    }
    cb_push_back(cb_w, 8);
    cb_pop_front(cb_sm, 1);
    cb_wait_front(cb_w, 8);
    for (uint32_t h = 0; h < 2; ++h) {
        const uint32_t cin = h ? cb_in1 : cb_in0;
        const uint32_t ct = h ? cb_tl1 : cb_tl0;
        cb_wait_front(cin, NS);
        cb_reserve_back(ct, NS);
        reconfig_data_format_srca(cin);
        pack_reconfig_data_format(ct);
#ifndef MSDA_MK_NOTILIZE
        tilize_init(cin, 2, ct);
        for (uint32_t bk = 0; bk < NS / 2; ++bk) {
            tilize_block(cin, 2, ct, 2 * bk, 2 * bk);
        }
        tilize_uninit(cin, ct);
#endif
        cb_push_back(ct, NS);
        cb_pop_front(cin, NS);
    }
    // ---- 2. the 4P (point, corner) steps
    cb_reserve_back(cb_acc, 1);
    reconfig_data_format_srca(cb_e);
    reconfig_data_format_srcb(cb_w);
    pack_reconfig_data_format(cb_acc);
    for (uint32_t h = 0; h < 2; ++h) {
        const uint32_t ci_ = h ? cb_tl1 : cb_tl0;
        cb_wait_front(ci_, NS);
#ifdef MSDA_MK_NOSTEPS
        for (uint32_t k = 0; k < 1; ++k) {
#else
        for (uint32_t k = 0; k < NS; ++k) {
#endif
            const uint32_t p = h * PH + (k >> 2);
            const uint32_t c = k & 3;
            tile_regs_acquire();
            matmul_init(cb_w, cb_e);
            matmul_tiles(cb_w, cb_e, h * 4 + c, p, 0);
            mul_reuse_dest_init<EltwiseBinaryReuseDestType::DEST_TO_SRCB>(ci_);
            mul_reuse_dest_tiles<EltwiseBinaryReuseDestType::DEST_TO_SRCB>(ci_, k, 0);
            tile_regs_commit();
            tile_regs_wait();
            pack_reconfig_l1_acc((h == 0 && k == 0) ? 0 : 1);
            pack_tile<true>(0, cb_acc, 0);
            tile_regs_release();
        }
        cb_pop_front(ci_, NS);
    }
    pack_reconfig_l1_acc(0);
    cb_push_back(cb_acc, 1);
    cb_pop_front(cb_w, 8);
    // ---- 3. fp32 sum -> bf16 output tile
    cb_wait_front(cb_acc, 1);
    cb_reserve_back(cb_out, 1);
#ifndef MSDA_MK_NOFINAL
    tile_regs_acquire();
    reconfig_data_format_srca(cb_acc);
    copy_init(cb_acc);
    copy_tile(cb_acc, 0, 0);
    tile_regs_commit();
    tile_regs_wait();
    pack_reconfig_data_format(cb_out);
    pack_tile(0, cb_out);
    tile_regs_release();
#endif
    cb_push_back(cb_out, 1);
    cb_pop_front(cb_acc, 1);
}

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    uint32_t nwork = NB * H;
    if constexpr (ACT) {
        cb_wait_front(cb_cnt, 1);
        uint32_t tot = 0;
        for (uint32_t c = 0; c < SUB; ++c) {
            tot += read_tile_value(cb_cnt, 0, c);
        }
        nwork = tot * H;
    }
    const uint32_t w0 = ci * nwork / NC;
    const uint32_t w1 = (ci + 1) * nwork / NC;
    if (w0 >= w1) {
        return;
    }
    const uint32_t bf = w0 / H;
    const uint32_t bl = (w1 - 1) / H;
    compute_kernel_hw_startup(cb_lg, cb_scmax, cb_exps);
    cb_wait_front(cb_scmax, 1);
    cb_wait_front(cb_scsum, 1);
    cb_wait_front(cb_msk, 2);
    cb_wait_front(cb_e, P);
    cb_wait_front(cb_wsa, KT * NOUT);
    cb_wait_front(cb_bias, NOUT);
    if constexpr (HAS_SHIFT) {
        cb_wait_front(cb_shift, LOC_TW);
    }
    if constexpr (RAW) {
        cb_wait_front(cb_aff, 2);
    }
    front();
    for (uint32_t b = bf; b <= bl; ++b) {
        if (b + 1 <= bl) {
            front();
        }
        const uint32_t j0 = b == bf ? w0 - bf * H : 0;
        const uint32_t j1 = b == bl ? w1 - bl * H : H;
        for (uint32_t j = j0; j < j1; ++j) {
            item();
        }
    }
}