File size: 8,079 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 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 | // SPDX-License-Identifier: Apache-2.0
// STEM=fold kernel (meteor-p150 OPT2 round 3 R3.1, OPT_PLAN A6, tt/kin.py expand_fold): the camera input uploaded as
// whole uint8 rows (one [W * 3] page per image row) written as the stem's space-to-depth (fold 2) input: the image
// zero-padded by P on every side, cut into 2 x 2 pixel blocks, one block = one folded pixel of 16 bf16 channels
// (32 B page): channel (a * 2 + b) * 3 + c = pixel (2Y + a - P, 2X + b - P), colour c, as bf16 (0..255 exact);
// channels 12..15 = +0.0. The 7 x 7 / s2 / p3 stem conv on the image is then a 4 x 4 / s1 / p0 conv on this map
// (weights re-indexed on the host). Output [N * HF * WF, 16] bf16 ROW_MAJOR, HF = (H + 2P) / 2, WF = (W + 2P) / 2.
// One data-movement kernel on one RISC / NOC (reads and writes from the same NOC, SYS-1419).
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns the contiguous folded-row
// range [i * B + min(i, E), + B + (i < E)) of the RF = N * HF folded rows, B = RF / NC, E = RF % NC.
// CT args: [0] RF, [1] NC, [2] GX, [3] H, [4] W, [5] HF, [6] WF, [7] P, [8] T4, then the TensorAccessorArgs of the
// input and the output.
// T4 = 1 (OPT2 round 6 R6.1, METEOR_STEM_IN=t4): the output is the 4-D TILE tensor [N, HF, WF, 16] (padded to
// [N, HF, TPR * 32, 32], TPR = ceil(WF / 32) tiles per folded row; page = tile (N * HF + Y) * TPR + X / 32). A folded
// pixel's 16 channels are exactly one 32 B face row: pixel r = X % 32 of a tile sits at face (r / 16) * 2 (faces 0 /
// 2), row r % 16; faces 1 / 3 (channels 16..31) and the pad pixels X >= WF stay +0.0. The conv's DRAM-path slices
// then read whole tiles instead of 32 B pages (the values are the same: bit-identical stem output). Common RT args: [in_addr, out_addr]. CB 0: two slots x two input rows (64 B stride), CB 1: two output rows, CB 2: the
// 256-entry uint8 -> bf16 table.
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t RF = 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 H = get_compile_time_arg_val(3);
constexpr uint32_t W = get_compile_time_arg_val(4);
constexpr uint32_t HF = get_compile_time_arg_val(5);
constexpr uint32_t WF = get_compile_time_arg_val(6);
constexpr uint32_t P = get_compile_time_arg_val(7);
constexpr uint32_t T4 = get_compile_time_arg_val(8);
constexpr auto in_args = TensorAccessorArgs<9>();
constexpr auto out_args = TensorAccessorArgs<in_args.next_compile_time_args_offset()>();
constexpr uint32_t IN_ROW = W * 3; // bytes of one uint8 image row
constexpr uint32_t IN_STRIDE = (IN_ROW + 63) / 64 * 64; // L1 stride of a staged row (DRAM reads: 64 B aligned)
constexpr uint32_t PIX = 32; // bytes of one folded pixel (16 bf16 channels)
constexpr uint32_t TPR = (WF + 31) / 32; // T4: tiles per folded row
constexpr uint32_t TILE_B = 2048; // T4: one bf16 32 x 32 tile
constexpr uint32_t OUT_ROW = T4 ? TPR * TILE_B : WF * PIX;
// L1 byte offset of folded pixel x in a staged output row
inline uint32_t pix_off(uint32_t x) {
if constexpr (T4) {
const uint32_t r = x & 31;
return (x >> 5) * TILE_B + (r >> 4) * 1024 + (r & 15) * PIX;
} else {
return x * PIX;
}
}
// the image rows (2Y - P, 2Y + 1 - P) of folded row Y of camera n into slot `ib` (rows outside the image: not read)
template <typename A>
inline void read_rows(const A& in, uint32_t f, uint32_t ib) {
const uint32_t n = f / HF, y = f % HF;
for (uint32_t a = 0; a < 2; ++a) {
const int32_t r = static_cast<int32_t>(2 * y + a) - static_cast<int32_t>(P);
if (r >= 0 && r < static_cast<int32_t>(H)) {
noc_async_read(in.get_noc_addr(n * H + static_cast<uint32_t>(r)), ib + a * IN_STRIDE, IN_ROW);
}
}
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = RF / NC, E = RF % NC;
const uint32_t f0 = core * B + (core < E ? core : E);
const uint32_t nf = B + (core < E ? 1 : 0);
const auto in = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0));
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(1));
// bf16 bits of the integers 0..255 (exact; the stock uint8 -> bf16 typecast)
volatile tt_l1_ptr uint32_t* lut = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(2)); // CB 2
lut[0] = 0;
for (uint32_t v = 1; v < 256; ++v) {
uint32_t e = 0;
while ((v >> (e + 1)) != 0) {
++e;
}
lut[v] = ((127 + e) << 7) | ((v << (7 - e)) & 0x7F);
}
const uint32_t ibuf = get_write_ptr(0);
const uint32_t obuf = get_write_ptr(1);
// both output staging rows zeroed once: channels 12..15 (words 6, 7 of a pixel) are never written again (T4: nor
// faces 1 / 3 and the pad pixels)
for (uint32_t k = 0; k < 2; ++k) {
volatile tt_l1_ptr uint32_t* o = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(obuf + k * OUT_ROW);
for (uint32_t i = 0; i < OUT_ROW / 4; ++i) {
o[i] = 0;
}
}
if (nf == 0) {
return;
}
read_rows(in, f0, ibuf);
for (uint32_t j = 0; j < nf; ++j) {
const uint32_t f = f0 + j;
const uint32_t k = j & 1;
const uint32_t ib = ibuf + k * 2 * IN_STRIDE;
const uint32_t ob = obuf + k * OUT_ROW;
noc_async_read_barrier();
if (j + 1 < nf) { // prefetch the next folded row's image rows into the other slot
read_rows(in, f + 1, ibuf + (k ^ 1) * 2 * IN_STRIDE);
}
if (j >= 2) {
noc_async_writes_flushed(); // the writes of row j - 2 (same output slot) have left L1
}
const uint32_t y = f % HF;
volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ob);
for (uint32_t a = 0; a < 2; ++a) {
const int32_t r = static_cast<int32_t>(2 * y + a) - static_cast<int32_t>(P);
const bool row_ok = r >= 0 && r < static_cast<int32_t>(H);
volatile tt_l1_ptr uint8_t* src = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(ib + a * IN_STRIDE);
for (uint32_t x = 0; x < WF; ++x) {
volatile tt_l1_ptr uint32_t* d = dst + pix_off(x) / 4 + a * 3;
const int32_t c0 = static_cast<int32_t>(2 * x) - static_cast<int32_t>(P);
if (!row_ok) {
d[0] = 0;
d[1] = 0;
d[2] = 0;
} else if (c0 >= 0 && c0 + 1 < static_cast<int32_t>(W)) { // both columns inside the image
volatile tt_l1_ptr uint8_t* s = src + 3 * c0;
d[0] = lut[s[0]] | (lut[s[1]] << 16);
d[1] = lut[s[2]] | (lut[s[3]] << 16);
d[2] = lut[s[4]] | (lut[s[5]] << 16);
} else { // an edge block: per column
uint32_t v[6];
for (uint32_t i = 0; i < 6; ++i) {
const int32_t col = c0 + static_cast<int32_t>(i / 3);
v[i] = (col >= 0 && col < static_cast<int32_t>(W)) ? lut[src[3 * col + i % 3]] : 0;
}
d[0] = v[0] | (v[1] << 16);
d[1] = v[2] | (v[3] << 16);
d[2] = v[4] | (v[5] << 16);
}
}
}
if constexpr (T4) {
const uint32_t t0 = f * TPR;
for (uint32_t t = 0; t < TPR; ++t) {
noc_async_write(ob + t * TILE_B, out.get_noc_addr(t0 + t), TILE_B);
}
} else {
const uint32_t p0 = f * WF;
for (uint32_t x = 0; x < WF; ++x) {
noc_async_write(ob + x * PIX, out.get_noc_addr(p0 + x), PIX);
}
}
}
noc_async_write_barrier();
}
|