File size: 7,952 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 | // SPDX-License-Identifier: Apache-2.0
// EPI reader (sceneseg-p150 OPT_PLAN EPI, tt/epi.py): the NP fp32 term outputs of one conv / linear (DRAM or L1
// interleaved TILE, one shape) into CBs 0 .. NP-1, and once per program the CT fp32 bias tiles (the bias row
// replicated over the 32 rows of a tile on the host: what binary_ng's reader builds with fill_tile_with_first_row)
// into CB 3, kept resident for the whole program.
// EPI-split (OPT_REPORT R2.3) operands: HAS_RES = a residual map of the output's shape (tile t) into CB 4, read with
// the parts; HAS_ROW = the CT tiles of a [1, 1, 1, C] fp32 row tensor (the SE gate, computed on the device) into
// CB 5, row 0 replicated over the 32 rows here (binary_ng's fill_tile_with_first_row), resident.
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores in row-major order owns the
// contiguous tile range [i * B + min(i, E), + B + (i < E)) of the T output tiles, B = T / NC, E = T % NC.
// A part may be a column range of a wider tensor (ST2: the [W_hi | W_lo] conv's [R, 2C] output read as its two
// products): output tile t = (r, c) = (t / CT, t % CT) reads tile r * W_i + O_i + c of part i.
// Resident tiles (R4.4 EPI-B, knob SCENESEG_EPI_BIAS_LOCAL): with BLOCAL a core reads only the bias / row tiles of the
// output columns its tile range touches (min(n, CT) tiles from column t0 % CT, wrapping) into their CB slots; the other
// slots stay unread (never used by its compute). Without it every core reads all CT tiles (CT x 4 KB per core: 17.7 MB
// per program on the 1152-channel maps).
// CT args: [0] T, [1] NC, [2] GX, [3] CT (output column tiles), [4] NP (0..3), [5] HAS_BIAS, [6] NB (tiles per read
// batch), [7 + 2i] / [8 + 2i] W_i / O_i (column tiles / column-tile offset of part i, i < 3), [13] HAS_RES,
// [14] HAS_ROW, [15] UNIT (tiles per work unit: the split hands out whole units; 1, or the tiles of one image
// row for the K8 class map), [16] BLOCAL, then the TensorAccessorArgs of part 0, part 1, part 2 (repeats of part 0 when
// NP < 3), the bias, the residual and the row (part 0 when absent). Common RT args: [p0, p1, p2, bias, res,
// row].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t T = 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 CT = get_compile_time_arg_val(3);
constexpr uint32_t NP = get_compile_time_arg_val(4);
constexpr uint32_t HAS_BIAS = get_compile_time_arg_val(5);
constexpr uint32_t NB = get_compile_time_arg_val(6);
constexpr uint32_t W0 = get_compile_time_arg_val(7), O0 = get_compile_time_arg_val(8);
constexpr uint32_t W1 = get_compile_time_arg_val(9), O1 = get_compile_time_arg_val(10);
constexpr uint32_t W2 = get_compile_time_arg_val(11), O2 = get_compile_time_arg_val(12);
constexpr uint32_t HAS_RES = get_compile_time_arg_val(13);
constexpr uint32_t HAS_ROW = get_compile_time_arg_val(14);
constexpr uint32_t UNIT = get_compile_time_arg_val(15);
constexpr uint32_t BLOCAL = get_compile_time_arg_val(16);
constexpr auto p0_args = TensorAccessorArgs<17>();
constexpr auto p1_args = TensorAccessorArgs<p0_args.next_compile_time_args_offset()>();
constexpr auto p2_args = TensorAccessorArgs<p1_args.next_compile_time_args_offset()>();
constexpr auto bs_args = TensorAccessorArgs<p2_args.next_compile_time_args_offset()>();
constexpr auto rs_args = TensorAccessorArgs<bs_args.next_compile_time_args_offset()>();
constexpr auto rw_args = TensorAccessorArgs<rs_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t cb_bias = 3, cb_res = 4, cb_row = 5;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = T / UNIT / NC, E = T / UNIT % NC;
const uint32_t t0 = (core * B + (core < E ? core : E)) * UNIT;
const uint32_t n = (B + (core < E ? 1 : 0)) * UNIT;
if (n == 0) {
return;
}
const auto a0 = TensorAccessor(p0_args, get_common_arg_val<uint32_t>(0), TB);
const auto a1 = TensorAccessor(p1_args, get_common_arg_val<uint32_t>(1), TB);
const auto a2 = TensorAccessor(p2_args, get_common_arg_val<uint32_t>(2), TB);
const auto ar = TensorAccessor(rs_args, get_common_arg_val<uint32_t>(4), TB);
// the resident column tiles this core reads: NR from column C0 (wrapping), or all CT
const uint32_t NR = (BLOCAL && n < CT) ? n : CT;
const uint32_t C0 = (BLOCAL && n < CT) ? t0 % CT : 0;
if constexpr (HAS_BIAS) {
const auto ab = TensorAccessor(bs_args, get_common_arg_val<uint32_t>(3), TB);
cb_reserve_back(cb_bias, CT);
const uint32_t base = get_write_ptr(cb_bias);
for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
noc_async_read(ab.get_noc_addr(c), base + c * TB, TB);
}
noc_async_read_barrier();
cb_push_back(cb_bias, CT);
}
if constexpr (HAS_ROW) {
const auto aw = TensorAccessor(rw_args, get_common_arg_val<uint32_t>(5), TB);
cb_reserve_back(cb_row, CT);
const uint32_t base = get_write_ptr(cb_row);
for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
noc_async_read(aw.get_noc_addr(c), base + c * TB, TB);
}
noc_async_read_barrier();
// row 0 of faces 0 / 1 down their 16 rows, then faces 2 / 3 = faces 0 / 1
for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
const uint32_t tl = base + c * TB;
for (uint32_t f = 0; f < 2; ++f) {
const uint64_t src = get_noc_addr(tl + f * 1024);
for (uint32_t r = 1; r < 16; ++r) {
noc_async_read(src, tl + f * 1024 + r * 64, 64);
}
}
}
noc_async_read_barrier();
for (uint32_t k = 0, c = C0; k < NR; ++k, c = (c + 1 == CT ? 0 : c + 1)) {
const uint32_t tl = base + c * TB;
noc_async_read(get_noc_addr(tl), tl + 2048, 2048);
}
noc_async_read_barrier();
cb_push_back(cb_row, CT);
}
for (uint32_t t = t0; t < t0 + n; t += NB) {
const uint32_t nb = (t + NB <= t0 + n) ? NB : t0 + n - t;
uint32_t w0 = 0, w1 = 0, w2 = 0, wr = 0;
if constexpr (NP > 0) {
cb_reserve_back(0, nb);
w0 = get_write_ptr(0);
}
if constexpr (NP > 1) {
cb_reserve_back(1, nb);
w1 = get_write_ptr(1);
}
if constexpr (NP > 2) {
cb_reserve_back(2, nb);
w2 = get_write_ptr(2);
}
if constexpr (HAS_RES) {
cb_reserve_back(cb_res, nb);
wr = get_write_ptr(cb_res);
}
for (uint32_t j = 0; j < nb; ++j) {
const uint32_t r = (t + j) / CT, c = (t + j) % CT;
if constexpr (NP > 0) {
noc_async_read(a0.get_noc_addr(r * W0 + O0 + c), w0, TB);
w0 += TB;
}
if constexpr (NP > 1) {
noc_async_read(a1.get_noc_addr(r * W1 + O1 + c), w1, TB);
w1 += TB;
}
if constexpr (NP > 2) {
noc_async_read(a2.get_noc_addr(r * W2 + O2 + c), w2, TB);
w2 += TB;
}
if constexpr (HAS_RES) {
noc_async_read(ar.get_noc_addr(t + j), wr, TB);
wr += TB;
}
}
noc_async_read_barrier();
if constexpr (NP > 0) {
cb_push_back(0, nb);
}
if constexpr (NP > 1) {
cb_push_back(1, nb);
}
if constexpr (NP > 2) {
cb_push_back(2, nb);
}
if constexpr (HAS_RES) {
cb_push_back(cb_res, nb);
}
}
}
|