File size: 2,871 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 | // SPDX-License-Identifier: Apache-2.0
// STEM_POOL reader (meteor-p150 OPT round 3 R3.2, tt/pool.py): the 3x3 / stride-2 / pad-1 max pool of a TILE fp32
// map [1, 1, N*H*W, C] (rows n, h, w; W % 32 == 0, so a tile row block is 32 consecutive w of one (n, h)).
// One work unit = (n, oy, ct): the input rows iy = 2oy-1, 2oy, 2oy+1, channel tile ct, all WT = W / 32 tiles of each
// row, into CBs 0 / 1 / 2 (WT tiles each). Row -1 (oy = 0) is the pool's zero padding: CB 0 is filled with +0.0
// (the stock chain pads with 0.0, tt.layers.maxpool3x3s2). H = 2 * OH, so row 2oy+1 always exists.
// Work split (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 = N * OH * CT units, B = NU / NC, E = NU % NC.
// CT args: [0] NU, [1] NC, [2] GX, [3] H, [4] WT, [5] CT, [6] OH, then the TensorAccessorArgs of the input.
// Common RT args: [in_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 H = get_compile_time_arg_val(3);
constexpr uint32_t WT = get_compile_time_arg_val(4);
constexpr uint32_t CT = get_compile_time_arg_val(5);
constexpr uint32_t OH = get_compile_time_arg_val(6);
constexpr auto in_args = TensorAccessorArgs<7>();
constexpr uint32_t TB = 4096;
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 in = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), TB);
for (uint32_t u = u0; u < u0 + nu; ++u) {
const uint32_t ct = u % CT;
const uint32_t q = u / CT;
const uint32_t oy = q % OH;
const uint32_t n = q / OH;
for (uint32_t k = 0; k < 3; ++k) {
cb_reserve_back(k, WT);
uint32_t w = get_write_ptr(k);
if (oy == 0 && k == 0) {
volatile tt_l1_ptr uint32_t* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(w);
for (uint32_t i = 0; i < WT * 1024; ++i) {
z[i] = 0;
}
} else {
const uint32_t iy = 2 * oy + k - 1;
const uint32_t base = ((n * H + iy) * WT) * CT + ct;
for (uint32_t tx = 0; tx < WT; ++tx) {
noc_async_read(in.get_noc_addr(base + tx * CT), w, TB);
w += TB;
}
}
}
noc_async_read_barrier();
cb_push_back(0, WT);
cb_push_back(1, WT);
cb_push_back(2, WT);
}
}
|