File size: 4,486 Bytes
8eeb550 | 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 | // SPDX-License-Identifier: Apache-2.0
// IMG_IN kernel (meteor-p150 OPT round 3 R3.3, tt/kin.py): the camera input uploaded as whole uint8 rows (one
// [W * 3] page per image row, 2,304 B: H2D of 3,456 pages instead of 2.65 M pages of 3 B) expanded on the device
// into the stem conv's ROW_MAJOR input [N*H*W, CO] bf16 (one CO * 2 B page per pixel; channels 0..2 = the pixel's
// R, G, B as bf16 (0..255 are exact in bf16: the stock typecast), channels 3..CO-1 = +0.0, the zero channels the conv
// pads its 3-channel input with anyway).
// One data-movement kernel on one RISC / NOC (reads and writes from the same NOC: no two-NOC DRAM traffic, SYS-1419).
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns the contiguous image-row
// range [i * B + min(i, E), + B + (i < E)) of the R rows, B = R / NC, E = R % NC.
// CT args: [0] R, [1] NC, [2] GX, [3] W, [4] CO, then the TensorAccessorArgs of the input and the output.
// Common RT args: [in_addr, out_addr]. CB 0: two input rows (staging), CB 1: two output rows (staging), CB 2: the
// 256-entry uint8 -> bf16 table (1 KB, in L1 rather than on the small RISC stack).
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t R = 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 W = get_compile_time_arg_val(3);
constexpr uint32_t CO = get_compile_time_arg_val(4);
constexpr auto in_args = TensorAccessorArgs<5>();
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 PIX = CO * 2; // bytes of one output pixel page
constexpr uint32_t OUT_ROW = W * PIX;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = R / NC, E = R % NC;
const uint32_t r0 = core * B + (core < E ? core : E);
const uint32_t nr = B + (core < E ? 1 : 0);
// the accessors' page strides are the buffers' aligned page sizes (from the accessor CT args)
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 (8 significant bits: 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);
// the zero channels, written once into both output staging rows
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 (nr == 0) {
return;
}
noc_async_read(in.get_noc_addr(r0), ibuf, IN_ROW);
for (uint32_t j = 0; j < nr; ++j) {
const uint32_t k = j & 1;
const uint32_t ib = ibuf + k * IN_ROW;
const uint32_t ob = obuf + k * OUT_ROW;
noc_async_read_barrier();
if (j + 1 < nr) { // prefetch the next row into the other input slot
noc_async_read(in.get_noc_addr(r0 + j + 1), ibuf + (k ^ 1) * IN_ROW, IN_ROW);
}
if (j >= 2) {
noc_async_writes_flushed(); // the writes of row j - 2 (same output slot) have left L1
}
volatile tt_l1_ptr uint8_t* src = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(ib);
volatile tt_l1_ptr uint32_t* dst = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ob);
for (uint32_t w = 0; w < W; ++w) {
dst[w * (PIX / 4)] = lut[src[3 * w]] | (lut[src[3 * w + 1]] << 16);
dst[w * (PIX / 4) + 1] = lut[src[3 * w + 2]]; // channel 3 = +0.0; channels 4.. stay +0.0
}
const uint32_t p0 = (r0 + j) * W;
for (uint32_t w = 0; w < W; ++w) {
noc_async_write(ob + w * PIX, out.get_noc_addr(p0 + w), PIX);
}
}
noc_async_write_barrier();
}
|