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 /scatter_k_compute.cpp
changh95's picture
changh95
tt-model push bevformer-p150 (container)
55474c5 verified
Raw History Blame Contribute Delete
4.53 kB
// SPDX-License-Identifier: Apache-2.0
// scatter_k (tt/scatter_k.py), compute. Per 32-row block (DT tiles), the steps of the stock ops in their order, each
// result packed to bf16 like the stock ops' bf16 outputs (16-bit DEST, fp32 dest accumulation off, as binary_ng runs
// bf16 operands):
// s1 = g0 + g1, s2 = s1 + g2: the FPU add (add_tiles);
// out = s2 * scale: the SFPU multiply (both operands copied to DEST, mul_binary_tile) with the scale column broadcast
// over the row by the reader (binary_ng's column-broadcast fill).
// Which unit each stock op used was measured against the stock ops (logs/bevformer/opt_r8/scatp_*.log): the SFPU add
// rounds ties to even where the stock add does not; the FPU multiply rounds differently from the stock multiply.
#include <cstdint>
#include "api/compute/bcast.h"
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/eltwise_binary.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/pack.h"
constexpr uint32_t DT = get_compile_time_arg_val(0);
constexpr uint32_t NWORK = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t cb_g0 = 0, cb_g1 = 1, cb_g2 = 2, cb_sc = 3, cb_s1 = 4, cb_s2 = 5, cb_out = 16;
constexpr uint32_t SB = 4; // tiles per DEST batch (operand pairs: 8 of the 16-bit DEST half)
static_assert(DT % SB == 0, "scatter_k: D a multiple of 128");
// s = a + b: binary_ng's FPU add of bf16 operands (add_tiles; measured: the SFPU add's RNE does not reproduce it)
static inline void add_step(uint32_t a, uint32_t b, uint32_t o) {
cb_wait_front(a, DT);
cb_wait_front(b, DT);
cb_reserve_back(o, DT);
reconfig_data_format(a, b);
pack_reconfig_data_format(o);
add_init(a, b);
for (uint32_t t0 = 0; t0 < DT; t0 += SB) {
tile_regs_acquire();
for (uint32_t i = 0; i < SB; ++i) {
add_tiles(a, b, t0 + i, t0 + i, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < SB; ++i) {
pack_tile<true>(i, o, t0 + i);
}
tile_regs_release();
}
cb_push_back(o, DT);
cb_pop_front(a, DT);
cb_pop_front(b, DT);
}
// o = a * bcast: binary_ng's SFPU multiply (copy both to DEST, mul_binary_tile; the broadcast tile filled by the
// reader; measured: the FPU multiply does not reproduce it)
static inline void mul_step(uint32_t a, uint32_t b, uint32_t o) {
cb_wait_front(a, DT);
cb_wait_front(b, 1);
cb_reserve_back(o, DT);
pack_reconfig_data_format(o);
for (uint32_t t0 = 0; t0 < DT; t0 += SB) {
tile_regs_acquire();
reconfig_data_format_srca(a);
copy_init(a);
for (uint32_t i = 0; i < SB; ++i) {
copy_tile(a, t0 + i, 2 * i);
}
reconfig_data_format_srca(b);
copy_init(b);
for (uint32_t i = 0; i < SB; ++i) {
copy_tile(b, 0, 2 * i + 1);
mul_binary_tile_init();
mul_binary_tile(2 * i, 2 * i + 1, 2 * i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < SB; ++i) {
pack_tile<true>(2 * i, o, t0 + i);
}
tile_regs_release();
}
cb_push_back(o, DT);
cb_pop_front(a, DT);
cb_pop_front(b, 1);
}
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;
}
compute_kernel_hw_startup(cb_g0, cb_g1, cb_out);
for (uint32_t w = w0; w < w1; ++w) {
#if defined(SCATTER_K_STOP) && SCATTER_K_STOP == 1 // debug: out = g0 + g1 (g2 / scale consumed unused)
add_step(cb_g0, cb_g1, cb_out);
cb_wait_front(cb_g2, DT);
cb_pop_front(cb_g2, DT);
cb_wait_front(cb_sc, 1);
cb_pop_front(cb_sc, 1);
#elif defined(SCATTER_K_STOP) && SCATTER_K_STOP == 2 // debug: out = (g0 + g1) + g2
add_step(cb_g0, cb_g1, cb_s1);
add_step(cb_s1, cb_g2, cb_out);
cb_wait_front(cb_sc, 1);
cb_pop_front(cb_sc, 1);
#else
add_step(cb_g0, cb_g1, cb_s1);
add_step(cb_s1, cb_g2, cb_s2);
mul_step(cb_s2, cb_sc, cb_out);
#endif
}
}