Download code/tt_sceneseg/tt/kernels/s1_reader.cpp from changh95/sceneseg-p150: direct link, hf CLI and curl.
- Browser
- Download file 13 kB
-
https://huggingface.co/changh95/sceneseg-p150/resolve/main/code/tt_sceneseg/tt/kernels/s1_reader.cpp
- Command line
-
hf download hf://changh95/sceneseg-p150/code/tt_sceneseg/tt/kernels/s1_reader.cpp
-
curl -L -o s1_reader.cpp https://huggingface.co/changh95/sceneseg-p150/resolve/main/code/tt_sceneseg/tt/kernels/s1_reader.cpp
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]. | |
| 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; | |
| } | |
| } | |