Hugging Face's logo

changh95
/
sceneseg-p150

Image Segmentation
ONNX
blackhole
p150
tt-dit-server
tt-model-cache
tt-model-catalog
tt-model-container
Model card Files
xet
 Community
sceneseg-p150 / code /tt_sceneseg /tt /kernels /dw_compute.cpp
changh95's picture
changh95
tt-model push sceneseg-p150 (container)
e0adbab verified
Raw History Blame Contribute Delete
7.62 kB
// SPDX-License-Identifier: Apache-2.0
// K11 depthwise compute (sceneseg-p150 OPT_PLAN K11, tt/dw.py): per output tile (32 pixels x 32 channels, float32)
// acc = x_0 * w_0; acc = x_t * w_t + acc (t = 1 .. KK-1, SFPMAD: one fp32 rounding per tap)
// acc = acc + bias (add_binary_tile, fp32 SFPU add, as EPI)
// acc = silu(acc) (silu_tile, ttnn.silu's LLK, as EPI) [DW_ACT 1; 0 = none]
// x_t: the gathered input page of tap t (CB 0, unpacked to DEST as fp32), w_t: weight tile t of the current channel
// tile (CB 1, every row = the 32 channel weights of tap t), bias: CB 2. DEST tiles: 0 = acc, 1 = x, 2 = w.
// Work split and order as dw_reader.cpp; the weight / bias CBs are popped when the channel tile changes.
// CT args: [0] KK, [1] CT, [2] PT, [3] NC, [4] GX, [5] CH, [6] NWT (weight tiles per channel tile). Defines: DW_ACT,
// DW_WPACK.
// R7.1 DW_WPACK=1: the NWT = ceil(KK / 8) weight tiles hold 8 taps each (rows 4l .. 4l+3 = tap 8j + l); per output
// tile they are unpacked to DEST tiles 2 / 3 once (5x5: again for taps 16 .. 24) and tap t's SFPMADs read their weight
// vectors from its 4-row group (k11_tap_tile). An SFPU vector is 4 DEST rows x 8 columns of one parity (sfpi unit
// 2k = rows 4k' .. +3 even columns, 2k + 1 the odd ones), so the x vector of unit d takes the group's unit of the same
// parity: the same operands per SFPMAD as the one-tap-per-tile form (bit-identical).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/tile_move_copy.h"
#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
// one face (8 row groups) of DEST tile 0 = tile 1 * tile 2 (+ tile 0)
inline void k11_mul_face(const std::uint32_t, const std::uint32_t, const std::uint32_t) {
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat x = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::dst_reg[0] = x * w;
sfpi::dst_reg++;
}
}
inline void k11_mad_face(const std::uint32_t, const std::uint32_t, const std::uint32_t) {
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat x = sfpi::dst_reg[32];
sfpi::vFloat w = sfpi::dst_reg[64];
sfpi::vFloat a = sfpi::dst_reg[0];
sfpi::dst_reg[0] = x * w + a;
sfpi::dst_reg++;
}
}
// R7.1: the whole tile in one call (VectorMode::None, base = DEST tile 0): tile 0 = tile 1 * w (+ tile 0), w = the
// tap's weight group at sfpi unit WR relative to the face base for faces 0 / 1, WR - 16 for faces 2 / 3.
template <bool FIRST>
inline void k11_tap_tile(const std::uint32_t, const std::uint32_t, const std::uint32_t, const std::uint32_t wr) {
#pragma GCC unroll 4
for (int f = 0; f < 4; f++) {
const int wo = (int)wr - (f >= 2 ? 16 : 0);
sfpi::vFloat we = sfpi::dst_reg[wo];
sfpi::vFloat wd = sfpi::dst_reg[wo + 1];
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat x = sfpi::dst_reg[32];
if constexpr (FIRST) {
sfpi::dst_reg[0] = x * ((d & 1) ? wd : we);
} else {
sfpi::vFloat a = sfpi::dst_reg[0];
sfpi::dst_reg[0] = x * ((d & 1) ? wd : we) + a;
}
sfpi::dst_reg++;
}
}
}
} // namespace ckernel::sfpu
#endif
using namespace ckernel;
constexpr uint32_t KK = get_compile_time_arg_val(0);
constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t PT = 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 CH = get_compile_time_arg_val(5);
constexpr uint32_t NWT = get_compile_time_arg_val(6);
constexpr uint32_t NW = CT * PT;
constexpr uint32_t cb_x = 0, cb_w = 1, cb_b = 2, cb_out = 16;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NW / NC, E = NW % NC;
const uint32_t w0 = core * B + (core < E ? core : E);
const uint32_t w1 = w0 + B + (core < E ? 1 : 0);
compute_kernel_hw_startup(cb_x, cb_out);
uint32_t cur_ct = 0xFFFFFFFF;
uint32_t w = w0;
while (w < w1) {
const uint32_t ct = w / PT;
const uint32_t pt0 = w - ct * PT;
uint32_t pt1 = pt0 + CH;
if (pt1 > PT) {
pt1 = PT;
}
if (pt1 - pt0 > w1 - w) {
pt1 = pt0 + (w1 - w);
}
if (ct != cur_ct) {
if (cur_ct != 0xFFFFFFFF) {
cb_pop_front(cb_w, NWT);
cb_pop_front(cb_b, 1);
}
cb_wait_front(cb_w, NWT);
cb_wait_front(cb_b, 1);
cur_ct = ct;
}
for (uint32_t pt = pt0; pt < pt1; ++pt) {
tile_regs_acquire();
#if DW_WPACK == 1
copy_init(cb_w);
copy_tile(cb_w, 0, 2);
if constexpr (NWT > 1) {
copy_tile(cb_w, 1, 3);
}
copy_init(cb_x);
add_binary_tile_init();
for (uint32_t t = 0; t < KK; ++t) {
if (t == 16) {
copy_init(cb_w);
copy_tile(cb_w, 2, 2);
if constexpr (NWT > 3) {
copy_tile(cb_w, 3, 3);
}
copy_init(cb_x);
add_binary_tile_init();
}
cb_wait_front(cb_x, 1);
copy_tile(cb_x, 0, 1);
const uint32_t wr = 64 + 32 * ((t >> 3) & 1) + 16 * ((t >> 2) & 1) + 2 * (t & 3);
if (t == 0) {
MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, k11_tap_tile, (true), 0, 1, 0,
VectorMode::None, wr)));
} else {
MATH((SFPU_BINARY_CALL(DST_SYNC_MODE, DST_ACCUM_MODE, k11_tap_tile, (false), 0, 1, 0,
VectorMode::None, wr)));
}
cb_pop_front(cb_x, 1);
}
#else
for (uint32_t t = 0; t < KK; ++t) {
cb_wait_front(cb_x, 1);
copy_init(cb_x);
copy_tile(cb_x, 0, 1);
copy_init(cb_w);
copy_tile(cb_w, t, 2);
add_binary_tile_init();
if (t == 0) {
MATH((SFPU_BINARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, k11_mul_face, 0, 1, 0,
VectorMode::RC)));
} else {
MATH((SFPU_BINARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, k11_mad_face, 0, 1, 0,
VectorMode::RC)));
}
cb_pop_front(cb_x, 1);
}
#endif
copy_init(cb_b);
copy_tile(cb_b, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
#if DW_ACT == 1
silu_tile_init();
silu_tile(0);
#endif
tile_regs_commit();
cb_reserve_back(cb_out, 1);
tile_regs_wait();
pack_tile(0, cb_out);
tile_regs_release();
cb_push_back(cb_out, 1);
}
w += pt1 - pt0;
}
if (cur_ct != 0xFFFFFFFF) {
cb_pop_front(cb_w, NWT);
cb_pop_front(cb_b, 1);
}
}