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_compute.cpp
changh95's picture
changh95
tt-model push bevformer-p150 (container)
55474c5 verified
Raw History Blame Contribute Delete
11.6 kB
// SPDX-License-Identifier: Apache-2.0
// msda_k (tt/msda_k.py), compute. Per work item (32 queries x one head block; see msda_k_reader.cpp):
// 1. for each half h (the points of RISC h): DEST0..2 <- cb_geo_h (fx, fy, attn; fp32, UnpackToDestFp32); SFPU:
// W_c = a * wx * wy for the 4 corners (fp32) -> cb_w pages 4h .. 4h + 3 (fp32).
// 2. for each step (point p, corner c): DEST0 = W_{h,c} @ E_p (broadcast of column p over the 32 channels, exact),
// DEST0 *= V (the gathered bf16 value tile; DEST -> SrcB reuse), packer L1 accumulation into cb_acc (fp32).
// 3. cb_acc (fp32, UnpackToDestFp32) -> bf16 -> cb_out.
// CBs: cb_geo_h 3 pages / item (6 total), cb_in_h 2P pages / item (4P total), cb_w 8 (once per item), cb_acc 1,
// cb_e P (waited once), cb_out 1 per item (2 total).
#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/eltwise_binary.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"
#include "api/compute/bcast.h"
#include "api/compute/reduce.h"
#include "api/compute/softmax.h"
#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_compute.hpp"
#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/chain.hpp"
#include "ttnn/cpp/ttnn/kernel_lib/eltwise/api/convenience.hpp"
#include "ttnn/cpp/ttnn/kernel_lib/eltwise/unary/math.hpp"
#if defined(TRISC_MATH)
#include "llk_math_eltwise_unary_sfpu_macros.h"
#include "sfpi.h"
namespace ckernel::sfpu {
// DEST tiles 0, 1, 2 = fx, fy, attn -> tiles 0..3 = the NW, NE, SW, SE weights (one face per call)
inline void msda_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++;
}
}
// SM: DEST 0, 1, 2, 3 = fx, fy, softmax, half mask -> attn = softmax * mask (exact: x 1 or x 0), then as above
inline void msda_k_w_face_sm() {
#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::dst_reg[96];
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
namespace ckl = compute_kernel_lib;
constexpr uint32_t P = 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 SKIP = get_compile_time_arg_val(5); // strided items (msda_k_reader.cpp)
// SM (BEVFORMER_MSDA_SOFTMAX=kernel): the softmax of the head's logit tile (cb_lg, from the reader) runs here, the
// same LLK sequence (and CB formats / unpack modes) as ttnn.softmax(numeric_stable) on one fp32 tile (attention
// softmax.cpp, Wt = 1, block 1): row max -> x - max, exp -> row sum, 1 / sum -> x * (1 / sum)
constexpr uint32_t SM = get_compile_time_arg_val(6);
// PAIRS (VSRC 2): the readers deliver each point's corners as two [32 x 64] row-major blocks (NW | NE, SW | SE);
// they are tilized (exact) into cb_tl0 / cb_tl1 in the step order (point, corner) before the steps
constexpr uint32_t PAIRS = get_compile_time_arg_val(7);
constexpr uint32_t cb_tl0 = 19, cb_tl1 = 20;
constexpr uint32_t cb_in0 = 0, cb_in1 = 1, cb_geo0 = 2, cb_geo1 = 3, cb_e = 4, cb_w = 5, cb_acc = 6, cb_out = 16;
constexpr uint32_t cb_lg = 9, cb_scmax = 10, cb_scsum = 11, cb_msk = 12, cb_sm = 13, cb_max = 14, cb_exps = 15,
cb_rsum = 17;
static inline void softmax_tile() {
reconfig_data_format(cb_lg, cb_lg);
pack_reconfig_data_format(cb_exps);
copy_init(cb_lg);
compute_kernel_lib::reduce<PoolType::MAX, ReduceDim::REDUCE_ROW, cb_lg, cb_scmax, cb_max,
compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop,
compute_kernel_lib::ReduceDataFormatReconfigMode::INPUT>(
compute_kernel_lib::ReduceInputBlockShape::row(1));
ckl::eltwise_chain(
ckl::IterationShape::tiles(1).block_size(1),
ckl::BinaryFpu<ckl::BinaryFpuOp::Sub,
ckl::input(cb_lg, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd, ckl::InputTileMapping::Block,
ckl::DataFormatReconfig::Disabled),
ckl::input(cb_max, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd)>{},
ckl::Exp<static_cast<ckl::Approx>(0), ckl::Dst::D0>{},
ckl::PackTile<ckl::output(cb_exps, ckl::ReservePolicy::PerBlockSize, ckl::PushPolicy::PerBlockSize,
ckl::DataFormatReconfig::Disabled)>{});
cb_wait_front(cb_exps, 1);
compute_kernel_lib::reduce<PoolType::SUM, ReduceDim::REDUCE_ROW, cb_exps, cb_scsum, cb_rsum,
compute_kernel_lib::ReduceInputPolicy::WaitUpfrontNoPop>(
compute_kernel_lib::ReduceInputBlockShape::row(1), compute_kernel_lib::ReduceInputMemoryLayout::contiguous(),
compute_kernel_lib::NoAccumulation{}, [](uint32_t) {
recip_tile_init<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>();
recip_tile<ReciprocalDestAcc::FP32, ReciprocalApproxMode::Precise>(0);
});
ckl::mul<ckl::input(cb_exps, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd, ckl::InputTileMapping::Block),
ckl::input(cb_rsum, ckl::BroadcastDim::Col, ckl::WaitPolicy::Upfront, ckl::PopPolicy::AtEnd),
ckl::output(cb_sm, ckl::ReservePolicy::PerBlockSize, ckl::PushPolicy::PerBlockSize)>(
ckl::IterationShape::tiles(1).block_size(1));
}
constexpr uint32_t PH = P / 2;
constexpr uint32_t NS = PH * 4;
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;
}
if constexpr (SM) {
compute_kernel_hw_startup(cb_lg, cb_scmax, cb_exps);
cb_wait_front(cb_scmax, 1);
cb_wait_front(cb_scsum, 1);
cb_wait_front(cb_msk, 2);
} else {
compute_kernel_hw_startup(cb_geo0, cb_geo0, cb_w);
}
cb_wait_front(cb_e, P);
for (uint32_t w = w0; w < w1; w += wstep) {
if constexpr (SM) {
softmax_tile();
cb_wait_front(cb_sm, 1);
}
// ---- 1. corner weights
cb_reserve_back(cb_w, 8);
for (uint32_t h = 0; h < 2; ++h) {
const uint32_t cg = h ? cb_geo1 : cb_geo0;
cb_wait_front(cg, 3);
tile_regs_acquire();
reconfig_data_format_srca(cg);
copy_init(cg);
copy_tile(cg, 0, 0);
copy_tile(cg, 1, 1);
if constexpr (SM) {
reconfig_data_format_srca(cb_sm);
copy_init(cb_sm);
copy_tile(cb_sm, 0, 2);
reconfig_data_format_srca(cb_msk);
copy_init(cb_msk);
copy_tile(cb_msk, h, 3);
MATH((SFPU_UNARY_INIT(unused)));
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, msda_k_w_face_sm, 0,
VectorMode::RC)));
} else {
copy_tile(cg, 2, 2);
MATH((SFPU_UNARY_INIT(unused)));
MATH((SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, msda_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<true>(c, cb_w, h * 4 + c);
}
tile_regs_release();
cb_pop_front(cg, 3);
}
cb_push_back(cb_w, 8);
if constexpr (SM) {
cb_pop_front(cb_sm, 1);
}
cb_wait_front(cb_w, 8);
if constexpr (PAIRS) {
for (uint32_t h = 0; h < 2; ++h) {
const uint32_t cin = h ? cb_in1 : cb_in0;
const uint32_t ct = h ? cb_tl1 : cb_tl0;
cb_wait_front(cin, NS);
cb_reserve_back(ct, NS);
reconfig_data_format_srca(cin);
pack_reconfig_data_format(ct);
#ifdef MSDA_K_FAST_TILIZE
// BEVFORMER_MSDA_TILIZE=fast: the BH fast tilize (a lossless bf16 reorder)
fast_tilize_init(cin, 2, ct);
for (uint32_t bk = 0; bk < NS / 2; ++bk) {
fast_tilize_block(cin, 2, ct, 2 * bk, 2 * bk);
}
fast_tilize_uninit(cin, ct, 2);
#else
tilize_init(cin, 2, ct);
for (uint32_t bk = 0; bk < NS / 2; ++bk) {
tilize_block(cin, 2, ct, 2 * bk, 2 * bk);
}
tilize_uninit(cin, ct);
#endif
cb_push_back(ct, NS);
cb_pop_front(cin, NS);
}
}
// ---- 2. the 4P (point, corner) steps
cb_reserve_back(cb_acc, 1);
reconfig_data_format_srca(cb_e);
reconfig_data_format_srcb(cb_w);
pack_reconfig_data_format(cb_acc);
for (uint32_t h = 0; h < 2; ++h) {
const uint32_t ci_ = PAIRS ? (h ? cb_tl1 : cb_tl0) : (h ? cb_in1 : cb_in0);
cb_wait_front(ci_, NS);
#ifdef MSDA_K_NOSTEPS
for (uint32_t k = 0; k < 1; ++k) {
#else
for (uint32_t k = 0; k < NS; ++k) {
#endif
const uint32_t p = h * PH + (k >> 2);
const uint32_t c = k & 3;
tile_regs_acquire();
matmul_init(cb_w, cb_e);
matmul_tiles(cb_w, cb_e, h * 4 + c, p, 0);
mul_reuse_dest_init<EltwiseBinaryReuseDestType::DEST_TO_SRCB>(ci_);
mul_reuse_dest_tiles<EltwiseBinaryReuseDestType::DEST_TO_SRCB>(ci_, k, 0);
tile_regs_commit();
tile_regs_wait();
pack_reconfig_l1_acc((h == 0 && k == 0) ? 0 : 1);
pack_tile<true>(0, cb_acc, 0);
tile_regs_release();
}
cb_pop_front(ci_, NS);
}
pack_reconfig_l1_acc(0);
cb_push_back(cb_acc, 1);
cb_pop_front(cb_w, 8);
// ---- 3. fp32 sum -> bf16 output tile
cb_wait_front(cb_acc, 1);
cb_reserve_back(cb_out, 1);
tile_regs_acquire();
reconfig_data_format_srca(cb_acc);
copy_init(cb_acc);
copy_tile(cb_acc, 0, 0);
tile_regs_commit();
tile_regs_wait();
pack_reconfig_data_format(cb_out);
pack_tile(0, cb_out);
tile_regs_release();
cb_push_back(cb_out, 1);
cb_pop_front(cb_acc, 1);
}
}