code/tt_frnet/tt/kernels/kcat_reader.cpp

// 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);
        }
    }
}