// SPDX-License-Identifier: Apache-2.0
// msda_k (tt/msda_k.py), compute. Per work item (32 queries x one head block; see msda_k_reader.cpp):
//   1. for each half h (the points of RISC h): DEST0..2 <- cb_geo_h (fx, fy, attn; fp32, UnpackToDestFp32); SFPU:
//      W_c = a * wx * wy for the 4 corners (fp32) -> cb_w pages 4h .. 4h + 3 (fp32).
//   2. for each step (point p, corner c): DEST0 = W_{h,c} @ E_p (broadcast of column p over the 32 channels, exact),
//      DEST0 *= V (the gathered bf16 value tile; DEST -> SrcB reuse), packer L1 accumulation into cb_acc (fp32).
//   3. cb_acc (fp32, UnpackToDestFp32) -> bf16 -> cb_out.
// CBs: cb_geo_h 3 pages / item (6 total), cb_in_h 2P pages / item (4P total), cb_w 8 (once per item), cb_acc 1,
// cb_e P (waited once), cb_out 1 per item (2 total).
#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.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/bcast.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 tiles 0, 1, 2 = fx, fy, attn -> tiles 0..3 = the NW, NE, SW, SE weights (one face per call)
inline void msda_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++;
    }
}
// SM: DEST 0, 1, 2, 3 = fx, fy, softmax, half mask -> attn = softmax * mask (exact: x 1 or x 0), then as above
inline void msda_k_w_face_sm() {
#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::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(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 SKIP = get_compile_time_arg_val(5);   // strided items (msda_k_reader.cpp)
// SM (BEVFORMER_MSDA_SOFTMAX=kernel): the softmax of the head's logit tile (cb_lg, from the reader) runs here, the
// same LLK sequence (and CB formats / unpack modes) as ttnn.softmax(numeric_stable) on one fp32 tile (attention
// softmax.cpp, Wt = 1, block 1): row max -> x - max, exp -> row sum, 1 / sum -> x * (1 / sum)
constexpr uint32_t SM = get_compile_time_arg_val(6);
// PAIRS (VSRC 2): the readers deliver each point's corners as two [32 x 64] row-major blocks (NW | NE, SW | SE);
// they are tilized (exact) into cb_tl0 / cb_tl1 in the step order (point, corner) before the steps
constexpr uint32_t PAIRS = get_compile_time_arg_val(7);
constexpr uint32_t cb_tl0 = 19, cb_tl1 = 20;
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;

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;

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t w0 = SKIP ? ci : ci * NWORK / NC;
    const uint32_t w1 = SKIP ? NWORK : (ci + 1) * NWORK / NC;
    const uint32_t wstep = SKIP ? NC : 1;
    if (w0 >= w1) {
        return;
    }
    if constexpr (SM) {
        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);
    } else {
        compute_kernel_hw_startup(cb_geo0, cb_geo0, cb_w);
    }
    cb_wait_front(cb_e, P);
    for (uint32_t w = w0; w < w1; w += wstep) {
        if constexpr (SM) {
            softmax_tile();
            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);
            tile_regs_acquire();
            reconfig_data_format_srca(cg);
            copy_init(cg);
            copy_tile(cg, 0, 0);
            copy_tile(cg, 1, 1);
            if constexpr (SM) {
                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_k_w_face_sm, 0,
                                                      VectorMode::RC)));
            } else {
                copy_tile(cg, 2, 2);
                MATH((SFPU_UNARY_INIT(unused)));
                MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, msda_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<true>(c, cb_w, h * 4 + c);
            }
            tile_regs_release();
            cb_pop_front(cg, 3);
        }
        cb_push_back(cb_w, 8);
        if constexpr (SM) {
            cb_pop_front(cb_sm, 1);
        }
        cb_wait_front(cb_w, 8);
        if constexpr (PAIRS) {
            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);
#ifdef MSDA_K_FAST_TILIZE
                // BEVFORMER_MSDA_TILIZE=fast: the BH fast tilize (a lossless bf16 reorder)
                fast_tilize_init(cin, 2, ct);
                for (uint32_t bk = 0; bk < NS / 2; ++bk) {
                    fast_tilize_block(cin, 2, ct, 2 * bk, 2 * bk);
                }
                fast_tilize_uninit(cin, ct, 2);
#else
                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_ = PAIRS ? (h ? cb_tl1 : cb_tl0) : (h ? cb_in1 : cb_in0);
            cb_wait_front(ci_, NS);
#ifdef MSDA_K_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);
        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();
        cb_push_back(cb_out, 1);
        cb_pop_front(cb_acc, 1);
    }
}