// SPDX-License-Identifier: Apache-2.0
// dcnpair_k (tt/dcn_k.py dcn_pair_map, BEVFORMER_DCN_PAIRS; D20 round 2): the DCN input map -> horizontal pixel-pair
// map for dcn_k's corner reads (one 2C-wide read per row and corner pair: NW + NE, or SW + SE).
//   in : bf16 ROW_MAJOR [1, 1, N HV WV, C], interleaved (one C-channel stick per page; L1)
//   out: bf16 ROW_MAJOR [1, 1, N HV (WV + 1), 2 C], interleaved: page (n, y, i) = the sticks of pixels x = i - 1 and
//        x = i of row y (zeros for x = -1 and x = WV)
// Work: pages in contiguous ranges over the cores, batches of BATCH pages alternating over the two data-movement
// RISC-Vs (SLOT 0 / 1, each its own scratch CB): 2 reads per page (or zeros), one barrier, BATCH writes, one write
// barrier. No semaphores, no multicast; common RT args only (P3).
#include <stdint.h>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t SLOT = get_compile_time_arg_val(0);
constexpr uint32_t cb_s = get_compile_time_arg_val(1);
constexpr uint32_t HV = get_compile_time_arg_val(2);
constexpr uint32_t WV = get_compile_time_arg_val(3);
constexpr uint32_t ROWB = get_compile_time_arg_val(4);    // C * 2 bytes
constexpr uint32_t NPAGES = get_compile_time_arg_val(5);
constexpr uint32_t NC = get_compile_time_arg_val(6);
constexpr uint32_t GX = get_compile_time_arg_val(7);
constexpr uint32_t BATCH = get_compile_time_arg_val(8);
constexpr auto i_args = TensorAccessorArgs<9>();
constexpr auto o_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();
constexpr uint32_t W1 = WV + 1;
static_assert(ROWB % MEM_ZEROS_SIZE == 0, "dcnpair_k: C * 2 a multiple of the zeros page");

static inline void zeros(uint32_t dst) {
    const uint64_t z = get_noc_addr(MEM_ZEROS_BASE);
    for (uint32_t o = 0; o < ROWB; o += MEM_ZEROS_SIZE) {
        noc_async_read(z, dst + o, MEM_ZEROS_SIZE);
    }
}

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t p0 = ci * NPAGES / NC;
    const uint32_t p1 = (ci + 1) * NPAGES / NC;
    const auto ia = TensorAccessor(i_args, get_common_arg_val<uint32_t>(0), ROWB);
    const auto oa = TensorAccessor(o_args, get_common_arg_val<uint32_t>(1), 2 * ROWB);
    cb_reserve_back(cb_s, 1);
    const uint32_t buf = get_write_ptr(cb_s);
    for (uint32_t b0 = p0 + SLOT * BATCH; b0 < p1; b0 += 2 * BATCH) {
        const uint32_t b1 = b0 + BATCH < p1 ? b0 + BATCH : p1;
        for (uint32_t p = b0; p < b1; ++p) {
            const uint32_t row = p / W1;                 // (n, y)
            const uint32_t i = p - row * W1;
            const uint32_t dst = buf + (p - b0) * 2 * ROWB;
            const uint32_t q = row * WV + i;             // input page of pixel x = i
            if (i == 0) {
                zeros(dst);
            } else {
                noc_async_read(ia.get_noc_addr(q - 1), dst, ROWB);
            }
            if (i == WV) {
                zeros(dst + ROWB);
            } else {
                noc_async_read(ia.get_noc_addr(q), dst + ROWB, ROWB);
            }
        }
        noc_async_read_barrier();
        for (uint32_t p = b0; p < b1; ++p) {
            noc_async_write(buf + (p - b0) * 2 * ROWB, oa.get_noc_addr(p), 2 * ROWB);
        }
        noc_async_write_barrier();
    }
}