// SPDX-License-Identifier: Apache-2.0
// pairmap_k (tt/pairmap_k.py, BEVFORMER_MSDA_PAIRS): the SCA value projection -> horizontal pixel-pair maps for
// msda_k's 128-byte corner reads (one read per point and row: the NW + NE or SW + SE sticks).
//
//   in : bf16 TILE [B, 1, SUB * HV * WV, HEADS * 32], L1 interleaved (32-byte reads at 32-byte offsets); rows =
//        camera-major pixels, head j in columns j * 32 ..: the stick of (map m = (b HEADS + j) SUB + s, y, x) = row
//        s * HV * WV + y * WV + x of batch b, tile column j (TPB tiles per batch): two 32-byte face rows (faces
//        2 (r / 16) and + 1, row r % 16, r = row % 32)
//   out: bf16 ROW_MAJOR [B * HEADS * SUB, HV, WV + 1, 64], interleaved, 128-byte pages: page (m, y, i) = the sticks of
//        pixels x = i - 1 and x = i of row y (zeros for x = -1 and x = WV)
// Work: pages split in contiguous ranges over the cores, then alternately over the two data-movement RISC-Vs (SLOT 0 /
// 1, each its own scratch CB of BATCH pages): per batch, 4 reads per page (zeros page outside the row) into the
// scratch, one barrier, BATCH page 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 SUB = get_compile_time_arg_val(4);
constexpr uint32_t HEADS = get_compile_time_arg_val(5);
constexpr uint32_t NPAGES = get_compile_time_arg_val(6);
constexpr uint32_t NC = get_compile_time_arg_val(7);
constexpr uint32_t GX = get_compile_time_arg_val(8);
constexpr uint32_t BATCH = get_compile_time_arg_val(9);
constexpr uint32_t TPB = get_compile_time_arg_val(10);     // input tiles per batch
constexpr auto i_args = TensorAccessorArgs<11>();
constexpr auto o_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();

constexpr uint32_t TB16 = 2048, PAGE = 128, W1 = WV + 1, MAP = HV * WV;

template <typename A>
static inline void stick(const A& ia, uint32_t m, uint32_t y, int32_t x, uint64_t zeros, uint32_t dst) {
    if (x < 0 || x >= static_cast<int32_t>(WV)) {
        noc_async_read(zeros, dst, 32);
        noc_async_read(zeros, dst + 32, 32);
        return;
    }
    const uint32_t bt = m / (HEADS * SUB), jm = m - bt * HEADS * SUB;
    const uint32_t j = jm / SUB, s = jm - j * SUB;
    const uint32_t row = s * MAP + y * WV + static_cast<uint32_t>(x);
    const uint32_t r = row & 31;
    const uint32_t off = ((r >> 4) << 10) + ((r & 15) << 5);
    const uint64_t a = ia.get_noc_addr(bt * TPB + (row >> 5) * HEADS + j, off);
    noc_async_read(a, dst, 32);           // channels 0..15 (face 2 (r / 16))
    noc_async_read(a + 512, dst + 32, 32);  // channels 16..31 (the next face)
}

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), TB16);
    const auto oa = TensorAccessor(o_args, get_common_arg_val<uint32_t>(1), PAGE);
    const uint64_t zeros = get_noc_addr(MEM_ZEROS_BASE);
    cb_reserve_back(cb_s, 1);
    const uint32_t buf = get_write_ptr(cb_s);
    // batches b0 = p0 + (2 SLOT) BATCH, step 2 BATCH: the two RISC-Vs alternate whole batches
    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 m = p / (HV * W1);
            const uint32_t rem = p - m * HV * W1;
            const uint32_t y = rem / W1;
            const int32_t i = static_cast<int32_t>(rem - y * W1);
            const uint32_t dst = buf + (p - b0) * PAGE;
            stick(ia, m, y, i - 1, zeros, dst);
            stick(ia, m, y, i, zeros, dst + 64);
        }
        noc_async_read_barrier();
        for (uint32_t p = b0; p < b1; ++p) {
            noc_async_write(buf + (p - b0) * PAGE, oa.get_noc_addr(p), PAGE);
        }
        noc_async_write_barrier();
    }
}