File size: 2,769 Bytes
e91e989 | 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 | // SPDX-License-Identifier: Apache-2.0
// TGATE reader (meteor-p150 OPT2 round 3 R3.2, tt/tgate.py): the merged temporal-fuse / gate conv output g
// [R, GT * 32] fp32 TILE (columns 0..95 = tfuse.0's W x, columns 96..98 = the gate logit differences l_j - l_0) per
// 32-row block b: tiles b * GT + 0..2 into CB 0, tile b * GT + 3 into CB 1; once: the ones tile (CB 2) and the 3
// tfuse.0 bias tiles (CB 3, every row the bias).
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns the contiguous block
// range [i * B + min(i, E), + B + (i < E)) of the NB blocks, B = NB / NC, E = NB % NC.
// CT args: [0] NB, [1] NC, [2] GX, [3] GT, then the TensorAccessorArgs of g, ones, bias.
// Common RT args: [g_addr, ones_addr, bias_addr].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t NB = 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 GT = get_compile_time_arg_val(3);
constexpr auto g_args = TensorAccessorArgs<4>();
constexpr auto o_args = TensorAccessorArgs<g_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;
constexpr uint32_t cb_t = 0, cb_l = 1, cb_one = 2, cb_bias = 3;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NB / NC, E = NB % NC;
const uint32_t b0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
const auto g = TensorAccessor(g_args, get_common_arg_val<uint32_t>(0), TB);
const auto ones = TensorAccessor(o_args, get_common_arg_val<uint32_t>(1), TB);
const auto bias = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2), TB);
cb_reserve_back(cb_one, 1);
noc_async_read(ones.get_noc_addr(0), get_write_ptr(cb_one), TB);
cb_reserve_back(cb_bias, 3);
for (uint32_t k = 0; k < 3; ++k) {
noc_async_read(bias.get_noc_addr(k), get_write_ptr(cb_bias) + k * TB, TB);
}
noc_async_read_barrier();
cb_push_back(cb_one, 1);
cb_push_back(cb_bias, 3);
for (uint32_t b = b0; b < b0 + n; ++b) {
cb_reserve_back(cb_l, 1);
cb_reserve_back(cb_t, 3);
noc_async_read(g.get_noc_addr(b * GT + 3), get_write_ptr(cb_l), TB);
const uint32_t wt = get_write_ptr(cb_t);
for (uint32_t k = 0; k < 3; ++k) {
noc_async_read(g.get_noc_addr(b * GT + k), wt + k * TB, TB);
}
noc_async_read_barrier();
cb_push_back(cb_l, 1);
cb_push_back(cb_t, 3);
}
}
|