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 /dcn_k_compute.cpp
changh95's picture
changh95
tt-model push bevformer-p150 (container)
55474c5 verified
Raw History Blame Contribute Delete
13.5 kB
// SPDX-License-Identifier: Apache-2.0
// dcn_k (tt/dcn_k.py), compute. Per work item (32 pixels x one tap; see dcn_k_reader.cpp):
// 1. DEST0..2 <- cb_geo (fx, fy, mask in column 0; fp32, UnpackToDestFp32); SFPU: the 4 corner weights
// mask * wx * wy (fp32) -> cb_w (4 fp32 tiles, column 0 valid, the other columns 0).
// 2. for each corner c: tilize its 32 gathered rows (cb_rm0 for NW / NE, cb_rm1 for SW / SE) into CT tiles (cb_xt),
// DEST = xt * bcast_col(W_c) (HiFi4; the fp32 weight as tf32), packer L1 accumulation into cb_acc (CT fp32).
// 3. cb_acc -> bf16 -> cb_out (CT tiles: the item's im2col tiles).
#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/matmul.h"
#include "api/compute/pack.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/tilize.h"
#ifdef DCN_K_GLUE
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/typecast.h"
#endif
#if defined(TRISC_MATH)
#include "llk_math_eltwise_unary_sfpu_macros.h"
#include "sfpi.h"
namespace ckernel::sfpu {
inline void dcn_k_w_face() {
#pragma GCC unroll 2
for (int d = 0; d < 8; d++) {
sfpi::vFloat fx = sfpi::dst_reg[0];
sfpi::vFloat fy = sfpi::dst_reg[32];
sfpi::vFloat a = sfpi::dst_reg[64];
sfpi::vFloat gx = 1.0f - fx;
sfpi::vFloat ay0 = a * (1.0f - fy);
sfpi::vFloat ay1 = a * fy;
sfpi::dst_reg[0] = gx * ay0;
sfpi::dst_reg[32] = fx * ay0;
sfpi::dst_reg[64] = gx * ay1;
sfpi::dst_reg[96] = fx * ay1;
sfpi::dst_reg++;
}
}
} // namespace ckernel::sfpu
#endif
constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t NWORK = get_compile_time_arg_val(2);
constexpr uint32_t NC = get_compile_time_arg_val(3);
constexpr uint32_t GX = get_compile_time_arg_val(4);
constexpr uint32_t cb_rm0 = 0, cb_rm1 = 1, cb_geo = 2, cb_xt = 3, cb_w = 4, cb_acc = 5, cb_out = 16;
constexpr uint32_t DB = 4; // tiles per DEST section (fp32)
#ifdef DCN_K_GLUE
// BEVFORMER_DCN_GLUE=kernel (D20 round 2): per new block, glue_k's MODE-1 steps (loc = int32(off * S + B) with the
// SFPU mul / RNE add / typecast, sig = sigmoid(off)) -> loc to both readers (cb 6 / 7), sig to PART 0 (cb 8)
constexpr uint32_t cb_l0 = 6, cb_l1 = 7, cb_sg = 8, cb_goff = 9, cb_gb = 10, cb_gs = 11;
constexpr uint32_t F32 = static_cast<uint32_t>(DataFormat::Float32);
constexpr uint32_t I32 = static_cast<uint32_t>(DataFormat::Int32);
inline void glue_block() {
cb_wait_front(cb_goff, 1);
cb_wait_front(cb_gb, 1);
tile_regs_acquire();
reconfig_data_format_srca(cb_goff);
copy_tile_to_dst_init_short(cb_goff);
copy_tile(cb_goff, 0, 0);
reconfig_data_format_srca(cb_gs);
copy_tile_to_dst_init_short(cb_gs);
copy_tile(cb_gs, 0, 1);
mul_binary_tile_init();
mul_binary_tile(0, 1, 0);
reconfig_data_format_srca(cb_gb);
copy_tile_to_dst_init_short(cb_gb);
copy_tile(cb_gb, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
reconfig_data_format_srca(cb_goff);
copy_tile_to_dst_init_short(cb_goff);
copy_tile(cb_goff, 0, 2);
sigmoid_tile_init<false>();
sigmoid_tile<VectorMode::RC, false>(2);
typecast_tile_init<F32, I32>();
typecast_tile<F32, I32>(0);
tile_regs_commit();
cb_reserve_back(cb_l0, 1);
cb_reserve_back(cb_l1, 1);
cb_reserve_back(cb_sg, 1);
tile_regs_wait();
pack_reconfig_data_format(cb_l0);
pack_tile(0, cb_l0);
pack_tile(0, cb_l1);
pack_reconfig_data_format(cb_sg);
pack_tile(2, cb_sg);
tile_regs_release();
cb_push_back(cb_l0, 1);
cb_push_back(cb_l1, 1);
cb_push_back(cb_sg, 1);
cb_pop_front(cb_goff, 1);
cb_pop_front(cb_gb, 1);
}
#endif
static_assert(CT % DB == 0, "dcn_k: C a multiple of 128");
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_geo, cb_geo, cb_w);
#ifdef DCN_K_GLUE
cb_wait_front(cb_gs, 1);
uint32_t cur_b = 0xFFFFFFFF;
#endif
for (uint32_t w = w0; w < w1; ++w) {
#ifdef DCN_K_GLUE
if (w / 9 != cur_b) {
cur_b = w / 9;
glue_block();
}
#endif
#ifdef DCN_K_NEAREST
// BEVFORMER_DCN_SAMPLE=nearest (D20 round 3): one gathered corner per item (even items from cb_rm0, odd from
// cb_rm1), tilized, times the mask (cb_geo tile, column 0; column-broadcast multiply), packed to bf16
{
const uint32_t src = ((w - w0) & 1) ? cb_rm1 : cb_rm0;
cb_wait_front(src, CT);
cb_reserve_back(cb_xt, CT);
reconfig_data_format_srca(src);
pack_reconfig_data_format(cb_xt);
#ifdef DCN_K_FAST_TILIZE
fast_tilize_init(src, CT, cb_xt);
fast_tilize_block(src, CT, cb_xt, 0, 0);
fast_tilize_uninit(src, cb_xt, CT);
#else
tilize_init(src, CT, cb_xt);
tilize_block(src, CT, cb_xt, 0, 0);
tilize_uninit(src, cb_xt);
#endif
cb_push_back(cb_xt, CT);
cb_pop_front(src, CT);
cb_wait_front(cb_xt, CT);
cb_wait_front(cb_geo, 1);
cb_reserve_back(cb_out, CT);
reconfig_data_format_srca(cb_xt);
reconfig_data_format_srcb(cb_geo);
pack_reconfig_data_format(cb_out);
mul_bcast_cols_init(cb_xt, cb_geo);
for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
tile_regs_acquire();
for (uint32_t i = 0; i < DB; ++i) {
mul_tiles_bcast_cols(cb_xt, cb_geo, t0 + i, 0, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < DB; ++i) {
pack_tile<true>(i, cb_out, t0 + i);
}
tile_regs_release();
}
cb_push_back(cb_out, CT);
cb_pop_front(cb_xt, CT);
cb_pop_front(cb_geo, 1);
}
continue;
#endif
// ---- 1. corner weights
cb_wait_front(cb_geo, 3);
cb_reserve_back(cb_w, 4);
tile_regs_acquire();
reconfig_data_format_srca(cb_geo);
copy_init(cb_geo);
copy_tile(cb_geo, 0, 0);
copy_tile(cb_geo, 1, 1);
copy_tile(cb_geo, 2, 2);
MATH((SFPU_UNARY_INIT(unused)));
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, dcn_k_w_face, 0, VectorMode::RC)));
tile_regs_commit();
tile_regs_wait();
pack_reconfig_data_format(cb_w);
for (uint32_t c = 0; c < 4; ++c) {
pack_tile(c, cb_w);
}
tile_regs_release();
cb_push_back(cb_w, 4);
cb_pop_front(cb_geo, 3);
cb_wait_front(cb_w, 4);
// ---- 2. the 4 corners
#ifdef DCN_K_BATCH
// BEVFORMER_DCN_BATCH=1: the 4 corners tilized back to back (one tilize init / uninit per item), then the
// 4 weighted multiplies with one init (the same products, the same accumulation order)
#ifndef DCN_K_MM
cb_reserve_back(cb_acc, CT);
#endif
cb_wait_front(cb_rm0, 2 * CT);
cb_wait_front(cb_rm1, 2 * CT);
cb_reserve_back(cb_xt, 4 * CT);
reconfig_data_format_srca(cb_rm0);
pack_reconfig_data_format(cb_xt);
#ifdef DCN_K_NOTILIZE
// (anatomy probe only: no tilize, garbage tiles)
#elif defined(DCN_K_FAST_TILIZE)
fast_tilize_init(cb_rm0, CT, cb_xt);
for (uint32_t c = 0; c < 4; ++c) {
fast_tilize_block(c < 2 ? cb_rm0 : cb_rm1, CT, cb_xt, (c & 1) * CT, c * CT);
}
fast_tilize_uninit(cb_rm0, cb_xt, CT);
#elif defined(DCN_K_PAIRS)
// BEVFORMER_DCN_PAIRS: each reader's block is 32 rows x 2C ([NW | NE] / [SW | SE]): 2 CT tiles wide
tilize_init(cb_rm0, 2 * CT, cb_xt);
tilize_block(cb_rm0, 2 * CT, cb_xt, 0, 0);
tilize_block(cb_rm1, 2 * CT, cb_xt, 0, 2 * CT);
tilize_uninit(cb_rm0, cb_xt);
#else
tilize_init(cb_rm0, CT, cb_xt);
for (uint32_t c = 0; c < 4; ++c) {
tilize_block(c < 2 ? cb_rm0 : cb_rm1, CT, cb_xt, (c & 1) * CT, c * CT);
}
tilize_uninit(cb_rm0, cb_xt);
#endif
cb_push_back(cb_xt, 4 * CT);
cb_pop_front(cb_rm0, 2 * CT);
cb_pop_front(cb_rm1, 2 * CT);
cb_wait_front(cb_xt, 4 * CT);
#ifdef DCN_K_MM
// BEVFORMER_DCN_SUM=mm (D20 round 2): the reader put fx / fy / mask on the tile diagonal, so cb_w holds the 4
// corner weights as diagonal matrices diag(w_c) (fp32); the weighted corner sum is one DEST accumulation of
// 4 matmuls per output tile, out = sum_c diag(w_c) @ X_c, packed once to bf16 (no packer-L1 passes, no copy)
cb_reserve_back(cb_out, CT);
reconfig_data_format_srca(cb_xt);
reconfig_data_format_srcb(cb_w);
matmul_init(cb_w, cb_xt);
pack_reconfig_data_format(cb_out);
for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
tile_regs_acquire();
#ifndef DCN_K_NOMUL
for (uint32_t c = 0; c < 4; ++c) {
for (uint32_t i = 0; i < DB; ++i) {
matmul_tiles(cb_w, cb_xt, c, c * CT + t0 + i, i);
}
}
#endif
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < DB; ++i) {
pack_tile<true>(i, cb_out, t0 + i);
}
tile_regs_release();
}
cb_push_back(cb_out, CT);
cb_pop_front(cb_xt, 4 * CT);
cb_pop_front(cb_w, 4);
continue;
#endif
reconfig_data_format_srca(cb_xt);
reconfig_data_format_srcb(cb_w);
pack_reconfig_data_format(cb_acc);
mul_bcast_cols_init(cb_xt, cb_w);
for (uint32_t c = 0; c < 4; ++c) {
pack_reconfig_l1_acc(c == 0 ? 0 : 1);
for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
tile_regs_acquire();
for (uint32_t i = 0; i < DB; ++i) {
mul_tiles_bcast_cols(cb_xt, cb_w, c * CT + t0 + i, c, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < DB; ++i) {
pack_tile<true>(i, cb_acc, t0 + i);
}
tile_regs_release();
}
}
pack_reconfig_l1_acc(0);
cb_pop_front(cb_xt, 4 * CT);
#else
cb_reserve_back(cb_acc, CT);
for (uint32_t c = 0; c < 4; ++c) {
const uint32_t src = c < 2 ? cb_rm0 : cb_rm1;
if ((c & 1) == 0) {
cb_wait_front(src, 2 * CT);
}
cb_reserve_back(cb_xt, CT);
reconfig_data_format_srca(src);
pack_reconfig_data_format(cb_xt);
#ifndef DCN_K_NOTILIZE
#ifdef DCN_K_FAST_TILIZE
// BEVFORMER_DCN_TILIZE=fast: the BH fast tilize (unpack row chunks of 4 tiles; a lossless bf16 reorder)
fast_tilize_init(src, CT, cb_xt);
fast_tilize_block(src, CT, cb_xt, (c & 1) * CT, 0);
fast_tilize_uninit(src, cb_xt, CT);
#else
tilize_init(src, CT, cb_xt);
tilize_block(src, CT, cb_xt, (c & 1) * CT, 0);
tilize_uninit(src, cb_xt);
#endif
#endif
cb_push_back(cb_xt, CT);
cb_wait_front(cb_xt, CT);
reconfig_data_format_srca(cb_xt);
reconfig_data_format_srcb(cb_w);
pack_reconfig_data_format(cb_acc);
mul_bcast_cols_init(cb_xt, cb_w);
pack_reconfig_l1_acc(c == 0 ? 0 : 1);
#ifdef DCN_K_NOMUL
for (uint32_t t0 = 0; t0 < 0; t0 += DB) {
#else
for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
#endif
tile_regs_acquire();
for (uint32_t i = 0; i < DB; ++i) {
mul_tiles_bcast_cols(cb_xt, cb_w, t0 + i, c, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < DB; ++i) {
pack_tile<true>(i, cb_acc, t0 + i);
}
tile_regs_release();
}
pack_reconfig_l1_acc(0);
cb_pop_front(cb_xt, CT);
if (c & 1) {
cb_pop_front(src, 2 * CT);
}
}
#endif
cb_push_back(cb_acc, CT);
cb_pop_front(cb_w, 4);
// ---- 3. fp32 -> bf16 im2col tiles
cb_wait_front(cb_acc, CT);
cb_reserve_back(cb_out, CT);
reconfig_data_format_srca(cb_acc);
copy_init(cb_acc);
pack_reconfig_data_format(cb_out);
for (uint32_t t0 = 0; t0 < CT; t0 += DB) {
tile_regs_acquire();
for (uint32_t i = 0; i < DB; ++i) {
copy_tile(cb_acc, t0 + i, i);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t i = 0; i < DB; ++i) {
pack_tile<true>(i, cb_out, t0 + i);
}
tile_regs_release();
}
cb_push_back(cb_out, CT);
cb_pop_front(cb_acc, CT);
}
}