// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// cmm compute: out[1 x NT] = sum_b a_b[1 x KB] @ w_b[KB x NT]  (whole K accumulated in DST, fp32),
// packed to an fp32 intermediate, + bias (row broadcast), packed into the sharded output.
// compile-time: NB, KB, NT
#include <cstdint>
#include "api/compute/matmul.h"
#include "api/compute/bcast.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/cb_api.h"
#include "api/compute/pack.h"
#include "api/compute/reg_api.h"
#ifdef CMM_PROF
#include "tools/profiler/kernel_profiler.hpp"
#define KZONE(n) DeviceZoneScopedN(n)
#else
#define KZONE(n)
#endif
#ifdef CMM_GELU
#include "api/compute/common.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/compute_kernel_api.h"
#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
// same erf-form GELU as kernels/eltwise/gelu_compute.cpp (the model's default standalone GELU since round 5)
inline void cmm_gelu_face() {
    for (size_t i = 0; i < 8; i += 2) {
        sfpi::vFloat x0 = sfpi::dst_reg[i];
        sfpi::vFloat x1 = sfpi::dst_reg[i + 1];
        sfpi::vFloat u0 = sfpi::min(x0 * x0, 16.0f);
        sfpi::vFloat u1 = sfpi::min(x1 * x1, 16.0f);
        sfpi::vFloat a0, a1;
        { const sfpi::vFloat c8 = 8.040229593e-11f; const sfpi::vFloat c7 = -6.988785106e-09f;
        a0 = u0 * c8 + c7;
        a1 = u1 * c8 + c7; }
        { const sfpi::vFloat c = 2.712364733e-07f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -6.289864814e-06f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = 9.886907355e-05f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -1.133757411e-03f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = 9.877118282e-03f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = -6.641025841e-02f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        { const sfpi::vFloat c = 3.989226222e-01f;
        a0 = a0 * u0 + c;
        a1 = a1 * u1 + c; }
        sfpi::vFloat p0 = sfpi::clamp(x0 * a0 + 0.5f, 0.0f, 1.0f);
        sfpi::vFloat p1 = sfpi::clamp(x1 * a1 + 0.5f, 0.0f, 1.0f);
        sfpi::dst_reg[i] = x0 * p0;
        sfpi::dst_reg[i + 1] = x1 * p1;
    }
}
}  // namespace ckernel::sfpu
#endif
inline void cmm_gelu_tile(uint32_t idst) {
    MATH(SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, cmm_gelu_face, idst, VectorMode::RC));
}
#endif

#ifdef CMM_MM_FID
// round 8: the cmm matmul at an explicit fidelity (CMM_MM_FID, e.g. 2 = HiFi2) when this file is the first phase of a
// kernel compiled at another MATH_FIDELITY (kernels/cmm/fused_compute.cpp: the add+LN epilogue needs HiFi4).  Same
// calls as api/compute/matmul.h matmul_block_init / matmul_block (incl. the Blackhole dynamic throttle).
constexpr ckernel::MathFidelity CMM_FIDV = static_cast<ckernel::MathFidelity>(CMM_MM_FID);
ALWI void cmm_mm_init(uint32_t a, uint32_t b, uint32_t ct, uint32_t rt, uint32_t kt) {
    state_configure(b, a, __builtin_LINE());
    UNPACK((llk_unpack_AB_matmul_init(a, b, 0, ct, rt, kt)));
    MATH((llk_math_matmul_init<CMM_FIDV, MM_THROTTLE>(a, b, 0, ct, rt)));
#ifdef ARCH_BLACKHOLE
    MATH((throttled_mop_status = 0));
#endif
}
ALWI void cmm_mm_block(uint32_t a, uint32_t b, uint32_t ia, uint32_t ib, uint32_t ct, uint32_t rt, uint32_t kt) {
    state_configure(b, a, __builtin_LINE());
    UNPACK((llk_unpack_AB_matmul(a, b, ia, ib, ct, rt, kt)));
#ifdef ARCH_BLACKHOLE
    MATH(({
        volatile uint32_t mm_throttle_en = *(throttle_ptr) % 2;
        if (mm_throttle_en) {
            if (throttled_mop_status != 1) {
                llk_math_matmul_init<CMM_FIDV, MM_THROTTLE_MAX>(a, b, 0, ct, rt);
                throttled_mop_status = 1;
            }
            llk_math_matmul<CMM_FIDV, MM_THROTTLE_MAX>(0, ct, rt);
        } else {
            if (throttled_mop_status != 0) {
                llk_math_matmul_init<CMM_FIDV, MM_THROTTLE>(a, b, 0, ct, rt);
                throttled_mop_status = 0;
            }
            llk_math_matmul<CMM_FIDV, MM_THROTTLE>(0, ct, rt);
        }
    }));
#else
    MATH((llk_math_matmul<CMM_FIDV, MM_THROTTLE>(0, ct, rt)));
#endif
}
#define CMM_MM_INIT(a, b, ct, rt, kt) cmm_mm_init(a, b, ct, rt, kt)
#define CMM_MM_BLOCK(a, b, ia, ib, ct, rt, kt) cmm_mm_block(a, b, ia, ib, ct, rt, kt)
#else
#define CMM_MM_INIT(a, b, ct, rt, kt) matmul_block_init(a, b, 0, ct, rt, kt)
#define CMM_MM_BLOCK(a, b, ia, ib, ct, rt, kt) matmul_block(a, b, ia, ib, 0, 0, ct, rt, kt)
#endif

void kernel_main() {
    constexpr uint32_t NB = get_compile_time_arg_val(0);
    constexpr uint32_t KB = get_compile_time_arg_val(1);
    constexpr uint32_t NT = get_compile_time_arg_val(2);
    constexpr uint32_t c_a = 0, c_w = 1, c_a1 = 5, c_w1 = 6, c_bias = 2, c_out = 16, c_int = 24;

#ifdef CMM_GELU
    // phase 0: GELU of this core's own activation blocks (raw c_araw -> staging c_ast, per parity), which
    // its movers then multicast; runtime args: own block count of parity 0, of parity 1
    {
        constexpr uint32_t c_raw0 = 11, c_raw1 = 12, c_st0 = 4, c_st1 = 8;
        const uint32_t n_own[2] = {get_arg_val<uint32_t>(0), get_arg_val<uint32_t>(1)};
        init_sfpu(c_raw0, c_st0);
        for (uint32_t par = 0; par < 2; ++par) {
            const uint32_t c_raw = par ? c_raw1 : c_raw0, c_st = par ? c_st1 : c_st0;
            for (uint32_t blk = 0; blk < n_own[par]; ++blk) {
                {KZONE("GWAIT"); cb_wait_front(c_raw, 1);}
                KZONE("GELU");
                for (uint32_t t = 0; t < KB; ++t) {
                    cb_wait_front(c_raw, 1);
                    tile_regs_acquire();
                    copy_tile(c_raw, 0, 0);
                    cmm_gelu_tile(0);
                    tile_regs_commit();
                    tile_regs_wait();
                    cb_reserve_back(c_st, 1);
                    pack_tile(0, c_st);
                    cb_push_back(c_st, 1);
                    cb_pop_front(c_raw, 1);
                    tile_regs_release();
                }
            }
        }
    }
    compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_int);
#else
    compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_int);
#endif
#ifdef CMM_KSPLIT
    // round 11 (layout k2): this core holds CMM_MPC = 2 token tile rows (own row first, then the partner's) of HALF of K;
    // the 2 x NT fp32 partials stay in DST (full-sync DST, 8 fp32 tiles).  The partner's half goes out through c_send
    // (the BRISC writes it into the partner's c_int); this core's own half is packed with packer L1 accumulation onto
    // the partner's partial, which the partner wrote into THIS core's c_int (NCRISC pops the c_gate token once it
    // landed: the PACK thread blocks in cb_reserve_back(c_gate) until then).
    constexpr uint32_t MPC = CMM_MPC, c_send = 19, c_gate = 22;
    cb_reserve_back(c_gate, 1);  // arm the gate (1-page CB, now full until the NCRISC pops it)
    cb_push_back(c_gate, 1);
#elif defined(CMM_MPC)
    constexpr uint32_t MPC = CMM_MPC;  // round 11 (xmm): MPC token rows x NT columns, all of K, in full-sync DST
#else
    constexpr uint32_t MPC = 1;
#endif
    CMM_MM_INIT(c_a, c_w, NT, MPC, KB);
    tile_regs_acquire();
    // even K blocks arrive in (c_a, c_w) (NCRISC / NOC_0), odd ones in (c_a1, c_w1) (BRISC / NOC_1)
    for (uint32_t b = 0; b < NB; ++b) {
        const uint32_t ca = (b & 1) ? c_a1 : c_a, cw = (b & 1) ? c_w1 : c_w;
        {
            KZONE("CWA");
            cb_wait_front(ca, KB * MPC);
        }
        {
            KZONE("CWW");
            cb_wait_front(cw, KB * NT);
        }
#ifndef CMM_NO_MM
        for (uint32_t k = 0; k < KB; ++k) {
            CMM_MM_BLOCK(ca, cw, k, k * NT, NT, MPC, KB);
        }
#endif
        cb_pop_front(ca, KB * MPC);
        cb_pop_front(cw, KB * NT);
    }
    tile_regs_commit();
#ifdef CMM_KSPLIT
    cb_reserve_back(c_send, NT);
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(NT + i, c_send);
    }
    cb_push_back(c_send, NT);
    {
        KZONE("CXWAIT");
        cb_reserve_back(c_gate, 1);  // the partner's partial is in c_int
    }
    cb_reserve_back(c_int, NT);
    PACK((llk_pack_reconfig_l1_acc(1)));
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_int);
    }
    PACK((llk_pack_reconfig_l1_acc(0)));
    tile_regs_release();
    cb_push_back(c_int, NT);
#else
    cb_reserve_back(c_int, MPC * NT);
    tile_regs_wait();
    for (uint32_t i = 0; i < MPC * NT; ++i) {
        pack_tile(i, c_int);
    }
    tile_regs_release();
    cb_push_back(c_int, MPC * NT);
#endif
#if defined(CMM_MPC) && !defined(CMM_KSPLIT)
    // xmm: + bias (the same NT bias tiles for every token row) -> c_out (MPC x NT tiles)
    cb_wait_front(c_int, MPC * NT);
    cb_wait_front(c_bias, NT);
    reconfig_data_format(c_int, c_bias);
    pack_reconfig_data_format(c_out);
    add_bcast_rows_init(c_int, c_bias);
    cb_reserve_back(c_out, MPC * NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < MPC * NT; ++i) {
        add_tiles_bcast_rows(c_int, c_bias, i, i % NT, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < MPC * NT; ++i) {
        pack_tile(i, c_out);
    }
    tile_regs_release();
    cb_push_back(c_out, MPC * NT);
    cb_pop_front(c_int, MPC * NT);
    cb_pop_front(c_bias, NT);
    return;
#endif

    cb_wait_front(c_int, NT);
    cb_wait_front(c_bias, NT);
    reconfig_data_format(c_int, c_bias);
    pack_reconfig_data_format(c_out);
    add_bcast_rows_init(c_int, c_bias);
    cb_reserve_back(c_out, NT);
    tile_regs_acquire();
    for (uint32_t i = 0; i < NT; ++i) {
        add_tiles_bcast_rows(c_int, c_bias, i, i, i);
    }
    tile_regs_commit();
    tile_regs_wait();
    for (uint32_t i = 0; i < NT; ++i) {
        pack_tile(i, c_out);
    }
    tile_regs_release();
    cb_push_back(c_out, NT);
    cb_pop_front(c_int, NT);
    cb_pop_front(c_bias, NT);
}