File size: 5,851 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 | // SPDX-License-Identifier: Apache-2.0
// DPOOL (meteor-p150 OPT2 round 4 R4.4, tt/dpool.py): the four 2x2-pool quadrant maps of an fp32 TILE map
// [N * H * W, C] (NHWC rows, H even, W % 64 == 0) in ONE data-movement program: map m = (dy, dx) = (m / 2, m % 2) holds
// at output row (n * H / 2 + i) * W / 2 + j the input pixel (n, 2 i + dy, 2 j + dx), all C channels. The depth head's
// avgpool2 then is ((q00 + q10) + (q01 + q11)) * 0.25 with stock fp32 adds: the same operands and order as the port's
// slice / tilize / reshape chain (bit-identical), without its RM round trips.
// Work unit = one output tile row u (32 output pixels of one output image row: W / 2 % 32 == 0): its input pixels are
// the 64 consecutive pixels from (2 q) * W + 2 j0 (top) and (2 q + 1) * W + 2 j0 (bottom), q = u * 32 / (W / 2),
// j0 = u * 32 % (W / 2): input tile rows T, T + 1 (top) and T + W / 32, T + W / 32 + 1 (bottom), CT tiles each.
// Output row k of map (dy, dx) <- input pixel 2 k + dx of the top / bottom span: tile (2 k + dx) / 32, row
// (2 k + dx) % 32; each (row, face column) is a 16-word face row copied by the RISC.
// Element (r, c) of a 32 x 32 fp32 tile: face (r / 16) * 2 + (c % 32) / 16, word (r % 16) * 16 + c % 16.
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns units
// [i * B + min(i, E), + B + (i < E)), B = NU / NC, E = NU % NC. Reads and writes on one RISC / NOC (SYS-1419).
// CT args: [0] NU, [1] NC, [2] GX, [3] CT (channel tiles), [4] WT (= W / 32 input tile rows per image row), then the
// TensorAccessorArgs of the input and of the outputs (all four alike). Common RT args: [in, out0, out1, out2, out3].
// CB 0: 2 x 4 * CT tiles (input slots); CB 1: 2 x 4 * CT tiles (output slots); plain L1 scratch, never pushed.
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NU = 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 CT = get_compile_time_arg_val(3);
constexpr uint32_t WT = 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 TB = 4096; // fp32 tile bytes
constexpr uint32_t TW = 1024; // fp32 tile words
constexpr uint32_t SLOT = 4 * CT * TB; // 4 input tile rows (top 2, bottom 2) / 4 output maps
FORCE_INLINE void cbar() { asm volatile("" ::: "memory"); }
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NU / NC, E = NU % NC;
const uint32_t u0 = core * B + (core < E ? core : E);
const uint32_t nu = B + (core < E ? 1 : 0);
if (nu == 0) {
return;
}
const auto in = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), TB);
const uint32_t ibuf = get_write_ptr(0);
const uint32_t obuf = get_write_ptr(1);
constexpr uint32_t HALF = WT / 2; // output tile rows per output image row (W / 2 / 32)
// input slot layout: [s][t][ct] tiles, s = 0 top / 1 bottom, t = 0 / 1 the tile row of the 64-pixel span
auto issue_reads = [&](uint32_t u, uint32_t buf) {
const uint32_t q = u / HALF;
const uint32_t jb = u - q * HALF;
const uint32_t top = 2 * q * WT + 2 * jb;
for (uint32_t s = 0; s < 2; ++s) {
const uint32_t tr = top + s * WT;
for (uint32_t t = 0; t < 2; ++t) {
for (uint32_t ct = 0; ct < CT; ++ct) {
noc_async_read(in.get_noc_addr((tr + t) * CT + ct), buf + ((s * 2 + t) * CT + ct) * TB, TB);
}
}
}
};
issue_reads(u0, ibuf);
for (uint32_t i = 0; i < nu; ++i) {
const uint32_t ib = ibuf + (i & 1) * SLOT;
const uint32_t ob = obuf + (i & 1) * SLOT;
noc_async_read_barrier();
if (i + 1 < nu) {
issue_reads(u0 + i + 1, ibuf + ((i + 1) & 1) * SLOT);
}
if (i >= 2) {
noc_async_writes_flushed(); // the writes of unit i - 2 (output slot i & 1) have left L1
}
cbar();
const uint32_t* src = reinterpret_cast<const uint32_t*>(ib);
uint32_t* dst = reinterpret_cast<uint32_t*>(ob);
// output slot layout: [m][ct] tiles, m = dy * 2 + dx
for (uint32_t m = 0; m < 4; ++m) {
const uint32_t dy = m >> 1, dx = m & 1;
for (uint32_t ct = 0; ct < CT; ++ct) {
uint32_t* o = dst + (m * CT + ct) * TW;
for (uint32_t k = 0; k < 32; ++k) {
const uint32_t p = 2 * k + dx; // pixel of the 64-pixel span
const uint32_t t = p >> 5, r = p & 31;
const uint32_t* it = src + ((dy * 2 + t) * CT + ct) * TW + (r >> 4) * 512 + (r & 15) * 16;
uint32_t* ot = o + (k >> 4) * 512 + (k & 15) * 16;
#pragma GCC unroll 16
for (uint32_t w = 0; w < 16; ++w) {
ot[w] = it[w]; // face column 0
}
#pragma GCC unroll 16
for (uint32_t w = 0; w < 16; ++w) {
ot[256 + w] = it[256 + w]; // face column 1
}
}
}
}
cbar();
const uint32_t up = u0 + i;
for (uint32_t m = 0; m < 4; ++m) {
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(1 + m), TB);
for (uint32_t ct = 0; ct < CT; ++ct) {
noc_async_write(ob + (m * CT + ct) * TB, out.get_noc_addr(up * CT + ct), TB);
}
}
}
noc_async_write_barrier();
}
|