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_compute.cpp
changh95's picture
changh95
tt-model push bevformer-p150 (container)
55474c5 verified
Raw History Blame Contribute Delete
7.53 kB
// SPDX-License-Identifier: Apache-2.0
// kpre (tt/kpre.py), compute. The C14 x0.8 INTER_LINEAR resize of the normalised node image on the SFPU, in the
// exact float32 steps of the host twin (host.fast; each lerp is fma(x1 - x0, w, x0): the difference rounded to
// fp32, one fused multiply-add rounded once; probe logs/bevformer/opt_r4/probe_sfpu: SFPU sub / SFPMAD are the
// correctly rounded fp32 results on this domain, 0 of 262144 differ):
// H row r (4 fp32 pseudo-tiles): H = (P1 - P0) * WX + P0 (P0 / P1 / WX from the reader)
// output row y (4 bf16 pseudo-tiles): V = (H[r0 + 1] - H[r0]) * WY[y % 4] + H[r0], rounded to bf16 (RNE, as
// torch's float32 -> bfloat16 cast of the host path) on the SFPU, then packed (exact)
// cb_h keeps all H rows of the core (no wrap of the 2-row window).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#include "kpre_common.h"
#if defined(TRISC_MATH)
#include "llk_math_eltwise_unary_sfpu_macros.h"
#include "sfpi.h"
namespace ckernel::sfpu {
// DEST tiles 0 / 1 / 2 = x0, x1, w -> tile 0 = fma(x1 - x0, w, x0)
inline void kpre_lerp() {
#pragma GCC unroll 1
for (int d = 0; d < 8; d++) {
sfpi::vFloat x0 = sfpi::dst_reg[0];
sfpi::vFloat x1 = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::vFloat dx = x1 - x0;
sfpi::dst_reg[0] = dx * w + x0;
sfpi::dst_reg++;
}
}
// U16 rows: DEST tiles 0 / 1 = the uint16 bytes of planes p / p + 1 (low 16 bits of each 32-bit word), 2 = the
// per-lane -mean_c, 3 = w -> tile 0 = fma(x1 - x0, w, x0) with x = f32(u) + (-mean_c) (exact int -> fp32 cast, then
// the single fp32 add of the host's normalisation)
inline void kpre_lerp_u16() {
#pragma GCC unroll 1
for (int d = 0; d < 8; d++) {
sfpi::vInt u0 = sfpi::dst_reg[0];
sfpi::vInt u1 = sfpi::dst_reg[32];
sfpi::vFloat beta = sfpi::dst_reg[64];
sfpi::vFloat w = sfpi::dst_reg[96];
sfpi::vFloat x0 = sfpi::convert<sfpi::vFloat>(u0 & 0xFFFF) + beta;
sfpi::vFloat x1 = sfpi::convert<sfpi::vFloat>(u1 & 0xFFFF) + beta;
sfpi::vFloat dx = x1 - x0;
sfpi::dst_reg[0] = dx * w + x0;
sfpi::dst_reg++;
}
}
// the same, then round to bf16 (RNE: + 0x7fff + lsb, low 16 bits cleared; the bf16 pack is then exact)
inline void kpre_lerp_bf16() {
#pragma GCC unroll 1
for (int d = 0; d < 8; d++) {
sfpi::vFloat x0 = sfpi::dst_reg[0];
sfpi::vFloat x1 = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::vFloat dx = x1 - x0;
sfpi::vFloat v = dx * w + x0;
sfpi::vUInt bits = sfpi::as<sfpi::vUInt>(v);
sfpi::vUInt lsb = (bits >> 16) & 1;
bits = bits + 0x7fffU + lsb;
bits = bits & 0xFFFF0000U;
sfpi::dst_reg[0] = sfpi::as<sfpi::vFloat>(bits);
sfpi::dst_reg++;
}
}
} // namespace ckernel::sfpu
#endif
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: odd rows (from rlo) come from the writer RISC in cb_p2
constexpr uint32_t U16 = get_compile_time_arg_val(10); // 1: the P rows are raw uint16 planes (+ cb_beta)
constexpr uint32_t cb_p = 0, cb_wx = 1, cb_wy = 2, cb_h = 3, cb_p2 = 5, cb_beta = 7, cb_v = 16;
constexpr uint32_t QT = 4; // pseudo-tiles per row
template <bool BF16>
inline void lerp_rows(uint32_t cb_x, uint32_t i0, uint32_t i1, uint32_t cb_w, uint32_t iw, uint32_t cb_o) {
tile_regs_acquire();
#ifndef KPRE_NOCOPY
copy_tile(cb_x, i0, 0);
copy_tile(cb_x, i1, 1);
copy_tile(cb_w, iw, 2);
#endif
#ifndef KPRE_NOSFPU
if constexpr (BF16) {
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp_bf16, 0, VectorMode::RC)));
} else {
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp, 0, VectorMode::RC)));
}
#endif
tile_regs_commit();
tile_regs_wait();
pack_tile(0, cb_o);
tile_regs_release();
}
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.rows) {
return;
}
compute_kernel_hw_startup(cb_p, cb_h);
copy_init(cb_p);
MATH((SFPU_UNARY_INIT(unused)));
cb_wait_front(cb_wx, QT);
cb_wait_front(cb_wy, QT);
if constexpr (U16) {
cb_wait_front(cb_beta, 1);
}
// cb_h holds every H row of the core (rlo .. rhi, sized on the host): rows are never popped before the end, so
// the two rows an output row reads are always contiguous pages (a FIFO window could wrap the ring)
uint32_t produced = g.rlo; // next H row to compute
for (uint32_t y = g.ya; y <= g.yv; ++y) {
const uint32_t r0 = kpre_r0<RA, RB, RD>(y);
while (produced <= r0 + 1) {
// H row `produced`: (P1 - P0) * WX + P0, fp32
// phase planes s = 0..4 of the node row: H plane p = lerp(plane p, plane p + 1, (2p + 1) / 8)
const uint32_t cbp = (SPLIT && ((produced - g.rlo) & 1)) ? cb_p2 : cb_p;
cb_wait_front(cbp, QT + 1);
cb_reserve_back(cb_h, QT);
pack_reconfig_data_format(cb_h);
if constexpr (U16) {
for (uint32_t t = 0; t < QT; ++t) {
tile_regs_acquire();
copy_tile_to_dst_init_short_with_dt(cb_wx, cbp);
copy_tile(cbp, t, 0);
copy_tile(cbp, t + 1, 1);
copy_tile_to_dst_init_short_with_dt(cbp, cb_wx);
copy_tile(cb_beta, 0, 2);
copy_tile(cb_wx, t, 3);
#ifndef KPRE_NOSFPU
MATH((SFPU_UNARY_INIT(unused)));
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, kpre_lerp_u16, 0,
VectorMode::RC)));
#endif
tile_regs_commit();
tile_regs_wait();
pack_tile(0, cb_h);
tile_regs_release();
}
} else {
for (uint32_t t = 0; t < QT; ++t) {
lerp_rows<false>(cbp, t, t + 1, cb_wx, t, cb_h);
}
}
cb_push_back(cb_h, QT);
cb_pop_front(cbp, QT + 1);
++produced;
}
// output row y: (H[r0 + 1] - H[r0]) * WY + H[r0] -> bf16
const uint32_t h0 = (r0 - g.rlo) * QT;
cb_wait_front(cb_h, h0 + 2 * QT);
cb_reserve_back(cb_v, QT);
pack_reconfig_data_format(cb_v);
for (uint32_t t = 0; t < QT; ++t) {
lerp_rows<true>(cb_h, h0 + t, h0 + QT + t, cb_wy, y % 4, cb_v);
}
cb_push_back(cb_v, QT);
}
cb_pop_front(cb_h, (produced - g.rlo) * QT);
if constexpr (U16) {
cb_pop_front(cb_beta, 1);
}
}