sceneseg-p150 / code /tt_sceneseg /tt /kernels /s1_reader.cpp
changh95's picture
tt-model push sceneseg-p150 (container)
e0adbab verified
Raw History Blame Contribute Delete
13 kB
// SPDX-License-Identifier: Apache-2.0
// S1 reader (sceneseg-p150 OPT_PLAN S1, tt/subpixel.py): the operands of the sub-pixel up-block epilogue, per output
// tile of the 2H x 2W map (32 pixels x 32 channels):
// CB 0 / CB 1 the Z conv's hi / lo products of the tile's 32 pixels, gathered by the shifted depth-to-space
// Y[2h + a, 2w + b] = Z[(a, b), h + a, w + b]. Z: ZSPLIT 1 = one [(H+1)(W+1), 8 Cout] fp32 TILE
// tensor, columns term-major then phase (2a + b) then Cout; ZSPLIT 2 = one [(H+1)(W+1), 4 Cout]
// tensor per term, columns phase-major; ZSPLIT 4 = one [(H+1)(W+1), 2 Cout] tensor per phase,
// columns term-major. Whole 4 KB Z tiles are read from DRAM into an L1 cache
// (2 terms x 4 phases x S slots, slot = Z row tile % S), then each output row (two 64 B face rows) is
// copied locally by the NoC; never 64 B DRAM reads (OPT_REPORT R2.2).
// CB 2 / CB 6 the composed skip conv's hi / lo products (S: [4HW, 2 Cout] fp32 TILE), when HAS_SKIP
// CB 3 the tile's bias: the interior bias tile (resident in CB 9, copied locally), or for a border row
// tile its pattern tile (DRAM), from the per-row-tile index (CB 10: 0 = interior, k = pattern k - 1)
// Work split (no per-core RT args, probe P3): units u = c * R + r (c = channel tile, r = output row tile), core i of
// the first NC (row-major) owns the contiguous range [i * B + min(i, E), + B + (i < E)), B = U / NC, E = U % NC.
// Units go in batches of up to NB tiles (one channel tile, no cache-slot conflict): the batch's Z misses are read
// with one barrier, then the batch's pages are filled (CB pages addressed with wrap-around) with one barrier and
// pushed one by one (a producer push must not cross the CB end).
// CT args: [0] R, [1] CT, [2] NC, [3] GX, [4] H, [5] W, [6] HAS_SKIP, [7] IDXB (bytes of the index page), [8] NB,
// [9] ZSPLIT (1 / 2 / 4), [10] ZSHIFT (1: Z on the (H+1) x (W+1) grid of the 2x2 conv, phase (a, b) of
// output (2h + a, 2w + b) at Z pixel (h + a, w + b); 0: Z on the H x W grid of the 3x3 form, at (h, w)),
// [11] ZTERMS (2: Z holds the hi and lo weight-term products; 1: one product, no CB 1, round 4 P2),
// [12] STERMS (the same for the skip conv: 2 = [Shi | Slo], 1 = S only, no CB 6),
// [13] BLOCAL (R4.4 EPI-B: read only the interior bias tiles of the channel tiles this core's units cover,
// c = u0 / R .. (u1 - 1) / R; 0 = all CT),
// then the TensorAccessorArgs of Z0..Z3 (Z0 repeated past ZSPLIT), S (Z0 when
// absent), the interior bias, the patterns, the index.
// Common RT args: [z0, z1, z2, z3, s, bias, patterns, index].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t R = get_compile_time_arg_val(0);
constexpr uint32_t CT = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t H = get_compile_time_arg_val(4);
constexpr uint32_t W = get_compile_time_arg_val(5);
constexpr uint32_t HAS_SKIP = get_compile_time_arg_val(6);
constexpr uint32_t IDXB = get_compile_time_arg_val(7);
constexpr uint32_t NB = get_compile_time_arg_val(8);
constexpr uint32_t ZSPLIT = get_compile_time_arg_val(9);
constexpr uint32_t ZSHIFT = get_compile_time_arg_val(10);
constexpr uint32_t ZTERMS = get_compile_time_arg_val(11);
constexpr uint32_t STERMS = get_compile_time_arg_val(12);
constexpr uint32_t BLOCAL = get_compile_time_arg_val(13);
constexpr auto z0_args = TensorAccessorArgs<14>();
constexpr auto z1_args = TensorAccessorArgs<z0_args.next_compile_time_args_offset()>();
constexpr auto z2_args = TensorAccessorArgs<z1_args.next_compile_time_args_offset()>();
constexpr auto z3_args = TensorAccessorArgs<z2_args.next_compile_time_args_offset()>();
constexpr auto s_args = TensorAccessorArgs<z3_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<s_args.next_compile_time_args_offset()>();
constexpr auto p_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr auto i_args = TensorAccessorArgs<p_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t W2 = 2 * W, WP = W + ZSHIFT, SW = STERMS * CT;
constexpr uint32_t ZW = ZTERMS * 4 * CT / ZSPLIT; // column tiles of a Z tensor
constexpr uint32_t ZTERM = ZSPLIT == 1 ? 4 * CT : (ZSPLIT == 4 ? CT : 0); // column-tile offset of the lo term
constexpr uint32_t ZPH = ZSPLIT == 4 ? 0 : CT; // column-tile offset per phase
constexpr uint32_t S = 4; // cache slots per (term, phase)
constexpr uint32_t NSLOT = 2 * 4 * S;
constexpr uint32_t cb_zh = 0, cb_zl = 1, cb_sh = 2, cb_b = 3, cb_sl = 6, cb_cache = 8, cb_bint = 9, cb_idx = 10;
constexpr uint32_t cb_src = 11; // L1 scratch: NB x 32 source-row offsets (uint32)
FORCE_INLINE uint32_t face_off(uint32_t row) { return (row >> 4) * 2048 + (row & 15) * 64; }
// address of page j after the write pointer of a CB, wrapped at its end
FORCE_INLINE uint32_t page_addr(uint32_t cb, uint32_t j) {
const auto& ci = get_local_cb_interface(cb);
uint32_t a = ci.fifo_wr_ptr + j * TB;
if (a >= ci.fifo_limit) {
a -= ci.fifo_size;
}
return a;
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t U = R * CT;
constexpr uint32_t B = U / NC, E = U % NC;
const uint32_t u0 = core * B + (core < E ? core : E);
const uint32_t u1 = u0 + B + (core < E ? 1 : 0);
if (u1 == u0) {
return;
}
const auto az0 = TensorAccessor(z0_args, get_common_arg_val<uint32_t>(0), TB);
const auto az1 = TensorAccessor(z1_args, get_common_arg_val<uint32_t>(1), TB);
const auto az2 = TensorAccessor(z2_args, get_common_arg_val<uint32_t>(2), TB);
const auto az3 = TensorAccessor(z3_args, get_common_arg_val<uint32_t>(3), TB);
const auto as = TensorAccessor(s_args, get_common_arg_val<uint32_t>(4), TB);
const auto ab = TensorAccessor(b_args, get_common_arg_val<uint32_t>(5), TB);
const auto ap = TensorAccessor(p_args, get_common_arg_val<uint32_t>(6), TB);
const auto ai = TensorAccessor(i_args, get_common_arg_val<uint32_t>(7), IDXB);
// NoC address of a Z tile of term tm (0 hi, 1 lo) and phase ph (tile = row tile * ZW + column tile)
auto zaddr = [&](uint32_t tm, uint32_t ph, uint32_t tile) -> uint64_t {
if constexpr (ZSPLIT == 1) {
return az0.get_noc_addr(tile);
} else if constexpr (ZSPLIT == 2) {
return tm == 0 ? az0.get_noc_addr(tile) : az1.get_noc_addr(tile);
} else {
return ph == 0 ? az0.get_noc_addr(tile)
: (ph == 1 ? az1.get_noc_addr(tile)
: (ph == 2 ? az2.get_noc_addr(tile) : az3.get_noc_addr(tile)));
}
};
const uint32_t idx_l1 = get_write_ptr(cb_idx);
noc_async_read(ai.get_noc_addr(0), idx_l1, IDXB);
const uint32_t bint = get_write_ptr(cb_bint);
const uint32_t c_lo = BLOCAL ? u0 / R : 0, c_hi = BLOCAL ? (u1 - 1) / R : CT - 1;
for (uint32_t c = c_lo; c <= c_hi; ++c) {
noc_async_read(ab.get_noc_addr(c), bint + c * TB, TB);
}
noc_async_read_barrier();
volatile tt_l1_ptr uint32_t* idx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(idx_l1);
const uint32_t cache = get_write_ptr(cb_cache);
const uint64_t local = get_noc_addr(cache);
uint32_t tag[NSLOT / 2]; // Z row tile held by slot (phase, k) (both terms alike)
uint32_t used[NSLOT / 2]; // batch that last used the slot
for (uint32_t i = 0; i < NSLOT / 2; ++i) {
tag[i] = 0xFFFFFFFFu;
used[i] = 0xFFFFFFFFu;
}
// per batch tile and row: cache offset of the source row (L1 scratch, not the RISC stack)
volatile tt_l1_ptr uint32_t* src = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_src));
uint32_t miss[NSLOT / 2], miss_rt[NSLOT / 2], miss_ph[NSLOT / 2];
uint32_t cur_c = 0xFFFFFFFFu;
uint32_t batch = 0;
uint32_t u = u0;
while (u < u1) {
const uint32_t c = u / R;
if (c != cur_c) {
cur_c = c;
for (uint32_t i = 0; i < NSLOT / 2; ++i) {
tag[i] = 0xFFFFFFFFu;
}
}
// ---- plan the batch: up to NB tiles of channel tile c without a slot conflict
uint32_t k = 0, nmiss = 0;
while (k < NB && u + k < u1 && (u + k) / R == c) {
const uint32_t r = (u + k) % R;
uint32_t p = r * 32;
uint32_t y = p / W2;
uint32_t x = p - y * W2;
volatile tt_l1_ptr uint32_t* rows = src + k * 32;
bool ok = true;
for (uint32_t i = 0; i < 32 && ok; ++i) {
const uint32_t a = y & 1, b = x & 1;
const uint32_t zp = ((y >> 1) + a * ZSHIFT) * WP + (x >> 1) + b * ZSHIFT;
const uint32_t zrt = zp >> 5;
const uint32_t ph = 2 * a + b;
const uint32_t sl = ph * S + (zrt % S);
if (tag[sl] != zrt) {
if (used[sl] == batch) {
ok = false; // the slot holds a row tile this batch still needs
break;
}
tag[sl] = zrt;
miss[nmiss] = sl;
miss_rt[nmiss] = zrt * ZW + ph * ZPH + c;
miss_ph[nmiss] = ph;
++nmiss;
}
used[sl] = batch;
rows[i] = sl * TB + face_off(zp & 31);
if (++x == W2) {
x = 0;
++y;
}
}
if (!ok) {
break;
}
++k;
}
// (k >= 1 always: a fresh batch has every slot free for one tile's <= 8 (phase, row tile) pairs: a tile
// spans at most two image rows, 2W >= 32 is checked on the host)
for (uint32_t m = 0; m < nmiss; ++m) {
const uint32_t dst = cache + miss[m] * TB;
noc_async_read(zaddr(0, miss_ph[m], miss_rt[m]), dst, TB); // hi term
if constexpr (ZTERMS == 2) {
noc_async_read(zaddr(1, miss_ph[m], miss_rt[m] + ZTERM), dst + (NSLOT / 2) * TB, TB); // lo term
}
}
noc_async_read_barrier();
// ---- fill the batch's pages
cb_reserve_back(cb_zh, k);
if constexpr (ZTERMS == 2) {
cb_reserve_back(cb_zl, k);
}
cb_reserve_back(cb_b, k);
if constexpr (HAS_SKIP) {
cb_reserve_back(cb_sh, k);
if constexpr (STERMS == 2) {
cb_reserve_back(cb_sl, k);
}
}
noc_async_read_set_state(local);
for (uint32_t j = 0; j < k; ++j) {
const uint32_t dh = page_addr(cb_zh, j);
const uint32_t dl = ZTERMS == 2 ? page_addr(cb_zl, j) : 0;
for (uint32_t i = 0; i < 32; ++i) {
const uint32_t so = cache + src[j * 32 + i];
const uint32_t fo = face_off(i);
noc_async_read_with_state(so, dh + fo, 64);
noc_async_read_with_state(so + 1024, dh + fo + 1024, 64);
if constexpr (ZTERMS == 2) {
noc_async_read_with_state(so + (NSLOT / 2) * TB, dl + fo, 64);
noc_async_read_with_state(so + (NSLOT / 2) * TB + 1024, dl + fo + 1024, 64);
}
}
}
for (uint32_t j = 0; j < k; ++j) {
const uint32_t r = (u + j) % R;
const uint32_t pi = idx[r];
const uint32_t db = page_addr(cb_b, j);
if (pi == 0) {
noc_async_read(get_noc_addr(bint + c * TB), db, TB);
} else {
noc_async_read(ap.get_noc_addr((pi - 1) * CT + c), db, TB);
}
if constexpr (HAS_SKIP) {
noc_async_read(as.get_noc_addr(r * SW + c), page_addr(cb_sh, j), TB);
if constexpr (STERMS == 2) {
noc_async_read(as.get_noc_addr(r * SW + CT + c), page_addr(cb_sl, j), TB);
}
}
}
noc_async_read_barrier();
for (uint32_t j = 0; j < k; ++j) {
cb_push_back(cb_zh, 1);
if constexpr (ZTERMS == 2) {
cb_push_back(cb_zl, 1);
}
cb_push_back(cb_b, 1);
if constexpr (HAS_SKIP) {
cb_push_back(cb_sh, 1);
if constexpr (STERMS == 2) {
cb_push_back(cb_sl, 1);
}
}
}
u += k;
++batch;
}
}