// SPDX-License-Identifier: Apache-2.0
// TFUSE writer (meteor-p150 OPT2 MK-TF, tt/tgate.py fuse_mm): the 3 bf16 output tiles of each 32-row block (CB 16)
// to the interleaved bf16 TILE output (tiles b * 3 + k; flat [R, 96] or the 4-D [1, H, W, 96] form with W padded to
// XB * 32 rows per image row). 4-D form: in the last block of every image row (b % XB == XB - 1) the rows VR..31
// (the width padding) are zeroed in L1 before the write, so the pad pixels hold zeros as a tilize would write them.
// Work split as the reader. CT args: [0] NB, [1] NC, [2] GX, [3] XB (0: flat, no padding), [4] VR, then the
// TensorAccessorArgs of the output. Common RT args: [out_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NB = 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 XB = get_compile_time_arg_val(3);
constexpr uint32_t VR = get_compile_time_arg_val(4);
constexpr auto out_args = TensorAccessorArgs<5>();
constexpr uint32_t TB16 = 2048;
constexpr uint32_t cb_out = 16;

// bf16 tile: faces 0 / 1 = rows 0..15 (columns 0..15 / 16..31), faces 2 / 3 = rows 16..31; 16 x 2 B = 32 B per face
// row, 512 B per face
inline void zero_rows(uint32_t tile_addr) {
    for (uint32_t r = VR; r < 32; ++r) {
        const uint32_t f0 = (r < 16) ? 0 : 2;
        const uint32_t fr = r & 15;
        for (uint32_t f = f0; f < f0 + 2; ++f) {
            volatile tt_l1_ptr uint32_t* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(tile_addr + f * 512 + fr * 32);
            for (uint32_t i = 0; i < 8; ++i) {
                p[i] = 0;
            }
        }
    }
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NB / NC, E = NB % NC;
    const uint32_t b0 = 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), TB16);
    for (uint32_t b = b0; b < b0 + n; ++b) {
        cb_wait_front(cb_out, 3);
        const uint32_t rp = get_read_ptr(cb_out);
        if constexpr (XB > 0 && VR < 32) {
            if (b % XB == XB - 1) {
                for (uint32_t k = 0; k < 3; ++k) {
                    zero_rows(rp + k * TB16);
                }
            }
        }
        for (uint32_t k = 0; k < 3; ++k) {
            noc_async_write(rp + k * TB16, out.get_noc_addr(b * 3 + k), TB16);
        }
        noc_async_writes_flushed();
        cb_pop_front(cb_out, 3);
    }
    noc_async_write_barrier();
}