File size: 3,149 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 | // SPDX-License-Identifier: Apache-2.0
// P2P (OPT round 4, FRNET_P2P): pixel2point gather straight into the KCAT operand, reader (RISCV_0).
// For this core's contiguous range of 32-row blocks: the block's 32 uint32 pixel indices (from the ROW_MAJOR [1, N]
// index tensor, one page of N * 4 bytes), then the 32 bf16 table rows (ROW_MAJOR [V, C], one page of C * 2 bytes per
// pixel) gathered into one row-major block of CB 0 (Kt pages of 2048 bytes = 32 rows x C bf16), which the compute
// kernel tilizes. BB blocks are read per barrier (CB 0 holds 2 BB blocks); the indices are read IB blocks at a time.
// CT args: [0] per-core RT-arg count (P3), [1] Kt (C / 32), [2] row bytes (C * 2), [3] BB, [4] IB, [5] index page
// bytes (N * 4), [6] DIAG (0; 1 = timing diagnostic without table reads), then the TensorAccessorArgs of the table and of the index.
// Common RT args: [table_addr, index_addr]. Per-core RT args: [b0, nb] (first block, block count).
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t Kt = get_compile_time_arg_val(1);
constexpr uint32_t RB = get_compile_time_arg_val(2);
constexpr uint32_t BB = get_compile_time_arg_val(3);
constexpr uint32_t IB = get_compile_time_arg_val(4);
constexpr uint32_t IPB = get_compile_time_arg_val(5);
constexpr uint32_t DIAG = get_compile_time_arg_val(6); // timing diagnostic: 1 = no table reads (garbage output)
constexpr auto t_args = TensorAccessorArgs<7>();
constexpr auto i_args = TensorAccessorArgs<t_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_rm = 0, cb_idx = 1;
constexpr uint32_t BLOCK_BYTES = 32 * RB;
void kernel_main() {
const uint32_t b0 = get_arg_val<uint32_t>(0);
const uint32_t nb = get_arg_val<uint32_t>(1);
if (nb == 0) {
return;
}
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);
const uint32_t ibuf = get_write_ptr(cb_idx);
volatile tt_l1_ptr uint32_t* idx = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(ibuf);
const uint32_t end = b0 + nb;
uint32_t ib_end = b0; // indices of blocks [ib_start, ib_end) are in ibuf
uint32_t ib_start = b0;
for (uint32_t b = b0; b < end; b += BB) {
const uint32_t k = (end - b) < BB ? (end - b) : BB;
if (b + k > ib_end) { // refill: IB blocks of indices from block b (BB divides IB: never straddles)
ib_start = b;
const uint32_t m = (end - b) < IB ? (end - b) : IB;
noc_async_read(index.get_noc_addr(0, b * 128), ibuf, m * 128);
noc_async_read_barrier();
ib_end = b + m;
}
cb_reserve_back(cb_rm, k * Kt);
uint32_t wp = get_write_ptr(cb_rm);
const uint32_t base = (b - ib_start) * 32;
if constexpr (DIAG != 1) {
for (uint32_t j = 0; j < 32 * k; ++j) {
noc_async_read(table.get_noc_addr(idx[base + j]), wp, RB);
wp += RB;
}
}
noc_async_read_barrier();
cb_push_back(cb_rm, k * Kt);
}
}
|