// SPDX-License-Identifier: Apache-2.0
// TERMS_EPI writer (meteor-p150 OPT_PLAN #2, tt/epi.py): the output tiles of this core's contiguous range (work split
// as the reader) to the interleaved outputs, flushed per tile, one barrier at the end: the fp32 main output (CB 16,
// HAS_MAIN) and, with SPLIT, the bf16 hi / lo terms (CBs 17 / 18).
// CT args: [0] T, [1] NC, [2] GX, [3] HAS_MAIN, [4] SPLIT, then the TensorAccessorArgs of the main, hi and lo outputs
// (repeats of the first present one when absent). Common RT args: [main_addr, hi_addr, lo_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 HAS_MAIN = get_compile_time_arg_val(3);
constexpr uint32_t SPLIT = get_compile_time_arg_val(4);
constexpr auto out_args = TensorAccessorArgs<5>();
constexpr auto hi_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
constexpr auto lo_args = TensorAccessorArgs<hi_args.next_compile_time_args_offset()>();
constexpr uint32_t cb_out = 16, cb_hi = 17, cb_lo = 18;
constexpr uint32_t 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 out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), TB32);
    const auto hi = TensorAccessor(hi_args, get_common_arg_val<uint32_t>(1), TB16);
    const auto lo = TensorAccessor(lo_args, get_common_arg_val<uint32_t>(2), TB16);
    for (uint32_t t = t0; t < t0 + n; ++t) {
        if constexpr (HAS_MAIN) {
            cb_wait_front(cb_out, 1);
            noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(t), TB32);
        }
        if constexpr (SPLIT) {
            cb_wait_front(cb_hi, 1);
            cb_wait_front(cb_lo, 1);
            noc_async_write(get_read_ptr(cb_hi), hi.get_noc_addr(t), TB16);
            noc_async_write(get_read_ptr(cb_lo), lo.get_noc_addr(t), TB16);
        }
        noc_async_writes_flushed();
        if constexpr (HAS_MAIN) {
            cb_pop_front(cb_out, 1);
        }
        if constexpr (SPLIT) {
            cb_pop_front(cb_hi, 1);
            cb_pop_front(cb_lo, 1);
        }
    }
    noc_async_write_barrier();
}