File size: 5,401 Bytes
55474c5 | 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 | // SPDX-License-Identifier: Apache-2.0
// kpre (tt/kpre.py, OPT_PLAN R4 device pre-processing), reader (NCRISC). One camera's uint8 node image rows
// (1600 x 3, RGB or BGR as uploaded) -> per H row r: the two lerp operands of the C14 x0.8 horizontal resize as fp32
// pseudo-tiles (1024 contiguous values; the compute kernel's unpack -> SFPU -> pack round trip keeps the order):
// P0[j] = LUT[c][u8[r][(5k + p) * 3 + s]], P1[j] = LUT[c][u8[r][(5k + p + 1) * 3 + s]], j = (4k + p) * 3 + c
// (output pixel x = 4k + p, phase p; s = 2 - c for RGB input) where LUT[c][u] = f32(u + (-mean_c)) (the C14
// normalisation, computed on the host: exact by construction). Once per program: the lerp weight pattern (4 fp32
// pseudo-tiles, w[j] = (2p + 1) / 8) and the 4 vertical weight tiles ((2q + 1) / 8 everywhere) into resident CBs.
// Data movement only; no semaphores, no multicast; common runtime args only (P3): the 4 buffer addresses.
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
#include "kpre_common.h"
constexpr uint32_t N = get_compile_time_arg_val(0);
constexpr uint32_t S = get_compile_time_arg_val(1);
constexpr uint32_t W = get_compile_time_arg_val(2);
constexpr uint32_t RH = get_compile_time_arg_val(3);
constexpr uint32_t NC = get_compile_time_arg_val(4);
constexpr uint32_t GX = get_compile_time_arg_val(5);
constexpr uint32_t RA = get_compile_time_arg_val(6);
constexpr uint32_t RB = get_compile_time_arg_val(7);
constexpr uint32_t RD = get_compile_time_arg_val(8);
constexpr uint32_t SWAP = get_compile_time_arg_val(9); // 1: the image is RGB (network channel c reads 2 - c)
constexpr uint32_t ROW_BYTES = get_compile_time_arg_val(10); // node row: 1600 x 3 bytes
constexpr uint32_t PERIODS = get_compile_time_arg_val(11); // W / 4 (4 output pixels from 5 node pixels)
constexpr uint32_t SPLIT = get_compile_time_arg_val(12); // 1: the writer RISC builds the odd rows (cb_p2)
constexpr uint32_t U16 = get_compile_time_arg_val(13); // 1: raw uint16 planes (normalised on the SFPU)
constexpr auto img_args = TensorAccessorArgs<14>();
constexpr auto lut_args = TensorAccessorArgs<img_args.next_compile_time_args_offset()>();
constexpr auto wx_args = TensorAccessorArgs<lut_args.next_compile_time_args_offset()>();
constexpr auto wy_args = TensorAccessorArgs<wx_args.next_compile_time_args_offset()>();
constexpr auto beta_args = TensorAccessorArgs<wy_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_p = 0, cb_wx = 1, cb_wy = 2, cb_scratch = 4, cb_beta = 7;
constexpr uint32_t TB32 = 4096, TILE_WORDS = 1024;
constexpr uint32_t PT = 5; // pseudo-tiles per node row: its 5 phase planes
static_assert(ROW_BYTES % 64 == 0, "node rows must be 64-byte multiples (DRAM read alignment)");
static_assert(PERIODS * 3 <= TILE_WORDS, "a phase plane fits one pseudo-tile");
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
const KpreGeom g = kpre_geom<N, S, W, RH, NC, RA, RB, RD>(core);
if (!g.rows) {
return;
}
const auto img = TensorAccessor(img_args, get_common_arg_val<uint32_t>(0), ROW_BYTES);
const auto lut_t = TensorAccessor(lut_args, get_common_arg_val<uint32_t>(1), 3 * 256 * 4);
const auto wx_t = TensorAccessor(wx_args, get_common_arg_val<uint32_t>(2), 4 * TB32);
const auto wy_t = TensorAccessor(wy_args, get_common_arg_val<uint32_t>(3), 4 * TB32);
// scratch: LUT (3 KB) then two node rows (double buffer)
const uint32_t base = (get_write_ptr(cb_scratch) + 63u) & ~63u;
const uint32_t lut_l1 = base;
const uint32_t row_l1[2] = {base + 3072, base + 3072 + ROW_BYTES};
noc_async_read(lut_t.get_noc_addr(0), lut_l1, 3 * 256 * 4);
cb_reserve_back(cb_wx, 4);
noc_async_read(wx_t.get_noc_addr(0), get_write_ptr(cb_wx), 4 * TB32);
cb_reserve_back(cb_wy, 4);
noc_async_read(wy_t.get_noc_addr(0), get_write_ptr(cb_wy), 4 * TB32);
if constexpr (U16) { // the per-lane -mean_c tile (normalisation on the SFPU)
const auto beta_t = TensorAccessor(beta_args, get_common_arg_val<uint32_t>(5), TB32);
cb_reserve_back(cb_beta, 1);
noc_async_read(beta_t.get_noc_addr(0), get_write_ptr(cb_beta), TB32);
}
noc_async_read(img.get_noc_addr(g.rlo), row_l1[0], ROW_BYTES);
noc_async_read_barrier();
cb_push_back(cb_wx, 4);
cb_push_back(cb_wy, 4);
if constexpr (U16) {
cb_push_back(cb_beta, 1);
}
const uint32_t* lut = reinterpret_cast<const uint32_t*>(lut_l1);
constexpr uint32_t STEP = SPLIT ? 2 : 1; // SPLIT: rows rlo, rlo + 2, ... here; the writer builds the others
uint32_t buf = 0;
for (uint32_t r = g.rlo; r <= g.rhi; r += STEP) {
if (r + STEP <= g.rhi) { // prefetch the next node row
noc_async_read(img.get_noc_addr(r + STEP), row_l1[buf ^ 1], ROW_BYTES);
}
cb_reserve_back(cb_p, PT);
uint32_t* d0 = reinterpret_cast<uint32_t*>(get_write_ptr(cb_p));
const uint8_t* src = reinterpret_cast<const uint8_t*>(row_l1[buf]);
asm volatile("" ::: "memory");
if constexpr (U16) {
kpre_build_row_u16<PERIODS, SWAP>(src, reinterpret_cast<uint16_t*>(d0));
} else {
kpre_build_row<PERIODS, SWAP>(src, lut, d0);
}
asm volatile("" ::: "memory");
cb_push_back(cb_p, PT);
noc_async_read_barrier();
buf ^= 1;
}
}
|