File size: 5,759 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 | // SPDX-License-Identifier: Apache-2.0
// Stem reader (sceneseg-p150 OPT_PLAN K-IN kernel form / MK0, tt/stem.py): the stem conv's im2col tiles straight from
// the uploaded uint8 squash (ROW_MAJOR, one page of WIN * 3 bytes per image row: B, G, R per pixel), as bf16
// (0..255 is exact in bf16). Output tile t = 32 consecutive output pixels of one output row (WOUT % 32 == 0); A row
// i = pixel i, column k = (dy * 3 + dx) * 3 + c (27 used, 5 zero) of the 3x3 stride-2 padding-1 window; taps outside
// the image are 0 (the zero padding of the uint8 input; the normalisation's mean term is in the bias map). Also the
// tile's fp32 bias-map tile (the stem's exact border-aware bias, tt/layers.py stem_fold) into CB 2, and once the two
// bf16 weight tiles (hi, lo: [k, cout]) into CB 1.
// Work split (no per-core RT args, probe P3): core i of the first NC (row-major) owns tiles [i * B + min(i, E), + B +
// (i < E)), B = T / NC, E = T % NC. The three input rows of an output row are read once per output row.
// CT args: [0] T, [1] NC, [2] GX, [3] HIN, [4] WIN, [5] WOUT, then the TensorAccessorArgs of the image, the weight
// tiles and the bias map. Common RT args: [image, weights, bias].
#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 HIN = get_compile_time_arg_val(3);
constexpr uint32_t WIN = get_compile_time_arg_val(4);
constexpr uint32_t WOUT = get_compile_time_arg_val(5);
constexpr auto im_args = TensorAccessorArgs<6>();
constexpr auto w_args = TensorAccessorArgs<im_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr uint32_t ROWB = WIN * 3;
constexpr uint32_t TPR = WOUT / 32; // output tiles per output row
constexpr uint32_t cb_a = 0, cb_w = 1, cb_b = 2, cb_rows = 8, cb_lut = 9;
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 / NC, E = T % NC;
const uint32_t t0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
const auto aim = TensorAccessor(im_args, get_common_arg_val<uint32_t>(0), ROWB);
const auto aw = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), 2048);
const auto ab = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), 4096);
cb_reserve_back(cb_w, 2);
noc_async_read(aw.get_noc_addr(0), get_write_ptr(cb_w), 2048);
noc_async_read(aw.get_noc_addr(1), get_write_ptr(cb_w) + 2048, 2048);
// bf16 bits of 0..255 (exact): v = 1.m * 2^e, 8 significant bits fit the 7-bit mantissa + the implicit 1
volatile tt_l1_ptr uint16_t* lut = reinterpret_cast<volatile tt_l1_ptr uint16_t*>(get_write_ptr(cb_lut));
lut[0] = 0;
for (uint32_t v = 1; v < 256; ++v) {
uint32_t e = 7;
while (((v >> e) & 1) == 0) {
--e;
}
const uint32_t bits = ((e + 127) << 23) | ((v << (23 - e)) & 0x7FFFFF);
lut[v] = (uint16_t)(bits >> 16);
}
const uint32_t rows = get_write_ptr(cb_rows);
volatile tt_l1_ptr uint8_t* rowp = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(rows);
noc_async_read_barrier();
cb_push_back(cb_w, 2);
uint32_t cur = 0xFFFFFFFFu;
bool valid[3] = {false, false, false};
for (uint32_t t = t0; t < t0 + n; ++t) {
const uint32_t oy = t / TPR, ox0 = (t % TPR) * 32;
if (oy != cur) {
cur = oy;
for (uint32_t dy = 0; dy < 3; ++dy) {
const int32_t iy = (int32_t)(2 * oy + dy) - 1;
valid[dy] = iy >= 0 && iy < (int32_t)HIN;
if (valid[dy]) {
noc_async_read(aim.get_noc_addr((uint32_t)iy), rows + dy * ROWB, ROWB);
}
}
noc_async_read_barrier();
}
cb_reserve_back(cb_b, 1);
noc_async_read(ab.get_noc_addr(t), get_write_ptr(cb_b), 4096);
cb_reserve_back(cb_a, 1);
volatile tt_l1_ptr uint32_t* a = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_a));
for (uint32_t i = 0; i < 32; ++i) {
const int32_t ixc = (int32_t)(2 * (ox0 + i)); // input column of dx = 1
uint16_t v[32];
uint32_t k = 0;
for (uint32_t dy = 0; dy < 3; ++dy) {
for (int32_t dx = -1; dx <= 1; ++dx) {
const int32_t ix = ixc + dx;
if (valid[dy] && ix >= 0 && ix < (int32_t)WIN) {
const uint32_t o = dy * ROWB + (uint32_t)ix * 3;
v[k] = lut[rowp[o]];
v[k + 1] = lut[rowp[o + 1]];
v[k + 2] = lut[rowp[o + 2]];
} else {
v[k] = v[k + 1] = v[k + 2] = 0;
}
k += 3;
}
}
for (; k < 32; ++k) {
v[k] = 0;
}
// row i of the tile: face (i / 16) * 2 + 0 holds k 0..15, face + 1 holds k 16..31; 16 bf16 per face row
const uint32_t base = ((i >> 4) * 2) * 128 + (i & 15) * 8; // in uint32 (2 bf16) units
for (uint32_t j = 0; j < 8; ++j) {
a[base + j] = (uint32_t)v[2 * j] | ((uint32_t)v[2 * j + 1] << 16);
a[base + 128 + j] = (uint32_t)v[16 + 2 * j] | ((uint32_t)v[16 + 2 * j + 1] << 16);
}
}
cb_push_back(cb_a, 1);
noc_async_read_barrier();
cb_push_back(cb_b, 1);
}
}
|