File size: 3,765 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 | // SPDX-License-Identifier: Apache-2.0
// Fused fp32 matmul attention (tt/fattn_kernel.py), writer (RISCV_1): per unit u (head h = u / Mt, query tile row
// r = u % Mt) the output tile as tile (r, h) of the merged-heads [1, 1, Sq, H * 32] tensor (what nlp_concat_heads
// makes of the [1, H, Sq, 32] P V output).
// With kcat_out (KCAT_EMIT) the output is instead the split operand [o_hi | o_hi | o_lo | 1 | 0..] of the next
// K-concatenated linear (row tiles Ktp, kernels/kcat_writer.cpp's layout): o_hi to (r, h) and (r, H + h), o_lo to
// (r, 2 H + h); the unit of head 0 also writes the ones tile (r, 3 H) and the zero pad tiles.
// CT args: [0] Mt, [1] H, [2] per-core RT-arg count (P3), [3] kcat_out, [4] Ktp, then the TensorAccessorArgs of out.
// Common RT args: [out_addr]. Per-core RT args: [u0, n].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t Mt = get_compile_time_arg_val(0);
constexpr uint32_t H = get_compile_time_arg_val(1);
constexpr uint32_t kcat_out = get_compile_time_arg_val(3);
constexpr uint32_t Ktp = get_compile_time_arg_val(4);
constexpr auto o_args = TensorAccessorArgs<5>();
constexpr uint32_t cb_o = 16, cb_lo = 17, cb_one = 18, cb_zero = 19;
constexpr uint32_t TB = 4096;
void kernel_main() {
const uint32_t o_addr = get_common_arg_val<uint32_t>(0);
const uint32_t u0 = get_arg_val<uint32_t>(0);
const uint32_t n = get_arg_val<uint32_t>(1);
const auto o = TensorAccessor(o_args, o_addr, TB);
if constexpr (!kcat_out) {
for (uint32_t u = u0; u < u0 + n; ++u) {
const uint32_t h = u / Mt;
const uint32_t r = u % Mt;
cb_wait_front(cb_o, 1);
noc_async_write(get_read_ptr(cb_o), o.get_noc_addr(r * H + h), TB);
noc_async_writes_flushed();
cb_pop_front(cb_o, 1);
}
} else {
// ones tile: element (i, c) at face (i / 16) * 2 + (c / 16), offset (i % 16) * 16 + c % 16
uint32_t one_ptr = 0, zero_ptr = 0;
if (u0 < Mt) {
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 * H + 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 u = u0; u < u0 + n; ++u) {
const uint32_t h = u / Mt;
const uint32_t r = u % Mt;
const uint32_t ob = r * Ktp;
cb_wait_front(cb_o, 1);
const uint32_t hp = get_read_ptr(cb_o);
noc_async_write(hp, o.get_noc_addr(ob + h), TB);
noc_async_write(hp, o.get_noc_addr(ob + H + h), TB);
cb_wait_front(cb_lo, 1);
noc_async_write(get_read_ptr(cb_lo), o.get_noc_addr(ob + 2 * H + h), TB);
if (h == 0) {
noc_async_write(one_ptr, o.get_noc_addr(ob + 3 * H), TB);
for (uint32_t z = 3 * H + 1; z < Ktp; ++z) {
noc_async_write(zero_ptr, o.get_noc_addr(ob + z), TB);
}
}
noc_async_writes_flushed();
cb_pop_front(cb_o, 1);
cb_pop_front(cb_lo, 1);
}
}
noc_async_write_barrier();
}
|