File size: 6,066 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 | // SPDX-License-Identifier: Apache-2.0
// FOUT reader (meteor-p150 OPT2 MK-FO, tt/fout.py): the fused BEV map fused = raw + t and its layouts in one program.
// Work units (no per-core RT args, probe P3: core i = y * GX + x of the first NC cores owns the contiguous unit range
// [i * B + min(i, E), + B + (i < E)) of the NU = NBF + NB4 units, B = NU / NC, E = NU % NC):
// u < NBF (flat unit b = u): pixels p0 = 32 b, n = 32; the other operand = raw's flat tiles b * 3 + 0..2
// u >= NBF (4-D unit v = u - NBF): image row y = v / XB, column block xb = v % XB, pixels p0 = y * W + 32 xb,
// n = min(32, W - 32 xb); the other operand = r4's tiles v * 3 + 0..2
// Per unit: the n ROW_MAJOR pixel rows of t (192 B pages p0..p0+n-1) read into an L1 staging buffer (CB 1; NBANK
// reads of up to 32 / NBANK consecutive same-bank pages, see SLOT), then
// tilized by 32 B local NOC copies into CB 0 (3 bf16 tiles: tile k, face 2 (i / 16) + h, face row i % 16 <- row i,
// bytes 64 k + 32 h); rows n..31 zeroed (only when n < 32); the 3 tiles of the other operand into CB 2.
// CT args: [0] NBF, [1] NB4, [2] NC, [3] GX, [4] XB, [5] W, [6] NBANK (DRAM banks; 32 % NBANK == 0), then the
// TensorAccessorArgs of t, raw, r4.
// Common RT args: [t_addr, raw_addr, r4_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NBF = get_compile_time_arg_val(0);
constexpr uint32_t NB4 = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t XB = get_compile_time_arg_val(4);
constexpr uint32_t W = get_compile_time_arg_val(5);
constexpr uint32_t NBANK = get_compile_time_arg_val(6);
constexpr auto t_args = TensorAccessorArgs<7>();
constexpr auto r_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr auto f_args = TensorAccessorArgs<r_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048, ROW = 192;
constexpr uint32_t cb_a = 0, cb_stage = 1, cb_b = 2;
inline void unit_rows(uint32_t u, uint32_t& p0, uint32_t& n) {
if (u < NBF) {
p0 = u * 32;
n = 32;
} else {
const uint32_t v = u - NBF;
const uint32_t y = v / XB, xb = v % XB;
p0 = y * W + xb * 32;
n = (W - xb * 32) < 32 ? (W - xb * 32) : 32;
}
}
// staging slot of pixel row i: rows i, i + 8, i + 16, i + 24 are consecutive pages of one DRAM bank (interleaved
// pages: bank p % NBANK, offset p / NBANK), read as ONE contiguous run of up to 4 rows
#define SLOT(i) ((((i) % NBANK) * (32 / NBANK)) + ((i) / NBANK))
template <typename TA>
inline void read_rows(const TA& t, uint32_t p0, uint32_t n, uint32_t stage) {
for (uint32_t j = 0; j < NBANK && j < n; ++j) {
const uint32_t cnt = (n - j + NBANK - 1) / NBANK;
noc_async_read(t.get_noc_addr(p0 + j), stage + j * (32 / NBANK) * ROW, cnt * ROW);
}
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t NU = NBF + NB4;
constexpr uint32_t B = NU / NC, E = NU % NC;
const uint32_t u0 = core * B + (core < E ? core : E);
const uint32_t n_units = B + (core < E ? 1 : 0);
if (n_units == 0) {
return;
}
const auto t = TensorAccessor(t_args, get_common_arg_val<uint32_t>(0), ROW);
const auto raw = TensorAccessor(r_args, get_common_arg_val<uint32_t>(1), TB16);
const auto r4 = TensorAccessor(f_args, get_common_arg_val<uint32_t>(2), TB16);
// two reader-private staging buffers (CB 1, 2 x 32 rows): reserved once, never pushed; unit u's rows are read
// into stage[u & 1] while unit u - 1's rows are tilized
cb_reserve_back(cb_stage, 1);
const uint32_t stage0 = get_write_ptr(cb_stage);
uint32_t p0, n;
unit_rows(u0, p0, n);
read_rows(t, p0, n, stage0);
noc_async_read_barrier();
const uint64_t here = get_noc_addr(stage0);
for (uint32_t u = u0; u < u0 + n_units; ++u) {
const uint32_t stage = stage0 + ((u - u0) & 1) * 32 * ROW;
cb_reserve_back(cb_a, 3);
cb_reserve_back(cb_b, 3);
const uint32_t wa = get_write_ptr(cb_a);
const uint32_t wb = get_write_ptr(cb_b);
if (n < 32) {
for (uint32_t i = n; i < 32; ++i) {
for (uint32_t k = 0; k < 3; ++k) {
for (uint32_t h = 0; h < 2; ++h) {
volatile tt_l1_ptr uint32_t* q = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(
wa + k * TB16 + ((i >> 4) * 2 + h) * 512 + (i & 15) * 32);
for (uint32_t w = 0; w < 8; ++w) {
q[w] = 0;
}
}
}
}
}
// tilize: 32 B face rows by local NOC copies (state set once per unit: the DRAM reads below reuse the buffer)
noc_async_read_one_packet_set_state(here, 32);
for (uint32_t i = 0; i < n; ++i) {
const uint32_t src = stage + SLOT(i) * ROW;
const uint32_t drow = wa + ((i >> 4) * 2) * 512 + (i & 15) * 32;
for (uint32_t k = 0; k < 3; ++k) {
noc_async_read_one_packet_with_state(src + k * 64, drow + k * TB16);
noc_async_read_one_packet_with_state(src + k * 64 + 32, drow + k * TB16 + 512);
}
}
// the other operand's tiles of this unit, then the next unit's rows into the other staging buffer
for (uint32_t k = 0; k < 3; ++k) {
noc_async_read(u < NBF ? raw.get_noc_addr(u * 3 + k) : r4.get_noc_addr((u - NBF) * 3 + k), wb + k * TB16,
TB16);
}
if (u + 1 < u0 + n_units) {
const uint32_t nstage = stage0 + ((u + 1 - u0) & 1) * 32 * ROW;
unit_rows(u + 1, p0, n);
read_rows(t, p0, n, nstage);
}
noc_async_read_barrier();
cb_push_back(cb_b, 3);
cb_push_back(cb_a, 3);
}
}
|