File size: 2,074 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 | // SPDX-License-Identifier: Apache-2.0
// KCAT operand build (tt/kcat.py), reader (RISCV_0): this core's contiguous range of input tiles, B tiles per
// read barrier.
// SKIP (OPT round 3, HEAD-SKIP): the same tile range of the fp32 skip tensor into CB 1 (the compute adds it to x).
// CT args: [0] per-core RT-arg count (P3), [1] input tile bytes, [2] B (tiles per batch; CBs 0 / 1 hold 2 B pages),
// [3] SKIP (0 / 1), then the TensorAccessorArgs of x and of skip (x's again when SKIP = 0).
// Common RT args: [x_addr, skip_addr]. Per-core RT args: [t0, n] (flattened input tile indices).
#include <cstdint>
#include "api/dataflow/dataflow_api.h"
constexpr uint32_t TB = get_compile_time_arg_val(1);
constexpr uint32_t B = get_compile_time_arg_val(2);
constexpr uint32_t SKIP = get_compile_time_arg_val(3);
constexpr auto x_args = TensorAccessorArgs<4>();
constexpr auto s_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_x = 0, cb_s = 1;
void kernel_main() {
const uint32_t x_addr = get_common_arg_val<uint32_t>(0);
const uint32_t t0 = get_arg_val<uint32_t>(0);
const uint32_t n = get_arg_val<uint32_t>(1);
const auto x = TensorAccessor(x_args, x_addr, TB);
const auto sk = TensorAccessor(s_args, get_common_arg_val<uint32_t>(1), TB);
const uint32_t end = t0 + n;
for (uint32_t t = t0; t < end; t += B) {
const uint32_t k = (end - t) < B ? (end - t) : B;
cb_reserve_back(cb_x, k);
uint32_t wp = get_write_ptr(cb_x);
for (uint32_t i = 0; i < k; ++i) {
noc_async_read(x.get_noc_addr(t + i), wp, TB);
wp += TB;
}
if constexpr (SKIP) {
cb_reserve_back(cb_s, k);
uint32_t sp = get_write_ptr(cb_s);
for (uint32_t i = 0; i < k; ++i) {
noc_async_read(sk.get_noc_addr(t + i), sp, TB);
sp += TB;
}
}
noc_async_read_barrier();
cb_push_back(cb_x, k);
if constexpr (SKIP) {
cb_push_back(cb_s, k);
}
}
}
|