// SPDX-License-Identifier: Apache-2.0
// KTILE reader (meteor-p150 OPT2, tt/ktile.py): a ROW_MAJOR map (logical [A, B, WD, C], one DRAM page per pixel of
// PAGE = C * ES bytes) tilized into the TILE form of the same logical shape, without a compute kernel.
// Work units (no per-core RT args, probe P3: core i = y * GX + x of the first NC cores owns the contiguous unit range
// [i * Q + min(i, E), + Q + (i < E)) of the NU = A * B * XB units, Q = NU / NC, E = NU % NC): unit v = 32-pixel row
// block xb = v % XB of plane v / XB, pixels p0 = (v / XB) * WD + 32 xb, n = min(32, WD - 32 xb) valid rows; its
// output = the CT tiles v * CT + 0..CT-1.
// Per unit: the n pixel rows read as NBANK runs of same-bank consecutive pages (rows j, j + NBANK, ... of one DRAM
// bank; staging slot of row i = (i % NBANK) * 32 / NBANK + i / NBANK) into a double-buffered L1 staging area (CB 1),
// the next unit's runs issued while this unit is tilized; tilize = 16-element face rows (16 * ES bytes) copied by
// stateful local NOC reads into the CT tiles of CB 0 (tile k, face 2 (i / 16) + h, face row i % 16 <- row i,
// elements 32 k + 16 h); the rows n..31 and the channel halves >= C zeroed (the tilize padding).
// Crop (a window of a larger 4-D map): the source page of output pixel j of plane q is SRC_OFF + q * SRC_PSTRIDE + j
// (no crop: SRC_OFF 0, SRC_PSTRIDE WD).
// CT args: [0] NU, [1] NC, [2] GX, [3] XB, [4] WD, [5] C, [6] ES, [7] NBANK, [8] CT, [9] SRC_OFF, [10] SRC_PSTRIDE,
// then the TensorAccessorArgs of the input. Common RT args: [in_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NU = 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 WD = get_compile_time_arg_val(4);
constexpr uint32_t CH = get_compile_time_arg_val(5);
constexpr uint32_t ES = get_compile_time_arg_val(6);
constexpr uint32_t NBANK = get_compile_time_arg_val(7);
constexpr uint32_t CT = get_compile_time_arg_val(8);
constexpr uint32_t SRC_OFF = get_compile_time_arg_val(9);
constexpr uint32_t SRC_PSTRIDE = get_compile_time_arg_val(10);
constexpr auto in_args = TensorAccessorArgs<11>();
constexpr uint32_t PAGE = CH * ES, TB = 1024 * ES, FACE = 256 * ES, FROW = 16 * ES;
constexpr uint32_t RUN = 32 / NBANK;
constexpr uint32_t cb_t = 0, cb_stage = 1;

inline void unit_rows(uint32_t v, uint32_t& p0, uint32_t& n) {
    const uint32_t xb = v % XB;
    p0 = SRC_OFF + (v / XB) * SRC_PSTRIDE + xb * 32;
    n = (WD - xb * 32) < 32 ? (WD - xb * 32) : 32;
}

template <typename TA>
inline void read_rows(const TA& src, uint32_t p0, uint32_t n, uint32_t stage) {
    for (uint32_t j = 0; j < NBANK && j < n; ++j) {
        const uint32_t cnt = (n - j + NBANK - 1) / NBANK;
        noc_async_read(src.get_noc_addr(p0 + j), stage + j * RUN * PAGE, cnt * PAGE);
    }
}

inline void zero(uint32_t addr, uint32_t bytes) {
    volatile tt_l1_ptr uint32_t* q = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(addr);
    for (uint32_t w = 0; w < bytes / 4; ++w) {
        q[w] = 0;
    }
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t Q = NU / NC, E = NU % NC;
    const uint32_t u0 = core * Q + (core < E ? core : E);
    const uint32_t nu = Q + (core < E ? 1 : 0);
    if (nu == 0) {
        return;
    }
    const auto src = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), PAGE);
    cb_reserve_back(cb_stage, 1);                      // reader-private, never pushed
    const uint32_t stage0 = get_write_ptr(cb_stage);
    const uint64_t here = get_noc_addr(stage0);
    uint32_t p0, n;
    unit_rows(u0, p0, n);
    read_rows(src, p0, n, stage0);
    noc_async_read_barrier();
    for (uint32_t v = u0; v < u0 + nu; ++v) {
        const uint32_t stage = stage0 + ((v - u0) & 1) * 32 * PAGE;
        cb_reserve_back(cb_t, CT);
        const uint32_t wt = get_write_ptr(cb_t);
        if constexpr (CH % 32 != 0) {                  // the last tile's upper 16 columns: channel padding
            zero(wt + (CT - 1) * TB + 1 * FACE, FACE);
            zero(wt + (CT - 1) * TB + 3 * FACE, FACE);
        }
        if (n < 32) {                                  // the width padding rows
            for (uint32_t i = n; i < 32; ++i) {
                for (uint32_t k = 0; k < CT; ++k) {
                    for (uint32_t h = 0; h < 2; ++h) {
                        zero(wt + k * TB + ((i >> 4) * 2 + h) * FACE + (i & 15) * FROW, FROW);
                    }
                }
            }
        }
        noc_async_read_one_packet_set_state(here, FROW);
        for (uint32_t i = 0; i < n; ++i) {
            const uint32_t s = stage + ((i % NBANK) * RUN + i / NBANK) * PAGE;
            const uint32_t d = wt + ((i >> 4) * 2) * FACE + (i & 15) * FROW;
            for (uint32_t k = 0; k < CT; ++k) {
                noc_async_read_one_packet_with_state(s + k * 32 * ES, d + k * TB);
                if (k * 32 + 16 < CH) {
                    noc_async_read_one_packet_with_state(s + (k * 32 + 16) * ES, d + k * TB + FACE);
                }
            }
        }
        if (v + 1 < u0 + nu) {
            unit_rows(v + 1, p0, n);
            read_rows(src, p0, n, stage0 + ((v + 1 - u0) & 1) * 32 * PAGE);
        }
        noc_async_read_barrier();
        cb_push_back(cb_t, CT);
    }
}