// 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 #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(); constexpr auto s_args = TensorAccessorArgs(); constexpr auto o_args = TensorAccessorArgs(); 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(); #endif constexpr uint32_t CT = C / 32; constexpr uint32_t ROWB = C * 2; constexpr uint32_t HW = static_cast(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(0), 2 * ROWB); #else const auto x_acc = TensorAccessor(x_args, get_common_arg_val(0), ROWB); #endif const auto l_acc = TensorAccessor(l_args, get_common_arg_val(1), TB32); const auto s_acc = TensorAccessor(s_args, get_common_arg_val(2), TB32); const auto o_acc = TensorAccessor(o_args, get_common_arg_val(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(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(lbuf); volatile tt_l1_ptr uint32_t* sig = reinterpret_cast(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(get_read_ptr(cb_loc)); if constexpr (WRITES_GEO) { cb_wait_front(cb_sig0, 1); sig = reinterpret_cast(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(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(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(g_addr); volatile tt_l1_ptr uint32_t* fy = reinterpret_cast(g_addr + TB32); volatile tt_l1_ptr uint32_t* ma = reinterpret_cast(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(0xFFFF); vy = (vy + 0x8000) & ~static_cast(0xFFFF); #endif const int32_t x0 = vx >> 16; const int32_t y0 = (vy >> 16) + static_cast(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(vx) & 0xFFFF); fy[e] = q16_to_f32(static_cast(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(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(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); } }