// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn2 reader (NCRISC), round 15: MANO-head 1-query cross-attention, one core per head h (D5).
//   q:  row 0 of tiles 2h, 2h+1 of the head-split query [1, H, 32, 64] (bf16, only row 0 meaningful) -> two [1, 32]
//       tiny tiles (c_qt, contiguous 64 B each: face-0 row 0 then face-1 row 0)
//   K: the 12 tiles (key tile j, dim tile d) of head h -> c_k, page j * NC + base + d (runtime NC / base: the head-split
//      [1, H, 192, 64] K is NC = 2, base = 12 h; round 15: straight from the K/V projection output [1, 1, 192, L*1024],
//      NC = its tile columns, base = 32 l + 2 h -- no head split op).  V: the writer, the same way.
//   the dense score tile (c_sd, fp32 32x32, which the compute fills with the 192 scores as six [1, 32] tiny tiles at its
//   start) is first filled with a large negative constant (a 4 KB fp32 tile in L1, read once per launch), so the reductions
//   over the whole tile see only the scores
//   reduce scaler: bf16 1.0 in row 0 of every face (c_scal); a bf16 tile of ones (c_one, the sum-of-e matmul operand)
// compile-time: TensorAccessorArgs(q), (k), (v), (consts: tile 0 = fp32 pad, 4 KB; tile 1 = bf16 ones, 2 KB)
// runtime: q_addr, k_addr, v_addr, consts_addr, h, k_nc, k_base
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    constexpr auto q_args = TensorAccessorArgs<0>();
    constexpr auto k_args = TensorAccessorArgs<q_args.next_compile_time_args_offset()>();
    constexpr auto v_args = TensorAccessorArgs<k_args.next_compile_time_args_offset()>();
    constexpr auto n_args = TensorAccessorArgs<v_args.next_compile_time_args_offset()>();
    constexpr uint32_t c_qt = 0, c_k = 1, c_v = 2, c_scal = 3, c_st = 4, c_one = 6;
    constexpr uint32_t MT = 6;           // key tiles (192 keys)
    constexpr uint32_t page = 2048, FB = 512;  // bf16 tile / face bytes
    const uint32_t q_addr = get_arg_val<uint32_t>(0), k_addr = get_arg_val<uint32_t>(1);
    const uint32_t v_addr = get_arg_val<uint32_t>(2), n_addr = get_arg_val<uint32_t>(3), h = get_arg_val<uint32_t>(4);
    const uint32_t k_nc = get_arg_val<uint32_t>(5), k_base = get_arg_val<uint32_t>(6);
    const auto q = TensorAccessor(q_args, q_addr, page);
    const auto k = TensorAccessor(k_args, k_addr, page);
    // consts: one 6 KB row-major uint32 L1 buffer (page = 6 KB): [0, 4096) the fp32 pad tile, [4096, 6144) the bf16 ones
    const auto cst = TensorAccessor(n_args, n_addr, 6144);
    const uint64_t cs = cst.get_noc_addr(0);
    // pad fill of the dense score tile (its memory = c_st's), before anything is pushed to the compute
    noc_async_read(cs, get_write_ptr(c_st), 4096);
    cb_reserve_back(c_one, 1);
    noc_async_read(cs + 4096, get_write_ptr(c_one), 2048);
    cb_reserve_back(c_k, 2 * MT);
    uint32_t kd = get_write_ptr(c_k);
    for (uint32_t t = 0; t < 2 * MT; ++t) {
        noc_async_read(k.get_noc_addr((t / 2) * k_nc + k_base + (t % 2)), kd, page);
        kd += page;
    }
    cb_reserve_back(c_scal, 1);
    {
        const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
        for (uint32_t off = 0; off < 2048; off += MEM_ZEROS_SIZE) {
            noc_async_read(zeros, get_write_ptr(c_scal) + off, MEM_ZEROS_SIZE);
        }
    }
#ifdef HA2_WAIT
    // round 15 (fused into the q GEMV program): the two q GEMV cores of head h write its query columns, then each
    // increments this core's semaphore 0 -- wait for both before reading q (K / pad / ones / scaler reads are in flight)
    {
        volatile tt_l1_ptr uint32_t* sem = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(0));
        noc_semaphore_wait(sem, 2);
    }
#endif
    cb_reserve_back(c_qt, 2);
    const uint32_t qd = get_write_ptr(c_qt);
    for (uint32_t d = 0; d < 2; ++d) {
        const uint64_t s = q.get_noc_addr(2 * h + d);
        noc_async_read(s, qd + d * 64, 32);
        noc_async_read(s + FB, qd + d * 64 + 32, 32);
    }
    noc_async_read_barrier();
    volatile tt_l1_ptr uint32_t* sc = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(c_scal));
    for (uint32_t f = 0; f < 4; ++f) {
        for (uint32_t j = 0; j < 8; ++j) {
            sc[f * 128 + j] = 0x3F803F80u;
        }
    }
    cb_push_back(c_qt, 2);
    cb_push_back(c_k, 2 * MT);
    cb_push_back(c_scal, 1);
    cb_push_back(c_one, 1);
}