// SPDX-License-Identifier: Apache-2.0
// FOUT writer (meteor-p150 OPT2 MK-FO, tt/fout.py): the 3 bf16 tiles of each unit (CB 16, units as the reader):
//   flat unit b: tiles -> fused (flat [1, 1, R, 96] TILE, tiles b * 3 + k); then untilized by 32 B local NOC copies
//                into an L1 staging buffer (CB 3: row i, bytes 64 k + 32 h <- tile k, face 2 (i / 16) + h, face row
//                i % 16) and written as the 32 ROW_MAJOR pixel rows p0..p0+31 of fused_rm (192 B pages)
//   4-D unit v:  tiles -> fused4 ([1, H, W, 96] TILE, tiles v * 3 + k), the rows n..31 (width padding) zeroed first
// The staging rows are kept in DRAM-bank order (row i at slot (i % NBANK) * 32 / NBANK + i / NBANK): the rows of one
// bank are consecutive pages there, written as one run of 32 / NBANK rows.
// CT args: [0] NBF, [1] NB4, [2] NC, [3] GX, [4] XB, [5] W, [6] NBANK, then the TensorAccessorArgs of fused,
// fused_rm, fused4.
// Common RT args: [fused_addr, fused_rm_addr, fused4_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NBF = get_compile_time_arg_val(0);
constexpr uint32_t NB4 = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t XB = get_compile_time_arg_val(4);
constexpr uint32_t W = get_compile_time_arg_val(5);
constexpr uint32_t NBANK = get_compile_time_arg_val(6);
constexpr auto f_args = TensorAccessorArgs<7>();
constexpr auto m_args = TensorAccessorArgs<f_args.next_compile_time_args_offset()>();
constexpr auto q_args = TensorAccessorArgs<m_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048, ROW = 192;
constexpr uint32_t cb_ustage = 3, cb_out = 16;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t NU = NBF + NB4;
    constexpr uint32_t B = NU / NC, E = NU % NC;
    const uint32_t u0 = core * B + (core < E ? core : E);
    const uint32_t n_units = B + (core < E ? 1 : 0);
    if (n_units == 0) {
        return;
    }
    const auto fused = TensorAccessor(f_args, get_common_arg_val<uint32_t>(0), TB16);
    const auto fused_rm = TensorAccessor(m_args, get_common_arg_val<uint32_t>(1), ROW);
    const auto fused4 = TensorAccessor(q_args, get_common_arg_val<uint32_t>(2), TB16);
    // two writer-private staging buffers (CB 3, 2 x 32 rows): reserved once, never pushed
    cb_reserve_back(cb_ustage, 1);
    const uint32_t stage0 = get_write_ptr(cb_ustage);
    const uint64_t here = get_noc_addr(stage0);
    for (uint32_t u = u0; u < u0 + n_units; ++u) {
        cb_wait_front(cb_out, 3);
        const uint32_t rp = get_read_ptr(cb_out);
        if (u < NBF) {
            for (uint32_t k = 0; k < 3; ++k) {
                noc_async_write(rp + k * TB16, fused.get_noc_addr(u * 3 + k), TB16);
            }
            const uint32_t stage = stage0 + ((u - u0) & 1) * 32 * ROW;
            noc_async_read_one_packet_set_state(here, 32);
            for (uint32_t i = 0; i < 32; ++i) {
                const uint32_t srow = rp + ((i >> 4) * 2) * 512 + (i & 15) * 32;
                const uint32_t drow = stage + ((i % NBANK) * (32 / NBANK) + i / NBANK) * ROW;
                for (uint32_t k = 0; k < 3; ++k) {
                    noc_async_read_one_packet_with_state(srow + k * TB16, drow + k * 64);
                    noc_async_read_one_packet_with_state(srow + k * TB16 + 512, drow + k * 64 + 32);
                }
            }
            noc_async_read_barrier();
            const uint32_t p0 = u * 32;
            for (uint32_t j = 0; j < NBANK; ++j) {
                noc_async_write(stage + j * (32 / NBANK) * ROW, fused_rm.get_noc_addr(p0 + j), (32 / NBANK) * ROW);
            }
        } else {
            const uint32_t v = u - NBF;
            const uint32_t xb = v % XB;
            const uint32_t n = (W - xb * 32) < 32 ? (W - xb * 32) : 32;
            if (n < 32) {
                for (uint32_t k = 0; k < 3; ++k) {
                    for (uint32_t i = n; i < 32; ++i) {
                        for (uint32_t h = 0; h < 2; ++h) {
                            volatile tt_l1_ptr uint32_t* q = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(
                                rp + k * TB16 + ((i >> 4) * 2 + h) * 512 + (i & 15) * 32);
                            for (uint32_t w = 0; w < 8; ++w) {
                                q[w] = 0;
                            }
                        }
                    }
                }
            }
            for (uint32_t k = 0; k < 3; ++k) {
                noc_async_write(rp + k * TB16, fused4.get_noc_addr(v * 3 + k), TB16);
            }
        }
        // the tiles (and the staging rows of unit u, reused at u + 2) have left L1
        noc_async_writes_flushed();
        cb_pop_front(cb_out, 3);
    }
    noc_async_write_barrier();
}