File size: 13,007 Bytes
e0adbab | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 225 226 227 228 229 230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 | // 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;
}
}
|