File size: 7,188 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 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 | // 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();
}
}
|