Hugging Face's logo

changh95
/
bevformer-p150

Object Detection
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-catalog
tt-model-container
Model card Files
xet
 Community
bevformer-p150 / code /tt_bevformer /tt /kernels /dcn_k_reader.cpp
changh95's picture
changh95
tt-model push bevformer-p150 (container)
55474c5 verified
Raw History Blame Contribute Delete
16.2 kB
// SPDX-License-Identifier: Apache-2.0
// dcn_k (tt/dcn_k.py, OPT_PLAN #3): the gather half of the DCNv2 im2col kernel. The same source runs on both
// data-movement RISC-Vs: PART 0 (NCRISC) gathers the corners NW, NE and writes the geometry tiles, PART 1 (BRISC)
// gathers SW, SE and writes the output tiles. Each RISC has its own CBs.
//
// Work item w (contiguous range per core; core index from the logical coordinates, common RT args only): block
// b = w / 9 (32 output pixels, rows b * 32 .. of the [N H W] map), tap k = w % 9.
// x bf16 ROW_MAJOR [1, 1, N H W, C], L1 interleaved: page q = the C-channel stick of pixel q;
// loc int32 TILE [1, 1, N H W, 32]: columns 2k, 2k + 1 = round(pixel * 2^16) of tap k's sampling point (x, y);
// sig fp32 TILE [1, 1, N H W, 32]: column 18 + k = sigmoid(mask logit of tap k);
// out bf16 TILE [1, 1, N H W, 9 C]: the modulated im2col, tiles b * 9 CT + k * CT + t (CT = C / 32).
// Per item: cb_rm (this RISC) <- 2 x 32 rows x C bf16 row-major (corner 2 PART + i at pages i CT .. i CT + CT - 1);
// PART 0 also cb_geo <- fx, fy, mask of the item's 32 rows in column 0 (fp32 bit patterns; the other columns zero).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t PART = get_compile_time_arg_val(1);
constexpr uint32_t cb_rm = get_compile_time_arg_val(2);
constexpr uint32_t cb_geo = get_compile_time_arg_val(3);
constexpr uint32_t cb_loc = get_compile_time_arg_val(4);
constexpr uint32_t cb_out = get_compile_time_arg_val(5);
constexpr int32_t HV = get_compile_time_arg_val(6);
constexpr int32_t WV = get_compile_time_arg_val(7);
constexpr uint32_t NQ = get_compile_time_arg_val(8);
constexpr uint32_t NB = get_compile_time_arg_val(9);
constexpr uint32_t C = get_compile_time_arg_val(10);
constexpr uint32_t NWORK = get_compile_time_arg_val(11);
constexpr uint32_t NC = get_compile_time_arg_val(12);
constexpr uint32_t GX = get_compile_time_arg_val(13);
constexpr uint32_t WRITES_OUT = get_compile_time_arg_val(14);
constexpr uint32_t WRITES_GEO = get_compile_time_arg_val(15);
// B0: the first 32-row block of this call (BEVFORMER_DCN_IM2COL=l1 runs the im2col in row chunks whose output is an
// L1 height-sharded tensor the GEMM reads in place): items cover blocks B0 .. B0 + NB - 1, loc / sig / x are read at
// the absolute rows, the output tiles are written at the chunk's own rows
constexpr uint32_t B0 = get_compile_time_arg_val(16);
constexpr auto x_args = TensorAccessorArgs<17>();
constexpr auto l_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto s_args = TensorAccessorArgs<l_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048, TB32 = 4096;
#ifdef DCN_K_GLUE
// BEVFORMER_DCN_GLUE=kernel (D20 round 2): the glue runs in the compute kernel. `loc` is the raw offset conv output
// (fp32 TILE [NHW, 32], read in place, e.g. the conv's L1 shards), `sig` the B table (fp32 TILE [NHW, 32]: base * S +
// O) and `sc` the S tile (fp32 TILE, row 0 replicated: tile 0 of the full scale map). PART 0 reads S once and each
// block's off / B tiles into cb 9 / 10, the compute kernel returns loc (int32) to both RISCs (cb 6 / 7) and sig
// (fp32) to PART 0 (cb 8): the same LLK steps as glue_k (bit-identical to the separate program).
constexpr uint32_t cb_goff = 9, cb_gb = 10, cb_gs = 11, cb_sig0 = 8;
constexpr auto sc_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
#endif
constexpr uint32_t CT = C / 32;
constexpr uint32_t ROWB = C * 2;
constexpr uint32_t HW = static_cast<uint32_t>(HV * WV);
constexpr int32_t LOC_LIM = 1 << 28;
static_assert(C % 32 == 0 && ROWB % MEM_ZEROS_SIZE == 0, "dcn_k: C a multiple of 256");
static inline uint32_t tile_elem(uint32_t r, uint32_t c) {
return (((r >> 4) << 1) + (c >> 4)) * 256 + ((r & 15) << 4) + (c & 15);
}
static inline uint32_t q16_to_f32(uint32_t q) {
if (q == 0) {
return 0;
}
const uint32_t n = 31 - __builtin_clz(q);
return ((127 - 16 + n) << 23) | ((q << (23 - n)) & 0x7FFFFF);
}
static inline void zero_l1(uint32_t addr, uint32_t bytes) {
const uint64_t z = get_noc_addr(MEM_ZEROS_BASE);
for (uint32_t o = 0; o < bytes; o += MEM_ZEROS_SIZE) {
noc_async_read(z, addr + o, MEM_ZEROS_SIZE);
}
}
void kernel_main() {
const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (ci >= NC) {
return;
}
const uint32_t w0 = ci * NWORK / NC;
const uint32_t w1 = (ci + 1) * NWORK / NC;
if (w0 == w1) {
return;
}
#ifdef DCN_K_PAIRS
// BEVFORMER_DCN_PAIRS (D20 round 2): x is the pixel-pair map [N HV (WV + 1), 2 C] (dcnpair_k): page (n, y, i) =
// the sticks of x = i - 1 and x = i; one 2C-wide read per row brings both corners of this RISC (NW + NE or
// SW + SE) into cb_rm as one row of a 32 x 2C row-major block (the compute kernel tilizes it 2 CT tiles wide)
const auto x_acc = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), 2 * ROWB);
#else
const auto x_acc = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), ROWB);
#endif
const auto l_acc = TensorAccessor(l_args, get_common_arg_val<uint32_t>(1), TB32);
const auto s_acc = TensorAccessor(s_args, get_common_arg_val<uint32_t>(2), TB32);
const auto o_acc = TensorAccessor(o_args, get_common_arg_val<uint32_t>(3), TB16);
if constexpr (WRITES_GEO) {
cb_reserve_back(cb_geo, 6);
zero_l1(get_write_ptr(cb_geo), 6 * TB32);
}
#ifdef DCN_K_GLUE
if constexpr (PART == 0) {
const auto sc_acc = TensorAccessor(sc_args, get_common_arg_val<uint32_t>(4), TB32);
cb_reserve_back(cb_gs, 1);
noc_async_read(sc_acc.get_noc_addr(0), get_write_ptr(cb_gs), TB32);
noc_async_read_barrier();
cb_push_back(cb_gs, 1);
}
volatile tt_l1_ptr int32_t* loc = nullptr;
volatile tt_l1_ptr uint32_t* sig = nullptr;
#else
cb_reserve_back(cb_loc, 2);
const uint32_t lbuf = get_write_ptr(cb_loc);
noc_async_read_barrier();
volatile tt_l1_ptr int32_t* loc = reinterpret_cast<volatile tt_l1_ptr int32_t*>(lbuf);
volatile tt_l1_ptr uint32_t* sig = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(lbuf + TB32);
#endif
uint32_t cur_b = 0xFFFFFFFF;
for (uint32_t w = w0; w < w1; ++w) {
const uint32_t b = w / 9 + B0; // absolute block
const uint32_t k = w - (b - B0) * 9;
if (b != cur_b) {
#ifdef DCN_K_GLUE
if (cur_b != 0xFFFFFFFF) { // the previous block's loc / sig pages
cb_pop_front(cb_loc, 1);
if constexpr (WRITES_GEO) {
cb_pop_front(cb_sig0, 1);
}
}
if constexpr (PART == 0) { // this block's glue inputs
cb_reserve_back(cb_goff, 1);
cb_reserve_back(cb_gb, 1);
noc_async_read(l_acc.get_noc_addr(b), get_write_ptr(cb_goff), TB32);
noc_async_read(s_acc.get_noc_addr(b), get_write_ptr(cb_gb), TB32);
noc_async_read_barrier();
cb_push_back(cb_goff, 1);
cb_push_back(cb_gb, 1);
}
cb_wait_front(cb_loc, 1);
loc = reinterpret_cast<volatile tt_l1_ptr int32_t*>(get_read_ptr(cb_loc));
if constexpr (WRITES_GEO) {
cb_wait_front(cb_sig0, 1);
sig = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_read_ptr(cb_sig0));
}
#else
noc_async_read(l_acc.get_noc_addr(b), lbuf, TB32);
if constexpr (WRITES_GEO) {
noc_async_read(s_acc.get_noc_addr(b), lbuf + TB32, TB32);
}
noc_async_read_barrier();
#endif
cur_b = b;
}
#ifdef DCN_K_NEAREST
// BEVFORMER_DCN_SAMPLE=nearest (D20 round 3): one corner per sampling point (the nearest pixel, round half
// up); the items alternate between the RISCs (item w0 + 2i on PART 0, w0 + 2i + 1 on PART 1), each gathering
// its item's 32 sticks into its own cb_rm (CT pages); PART 0 writes the mask tile (column 0) of every item
{
const bool mine = ((w - w0) & 1) == PART;
uint32_t rm = 0;
if (mine) {
cb_reserve_back(cb_rm, CT);
rm = get_write_ptr(cb_rm);
}
uint32_t g_addr = 0;
if constexpr (WRITES_GEO) {
cb_reserve_back(cb_geo, 1);
g_addr = get_write_ptr(cb_geo);
}
volatile tt_l1_ptr uint32_t* ma = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr);
if (mine || WRITES_GEO) {
for (uint32_t r = 0; r < 32; ++r) {
const uint32_t q = b * 32 + r;
const uint32_t e = tile_elem(r, 0);
if (q >= NQ) {
if constexpr (WRITES_GEO) {
ma[e] = 0;
}
if (mine) {
zero_l1(rm + r * ROWB, ROWB);
}
continue;
}
if constexpr (WRITES_GEO) {
ma[e] = sig[tile_elem(r, 18 + k)];
}
if (!mine) {
continue;
}
const uint32_t el = tile_elem(r, 2 * k);
int32_t vx = loc[el];
int32_t vy = loc[el + 1];
vx = vx > LOC_LIM ? LOC_LIM : (vx < -LOC_LIM ? -LOC_LIM : vx);
vy = vy > LOC_LIM ? LOC_LIM : (vy < -LOC_LIM ? -LOC_LIM : vy);
const int32_t x0 = (vx + 0x8000) >> 16;
const int32_t y0 = (vy + 0x8000) >> 16;
if (y0 >= 0 && y0 < HV && x0 >= 0 && x0 < WV) {
const uint32_t pos = (q / HW) * HW + static_cast<uint32_t>(y0 * WV + x0);
noc_async_read(x_acc.get_noc_addr(pos), rm + r * ROWB, ROWB);
} else {
zero_l1(rm + r * ROWB, ROWB);
}
}
}
noc_async_read_barrier();
if constexpr (WRITES_GEO) {
cb_push_back(cb_geo, 1);
}
if (mine) {
cb_push_back(cb_rm, CT);
}
}
#else
cb_reserve_back(cb_rm, 2 * CT);
const uint32_t rm = get_write_ptr(cb_rm);
uint32_t g_addr = 0;
if constexpr (WRITES_GEO) {
cb_reserve_back(cb_geo, 3);
g_addr = get_write_ptr(cb_geo);
}
volatile tt_l1_ptr uint32_t* fx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr);
volatile tt_l1_ptr uint32_t* fy = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + TB32);
volatile tt_l1_ptr uint32_t* ma = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + 2 * TB32);
for (uint32_t r = 0; r < 32; ++r) {
const uint32_t q = b * 32 + r;
const uint32_t dst0 = rm + r * ROWB;
const uint32_t dst1 = dst0 + CT * TB16;
#ifdef DCN_K_MM
constexpr bool DIAG = true; // BEVFORMER_DCN_SUM=mm: fx / fy / mask on the diagonal (diag(w_c) tiles)
#else
constexpr bool DIAG = false;
#endif
if (q >= NQ) {
if constexpr (WRITES_GEO) {
const uint32_t e = tile_elem(r, DIAG ? r : 0);
fx[e] = 0;
fy[e] = 0;
ma[e] = 0;
}
#ifdef DCN_K_PAIRS
zero_l1(rm + r * 2 * ROWB, 2 * ROWB);
#else
zero_l1(dst0, ROWB);
zero_l1(dst1, ROWB);
#endif
continue;
}
const uint32_t el = tile_elem(r, 2 * k);
int32_t vx = loc[el];
int32_t vy = loc[el + 1];
vx = vx > LOC_LIM ? LOC_LIM : (vx < -LOC_LIM ? -LOC_LIM : vx);
vy = vy > LOC_LIM ? LOC_LIM : (vy < -LOC_LIM ? -LOC_LIM : vy);
#ifdef DCN_K_NEAREST_EMU
// accuracy probe of DCN_K_NEAREST on the bilinear path: zero fractions -> weight 1 on NW only
vx = (vx + 0x8000) & ~static_cast<int32_t>(0xFFFF);
vy = (vy + 0x8000) & ~static_cast<int32_t>(0xFFFF);
#endif
const int32_t x0 = vx >> 16;
const int32_t y0 = (vy >> 16) + static_cast<int32_t>(PART); // PART 1: the lower row (SW, SE)
#ifndef DCN_K_NOGEO
if constexpr (WRITES_GEO) {
const uint32_t e = tile_elem(r, DIAG ? r : 0);
fx[e] = q16_to_f32(static_cast<uint32_t>(vx) & 0xFFFF);
fy[e] = q16_to_f32(static_cast<uint32_t>(vy) & 0xFFFF);
ma[e] = sig[tile_elem(r, 18 + k)];
}
#endif
const bool yv = y0 >= 0 && y0 < HV;
#ifdef DCN_K_PAIRS
{
const uint32_t pd = rm + r * 2 * ROWB;
#ifdef DCN_K_NOREAD
(void)pd;
continue;
#endif
if (yv && x0 >= -1 && x0 < WV) {
const uint32_t pp = (q / HW) * (HV * (WV + 1)) + static_cast<uint32_t>(y0 * (WV + 1) + x0 + 1);
noc_async_read(x_acc.get_noc_addr(pp), pd, 2 * ROWB);
} else {
zero_l1(pd, 2 * ROWB);
}
continue;
}
#endif
const uint32_t nbase = (q / HW) * HW;
const bool xv0 = x0 >= 0 && x0 < WV, xv1 = x0 + 1 >= 0 && x0 + 1 < WV;
const uint32_t pos = nbase + static_cast<uint32_t>(y0 * WV + x0);
#ifndef DCN_K_NOREAD
if (yv && xv0) {
#if defined(DCN_K_RLOCAL)
noc_async_read(get_noc_addr(dst1), dst0, ROWB); // (anatomy probe: own L1, no mesh hops)
#elif defined(DCN_K_RSMALL)
noc_async_read(x_acc.get_noc_addr(pos), dst0, 32); // (anatomy probe: same requests, 1/16 bytes)
#elif defined(DCN_K_RSELF)
noc_async_read(x_acc.get_noc_addr(q), dst0, ROWB); // (anatomy probe: the pixel's own page)
#else
noc_async_read(x_acc.get_noc_addr(pos), dst0, ROWB);
#endif
} else {
zero_l1(dst0, ROWB);
}
if (yv && xv1) {
#if defined(DCN_K_RLOCAL)
noc_async_read(get_noc_addr(dst0), dst1, ROWB);
#elif defined(DCN_K_RSMALL)
noc_async_read(x_acc.get_noc_addr(pos + 1), dst1, 32);
#elif defined(DCN_K_RSELF)
noc_async_read(x_acc.get_noc_addr(q), dst1, ROWB);
#else
noc_async_read(x_acc.get_noc_addr(pos + 1), dst1, ROWB);
#endif
} else {
zero_l1(dst1, ROWB);
}
#else
(void)pos;
#endif
}
noc_async_read_barrier();
if constexpr (WRITES_GEO) {
cb_push_back(cb_geo, 3);
}
cb_push_back(cb_rm, 2 * CT);
#endif // DCN_K_NEAREST
if constexpr (WRITES_OUT) {
if (w > w0) {
const uint32_t bp = (w - 1) / 9;
const uint32_t kp = (w - 1) - bp * 9;
cb_wait_front(cb_out, CT);
uint32_t rp = get_read_ptr(cb_out);
const uint32_t t0 = bp * 9 * CT + kp * CT;
#ifndef DCN_K_NOWRITE
for (uint32_t t = 0; t < CT; ++t) {
noc_async_write(rp + t * TB16, o_acc.get_noc_addr(t0 + t), TB16);
}
#endif
noc_async_writes_flushed();
cb_pop_front(cb_out, CT);
}
}
}
if constexpr (WRITES_OUT) {
const uint32_t bp = (w1 - 1) / 9;
const uint32_t kp = (w1 - 1) - bp * 9;
cb_wait_front(cb_out, CT);
uint32_t rp = get_read_ptr(cb_out);
const uint32_t t0 = bp * 9 * CT + kp * CT;
for (uint32_t t = 0; t < CT; ++t) {
noc_async_write(rp + t * TB16, o_acc.get_noc_addr(t0 + t), TB16);
}
noc_async_write_barrier();
cb_pop_front(cb_out, CT);
}
}