File size: 3,642 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 | // SPDX-License-Identifier: Apache-2.0
// KUNTILE reader (meteor-p150 OPT2, tt/ktile.py untilize): a TILE map (logical [A, B, WD, C], CT tile columns) as
// ROW_MAJOR (one DRAM page per pixel of PAGE = C * ES bytes), without a compute kernel. Units as KTILE (unit v = row
// block v % XB of plane v / XB, its CT input tiles v * CT + k, n = min(32, WD - 32 xb) valid rows). Per unit: the
// tiles read into a double-buffered reader-private area (CB 0; the next unit's tiles issued while this unit is
// copied), then the n rows untilized by stateful local NOC copies of 16-element face rows into a staging page (CB 1,
// pushed to the writer; row i at slot (i % NBANK) * 32 / NBANK + i / NBANK, so each DRAM bank's rows are one run).
// CT args: [0] NU, [1] NC, [2] GX, [3] XB, [4] WD, [5] C, [6] ES, [7] NBANK, [8] CT, 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 auto in_args = TensorAccessorArgs<9>();
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_tiles = 0, cb_stage = 1;
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), TB);
cb_reserve_back(cb_tiles, 1); // reader-private (2 x CT tiles), never pushed
const uint32_t tiles0 = get_write_ptr(cb_tiles);
for (uint32_t k = 0; k < CT; ++k) {
noc_async_read(src.get_noc_addr(u0 * CT + k), tiles0 + k * TB, TB);
}
noc_async_read_barrier();
const uint64_t here = get_noc_addr(tiles0);
for (uint32_t v = u0; v < u0 + nu; ++v) {
const uint32_t tl = tiles0 + ((v - u0) & 1) * CT * TB;
const uint32_t xb = v % XB;
const uint32_t n = (WD - xb * 32) < 32 ? (WD - xb * 32) : 32;
cb_reserve_back(cb_stage, 1);
const uint32_t st = get_write_ptr(cb_stage);
noc_async_read_one_packet_set_state(here, FROW);
for (uint32_t i = 0; i < n; ++i) {
const uint32_t s = tl + ((i >> 4) * 2) * FACE + (i & 15) * FROW;
const uint32_t d = st + ((i % NBANK) * RUN + i / NBANK) * PAGE;
for (uint32_t k = 0; k < CT; ++k) {
noc_async_read_one_packet_with_state(s + k * TB, d + k * 32 * ES);
if (k * 32 + 16 < CH) {
noc_async_read_one_packet_with_state(s + k * TB + FACE, d + (k * 32 + 16) * ES);
}
}
}
if (v + 1 < u0 + nu) {
const uint32_t nt = tiles0 + ((v + 1 - u0) & 1) * CT * TB;
for (uint32_t k = 0; k < CT; ++k) {
noc_async_read(src.get_noc_addr((v + 1) * CT + k), nt + k * TB, TB);
}
}
noc_async_read_barrier();
cb_push_back(cb_stage, 1);
}
}
|