File size: 5,762 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 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 108 109 110 | // SPDX-License-Identifier: Apache-2.0
// MK-H (OPT round 5, FRNET_POINT_MK=head; OPT_PLAN 3.1 / D19): FRHead as ONE row-streamed module megakernel, reader
// (RISCV_0). Once: every layer's K-concatenated two-term weight W' (bf16 TILE, K' x Nt tiles, row-major tile order)
// into CB 2 (layer after layer, never popped) and every bias row (fp32 TILE, Nt tiles) into CB 3. Then per tile row
// r of this core's contiguous range:
// - the 32 pixel indices of the rows (ROW_MAJOR uint32 [1, N], one page) -- read IB blocks at a time -- and the
// 32 bf16 rows of the gather table (ROW_MAJOR [V, C], one C * 2 byte page per pixel) into CB 0 (KT0 pages =
// one row-major 32-row block; the compute tilizes it): the pixel2point(fuse2) gather;
// - for each layer with a skip, its Nt fp32 tiles of row r (q, pf2, pf1, pf0) into CB 4, in layer order.
// CT args: [0] per-core RT-arg count (P3), [1] NL, [2] KT0 (C / 32), [3] index page bytes (N * 4), [4] IB,
// [5 + 3 l ..] per layer l < MAXL: K' tiles, Nt, SKIP; then the TensorAccessorArgs of the table, of the index, and one
// shared by all W', one by all biases, one by all skips (DRAM interleaved; tt/mk_head.py asserts it).
// Common RT args: [table, index, W_0 .. W_4, B_0 .. B_4, S_0 .. S_4] (unused slots 0).
// Per-core RT args: [r0, nr] (first tile row, tile-row count).
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t MAXL = 5;
constexpr uint32_t NL = get_compile_time_arg_val(1);
constexpr uint32_t KT0 = get_compile_time_arg_val(2);
constexpr uint32_t IPB = get_compile_time_arg_val(3);
constexpr uint32_t IB = get_compile_time_arg_val(4);
constexpr uint32_t KP[MAXL] = {get_compile_time_arg_val(5), get_compile_time_arg_val(8), get_compile_time_arg_val(11),
get_compile_time_arg_val(14), get_compile_time_arg_val(17)};
constexpr uint32_t NT[MAXL] = {get_compile_time_arg_val(6), get_compile_time_arg_val(9), get_compile_time_arg_val(12),
get_compile_time_arg_val(15), get_compile_time_arg_val(18)};
constexpr uint32_t SK[MAXL] = {get_compile_time_arg_val(7), get_compile_time_arg_val(10), get_compile_time_arg_val(13),
get_compile_time_arg_val(16), get_compile_time_arg_val(19)};
constexpr uint32_t ST = SK[0] * NT[0] + SK[1] * NT[1] + SK[2] * NT[2] + SK[3] * NT[3] + SK[4] * NT[4];
constexpr auto t_args = TensorAccessorArgs<20>();
constexpr auto i_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr auto w_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr auto s_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_rm = 0, cb_idx = 1, cb_w = 2, cb_b = 3, cb_s = 4;
constexpr uint32_t RB = KT0 * 64; // table row bytes (C bf16)
constexpr uint32_t WTB = 2048, FTB = 4096;
void kernel_main() {
const uint32_t r0 = get_arg_val<uint32_t>(0);
const uint32_t nr = get_arg_val<uint32_t>(1);
if (nr == 0) {
return;
}
// weights and biases, once
for (uint32_t l = 0; l < NL; ++l) {
const uint32_t nw = KP[l] * NT[l];
const auto w = TensorAccessor(w_args, get_common_arg_val<uint32_t>(2 + l), WTB);
cb_reserve_back(cb_w, nw);
uint32_t wp = get_write_ptr(cb_w);
for (uint32_t t = 0; t < nw; ++t) {
noc_async_read(w.get_noc_addr(t), wp, WTB);
wp += WTB;
}
const auto b = TensorAccessor(b_args, get_common_arg_val<uint32_t>(2 + MAXL + l), FTB);
cb_reserve_back(cb_b, NT[l]);
uint32_t bp = get_write_ptr(cb_b);
for (uint32_t t = 0; t < NT[l]; ++t) {
noc_async_read(b.get_noc_addr(t), bp, FTB);
bp += FTB;
}
noc_async_read_barrier();
cb_push_back(cb_w, nw);
cb_push_back(cb_b, NT[l]);
}
const auto table = TensorAccessor(t_args, get_common_arg_val<uint32_t>(0), RB);
const auto index = TensorAccessor(i_args, get_common_arg_val<uint32_t>(1), IPB);
cb_reserve_back(cb_idx, 1);
volatile tt_l1_ptr uint32_t* idx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_write_ptr(cb_idx));
const uint32_t ibuf = get_write_ptr(cb_idx);
const uint32_t end = r0 + nr;
uint32_t ib_start = r0, ib_end = r0;
for (uint32_t r = r0; r < end; ++r) {
if (r >= ib_end) {
ib_start = r;
const uint32_t m = (end - r) < IB ? (end - r) : IB;
noc_async_read(index.get_noc_addr(0, r * 128), ibuf, m * 128);
noc_async_read_barrier();
ib_end = r + m;
}
cb_reserve_back(cb_rm, KT0);
uint32_t wp = get_write_ptr(cb_rm);
const uint32_t base = (r - ib_start) * 32;
for (uint32_t j = 0; j < 32; ++j) {
noc_async_read(table.get_noc_addr(idx[base + j]), wp, RB);
wp += RB;
}
if constexpr (ST > 0) {
cb_reserve_back(cb_s, ST); // the row's skip tiles of every layer (CB 4 holds 2 ST: no wrap inside)
uint32_t sp = get_write_ptr(cb_s);
for (uint32_t l = 0; l < NL; ++l) {
if (SK[l] == 0) {
continue;
}
const auto s = TensorAccessor(s_args, get_common_arg_val<uint32_t>(2 + 2 * MAXL + l), FTB);
for (uint32_t j = 0; j < NT[l]; ++j) {
noc_async_read(s.get_noc_addr(r * NT[l] + j), sp, FTB);
sp += FTB;
}
}
}
noc_async_read_barrier();
cb_push_back(cb_rm, KT0);
if constexpr (ST > 0) {
cb_push_back(cb_s, ST);
}
}
}
|