File size: 3,351 Bytes
7a24154 | 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 | // SPDX-License-Identifier: Apache-2.0
// EPI reader (frnet-p150 OPT round 2, tt/epi.py; after sceneseg-p150's EPI): the NP (1..5) fp32 TILE operands of one
// image-branch epilogue (K-split partial convs, or a conv output and its residual), interleaved, one padded shape,
// into CBs 0 .. NP-1, NB tiles per batch.
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores in row-major order owns the
// contiguous tile range [i * B + min(i, E), + B + (i < E)) of the T tiles, B = T / NC, E = T % NC.
// CT args: [0] T, [1] NC, [2] GX, [3] NP, [4] NB, then the TensorAccessorArgs of parts 0..4 (repeats of part 0 when
// NP < 5). Common RT args: [p0, p1, p2, p3, p4] buffer addresses.
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t T = 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 NP = get_compile_time_arg_val(3);
constexpr uint32_t NB = get_compile_time_arg_val(4);
constexpr auto p0_args = TensorAccessorArgs<5>();
constexpr auto p1_args = TensorAccessorArgs<p0_args.next_compile_time_args_offset()>();
constexpr auto p2_args = TensorAccessorArgs<p1_args.next_compile_time_args_offset()>();
constexpr auto p3_args = TensorAccessorArgs<p2_args.next_compile_time_args_offset()>();
constexpr auto p4_args = TensorAccessorArgs<p3_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
template <typename A>
FORCE_INLINE void read_part(const A& acc, uint32_t cb, uint32_t t, uint32_t nb) {
cb_reserve_back(cb, nb);
uint32_t w = get_write_ptr(cb);
for (uint32_t j = 0; j < nb; ++j) {
noc_async_read(acc.get_noc_addr(t + j), w, TB);
w += TB;
}
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = T / NC, E = T % NC;
const uint32_t t0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
const auto a0 = TensorAccessor(p0_args, get_common_arg_val<uint32_t>(0), TB);
const auto a1 = TensorAccessor(p1_args, get_common_arg_val<uint32_t>(1), TB);
const auto a2 = TensorAccessor(p2_args, get_common_arg_val<uint32_t>(2), TB);
const auto a3 = TensorAccessor(p3_args, get_common_arg_val<uint32_t>(3), TB);
const auto a4 = TensorAccessor(p4_args, get_common_arg_val<uint32_t>(4), TB);
for (uint32_t t = t0; t < t0 + n; t += NB) {
const uint32_t nb = (t + NB <= t0 + n) ? NB : t0 + n - t;
read_part(a0, 0, t, nb);
if constexpr (NP > 1) {
read_part(a1, 1, t, nb);
}
if constexpr (NP > 2) {
read_part(a2, 2, t, nb);
}
if constexpr (NP > 3) {
read_part(a3, 3, t, nb);
}
if constexpr (NP > 4) {
read_part(a4, 4, t, nb);
}
noc_async_read_barrier();
cb_push_back(0, nb);
if constexpr (NP > 1) {
cb_push_back(1, nb);
}
if constexpr (NP > 2) {
cb_push_back(2, nb);
}
if constexpr (NP > 3) {
cb_push_back(3, nb);
}
if constexpr (NP > 4) {
cb_push_back(4, nb);
}
}
}
|