// SPDX-FileCopyrightText: © 2026 Tenstorrent USA, Inc.
// SPDX-License-Identifier: Apache-2.0
//
// SuperPoint keypoint list on device (data movement only, one core, ttnn.generic_op).
// Concatenates the per-RISC candidate slots of nms_unfold_kp.cpp (raster order) into
//   HDR [1, 16 + 4*KMAX] uint32: [0] total candidates, [1] overflow (a slot exceeded CAP), [2] n =
//        min(total, KMAX); keypoint k: [16 + 4k] = (y << 16) | x, [17 + 4k] = bf16 score bits,
//        [18 + 4k] = (y0c*wc << 16) | y1c*wc, [19 + 4k] = (x0c << 16) | x1c (bilinear cell rows /
//        columns from TAB [1, 2W + 2H] = per x (x0c, x1c), per y (y0c*wc, y1c*wc)).
// Entries k >= n are left stale (never read: consumers stop at n / skip whole tile rows >= n,
// and the entries n .. roundup(n, 32) are zero-filled so partial tile rows see pixel (0, 0)).
#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 a core
    constexpr uint32_t SLOT_BYTES = (CAP + 1) * 4;
    constexpr uint32_t HDR_WORDS = 16 + 4 * KMAX;
    constexpr uint32_t H = get_compile_time_arg_val(7);
    constexpr uint32_t TAB_BYTES = (2 * W + 2 * H) * 4;
    constexpr bool LOCAL = get_compile_time_arg_val(8);  // C sharded in this core's L1 (read in place)
    constexpr auto c_args = TensorAccessorArgs<9>();
    constexpr auto hdr_args = TensorAccessorArgs<c_args.next_compile_time_args_offset()>();
    constexpr auto tab_args = TensorAccessorArgs<hdr_args.next_compile_time_args_offset()>();
    const auto cacc = TensorAccessor(c_args, c_addr, SLOT_BYTES);
    const auto hdracc = TensorAccessor(hdr_args, hdr_addr, HDR_WORDS * 4);
    const auto tabacc = TensorAccessor(tab_args, tab_addr, TAB_BYTES);

    const uint32_t base = get_write_ptr(cb_scratch);
    const uint32_t slots_l1 = LOCAL ? c_addr : base;
    const uint32_t tab_l1 = base + NSLOT * SLOT_BYTES;
    const uint32_t hdr_l1 = tab_l1 + ((TAB_BYTES + 63) & ~63u);

    if constexpr (!LOCAL) {
        for (uint32_t s = 0; s < NSLOT; ++s) {
            noc_async_read(cacc.get_noc_addr(s), slots_l1 + s * SLOT_BYTES, SLOT_BYTES);
        }
    }
    noc_async_read(tabacc.get_noc_addr(0), tab_l1, TAB_BYTES);
    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;
    uint32_t total = 0, overflow = 0, n = 0;
    for (uint32_t s = 0; s < NSLOT; ++s) {
        const uint32_t* slot = reinterpret_cast<const uint32_t*>(slots_l1 + s * SLOT_BYTES);
        uint32_t cnt = slot[0];
        if (cnt == 0) {
            continue;
        }
        total += cnt;
        if (cnt > CAP) {
            overflow = 1;
            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) {  // a slot spans only a few image rows: no division
                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;
    }
    const uint32_t nr = (n + 31) & ~31u;
    for (uint32_t k = 4 * n; k < 4 * nr; ++k) {
        out[k] = 0;
    }
    hdr[0] = total;
    hdr[1] = overflow;
    hdr[2] = n;
    noc_async_write(hdr_l1, hdracc.get_noc_addr(0), (16 + 4 * nr) * 4);
    noc_async_write_barrier();
}