File size: 4,346 Bytes
c699c4c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 | // 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();
}
|