File size: 3,412 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 | // SPDX-License-Identifier: Apache-2.0
// P2P (OPT round 4, FRNET_P2P): pixel2point gather straight into the KCAT operand, writer (RISCV_1): the Kt tiles of
// block (tile row) r go to X' [.., N, 32 OKt] at tiles (r, off + j) and (r, off + Kt + j): the [x | x] columns of a
// bf16 part (as kcat_writer.cpp mode 1).
// CT args: [0] per-core RT-arg count (P3), [1] Kt, [2] OKt, [3] off, [4] DIAG (0; 2 = timing diagnostic without writes),
// [5] BCOL (tile column of the bias block), [6] BT (bias block tiles; 0 = none; kcat_writer.cpp), then the
// TensorAccessorArgs of X'.
// Common RT args: [out_addr]. Per-core RT args: [b0, nb].
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t Kt = get_compile_time_arg_val(1);
constexpr uint32_t OKt = get_compile_time_arg_val(2);
constexpr uint32_t off = get_compile_time_arg_val(3);
constexpr uint32_t DIAG = get_compile_time_arg_val(4); // timing diagnostic: 2 = no DRAM writes
constexpr uint32_t BCOL = get_compile_time_arg_val(5);
constexpr uint32_t BT = get_compile_time_arg_val(6);
constexpr auto o_args = TensorAccessorArgs<7>();
constexpr uint32_t cb_t = 16, cb_c = 18;
constexpr uint32_t TB = 2048;
// The bias block (FRNET_KCAT_BIAS=col; BT > 0 only in the writer of part 0): tile 0 of the block is 1.0 in columns
// 0..2 of every row (x the W' rows b_hi, b_lo, b_lo2), tiles 1..BT-1 are zero. Built once in CB 18 (2 pages; never
// pushed) with 32-bit L1 stores: faces of 16 x 16 bf16, columns 0..3 of row i at face (i / 16) * 2.
FORCE_INLINE uint32_t fill_bias_tiles(uint32_t cb_c) {
cb_reserve_back(cb_c, 2);
const uint32_t base = get_write_ptr(cb_c);
volatile tt_l1_ptr uint32_t* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base);
for (uint32_t i = 0; i < 2 * 2048 / 4; ++i) {
p[i] = 0;
}
for (uint32_t i = 0; i < 32; ++i) {
const uint32_t e = (i >> 4) * 512 + (i & 15) * 16; // bf16 element index of (row i, column 0)
p[e / 2] = 0x3F803F80u; // columns 0, 1 = 1.0
p[e / 2 + 1] = 0x00003F80u; // column 2 = 1.0, column 3 = 0
}
return base;
}
template <typename A>
FORCE_INLINE void write_bias_block(uint32_t ones, uint32_t row_tile0, uint32_t bt, const A& o) {
noc_async_write(ones, o.get_noc_addr(row_tile0), 2048);
for (uint32_t b = 1; b < bt; ++b) {
noc_async_write(ones + 2048, o.get_noc_addr(row_tile0 + b), 2048);
}
}
void kernel_main() {
const uint32_t b0 = get_arg_val<uint32_t>(0);
const uint32_t nb = get_arg_val<uint32_t>(1);
const auto o = TensorAccessor(o_args, get_common_arg_val<uint32_t>(0), TB);
uint32_t ones = 0;
if constexpr (BT > 0) {
ones = fill_bias_tiles(cb_c);
}
for (uint32_t r = b0; r < b0 + nb; ++r) {
if constexpr (BT > 0 && DIAG != 2) {
write_bias_block(ones, r * OKt + BCOL, BT, o);
}
cb_wait_front(cb_t, Kt);
uint32_t rp = get_read_ptr(cb_t);
const uint32_t ob = r * OKt + off;
if constexpr (DIAG != 2) {
for (uint32_t j = 0; j < Kt; ++j) {
noc_async_write(rp, o.get_noc_addr(ob + j), TB);
noc_async_write(rp, o.get_noc_addr(ob + Kt + j), TB);
rp += TB;
}
}
noc_async_writes_flushed();
cb_pop_front(cb_t, Kt);
}
noc_async_write_barrier();
}
|