File size: 3,767 Bytes
e0adbab | 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 | // SPDX-License-Identifier: Apache-2.0
// SE-K reader (sceneseg-p150 OPT_PLAN SE-K, tt/sek.py): the operands of one squeeze-excitation MLP, in the order the
// compute kernel consumes them: the fp32 row sum x (KT1 tiles, CB 0), the fc1 weight tiles (2 terms x KT1 x NT1
// bf16, CB 1), the fc1 bias tiles (NT1 fp32, CB 2), then this core's fc2 weight tiles (2 terms x NT1 x its output
// tiles, CB 3) and fc2 bias tiles (CB 4). Every core computes fc1 in full (no inter-core sync); core i of the first
// NC (row-major) owns fc2 output tiles [i * B + min(i, E), + B + (i < E)), B = NT2 / NC, E = NT2 % NC.
// CT args: [0] KT1, [1] NT1, [2] NT2, [3] NC, [4] GX, then the TensorAccessorArgs of x, W1, b1, W2, b2.
// Common RT args: [x, w1, b1, w2, b2].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t KT1 = get_compile_time_arg_val(0);
constexpr uint32_t NT1 = get_compile_time_arg_val(1);
constexpr uint32_t NT2 = get_compile_time_arg_val(2);
constexpr uint32_t NC = get_compile_time_arg_val(3);
constexpr uint32_t GX = get_compile_time_arg_val(4);
constexpr auto x_args = TensorAccessorArgs<5>();
constexpr auto w1_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto b1_args = TensorAccessorArgs<w1_args.next_compile_time_args_offset()>();
constexpr auto w2_args = TensorAccessorArgs<b1_args.next_compile_time_args_offset()>();
constexpr auto b2_args = TensorAccessorArgs<w2_args.next_compile_time_args_offset()>();
constexpr uint32_t TB32 = 4096, TB16 = 2048;
constexpr uint32_t cb_x = 0, cb_w1 = 1, cb_b1 = 2, cb_w2 = 3, cb_b2 = 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 = NT2 / NC, E = NT2 % NC;
const uint32_t n0 = core * B + (core < E ? core : E);
const uint32_t nn = B + (core < E ? 1 : 0);
const auto ax = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB32);
const auto aw1 = TensorAccessor(w1_args, get_common_arg_val<uint32_t>(1), TB16);
const auto ab1 = TensorAccessor(b1_args, get_common_arg_val<uint32_t>(2), TB32);
const auto aw2 = TensorAccessor(w2_args, get_common_arg_val<uint32_t>(3), TB16);
const auto ab2 = TensorAccessor(b2_args, get_common_arg_val<uint32_t>(4), TB32);
cb_reserve_back(cb_x, KT1);
uint32_t l1 = get_write_ptr(cb_x);
for (uint32_t k = 0; k < KT1; ++k, l1 += TB32) {
noc_async_read(ax.get_noc_addr(k), l1, TB32);
}
noc_async_read_barrier();
cb_push_back(cb_x, KT1);
cb_reserve_back(cb_w1, 2 * KT1 * NT1);
l1 = get_write_ptr(cb_w1);
for (uint32_t i = 0; i < 2 * KT1 * NT1; ++i, l1 += TB16) {
noc_async_read(aw1.get_noc_addr(i), l1, TB16);
}
cb_reserve_back(cb_b1, NT1);
l1 = get_write_ptr(cb_b1);
for (uint32_t n = 0; n < NT1; ++n, l1 += TB32) {
noc_async_read(ab1.get_noc_addr(n), l1, TB32);
}
noc_async_read_barrier();
cb_push_back(cb_w1, 2 * KT1 * NT1);
cb_push_back(cb_b1, NT1);
// fc2: tile (t * NT1 + k2) * NT2 + n of W2, pushed term-major, then k2, then this core's n
cb_reserve_back(cb_w2, 2 * NT1 * nn);
l1 = get_write_ptr(cb_w2);
for (uint32_t t = 0; t < 2; ++t) {
for (uint32_t k2 = 0; k2 < NT1; ++k2) {
for (uint32_t j = 0; j < nn; ++j, l1 += TB16) {
noc_async_read(aw2.get_noc_addr((t * NT1 + k2) * NT2 + n0 + j), l1, TB16);
}
}
}
cb_reserve_back(cb_b2, nn);
l1 = get_write_ptr(cb_b2);
for (uint32_t j = 0; j < nn; ++j, l1 += TB32) {
noc_async_read(ab2.get_noc_addr(n0 + j), l1, TB32);
}
noc_async_read_barrier();
cb_push_back(cb_w2, 2 * NT1 * nn);
cb_push_back(cb_b2, nn);
}
|