File size: 4,842 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 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 | // SPDX-License-Identifier: Apache-2.0
// TERMS_EPI reader (meteor-p150 OPT_PLAN #2, tt/epi.py): the NP fp32 term outputs of one TermsConv (interleaved TILE,
// one shape) into CBs 0 .. NP-1 and, with HAS_RES, the residual map (same shape) into CB 4, 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 output tiles, B = T / NC, E = T % NC.
// A part may be a column range of a wider tensor (ST2, METEOR_TERMS_ST2: the stacked [W_hi | W_lo] conv's [R, 2C]
// output read as its two products): output tile t = (r, c) = (t / CT, t % CT) reads tile r * W_i + O_i + c of part i
// (W_i = CT, O_i = 0 for a plain part: tile t). The residual has the output's shape (tile t).
// CT args: [0] T, [1] NC, [2] GX, [3] NP (1..3), [4] HAS_RES, [5] NB, [6] CT (output column tiles), [7 + 2i] /
// [8 + 2i] W_i / O_i (column tiles / column-tile offset of part i), then the TensorAccessorArgs of part 0, 1, 2
// (repeats of part 0 when NP < 3) and the residual (part 0 when absent). Common RT args: [p0, p1, p2, res].
#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 HAS_RES = get_compile_time_arg_val(4);
constexpr uint32_t NB = get_compile_time_arg_val(5);
constexpr uint32_t CT = get_compile_time_arg_val(6);
constexpr uint32_t W0 = get_compile_time_arg_val(7), O0 = get_compile_time_arg_val(8);
constexpr uint32_t W1 = get_compile_time_arg_val(9), O1 = get_compile_time_arg_val(10);
constexpr uint32_t W2 = get_compile_time_arg_val(11), O2 = get_compile_time_arg_val(12);
constexpr bool PLAIN = (W0 == CT && O0 == 0 && W1 == CT && O1 == 0 && W2 == CT && O2 == 0);
constexpr auto p0_args = TensorAccessorArgs<13>();
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 rs_args = TensorAccessorArgs<p2_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t cb_res = 4;
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 ar = TensorAccessor(rs_args, get_common_arg_val<uint32_t>(3), TB);
// (r, c) of the next tile for the column-range parts
uint32_t r = t0 / CT, c = t0 % CT;
for (uint32_t t = t0; t < t0 + n; t += NB) {
const uint32_t nb = (t + NB <= t0 + n) ? NB : t0 + n - t;
uint32_t w0 = 0, w1 = 0, w2 = 0, wr = 0;
cb_reserve_back(0, nb);
w0 = get_write_ptr(0);
if constexpr (NP > 1) {
cb_reserve_back(1, nb);
w1 = get_write_ptr(1);
}
if constexpr (NP > 2) {
cb_reserve_back(2, nb);
w2 = get_write_ptr(2);
}
if constexpr (HAS_RES) {
cb_reserve_back(cb_res, nb);
wr = get_write_ptr(cb_res);
}
for (uint32_t j = 0; j < nb; ++j) {
uint32_t i0 = t + j, i1 = t + j, i2 = t + j;
if constexpr (!PLAIN) {
i0 = r * W0 + O0 + c;
i1 = r * W1 + O1 + c;
i2 = r * W2 + O2 + c;
if (++c == CT) {
c = 0;
++r;
}
}
noc_async_read(a0.get_noc_addr(i0), w0, TB);
w0 += TB;
if constexpr (NP > 1) {
noc_async_read(a1.get_noc_addr(i1), w1, TB);
w1 += TB;
}
if constexpr (NP > 2) {
noc_async_read(a2.get_noc_addr(i2), w2, TB);
w2 += TB;
}
if constexpr (HAS_RES) {
noc_async_read(ar.get_noc_addr(t + j), wr, TB);
wr += TB;
}
}
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 (HAS_RES) {
cb_push_back(cb_res, nb);
}
}
}
|