File size: 5,506 Bytes
e91e989 | 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 | // SPDX-License-Identifier: Apache-2.0
// KTILE reader (meteor-p150 OPT2, tt/ktile.py): a ROW_MAJOR map (logical [A, B, WD, C], one DRAM page per pixel of
// PAGE = C * ES bytes) tilized into the TILE form of the same logical shape, without a compute kernel.
// Work units (no per-core RT args, probe P3: core i = y * GX + x of the first NC cores owns the contiguous unit range
// [i * Q + min(i, E), + Q + (i < E)) of the NU = A * B * XB units, Q = NU / NC, E = NU % NC): unit v = 32-pixel row
// block xb = v % XB of plane v / XB, pixels p0 = (v / XB) * WD + 32 xb, n = min(32, WD - 32 xb) valid rows; its
// output = the CT tiles v * CT + 0..CT-1.
// Per unit: the n pixel rows read as NBANK runs of same-bank consecutive pages (rows j, j + NBANK, ... of one DRAM
// bank; staging slot of row i = (i % NBANK) * 32 / NBANK + i / NBANK) into a double-buffered L1 staging area (CB 1),
// the next unit's runs issued while this unit is tilized; tilize = 16-element face rows (16 * ES bytes) copied by
// stateful local NOC reads into the CT tiles of CB 0 (tile k, face 2 (i / 16) + h, face row i % 16 <- row i,
// elements 32 k + 16 h); the rows n..31 and the channel halves >= C zeroed (the tilize padding).
// Crop (a window of a larger 4-D map): the source page of output pixel j of plane q is SRC_OFF + q * SRC_PSTRIDE + j
// (no crop: SRC_OFF 0, SRC_PSTRIDE WD).
// CT args: [0] NU, [1] NC, [2] GX, [3] XB, [4] WD, [5] C, [6] ES, [7] NBANK, [8] CT, [9] SRC_OFF, [10] SRC_PSTRIDE,
// then the TensorAccessorArgs of the input. Common RT args: [in_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NU = 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 XB = get_compile_time_arg_val(3);
constexpr uint32_t WD = get_compile_time_arg_val(4);
constexpr uint32_t CH = get_compile_time_arg_val(5);
constexpr uint32_t ES = get_compile_time_arg_val(6);
constexpr uint32_t NBANK = get_compile_time_arg_val(7);
constexpr uint32_t CT = get_compile_time_arg_val(8);
constexpr uint32_t SRC_OFF = get_compile_time_arg_val(9);
constexpr uint32_t SRC_PSTRIDE = get_compile_time_arg_val(10);
constexpr auto in_args = TensorAccessorArgs<11>();
constexpr uint32_t PAGE = CH * ES, TB = 1024 * ES, FACE = 256 * ES, FROW = 16 * ES;
constexpr uint32_t RUN = 32 / NBANK;
constexpr uint32_t cb_t = 0, cb_stage = 1;
inline void unit_rows(uint32_t v, uint32_t& p0, uint32_t& n) {
const uint32_t xb = v % XB;
p0 = SRC_OFF + (v / XB) * SRC_PSTRIDE + xb * 32;
n = (WD - xb * 32) < 32 ? (WD - xb * 32) : 32;
}
template <typename TA>
inline void read_rows(const TA& src, uint32_t p0, uint32_t n, uint32_t stage) {
for (uint32_t j = 0; j < NBANK && j < n; ++j) {
const uint32_t cnt = (n - j + NBANK - 1) / NBANK;
noc_async_read(src.get_noc_addr(p0 + j), stage + j * RUN * PAGE, cnt * PAGE);
}
}
inline void zero(uint32_t addr, uint32_t bytes) {
volatile tt_l1_ptr uint32_t* q = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(addr);
for (uint32_t w = 0; w < bytes / 4; ++w) {
q[w] = 0;
}
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t Q = NU / NC, E = NU % NC;
const uint32_t u0 = core * Q + (core < E ? core : E);
const uint32_t nu = Q + (core < E ? 1 : 0);
if (nu == 0) {
return;
}
const auto src = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), PAGE);
cb_reserve_back(cb_stage, 1); // reader-private, never pushed
const uint32_t stage0 = get_write_ptr(cb_stage);
const uint64_t here = get_noc_addr(stage0);
uint32_t p0, n;
unit_rows(u0, p0, n);
read_rows(src, p0, n, stage0);
noc_async_read_barrier();
for (uint32_t v = u0; v < u0 + nu; ++v) {
const uint32_t stage = stage0 + ((v - u0) & 1) * 32 * PAGE;
cb_reserve_back(cb_t, CT);
const uint32_t wt = get_write_ptr(cb_t);
if constexpr (CH % 32 != 0) { // the last tile's upper 16 columns: channel padding
zero(wt + (CT - 1) * TB + 1 * FACE, FACE);
zero(wt + (CT - 1) * TB + 3 * FACE, FACE);
}
if (n < 32) { // the width padding rows
for (uint32_t i = n; i < 32; ++i) {
for (uint32_t k = 0; k < CT; ++k) {
for (uint32_t h = 0; h < 2; ++h) {
zero(wt + k * TB + ((i >> 4) * 2 + h) * FACE + (i & 15) * FROW, FROW);
}
}
}
}
noc_async_read_one_packet_set_state(here, FROW);
for (uint32_t i = 0; i < n; ++i) {
const uint32_t s = stage + ((i % NBANK) * RUN + i / NBANK) * PAGE;
const uint32_t d = wt + ((i >> 4) * 2) * FACE + (i & 15) * FROW;
for (uint32_t k = 0; k < CT; ++k) {
noc_async_read_one_packet_with_state(s + k * 32 * ES, d + k * TB);
if (k * 32 + 16 < CH) {
noc_async_read_one_packet_with_state(s + (k * 32 + 16) * ES, d + k * TB + FACE);
}
}
}
if (v + 1 < u0 + nu) {
unit_rows(v + 1, p0, n);
read_rows(src, p0, n, stage0 + ((v + 1 - u0) & 1) * 32 * PAGE);
}
noc_async_read_barrier();
cb_push_back(cb_t, CT);
}
}
|