// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// Round 11: cmm compute for the k2 K-split layout with a WIDE output block (NT > 4 tile columns per core, e.g. qkv:
// 2 token tile rows x 12 tile columns).  The 2 x NT fp32 accumulators do not fit DST, so the matmul runs in subblocks of
// SB = 4 columns (2 x 4 = 8 fp32 DST tiles, full-sync DST) and accumulates over the K blocks with packer L1
// accumulation into c_part (fp32, [row i][col n] = tile i * NT + n; block 0 overwrites).  Then:
//   - the partner row's partial (c_part tiles NT .. 2 NT - 1) is written by the BRISC into the partner's c_int;
//   - this core's own partial (tiles 0 .. NT - 1) is copied into DST exactly (c_part is unpacked to DST in fp32,
//     UnpackToDestFp32) and packed with L1 accumulation onto the partner's partial in c_int (exact fp32 sum) once the
//     NCRISC released the c_gate token (partner partial landed);
//   - + bias (row broadcast) -> c_out (the output dtype, e.g. bfp8), in chunks of 4 tiles.
// c_part is never popped (the BRISC reads its partner half after the compute is done with it; CB state resets per launch).
// compile-time: NB (K blocks of this core), 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

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 MPC = CMM_MPC;
#ifdef CMM_K2W_SB
    constexpr uint32_t SB = CMM_K2W_SB;  // 2: half-sync DST (4 fp32 tiles, math / pack overlap); 4: full-sync DST
#else
    constexpr uint32_t SB = 4;
#endif
    constexpr uint32_t DSTC = MPC * SB;  // fp32 DST tiles per acquire
    constexpr uint32_t CH = 4;  // tiles per DST acquire in the sum / bias passes (fits half- and full-sync fp32 DST)
    static_assert(NT % SB == 0 && NT % CH == 0 && DSTC <= 8, "subblocks");
    constexpr uint32_t NSB = NT / SB;
    constexpr uint32_t c_a = 0, c_w = 1, c_a1 = 5, c_w1 = 6, c_bias = 2, c_out = 16, c_int = 24, c_part = 26, c_gate = 22;

    compute_kernel_hw_startup<SrcOrder::Reverse>(c_a, c_w, c_part);
    cb_reserve_back(c_gate, 1);  // arm the gate (1-page CB, full until the NCRISC pops it)
    cb_push_back(c_gate, 1);
    matmul_block_init(c_a, c_w, 0, SB, MPC, KB);
    cb_reserve_back(c_part, MPC * NT);
    for (uint32_t b = 0; b < NB; ++b) {
        KZONE("CBLK");
        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);
        }
        for (uint32_t sb = 0; sb < NSB; ++sb) {
            tile_regs_acquire();
#ifndef CMM_NO_MM
            for (uint32_t k = 0; k < KB; ++k) {
                matmul_block(ca, cw, k, k * NT + sb * SB, 0, 0, SB, MPC, KB);
            }
#endif
            tile_regs_commit();
            tile_regs_wait();
            if (b == 1 && sb == 0) {
                PACK((llk_pack_reconfig_l1_acc(1)));  // block 0 overwrote, every later block accumulates
            }
            for (uint32_t i = 0; i < MPC; ++i) {
                for (uint32_t c = 0; c < SB; ++c) {
                    pack_tile<true>(i * SB + c, c_part, i * NT + sb * SB + c);
                }
            }
            tile_regs_release();
        }
        cb_pop_front(ca, KB * MPC);
        cb_pop_front(cw, KB * NT);
    }
    PACK((llk_pack_reconfig_l1_acc(0)));
    cb_push_back(c_part, MPC * NT);

    // own partial + the partner's partial (exact fp32: unpack-to-dest copy, packer L1 accumulation)
    cb_wait_front(c_part, MPC * NT);
    {
    KZONE("CSUM");
    reconfig_data_format_srca(c_part);
    copy_tile_to_dst_init_short(c_part);
    pack_reconfig_data_format(c_int);
    {
        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 q = 0; q < NT; q += CH) {
        constexpr uint32_t n = CH;
        tile_regs_acquire();
        for (uint32_t t = 0; t < n; ++t) {
            copy_tile(c_part, q + t, t);
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t t = 0; t < n; ++t) {
            pack_tile<true>(t, c_int, q + t);
        }
        tile_regs_release();
    }
    PACK((llk_pack_reconfig_l1_acc(0)));
    cb_push_back(c_int, NT);
    }

    // + bias -> output
    KZONE("CBIAS");
    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);
    for (uint32_t q = 0; q < NT; q += CH) {
        tile_regs_acquire();
        for (uint32_t i = 0; i < CH; ++i) {
            add_tiles_bcast_rows(c_int, c_bias, q + i, q + i, i);
        }
        tile_regs_commit();
        tile_regs_wait();
        for (uint32_t i = 0; i < CH; ++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);
}