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_reader.cpp
changh95's picture
changh95
tt-model push sceneseg-p150 (container)
e0adbab verified
Raw History Blame Contribute Delete
12.6 kB
// SPDX-License-Identifier: Apache-2.0
// K11 depthwise reader (sceneseg-p150 OPT_PLAN K11, tt/dw.py): the float32 TILE input map [1, 1, Hin*Win, C] (DRAM
// interleaved) of one depthwise kxk conv, gathered into one TILE page per (output pixel tile, channel tile, tap).
//
// Work: the NW = CT * PT items (channel tile ct, output pixel tile pt; ct-major) split over the first NC cores in
// row-major order, core i owning [i * B + min(i, E), + B + (i < E)) (no per-core RT args: probe P3). A core walks
// its items in chunks of at most CH pixel tiles of one ct. Per chunk:
// 1. the input rows the chunk needs are loaded into an L1 band: two planes per row (channel half h = 0 / 1, 64 B
// per pixel = 16 fp32 channels), P zero columns on each side; for stride 2 each plane is stored column-parity
// major (even padded columns, then odd), so every tap reads consecutive output pixels from consecutive band
// pixels. A 16-pixel face half of an input tile is one contiguous DRAM run: stride 1 reads it straight into
// the band; stride 2 reads the row into a staging row (CB 5, two rows ping-pong) and scatters its pixels to
// the parity planes with 64 B local copies (64 B DRAM reads were latency-bound: R2 profile).
// 2. per output tile, the KK TILE pages (one per tap) are assembled from the band: face (rg, h) = 16 output
// pixels of channel half h, copied in runs (split at output-row ends); out-of-image rows and pixels past
// Hout*Wout read a zero block. Every copy is a local NoC read; ONE barrier per output tile (KK pages pushed
// together, CB 0 holds 2 * KK pages), not per page.
// When ct changes the NWT weight tiles (tile ct * NWT + j; NWT = KK: every row of tile j = w[ct*32 .. +32, tap j];
// R7.1 DW_WPACK: NWT = ceil(KK / 8), rows 4l .. 4l+3 of tile j = tap 8j + l) and the bias tile ct are read into
// CBs 1 / 2 (resident until the next ct).
// CT args: [0] HIN, [1] WIN, [2] HOUT, [3] WOUT, [4] K, [5] S, [6] P, [7] CT, [8] PT, [9] NC, [10] GX, [11] CH,
// [12] NRMAX (band rows), [13] STAGE (stride 2: 1 = staging rows, 0 = 64 B DRAM reads), [14] NWT (weight
// tiles per channel tile), then the TensorAccessorArgs of x, w, b.
// Common RT args: [x_addr, w_addr, b_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t HIN = get_compile_time_arg_val(0);
constexpr uint32_t WIN = get_compile_time_arg_val(1);
constexpr uint32_t HOUT = get_compile_time_arg_val(2);
constexpr uint32_t WOUT = get_compile_time_arg_val(3);
constexpr uint32_t K = get_compile_time_arg_val(4);
constexpr uint32_t S = get_compile_time_arg_val(5);
constexpr uint32_t P = get_compile_time_arg_val(6);
constexpr uint32_t CT = get_compile_time_arg_val(7);
constexpr uint32_t PT = get_compile_time_arg_val(8);
constexpr uint32_t NC = get_compile_time_arg_val(9);
constexpr uint32_t GX = get_compile_time_arg_val(10);
constexpr uint32_t CH = get_compile_time_arg_val(11);
constexpr uint32_t NRMAX = get_compile_time_arg_val(12);
constexpr uint32_t STAGE = get_compile_time_arg_val(13);
constexpr uint32_t NWT = get_compile_time_arg_val(14);
constexpr auto x_args = TensorAccessorArgs<15>();
constexpr auto w_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr uint32_t KK = K * K;
constexpr uint32_t HWO = HOUT * WOUT;
constexpr uint32_t NW = CT * PT;
constexpr uint32_t TB = 4096;
constexpr uint32_t WP = WIN + 2 * P; // padded input width (pixels)
constexpr uint32_t PH = (WP + 1) / 2; // stride 2: pixels of the even-column half of a plane row
constexpr uint32_t RS = (S == 1) ? WP : 2 * PH; // pixels per plane row
constexpr uint32_t cb_x = 0, cb_w = 1, cb_b = 2, cb_band = 3, cb_zero = 4, cb_stage = 5;
constexpr uint32_t STAGE_ROW = WIN * 128; // one staging row: two planes of WIN pixels x 64 B
// L1 byte address of padded column j of band row r, channel half h
FORCE_INLINE uint32_t band_addr(uint32_t base, uint32_t r, uint32_t h, uint32_t j) {
uint32_t pos;
if constexpr (S == 1) {
pos = j;
} else {
pos = (j & 1) * PH + (j >> 1);
}
return base + (((r * 2 + h) * RS) + pos) * 64;
}
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);
const auto xa = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB);
const auto wa = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), TB);
const auto ba = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), TB);
// the 1 KB zero block (16 pixels x 64 B)
const uint32_t zero = get_write_ptr(cb_zero);
{
volatile tt_l1_ptr uint32_t* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero);
for (uint32_t i = 0; i < 256; ++i) {
z[i] = 0;
}
}
const uint64_t zero_noc = get_noc_addr(zero);
const uint32_t band = get_write_ptr(cb_band);
const uint32_t stage = (S == 2 && STAGE) ? get_write_ptr(cb_stage) : 0;
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) {
cb_reserve_back(cb_w, NWT);
uint32_t l1 = get_write_ptr(cb_w);
for (uint32_t t = 0; t < NWT; ++t) {
noc_async_read(wa.get_noc_addr(ct * NWT + t), l1, TB);
l1 += TB;
}
cb_reserve_back(cb_b, 1);
noc_async_read(ba.get_noc_addr(ct), get_write_ptr(cb_b), TB);
noc_async_read_barrier();
cb_push_back(cb_w, NWT);
cb_push_back(cb_b, 1);
cur_ct = ct;
}
// ---- 1. the band: padded input rows R0 .. R0 + NR - 1 (input row iy = R0 + r - P)
const uint32_t p_first = pt0 * 32;
uint32_t p_last = pt1 * 32 - 1;
if (p_last >= HWO) {
p_last = HWO - 1;
}
const uint32_t ya = p_first / WOUT, yb = p_last / WOUT;
const uint32_t R0 = ya * S;
const uint32_t NR = (yb - ya) * S + K; // <= NRMAX by construction (host)
for (uint32_t r = 0; r < NR; ++r) {
const int32_t iy = (int32_t)(R0 + r) - (int32_t)P;
if (iy < 0 || iy >= (int32_t)HIN) {
continue;
}
for (uint32_t h = 0; h < 2; ++h) {
for (uint32_t j = 0; j < P; ++j) {
noc_async_read(zero_noc, band_addr(band, r, h, j), 64);
noc_async_read(zero_noc, band_addr(band, r, h, P + WIN + j), 64);
}
}
const uint32_t qa = (uint32_t)iy * WIN;
const uint32_t qb = qa + WIN;
if constexpr (S == 1) {
uint32_t q = qa;
while (q < qb) {
uint32_t ge = (q | 15) + 1;
if (ge > qb) {
ge = qb;
}
const uint32_t n = ge - q;
const uint32_t tile = (q >> 5) * CT + ct;
const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
for (uint32_t h = 0; h < 2; ++h) {
noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), band_addr(band, r, h, P + q - qa),
n * 64);
}
q = ge;
}
} else if constexpr (STAGE == 0) {
// stride 2 without staging (R2.1 form, knob DW_STAGE=0): one 64 B DRAM read per pixel and half
for (uint32_t q = qa; q < qb; ++q) {
const uint32_t tile = (q >> 5) * CT + ct;
const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
for (uint32_t h = 0; h < 2; ++h) {
noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), band_addr(band, r, h, P + q - qa), 64);
}
}
} else {
const uint32_t st = stage + (r & 1) * STAGE_ROW;
uint32_t q = qa;
while (q < qb) {
uint32_t ge = (q | 15) + 1;
if (ge > qb) {
ge = qb;
}
const uint32_t n = ge - q;
const uint32_t tile = (q >> 5) * CT + ct;
const uint32_t fo = ((q >> 4) & 1) * 2048 + (q & 15) * 64;
for (uint32_t h = 0; h < 2; ++h) {
noc_async_read(xa.get_noc_addr(tile, fo + h * 1024), st + h * (WIN * 64) + (q - qa) * 64,
n * 64);
}
q = ge;
}
// the staging row is complete (and the copies out of the other staging row, issued one row ago,
// are done before that row is refilled: this barrier covers them too)
noc_async_read_barrier();
noc_async_read_set_state(zero_noc);
for (uint32_t h = 0; h < 2; ++h) {
for (uint32_t px = 0; px < WIN; ++px) {
noc_async_read_with_state(st + h * (WIN * 64) + px * 64, band_addr(band, r, h, P + px), 64);
}
}
}
}
noc_async_read_barrier();
// ---- 2. KK TILE pages per output tile. The output pixels of the tile split into at most 8 segments (per
// 16-pixel row group: runs inside one output row, or past Hout*Wout); a tap only shifts each segment's band
// source, so the segments are computed once per tile. All copies are local: NoC state set once.
noc_async_read_set_state(zero_noc);
for (uint32_t pt = pt0; pt < pt1; ++pt) {
uint32_t seg_d[8], seg_n[8], seg_y[8], seg_x[8], ns = 0;
for (uint32_t rg = 0; rg < 2; ++rg) {
const uint32_t p = pt * 32 + rg * 16;
uint32_t y = p / WOUT;
uint32_t x = p - y * WOUT;
uint32_t i = 0;
while (i < 16) {
seg_d[ns] = rg * 2048 + i * 64;
if (p + i >= HWO) {
seg_n[ns] = (16 - i) * 64;
seg_y[ns] = 0xFFFF; // zero segment
seg_x[ns] = 0;
++ns;
break;
}
uint32_t n = WOUT - x;
if (n > 16 - i) {
n = 16 - i;
}
seg_n[ns] = n * 64;
seg_y[ns] = y;
seg_x[ns] = x;
++ns;
i += n;
x += n;
if (x == WOUT) {
x = 0;
++y;
}
}
}
cb_reserve_back(cb_x, KK);
uint32_t dst = get_write_ptr(cb_x);
for (uint32_t dy = 0; dy < K; ++dy) {
for (uint32_t dx = 0; dx < K; ++dx, dst += TB) {
for (uint32_t g = 0; g < ns; ++g) {
const uint32_t d0 = dst + seg_d[g];
const uint32_t rr = seg_y[g] * S + dy; // padded row (huge for a zero segment)
if (rr >= P && rr < P + HIN) {
const uint32_t src = band_addr(band, rr - R0, 0, seg_x[g] * S + dx);
noc_async_read_with_state(src, d0, seg_n[g]);
noc_async_read_with_state(src + RS * 64, d0 + 1024, seg_n[g]);
} else {
noc_async_read_with_state(zero, d0, seg_n[g]);
noc_async_read_with_state(zero, d0 + 1024, seg_n[g]);
}
}
}
}
noc_async_read_barrier();
cb_push_back(cb_x, KK);
}
w += pt1 - pt0;
}
}