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 /msda_k_reader.cpp
changh95's picture
changh95
tt-model push bevformer-p150 (container)
55474c5 verified
Raw History Blame Contribute Delete
18.5 kB
// SPDX-License-Identifier: Apache-2.0
// msda_k (tt/msda_k.py, OPT_PLAN #1 / #2 / #10): the data-movement half of the fixed-point MSDA kernel. The same
// source runs on both data-movement RISC-Vs; PART 0 (NCRISC, "reader") owns the points [0, P/2), PART 1 (BRISC,
// "writer") the points [P/2, P) and also writes the output tiles. Each RISC has its own CBs (cb_geo, cb_in, cb_loc),
// so there is no cross-RISC handshake.
//
// Work item w (contiguous range per core, core index from the logical coordinates: common RT args only, probe P3):
// j = w / NB (head block), b = w % NB (32-query block); value map m = j * SUB + b / (NB / SUB);
// loc int32 TILE [1, H, NB * 32, 32], tile j * NB + b: row r, column 2p / 2p + 1 = x / y of point p as
// round(pixel * 2^16) (pixel = (g + 1) size / 2 - 0.5, align_corners=False);
// attn fp32 TILE [1, H, NB * 32, 32], tile j * NB + b: the softmax weight of point p in column p;
// out bf16 TILE [1, 1, NB * 32, H * 32], tile b * H + j.
// value: VSRC 0 = (N, h, w, 32) ROW_MAJOR, page m * h * w + y * w + x; VSRC 1 = the head split [B, heads, rows, 32]
// TILE as nlp_create_qkv_heads writes it (map m = j * SUB + s: slice j, rows s * h * w ..; TPH tiles per slice), the
// stick of a pixel = row r of its tile: two 32-byte face rows (faces 2 (r / 16) and + 1), no untilize needed.
// Per item and RISC: cb_geo <- 3 fp32 tiles (fx, fy, attn with the other RISC's point columns zeroed) and
// cb_in <- 2P bf16 tiles (step k = (p - p0) * 4 + corner; corners NW, NE, SW, SE): row r of step k is the 32-channel
// value stick of that corner, read as two 32-byte halves straight into the tile's face rows; corners outside the map
// (and padding rows) read the zeros page. The compute kernel turns geo into the 4 corner weights on the SFPU.
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"
#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_dataflow.hpp"
constexpr uint32_t PART = get_compile_time_arg_val(1);
constexpr uint32_t cb_in = 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_e = get_compile_time_arg_val(5);
constexpr uint32_t cb_out = get_compile_time_arg_val(6);
constexpr int32_t HV = get_compile_time_arg_val(7);
constexpr int32_t WV = get_compile_time_arg_val(8);
constexpr uint32_t NQ = get_compile_time_arg_val(9);
constexpr uint32_t NB = get_compile_time_arg_val(10);
constexpr uint32_t H = get_compile_time_arg_val(11);
constexpr uint32_t SUB = get_compile_time_arg_val(12);
constexpr uint32_t P = get_compile_time_arg_val(13);
constexpr uint32_t NWORK = get_compile_time_arg_val(14);
constexpr uint32_t NC = get_compile_time_arg_val(15);
constexpr uint32_t GX = get_compile_time_arg_val(16);
constexpr uint32_t WRITES_OUT = get_compile_time_arg_val(17);
constexpr uint32_t MAKES_E = get_compile_time_arg_val(18);
// 0: ROW_MAJOR maps (64-byte pages); 1: TILE heads; 2 (BEVFORMER_MSDA_PAIRS): pixel-pair maps [N, h, w + 1, 64]
// (tt/pairmap_k.py: 128-byte page (m, y, i) = the sticks of x = i - 1 and x = i): ONE read per point and row lands
// the NW + NE (or SW + SE) sticks as a row of a [32 x 64] row-major staging block (tilized by the compute kernel)
constexpr uint32_t VSRC = get_compile_time_arg_val(19);
constexpr uint32_t TPH = get_compile_time_arg_val(20); // VSRC 1: tiles per head slice (SUB maps of HV x WV rows)
// SKIP: items strided over the cores (w = core, core + NC, ..: padding blocks spread evenly) and points with all four
// corners outside the map (or rows past NQ) issue no reads: their weight is set to 0, and their tile rows keep
// earlier (finite) gathers; the step tiles are zeroed once. 0: contiguous item ranges, zeros-page reads.
constexpr uint32_t SKIP = get_compile_time_arg_val(21);
// LOC_HS 0: loc is the head split [1, H, NB * 32, 32] (tile j * NB + b, columns 2p / 2p + 1); LOC_HS > 0
// (BEVFORMER_MSDA_LOC=compact): loc is the merged [1, 1, NB * 32, H * LOC_HS] map (the sampling projection's own
// layout: head j in columns j * LOC_HS .. + 2P), LOC_TW tiles per row block: tile b * LOC_TW + j * LOC_HS / 32,
// column (j * LOC_HS) % 32 + 2p
constexpr uint32_t LOC_HS = get_compile_time_arg_val(22);
constexpr uint32_t LOC_TW = get_compile_time_arg_val(23);
// SM 1 (BEVFORMER_MSDA_SOFTMAX=kernel): attn is the compact attention-logit map [1, 1, NB * 32, H * P] (head j in
// columns j * P ..; AT_TW tiles per row block); the PART-0 RISC rebuilds the head's logit tile (columns 0..P-1 the
// logits, the rest PAD_LOGIT: what the head-block projection writes) into cb_lg for the compute kernel's softmax,
// and makes the reduce scalers and the two half masks once. Rows past NQ get logits 0 (finite; their values are 0).
constexpr uint32_t SM = get_compile_time_arg_val(24);
constexpr uint32_t AT_TW = get_compile_time_arg_val(25);
constexpr uint32_t PAD_BITS = get_compile_time_arg_val(26); // fp32 bits of PAD_LOGIT
constexpr uint32_t cb_lg = 9, cb_scmax = 10, cb_scsum = 11, cb_msk = 12, cb_scr = 18;
constexpr auto v_args = TensorAccessorArgs<27>();
constexpr auto l_args = TensorAccessorArgs<v_args.next_compile_time_args_offset()>();
constexpr auto a_args = TensorAccessorArgs<l_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<a_args.next_compile_time_args_offset()>();
constexpr uint32_t PH = P / 2;
constexpr uint32_t P0 = PART * PH;
constexpr uint32_t NS = PH * 4; // steps (tiles) per item and RISC
constexpr uint32_t TB16 = 2048, TB32 = 4096;
constexpr uint32_t VPAGE = VSRC == 0 ? 64 : (VSRC == 2 ? 128 : 2048); // RM stick; pixel pair; TILE: a bf16 tile
static_assert(VSRC != 2 || !SKIP, "msda_k: pair maps and SKIP are exclusive");
constexpr uint32_t BPS = NB / SUB;
constexpr uint32_t MAP = static_cast<uint32_t>(HV * WV);
constexpr int32_t LOC_LIM = 1 << 28; // |pixel| < 4096 after the clamp: no overflow below
static_assert(P % 2 == 0 && P <= 8, "msda_k: P even, <= 8 (x / y of a point share a face)");
static_assert(NB % SUB == 0, "msda_k: blocks split evenly over the value maps of a head");
static_assert(LOC_HS == 0 || (LOC_HS >= 2 * P && 32 % LOC_HS == 0), "msda_k: LOC_HS holds 2P columns, divides 32");
static inline uint32_t tile_elem(uint32_t r, uint32_t c) {
return (((r >> 4) << 1) + (c >> 4)) * 256 + ((r & 15) << 4) + (c & 15);
}
// q in [0, 65536): the fp32 bit pattern of q / 65536 (exact)
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);
}
// one corner stick (32 channels) -> the face rows at dst (channels 0..15) and dst + 512 (16..31). pos: RM page
// (VSRC 0) or slice * TPH * 32 + row (VSRC 1).
// VSRC 3 (BEVFORMER_DEC_VALUE=fused, D20 round 3): the value is the projection itself, TILE [1, 1, SUB h w, H * 32]
// (no head split): the head block j of map row pos is tile (pos / 32) * H + j; g_hj = the current item's head
static uint32_t g_hj = 0;
template <typename A>
static inline void gather(const A& v_acc, uint32_t pos, bool ok, uint64_t zeros, uint32_t dst) {
if (!ok) {
noc_async_read(zeros, dst, 32);
noc_async_read(zeros, dst + 512, 32);
return;
}
uint64_t s;
uint32_t hi;
if constexpr (VSRC == 0) {
s = v_acc.get_noc_addr(pos);
hi = 32;
} else {
const uint32_t r = pos & 31;
const uint32_t t = VSRC == 3 ? (pos >> 5) * H + g_hj : pos >> 5;
s = v_acc.get_noc_addr(t, ((r >> 4) << 10) + ((r & 15) << 5));
hi = 512;
}
#ifndef MSDA_K_NOREAD
noc_async_read(s, dst, 32);
noc_async_read(s + hi, dst + 512, 32);
#endif
}
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 = SKIP ? ci : ci * NWORK / NC;
const uint32_t w1 = SKIP ? NWORK : (ci + 1) * NWORK / NC;
const uint32_t wstep = SKIP ? NC : 1;
if (w0 >= w1) {
return;
}
const auto v_acc = TensorAccessor(v_args, get_common_arg_val<uint32_t>(0), VPAGE);
const auto l_acc = TensorAccessor(l_args, get_common_arg_val<uint32_t>(1), TB32);
const auto a_acc = TensorAccessor(a_args, get_common_arg_val<uint32_t>(2), TB32);
const auto o_acc = TensorAccessor(o_args, get_common_arg_val<uint32_t>(3), TB16);
const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
if constexpr (MAKES_E) {
// E_p (bf16): row p all ones, so W @ E_p broadcasts column p of W over the 32 channels
cb_reserve_back(cb_e, P);
const uint32_t e_addr = get_write_ptr(cb_e);
zero_l1(e_addr, P * TB16);
noc_async_read_barrier();
volatile tt_l1_ptr uint16_t* e = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(e_addr);
for (uint32_t p = 0; p < P; ++p) {
for (uint32_t c = 0; c < 32; ++c) {
e[p * (TB16 / 2) + tile_elem(p, c)] = 0x3F80;
}
}
cb_push_back(cb_e, P);
}
if constexpr (SM && MAKES_E) {
// the softmax's reduce scalers, as ttnn's softmax reader makes them; the half masks (fp32 1 / 0 per column:
// RISC h owns the points [h P / 2, (h + 1) P / 2)); the logit pages pre-filled with PAD_LOGIT
dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scmax, ckernel::PoolType::MAX,
ckernel::ReduceDim::REDUCE_ROW>();
dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<cb_scsum, ckernel::PoolType::SUM,
ckernel::ReduceDim::REDUCE_ROW>();
cb_reserve_back(cb_msk, 2);
volatile tt_l1_ptr uint32_t* m = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_msk));
for (uint32_t h = 0; h < 2; ++h) {
for (uint32_t r = 0; r < 32; ++r) {
for (uint32_t c = 0; c < 32; ++c) {
m[h * 1024 + tile_elem(r, c)] = (c >= h * PH && c < (h + 1) * PH) ? 0x3F800000u : 0u;
}
}
}
cb_push_back(cb_msk, 2);
cb_reserve_back(cb_lg, 2);
volatile tt_l1_ptr uint32_t* lg = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
for (uint32_t i = 0; i < 2048; ++i) {
lg[i] = PAD_BITS;
}
}
// the geo pages: zero once (columns >= P and the right faces stay zero / finite forever)
cb_reserve_back(cb_geo, 6);
zero_l1(get_write_ptr(cb_geo), 6 * TB32);
if constexpr (SKIP) {
cb_reserve_back(cb_in, 2 * NS);
zero_l1(get_write_ptr(cb_in), 2 * NS * TB16);
}
cb_reserve_back(cb_loc, 1);
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);
uint32_t prev = 0xFFFFFFFF;
for (uint32_t w = w0; w < w1; w += wstep) {
const uint32_t j = w / NB;
const uint32_t b = w - j * NB;
// VSRC 0: page of the map's pixel 0; VSRC 1: the slice's first tile row + the map's first row in it
const uint32_t vbase = VSRC == 0 ? (j * SUB + b / BPS) * MAP
: VSRC == 2 ? (j * SUB + b / BPS) * HV * (WV + 1)
: VSRC == 3 ? (b / BPS) * MAP
: j * TPH * 32 + (b / BPS) * MAP;
g_hj = j;
cb_reserve_back(cb_geo, 3);
cb_reserve_back(cb_in, NS);
const uint32_t g_addr = get_write_ptr(cb_geo);
const uint32_t in_base = get_write_ptr(cb_in);
const uint32_t ltile = LOC_HS == 0 ? j * NB + b : b * LOC_TW + (j * LOC_HS) / 32;
const uint32_t lcol = LOC_HS == 0 ? 0 : (j * LOC_HS) % 32;
noc_async_read(l_acc.get_noc_addr(ltile), lbuf, TB32);
if constexpr (SM) {
if constexpr (MAKES_E) {
noc_async_read(a_acc.get_noc_addr(b * AT_TW + (j * P) / 32), get_write_ptr(cb_scr), TB32);
}
} else {
noc_async_read(a_acc.get_noc_addr(j * NB + b), g_addr + 2 * TB32, TB32);
}
noc_async_read_barrier();
if constexpr (SM && MAKES_E) {
// the head's logit tile: columns 0..P-1 <- compact columns (j P) % 32 + p; columns >= P keep PAD
cb_reserve_back(cb_lg, 1);
volatile tt_l1_ptr uint32_t* lg = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_lg));
volatile tt_l1_ptr uint32_t* src = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_scr));
const uint32_t acol = (j * P) % 32;
for (uint32_t r = 0; r < 32; ++r) {
const bool ok = b * 32 + r < NQ;
for (uint32_t p = 0; p < P; ++p) {
lg[tile_elem(r, p)] = ok ? src[tile_elem(r, acol + p)] : 0u;
}
}
cb_push_back(cb_lg, 1);
}
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* at = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(g_addr + 2 * TB32);
// the other RISC's points carry zero weight in this copy (their fx / fy here are stale but finite); SM: the
// compute kernel applies the half mask to the softmax instead
if constexpr (!SM) {
for (uint32_t r = 0; r < 32; ++r) {
for (uint32_t p = 0; p < P; ++p) {
if (p < P0 || p >= P0 + PH) {
at[tile_elem(r, p)] = 0;
}
}
}
}
for (uint32_t r = 0; r < 32; ++r) {
const uint32_t roff = (((r >> 4) << 1) << 9) + ((r & 15) << 5); // face (r / 16, 0), row r % 16
const bool row_ok = b * 32 + r < NQ;
for (uint32_t p = P0; p < P0 + PH; ++p) {
const uint32_t dst = in_base + (p - P0) * 4 * TB16 + roff;
const uint32_t ew = tile_elem(r, p);
if (!row_ok) {
fx[ew] = 0;
fy[ew] = 0;
if constexpr (!SM) {
at[ew] = 0;
}
if constexpr (SKIP) {
continue;
}
if constexpr (VSRC == 2) {
const uint32_t blk = in_base + (p - P0) * 4 * TB16 + r * 128;
noc_async_read(zeros, blk, 128);
noc_async_read(zeros, blk + 2 * TB16, 128);
continue;
}
for (uint32_t c = 0; c < 4; ++c) {
noc_async_read(zeros, dst + c * TB16, 32);
noc_async_read(zeros, dst + c * TB16 + 512, 32);
}
continue;
}
int32_t vx = loc[tile_elem(r, lcol + 2 * p)];
int32_t vy = loc[tile_elem(r, lcol + 2 * p + 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 >> 16;
const int32_t y0 = vy >> 16;
fx[ew] = q16_to_f32(static_cast<uint32_t>(vx) & 0xFFFF);
fy[ew] = q16_to_f32(static_cast<uint32_t>(vy) & 0xFFFF);
const bool xv0 = x0 >= 0 && x0 < WV, xv1 = x0 + 1 >= 0 && x0 + 1 < WV;
const bool yv0 = y0 >= 0 && y0 < HV, yv1 = y0 + 1 >= 0 && y0 + 1 < HV;
if constexpr (SKIP) {
static_assert(!(SKIP && SM), "msda_k: SKIP and SM are exclusive");
if (!((yv0 || yv1) && (xv0 || xv1))) {
at[ew] = 0;
continue;
}
}
if constexpr (VSRC == 2) {
// staging blocks of point p: top (NW | NE) at + 0, bottom (SW | SE) at + 2 tiles; row r: 128 B
const uint32_t blk = in_base + (p - P0) * 4 * TB16 + r * 128;
const bool xp = x0 >= -1 && x0 < WV; // page i = x0 + 1 in [0, WV]
const uint32_t pg = vbase + static_cast<uint32_t>(y0 * (WV + 1) + x0 + 1);
if (xp && yv0) {
noc_async_read(v_acc.get_noc_addr(pg), blk, 128);
} else {
noc_async_read(zeros, blk, 128);
}
if (xp && yv1) {
noc_async_read(v_acc.get_noc_addr(pg + WV + 1), blk + 2 * TB16, 128);
} else {
noc_async_read(zeros, blk + 2 * TB16, 128);
}
continue;
}
const uint32_t pos = vbase + static_cast<uint32_t>(y0 * WV + x0);
gather(v_acc, pos, yv0 && xv0, zeros, dst);
gather(v_acc, pos + 1, yv0 && xv1, zeros, dst + TB16);
gather(v_acc, pos + WV, yv1 && xv0, zeros, dst + 2 * TB16);
gather(v_acc, pos + WV + 1, yv1 && xv1, zeros, dst + 3 * TB16);
}
}
noc_async_read_barrier();
cb_push_back(cb_geo, 3);
cb_push_back(cb_in, NS);
if constexpr (WRITES_OUT) {
if (prev != 0xFFFFFFFF) {
const uint32_t jp = prev / NB;
const uint32_t bp = prev - jp * NB;
cb_wait_front(cb_out, 1);
noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(bp * H + jp), TB16);
noc_async_writes_flushed();
cb_pop_front(cb_out, 1);
}
}
prev = w;
}
if constexpr (WRITES_OUT) {
const uint32_t jp = prev / NB;
const uint32_t bp = prev - jp * NB;
cb_wait_front(cb_out, 1);
noc_async_write(get_read_ptr(cb_out), o_acc.get_noc_addr(bp * H + jp), TB16);
noc_async_write_barrier();
cb_pop_front(cb_out, 1);
}
}