code/tt_frnet/tt/kernels/epi_writer.cpp

// SPDX-License-Identifier: Apache-2.0
// EPI writer (frnet-p150 OPT round 2, tt/epi.py): this core's tiles (work split as the reader) to the fp32 map
// (CB 16, OUT32) and / or its bf16 copy (CB 17, OUT16), flushed per tile, one barrier at the end.
// CT args: [0] T, [1] NC, [2] GX, [3] OUT32, [4] OUT16, then the TensorAccessorArgs of the fp32 and the bf16 output
// (repeats of the present one when one is absent). Common RT args: [out32_addr, out16_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t T = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t OUT32 = get_compile_time_arg_val(3);
constexpr uint32_t OUT16 = get_compile_time_arg_val(4);
constexpr auto o32_args = TensorAccessorArgs<5>();
constexpr auto o16_args = TensorAccessorArgs<o32_args.next_compile_time_args_offset()>();
constexpr uint32_t cb32 = 16, cb16 = 17, TB32 = 4096, TB16 = 2048;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = T / NC, E = T % NC;
    const uint32_t t0 = core * B + (core < E ? core : E);
    const uint32_t n = B + (core < E ? 1 : 0);
    const auto o32 = TensorAccessor(o32_args, get_common_arg_val<uint32_t>(0), TB32);
    const auto o16 = TensorAccessor(o16_args, get_common_arg_val<uint32_t>(1), TB16);
    for (uint32_t t = t0; t < t0 + n; ++t) {
        if constexpr (OUT32) {
            cb_wait_front(cb32, 1);
            noc_async_write(get_read_ptr(cb32), o32.get_noc_addr(t), TB32);
        }
        if constexpr (OUT16) {
            cb_wait_front(cb16, 1);
            noc_async_write(get_read_ptr(cb16), o16.get_noc_addr(t), TB16);
        }
        noc_async_writes_flushed();
        if constexpr (OUT32) {
            cb_pop_front(cb32, 1);
        }
        if constexpr (OUT16) {
            cb_pop_front(cb16, 1);
        }
    }
    noc_async_write_barrier();
}