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 /kpre_writer.cpp
changh95's picture
changh95
tt-model push bevformer-p150 (container)
55474c5 verified
Raw History Blame Contribute Delete
4.83 kB
// SPDX-License-Identifier: Apache-2.0
// kpre (tt/kpre.py), writer (BRISC). Per output row y of the core's pixels: the computed bf16 row (4 pseudo-tiles,
// 1280 x 3 values in pixel-major order, from the compute kernel) or zeros (the pad rows y >= RH) -> the core's own
// shard of the stem conv input [H W, 8] (3 channels + 5 zeros per 16-byte pixel; local L1 stores, as kin_rows).
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
#include "kpre_common.h"
constexpr uint32_t N = get_compile_time_arg_val(0);
constexpr uint32_t S = get_compile_time_arg_val(1);
constexpr uint32_t W = get_compile_time_arg_val(2);
constexpr uint32_t RH = get_compile_time_arg_val(3);
constexpr uint32_t NC = get_compile_time_arg_val(4);
constexpr uint32_t GX = get_compile_time_arg_val(5);
constexpr uint32_t RA = get_compile_time_arg_val(6);
constexpr uint32_t RB = get_compile_time_arg_val(7);
constexpr uint32_t RD = get_compile_time_arg_val(8);
constexpr uint32_t SPLIT = get_compile_time_arg_val(9); // 1: this RISC also builds the odd P rows (cb_p2)
constexpr uint32_t SWAP = get_compile_time_arg_val(10);
constexpr uint32_t ROW_BYTES = get_compile_time_arg_val(11);
constexpr uint32_t PERIODS = get_compile_time_arg_val(12);
constexpr uint32_t U16 = get_compile_time_arg_val(13);
constexpr auto img_args = TensorAccessorArgs<14>();
constexpr auto lut_args = TensorAccessorArgs<img_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_v = 16, cb_p2 = 5, cb_scratch2 = 6;
constexpr uint32_t TILE_WORDS = 1024;
constexpr uint32_t QT = 4;
constexpr uint32_t C_PAD = 8;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
const KpreGeom g = kpre_geom<N, S, W, RH, NC, RA, RB, RD>(core);
if (!g.any) {
return;
}
const uint32_t out_addr = get_common_arg_val<uint32_t>(4);
const uint32_t shard0 = core * S;
// SPLIT: the odd P rows (rlo + 1, rlo + 3, ...), each built before the output row that needs it is awaited
uint32_t next = g.rlo + 1;
uint32_t lut_l1 = 0, row_l1 = 0;
const auto img = TensorAccessor(img_args, get_common_arg_val<uint32_t>(0), ROW_BYTES);
if constexpr (SPLIT) {
if (g.rows) {
const auto lut_t = TensorAccessor(lut_args, get_common_arg_val<uint32_t>(1), 3 * 256 * 4);
const uint32_t base = (get_write_ptr(cb_scratch2) + 63u) & ~63u;
lut_l1 = base;
row_l1 = base + 3072;
noc_async_read(lut_t.get_noc_addr(0), lut_l1, 3 * 256 * 4);
noc_async_read_barrier();
}
}
for (uint32_t y = g.ya; y <= g.yb; ++y) {
if constexpr (SPLIT) {
if (y < RH) {
const uint32_t need = kpre_r0<RA, RB, RD>(y) + 1;
for (; next <= need && next <= g.rhi; next += 2) {
noc_async_read(img.get_noc_addr(next), row_l1, ROW_BYTES);
cb_reserve_back(cb_p2, 5);
noc_async_read_barrier();
uint32_t* d0 = reinterpret_cast<uint32_t*>(get_write_ptr(cb_p2));
asm volatile("" ::: "memory");
if constexpr (U16) {
kpre_build_row_u16<PERIODS, SWAP>(reinterpret_cast<const uint8_t*>(row_l1),
reinterpret_cast<uint16_t*>(d0));
} else {
kpre_build_row<PERIODS, SWAP>(reinterpret_cast<const uint8_t*>(row_l1),
reinterpret_cast<const uint32_t*>(lut_l1), d0);
}
asm volatile("" ::: "memory");
cb_push_back(cb_p2, 5);
}
}
}
const uint32_t a = y * W > g.p0 ? y * W : g.p0;
const uint32_t e = (y + 1) * W < g.p1 ? (y + 1) * W : g.p1;
uint32_t* dst = reinterpret_cast<uint32_t*>(out_addr + (a - shard0) * C_PAD * 2);
uint32_t n = e - a;
if (y >= RH) {
for (uint32_t i = 0; i < n * 4; ++i) {
dst[i] = 0;
}
continue;
}
cb_wait_front(cb_v, QT);
// V row in phase-planar order: pixel x = 4k + p, channel c at [p * 1024 + 3k + c] (bf16)
const uint16_t* v = reinterpret_cast<const uint16_t*>(get_read_ptr(cb_v));
asm volatile("" ::: "memory");
#ifdef KPRE_NOWRITE
n = 0;
#endif
for (uint32_t x = a - y * W; n; ++x, --n) {
const uint16_t* sp = v + (x & 3u) * 1024 + (x >> 2) * 3;
dst[0] = static_cast<uint32_t>(sp[0]) | (static_cast<uint32_t>(sp[1]) << 16);
dst[1] = sp[2];
dst[2] = 0;
dst[3] = 0;
dst += 4;
}
asm volatile("" ::: "memory");
cb_pop_front(cb_v, QT);
}
}