// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// hattn reader (NCRISC), round 6: 1-token cross-attention of the MANO head, one core per head h.
//   q: row 0 of tiles 2h, 2h+1 of the row-form query [1, 1, 32, H*64] (bf16; already scaled by 1/sqrt(64) on the host)
//   K / V: straight from the batched K/V projection output kv [1, 1, MT*32, L*1024] (TILE, interleaved): layer `l`,
//          head h: K tile (j, d) = page j * NC + l*32 + 2h + d, V tile = page j * NC + l*32 + 16 + 2h + d  (j < MT)
// compile-time: MT, NC (kv tile columns), TensorAccessorArgs(q), TensorAccessorArgs(kv)
// runtime: q_addr, kv_addr, h, l
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    constexpr uint32_t MT = get_compile_time_arg_val(0);
    constexpr uint32_t NC = get_compile_time_arg_val(1);
    constexpr auto q_args = TensorAccessorArgs<2>();
    constexpr auto kv_args = TensorAccessorArgs<q_args.next_compile_time_args_offset()>();
    constexpr uint32_t c_q = 0, c_k = 1, c_v = 2, c_scal = 3;
    const uint32_t q_addr = get_arg_val<uint32_t>(0), kv_addr = get_arg_val<uint32_t>(1);
    const uint32_t h = get_arg_val<uint32_t>(2), l = get_arg_val<uint32_t>(3);
    constexpr uint32_t page = 2048, FB = 512;  // bf16 tile / face bytes
    const auto q = TensorAccessor(q_args, q_addr, page);
    const auto kv = TensorAccessor(kv_args, kv_addr, page);
    cb_reserve_back(c_q, 2);
    const uint32_t qd = get_write_ptr(c_q);
    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 * page, 32);
        noc_async_read(s + FB, qd + d * page + FB, 32);
    }
    cb_reserve_back(c_k, 2 * MT);
    cb_reserve_back(c_v, 2 * MT);
    uint32_t kd = get_write_ptr(c_k), vd = get_write_ptr(c_v);
    for (uint32_t j = 0; j < MT; ++j) {
        for (uint32_t d = 0; d < 2; ++d) {
            noc_async_read(kv.get_noc_addr(j * NC + l * 32 + 2 * h + d), kd, page);
            noc_async_read(kv.get_noc_addr(j * NC + l * 32 + 16 + 2 * h + d), vd, page);
            kd += page;
            vd += page;
        }
    }
    // reduce scaler: bf16 1.0 in row 0 of every face
    cb_reserve_back(c_scal, 1);
    volatile tt_l1_ptr uint32_t* sc = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(c_scal));
    for (uint32_t i = 0; i < 512; ++i) {
        sc[i] = 0;
    }
    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_scal, 1);
    noc_async_read_barrier();
    cb_push_back(c_q, 2);
    cb_push_back(c_k, 2 * MT);
    cb_push_back(c_v, 2 * MT);
}