File size: 4,520 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 | // SPDX-License-Identifier: Apache-2.0
// STEM_POOL writer (meteor-p150 OPT round 3 R3.2, tt/pool.py): the horizontal half of the 3x3 / s2 max pool and the
// output. Per unit (work split as the reader), the WT vertical-max tiles of CB 16 hold one input row of 32 * WT
// columns w (tile tx = w / 32, tile row w % 32) for one channel tile. Output column ox takes max(v[2ox-1], v[2ox],
// v[2ox+1]); column -1 is the zero padding (+0.0; W = 2 * OW, so 2ox+1 always exists). The OWT = WT / 2 output tiles
// are built in a scratch CB (24) on this RISC and written to the TILE fp32 output [1, 1, N*OH*OW, C].
// Compare: the fp32 bits as int32. The inputs are post-ReLU (>= +0.0, or -0.0), where the int32 order is the float
// order with -0.0 < +0.0 (that of the SFPU min / max swap); fp32 is never rounded here.
// Tile element (row r, column c): face ((r / 16) * 2 + c / 16) * 256 + (r % 16) * 16 + c % 16.
// CT args: [0] NU, [1] NC, [2] GX, [3] WT, [4] CT, [5] OH, then the TensorAccessorArgs of the output.
// Common RT args: [out_addr].
#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 WT = get_compile_time_arg_val(3);
constexpr uint32_t CT = get_compile_time_arg_val(4);
constexpr uint32_t OH = get_compile_time_arg_val(5);
constexpr auto out_args = TensorAccessorArgs<6>();
constexpr uint32_t OWT = WT / 2;
constexpr uint32_t TB = 4096;
constexpr uint32_t cb_v = 16, cb_s = 24;
static inline uint32_t row_off(uint32_t r) { return (r >> 4) * 512 + (r & 15) * 16; } // + (c / 16) * 256 + c % 16
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);
const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), TB);
volatile tt_l1_ptr int32_t* s = reinterpret_cast<volatile tt_l1_ptr int32_t*>(get_write_ptr(cb_s));
for (uint32_t u = u0; u < u0 + nu; ++u) {
const uint32_t ct = u % CT;
const uint32_t q = u / CT; // n * OH + oy
cb_wait_front(cb_v, WT);
volatile tt_l1_ptr int32_t* v = reinterpret_cast<volatile tt_l1_ptr int32_t*>(get_read_ptr(cb_v));
noc_async_write_barrier(); // the scratch tiles of the previous unit are out
for (uint32_t ot = 0; ot < OWT; ++ot) {
volatile tt_l1_ptr int32_t* o = s + ot * 1024;
for (uint32_t r = 0; r < 32; ++r) {
const uint32_t w1 = (ot * 32 + r) * 2; // 2 ox
const uint32_t w2 = w1 + 1;
volatile tt_l1_ptr int32_t* p1 = v + (w1 >> 5) * 1024 + row_off(w1 & 31);
volatile tt_l1_ptr int32_t* p2 = v + (w2 >> 5) * 1024 + row_off(w2 & 31);
volatile tt_l1_ptr int32_t* q0 = o + row_off(r);
if (w1 == 0) { // column -1: the +0.0 padding
for (uint32_t f = 0; f < 2; ++f) {
for (uint32_t c = 0; c < 16; ++c) {
int32_t m = p1[f * 256 + c];
const int32_t b = p2[f * 256 + c];
m = b > m ? b : m;
q0[f * 256 + c] = m > 0 ? m : 0;
}
}
} else {
const uint32_t w0 = w1 - 1;
volatile tt_l1_ptr int32_t* p0 = v + (w0 >> 5) * 1024 + row_off(w0 & 31);
for (uint32_t f = 0; f < 2; ++f) {
for (uint32_t c = 0; c < 16; ++c) {
int32_t m = p0[f * 256 + c];
const int32_t a = p1[f * 256 + c];
const int32_t b = p2[f * 256 + c];
m = a > m ? a : m;
q0[f * 256 + c] = b > m ? b : m;
}
}
}
}
}
cb_pop_front(cb_v, WT);
const uint32_t base = (q * OWT) * CT + ct;
for (uint32_t ot = 0; ot < OWT; ++ot) {
noc_async_write(reinterpret_cast<uint32_t>(s + ot * 1024), out.get_noc_addr(base + ot * CT), TB);
}
}
noc_async_write_barrier();
}
|