File size: 7,068 Bytes
be62f78 | 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 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 | // SPDX-License-Identifier: Apache-2.0
// Split-row fused fp32 LayerNorm (tt/ln_kernel.py, LN_SPLIT), writer (RISCV_1). Per row: (write_h) the new stream
// tile; sends this member's tile for the mean fold (c_11, from the compute) to page j of the root's gather CB c_10
// and raises the root's semaphore 0; on the root, broadcasts the packed mean (c_4, column 0) over the 32 columns and
// sends it to page 0 of every member's c_5 (semaphore 1 raised); the same for the squared tile and rstd (page 1);
// then writes the output tile.
// With kcat_out (KCAT_EMIT) out is the split operand [y_hi | y_hi | y_lo | 1 | 0..] of the next K-concatenated
// linear (row tiles Ktp, the kcat_writer layout): y_hi to (r, j) and (r, Wt + j), y_lo to (r, 2 Wt + j); the root
// also writes the ones tile (r, 3 Wt) and the zero pad tiles.
// CT args: [0] Wt, [1] per-core RT-arg count (P3), [2] write_h, [3] sfpu_bcast, [4] kcat_out, [5] Ktp, then the
// TensorAccessorArgs of out and h.
// Common RT args: [out_addr, h_addr]. Per-core RT args: as the reader's.
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t Wt = get_compile_time_arg_val(0);
constexpr uint32_t write_h = get_compile_time_arg_val(2);
constexpr uint32_t sfpu_bcast = get_compile_time_arg_val(3); // the root's compute packs a broadcast tile
constexpr uint32_t kcat_out = get_compile_time_arg_val(4);
constexpr uint32_t Ktp = get_compile_time_arg_val(5);
constexpr auto out_args = TensorAccessorArgs<6>();
constexpr auto h_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_stat = 4, cb_bc = 5, cb_gat = 10, cb_snd = 11, cb_scr = 12, cb_out = 16, cb_hout = 17,
cb_lo = 19, cb_one = 20, cb_zero = 21;
constexpr uint32_t TB = 4096;
FORCE_INLINE void bcast_col(uint32_t src, uint32_t dst) {
auto* s = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(src);
auto* d = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(dst);
for (uint32_t fo = 0; fo < 1024; fo += 512) {
for (uint32_t ro = 0; ro < 256; ro += 16) {
const uint32_t v = s[fo + ro];
volatile tt_l1_ptr uint32_t* l = d + fo + ro;
volatile tt_l1_ptr uint32_t* r = l + 256;
for (uint32_t c = 0; c < 16; ++c) {
l[c] = v;
r[c] = v;
}
}
}
}
void kernel_main() {
const uint32_t out_addr = get_common_arg_val<uint32_t>(0);
const uint32_t h_addr = get_common_arg_val<uint32_t>(1);
const uint32_t row0 = get_arg_val<uint32_t>(0);
const uint32_t n_rows = get_arg_val<uint32_t>(1);
const uint32_t stride = get_arg_val<uint32_t>(2);
const uint32_t j = get_arg_val<uint32_t>(3);
const uint32_t root_x = get_arg_val<uint32_t>(4);
const uint32_t root_y = get_arg_val<uint32_t>(5);
if (n_rows == 0) {
return;
}
const bool root = j == 0;
const auto out = TensorAccessor(out_args, out_addr, TB);
const auto hacc = TensorAccessor(h_args, h_addr, TB);
const uint32_t gat_base = get_write_ptr(cb_gat); // the same L1 address on every core of the program
const uint32_t bc_base = get_write_ptr(cb_bc);
const uint64_t root_sem = get_noc_addr(root_x, root_y, get_semaphore(0));
const uint32_t sem_b_addr = get_semaphore(1);
cb_reserve_back(cb_scr, 1);
const uint32_t scr = get_write_ptr(cb_scr);
auto send_to_root = [&]() {
cb_wait_front(cb_snd, 1);
noc_async_write(get_read_ptr(cb_snd), get_noc_addr(root_x, root_y, gat_base + j * TB), TB);
noc_async_write_barrier();
noc_semaphore_inc(root_sem, 1);
cb_pop_front(cb_snd, 1);
};
auto bcast_out = [&](uint32_t page) {
cb_wait_front(cb_stat, 1);
uint32_t src = get_read_ptr(cb_stat);
if constexpr (!sfpu_bcast) {
bcast_col(src, scr);
src = scr;
}
for (uint32_t m = 0; m < Wt; ++m) {
const uint32_t mx = get_arg_val<uint32_t>(6 + 2 * m), my = get_arg_val<uint32_t>(7 + 2 * m);
noc_async_write(src, get_noc_addr(mx, my, bc_base + page * TB), TB);
}
noc_async_write_barrier();
cb_pop_front(cb_stat, 1);
for (uint32_t m = 0; m < Wt; ++m) {
const uint32_t mx = get_arg_val<uint32_t>(6 + 2 * m), my = get_arg_val<uint32_t>(7 + 2 * m);
noc_semaphore_inc(get_noc_addr(mx, my, sem_b_addr), 1);
}
};
uint32_t one_ptr = 0, zero_ptr = 0;
if constexpr (kcat_out) {
if (root) {
// ones tile: element (i, c) at face (i / 16) * 2 + (c / 16), offset (i % 16) * 16 + c % 16
cb_reserve_back(cb_one, 1);
one_ptr = get_write_ptr(cb_one);
auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(one_ptr);
for (uint32_t i = 0; i < 1024; ++i) {
p[i] = 0;
}
for (uint32_t i = 0; i < 32; ++i) {
const uint32_t base = (i / 16) * 512 + (i % 16) * 16;
p[base] = 0x3F800000u;
p[base + 1] = 0x3F800000u;
}
if constexpr (Ktp > 3 * Wt + 1) {
cb_reserve_back(cb_zero, 1);
zero_ptr = get_write_ptr(cb_zero);
auto* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero_ptr);
for (uint32_t i = 0; i < 1024; ++i) {
z[i] = 0;
}
}
}
}
for (uint32_t i = 0; i < n_rows; ++i) {
const uint32_t r = row0 + i * stride;
if constexpr (write_h) {
cb_wait_front(cb_hout, 1);
noc_async_write(get_read_ptr(cb_hout), hacc.get_noc_addr(r * Wt + j), TB);
noc_async_writes_flushed();
cb_pop_front(cb_hout, 1);
}
send_to_root();
if (root) {
bcast_out(0);
}
send_to_root();
if (root) {
bcast_out(1);
}
if constexpr (kcat_out) {
const uint32_t ob = r * Ktp;
cb_wait_front(cb_out, 1);
const uint32_t hp = get_read_ptr(cb_out);
noc_async_write(hp, out.get_noc_addr(ob + j), TB);
noc_async_write(hp, out.get_noc_addr(ob + Wt + j), TB);
cb_wait_front(cb_lo, 1);
noc_async_write(get_read_ptr(cb_lo), out.get_noc_addr(ob + 2 * Wt + j), TB);
if (root) {
noc_async_write(one_ptr, out.get_noc_addr(ob + 3 * Wt), TB);
for (uint32_t z = 3 * Wt + 1; z < Ktp; ++z) {
noc_async_write(zero_ptr, out.get_noc_addr(ob + z), TB);
}
}
noc_async_writes_flushed();
cb_pop_front(cb_out, 1);
cb_pop_front(cb_lo, 1);
} else {
cb_wait_front(cb_out, 1);
noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(r * Wt + j), TB);
noc_async_writes_flushed();
cb_pop_front(cb_out, 1);
}
}
noc_async_write_barrier();
noc_async_atomic_barrier();
}
|