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 /epi_writer.cpp
changh95's picture
changh95
tt-model push sceneseg-p150 (container)
e0adbab verified
Raw History Blame Contribute Delete
5.39 kB
// SPDX-License-Identifier: Apache-2.0
// EPI writer (sceneseg-p150 OPT_PLAN EPI, tt/epi.py): the output tiles of this core's contiguous range (work split as
// the reader) to the interleaved output(s), flushed per tile, one barrier at the end: the main output (CB 16, when
// HAS_MAIN) and, with SPLIT (OPT_REPORT R2.3), the bf16 hi / lo terms (CBs 17 / 18).
// CLASS (K8, OPT_REPORT R2.4): the main output is the fp32 class-margin tile [m0, m1, m2, pad...] of 32 pixels; it is
// not written: the class of each pixel is decided here, exactly as ttnn.argmax(dim=1) over (m0, m1, m2) on the SFPU
// (argmax_nc_compute.cpp: running max with gt_binary_tile, strict '>', lowest index on ties; gt(a, b) = a > b in the
// sign-magnitude total order, false when both are zero / subnormal: ckernel_sfpu_binary_comp.h), and the uint8 class
// of the UNIT * 32 pixels of one image row is written as one row page of the uint8 ROW_MAJOR class map.
// CT args: [0] T, [1] NC, [2] GX, [3] OTB (main output tile bytes), [4] HAS_MAIN, [5] SPLIT, [6] UNIT, [7] CLASS,
// then the TensorAccessorArgs of the main (CLASS: the class map), hi and lo outputs (repeats of the first present
// one when absent). Common RT args: [main_addr, hi_addr, lo_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t T = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t OTB = get_compile_time_arg_val(3);
constexpr uint32_t HAS_MAIN = get_compile_time_arg_val(4);
constexpr uint32_t SPLIT = get_compile_time_arg_val(5);
constexpr uint32_t UNIT = get_compile_time_arg_val(6);
constexpr uint32_t CLASS = get_compile_time_arg_val(7);
constexpr auto out_args = TensorAccessorArgs<8>();
constexpr auto hi_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
constexpr auto lo_args = TensorAccessorArgs<hi_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_out = 16, cb_hi = 17, cb_lo = 18, cb_cls = 19;
constexpr uint32_t TB16 = 2048;
constexpr uint32_t ROWB = UNIT * 32; // CLASS: bytes (pixels) per class-map row page
// SFPU gt(a, b) on fp32 bit patterns (no NaN in the margins)
FORCE_INLINE bool sfpu_gt(uint32_t a, uint32_t b) {
if ((a & 0x7F800000u) == 0 && (b & 0x7F800000u) == 0) {
return false; // both zero / subnormal: equal
}
const int32_t ka = (a & 0x80000000u) ? -(int32_t)(a & 0x7FFFFFFFu) : (int32_t)a;
const int32_t kb = (b & 0x80000000u) ? -(int32_t)(b & 0x7FFFFFFFu) : (int32_t)b;
return ka > kb;
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = T / UNIT / NC, E = T / UNIT % NC;
const uint32_t t0 = (core * B + (core < E ? core : E)) * UNIT;
const uint32_t n = (B + (core < E ? 1 : 0)) * UNIT;
if constexpr (CLASS) {
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), ROWB);
const uint32_t row = get_write_ptr(cb_cls);
volatile tt_l1_ptr uint8_t* cls = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(row);
for (uint32_t t = t0; t < t0 + n; ++t) {
const uint32_t j = t % UNIT;
if (j == 0 && t != t0) {
noc_async_write_barrier(); // the previous row page left the buffer
}
cb_wait_front(cb_out, 1);
volatile tt_l1_ptr uint32_t* m = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_read_ptr(cb_out));
for (uint32_t r = 0; r < 32; ++r) {
const uint32_t w = (r >> 4) * 512 + (r & 15) * 16; // face (r / 16) * 2, row r % 16, column 0
uint32_t best = m[w];
uint8_t c = 0;
const uint32_t m1 = m[w + 1], m2 = m[w + 2];
if (sfpu_gt(m1, best)) {
best = m1;
c = 1;
}
if (sfpu_gt(m2, best)) {
c = 2;
}
cls[j * 32 + r] = c;
}
cb_pop_front(cb_out, 1);
if (j == UNIT - 1) {
noc_async_write(row, out.get_noc_addr(t / UNIT), ROWB);
}
}
noc_async_write_barrier();
return;
}
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), OTB);
const auto hi = TensorAccessor(hi_args, get_common_arg_val<uint32_t>(1), TB16);
const auto lo = TensorAccessor(lo_args, get_common_arg_val<uint32_t>(2), TB16);
for (uint32_t t = t0; t < t0 + n; ++t) {
if constexpr (HAS_MAIN) {
cb_wait_front(cb_out, 1);
noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(t), OTB);
}
if constexpr (SPLIT) {
cb_wait_front(cb_hi, 1);
cb_wait_front(cb_lo, 1);
noc_async_write(get_read_ptr(cb_hi), hi.get_noc_addr(t), TB16);
noc_async_write(get_read_ptr(cb_lo), lo.get_noc_addr(t), TB16);
}
noc_async_writes_flushed();
if constexpr (HAS_MAIN) {
cb_pop_front(cb_out, 1);
}
if constexpr (SPLIT) {
cb_pop_front(cb_hi, 1);
cb_pop_front(cb_lo, 1);
}
}
noc_async_write_barrier();
}