// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
// SPDX-License-Identifier: Apache-2.0
//
// SuperPoint keypoint list on device, two-RISC version of kp_compact.cpp (same HDR output format,
// see there). The candidate slots C [NSLOT, CAP+1] live in THIS core's L1 (sharded on core (0,0)),
// so they are read in place. Both data-movement RISCs run this kernel (PROC 0 / 1): each computes
// the kept count of every slot, the slots are split where the kept prefix reaches half of the total,
// PROC 0 writes entries of slots [0, split), PROC 1 of [split, NSLOT) into the shared L1 header
// image; PROC 1 then raises semaphore 0 and PROC 0 finishes the header and writes it to DRAM.
// KPC_HR (single-D2H mode): PROC 0 also picks the descriptor bucket b = max(spec, ceil(n/BSTEP) - 1)
// (spec = word 2 of the 64 B parameter tensor, the host's speculative bucket), stores b in HDR[3] (the
// sampler writes the rows to the same bucket) and copies the header into the first KPC_HR rows of
// bucket b, so the host reads header + descriptor rows in ONE transfer.
// Runtime args (KPC_HR): c_addr, hdr_addr, tab_addr, prm_addr, then the NB bucket addresses.
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    const uint32_t c_addr = get_arg_val<uint32_t>(0);
    const uint32_t hdr_addr = get_arg_val<uint32_t>(1);
    const uint32_t tab_addr = get_arg_val<uint32_t>(2);

    constexpr uint32_t cb_scratch = get_compile_time_arg_val(0);
    constexpr uint32_t NSLOT = get_compile_time_arg_val(1);
    constexpr uint32_t CAP = get_compile_time_arg_val(2);
    constexpr uint32_t W = get_compile_time_arg_val(3);
    constexpr uint32_t KMAX = get_compile_time_arg_val(4);
    constexpr uint32_t ROWS_PER_CORE = get_compile_time_arg_val(5);
    constexpr uint32_t SPLIT = get_compile_time_arg_val(6);  // BRISC rows of an NMS core
    constexpr uint32_t H = get_compile_time_arg_val(7);
    constexpr uint32_t PROC = get_compile_time_arg_val(8);
    constexpr uint32_t SLOT_BYTES = (CAP + 1) * 4;
    constexpr uint32_t HDR_WORDS = 16 + 4 * KMAX;
    constexpr uint32_t TAB_BYTES = (2 * W + 2 * H) * 4;
    constexpr uint32_t TAB_REGION = (TAB_BYTES + 63) & ~63u;
    constexpr auto hdr_args = TensorAccessorArgs<9>();
    constexpr auto tab_args = TensorAccessorArgs<hdr_args.next_compile_time_args_offset()>();
    const auto hdracc = TensorAccessor(hdr_args, hdr_addr, HDR_WORDS * 4);
    const auto tabacc = TensorAccessor(tab_args, tab_addr, TAB_BYTES);
#ifdef KPC_HR
    constexpr auto prm_args = TensorAccessorArgs<tab_args.next_compile_time_args_offset()>();
    constexpr auto bk_args = TensorAccessorArgs<prm_args.next_compile_time_args_offset()>();
    constexpr uint32_t ROWB = KPC_C * 4;  // one bucket row (page) in bytes
#endif

    const uint32_t base = get_write_ptr(cb_scratch);
    const uint32_t tab_l1 = base + PROC * TAB_REGION;
    const uint32_t hdr_l1 = base + 2 * TAB_REGION;
    noc_async_read(tabacc.get_noc_addr(0), tab_l1, TAB_BYTES);
#ifdef KPC_HR
    const uint32_t prm_l1 = hdr_l1 + HDR_WORDS * 4;
    if constexpr (PROC == 0) {
        const auto prmacc = TensorAccessor(prm_args, get_arg_val<uint32_t>(3), 64);
        noc_async_read(prmacc.get_noc_addr(0), prm_l1, 64);
    }
#endif

    // kept count per slot (capped by CAP and by the global KMAX budget), split point
    uint32_t total = 0, overflow = 0, kept = 0;
    for (uint32_t s = 0; s < NSLOT; ++s) {
        uint32_t cnt = reinterpret_cast<volatile uint32_t*>(c_addr + s * SLOT_BYTES)[0];
        total += cnt;
        if (cnt > CAP) {
            overflow = 1;
            cnt = CAP;
        }
        const uint32_t room = KMAX - kept;
        kept += cnt < room ? cnt : room;
    }
    const uint32_t half = kept / 2;
    uint32_t split = NSLOT, pre = 0, n0 = 0;
    for (uint32_t s = 0; s < NSLOT; ++s) {
        if (pre >= half) {
            split = s;
            n0 = pre;
            break;
        }
        uint32_t cnt = reinterpret_cast<volatile uint32_t*>(c_addr + s * SLOT_BYTES)[0];
        if (cnt > CAP) {
            cnt = CAP;
        }
        const uint32_t room = KMAX - pre;
        pre += cnt < room ? cnt : room;
    }
    if (split == NSLOT) {
        n0 = pre;
    }
    noc_async_read_barrier();
    const uint32_t* tx = reinterpret_cast<const uint32_t*>(tab_l1);
    const uint32_t* ty = tx + 2 * W;

    uint32_t* hdr = reinterpret_cast<uint32_t*>(hdr_l1);
    uint32_t* out = hdr + 16;
    const uint32_t s_begin = PROC == 0 ? 0 : split, s_end = PROC == 0 ? split : NSLOT;
    uint32_t n = PROC == 0 ? 0 : n0;
    for (uint32_t s = s_begin; s < s_end && n < KMAX; ++s) {
        const uint32_t* slot = reinterpret_cast<const uint32_t*>(c_addr + s * SLOT_BYTES);
        uint32_t cnt = slot[0];
        if (cnt == 0) {
            continue;
        }
        if (cnt > CAP) {
            cnt = CAP;
        }
        const uint32_t y_first = (s >> 1) * ROWS_PER_CORE + ((s & 1) ? SPLIT : 0);
        uint32_t m = KMAX - n;
        if (cnt < m) {
            m = cnt;
        }
        uint32_t* o = out + 4 * n;
        for (uint32_t j = 0; j < m; ++j, o += 4) {
            const uint32_t e = slot[1 + j];
            uint32_t x = e & 0xFFFF, y = y_first;
            while (x >= W) {
                x -= W;
                ++y;
            }
            o[0] = (y << 16) | x;
            o[1] = e >> 16;
            o[2] = (ty[2 * y] << 16) | ty[2 * y + 1];
            o[3] = (tx[2 * x] << 16) | tx[2 * x + 1];
        }
        n += m;
    }
    volatile tt_l1_ptr uint32_t* sem = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(0));
    if constexpr (PROC == 1) {
        *sem = 1;
        return;
    } else {
        noc_semaphore_wait(sem, 1);
        *sem = 0;
        const uint32_t nr = (kept + 31) & ~31u;
        for (uint32_t k = 4 * kept; k < 4 * nr; ++k) {
            out[k] = 0;
        }
        hdr[0] = total;
        hdr[1] = overflow;
        hdr[2] = kept;
#ifdef KPC_HR
        constexpr uint32_t NB = KMAX / KPC_BSTEP;
        uint32_t b = kept == 0 ? 0 : (kept + KPC_BSTEP - 1) / KPC_BSTEP - 1;
        const uint32_t spec = reinterpret_cast<volatile uint32_t*>(prm_l1)[2];
        if (spec > b) {
            b = spec;
        }
        if (b > NB - 1) {
            b = NB - 1;
        }
        hdr[3] = b;
#endif
        noc_async_write(hdr_l1, hdracc.get_noc_addr(0), (16 + 4 * nr) * 4);
#ifdef KPC_HR
        {
            const auto bacc = TensorAccessor(bk_args, get_arg_val<uint32_t>(4 + b), ROWB);
            const uint32_t bytes = (16 + 4 * nr) * 4;
            for (uint32_t p = 0, off = 0; off < bytes; ++p, off += ROWB) {
                const uint32_t sz = bytes - off < ROWB ? bytes - off : ROWB;
                noc_async_write(hdr_l1 + off, bacc.get_noc_addr(p), sz);
            }
            if (b != spec && spec < NB) {
                // the host reads bucket spec first: give it the counts and the bucket actually used
                const auto sacc = TensorAccessor(bk_args, get_arg_val<uint32_t>(4 + spec), ROWB);
                noc_async_write(hdr_l1, sacc.get_noc_addr(0), 64);
            }
        }
#endif
        noc_async_write_barrier();
    }
}