// SPDX-License-Identifier: Apache-2.0
// pairmap2 (tt/pairmap_k.py pair_maps, BEVFORMER_PAIRMAP=v2; D20 round 2), data movement: the same output as
// pairmap_k.cpp (horizontal pixel-pair maps, 128-byte pages (m, y, i) = the sticks of x = i - 1 and x = i, zeros
// outside the row) from whole tiles instead of 32-byte face-row reads: items (batch, tile row, head) in contiguous
// ranges over the cores; per item ROLE 0 (NCRISC) reads the 2 KB tile (-> the compute kernel's pack_untilize: 32
// pixel sticks of 64 B, consecutive in cb_rm) and the previous pixel's stick (the last row of the tile above, two
// 32-byte face rows); ROLE 1 (BRISC) writes each page as ONE 128-byte write of two consecutive untilized sticks (two
// 64-byte writes at row starts / tile starts, and the extra end-of-row page). About 1.1 requests per page instead
// of 5. No semaphores, no multicast; common RT args only (P3).
#include <stdint.h>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t ROLE = get_compile_time_arg_val(0);
constexpr uint32_t HV = get_compile_time_arg_val(1);
constexpr uint32_t WV = get_compile_time_arg_val(2);
constexpr uint32_t SUB = get_compile_time_arg_val(3);
constexpr uint32_t HEADS = get_compile_time_arg_val(4);
constexpr uint32_t NTR = get_compile_time_arg_val(5);      // tile rows per batch
constexpr uint32_t NITEMS = 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 auto i_args = TensorAccessorArgs<9>();
constexpr auto o_args = TensorAccessorArgs<i_args.next_compile_time_args_offset()>();
constexpr uint32_t TB16 = 2048, PAGE = 128, STK = 64, W1 = WV + 1, MAP = HV * WV;
constexpr uint32_t cb_in = 0, cb_prev = 1, cb_rm = 16;

void kernel_main() {
    const uint32_t ci = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (ci >= NC) {
        return;
    }
    const uint32_t w0 = ci * NITEMS / NC;
    const uint32_t w1 = (ci + 1) * NITEMS / 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);
    for (uint32_t w = w0; w < w1; ++w) {
        const uint32_t j = w % HEADS;
        const uint32_t bt_t = w / HEADS;                 // bt * NTR + t
        const uint32_t t = bt_t % NTR;
        const uint32_t bt = bt_t / NTR;
        if constexpr (ROLE == 0) {
            cb_reserve_back(cb_in, 1);
            cb_reserve_back(cb_prev, 1);
            noc_async_read(ia.get_noc_addr(bt_t * HEADS + j), get_write_ptr(cb_in), TB16);
            if (t > 0) {   // the previous pixel = row 31 of the tile above: faces 2 / 3, row 15
                const uint64_t a = ia.get_noc_addr((bt_t - 1) * HEADS + j, 1024 + 15 * 32);
                noc_async_read(a, get_write_ptr(cb_prev), 32);
                noc_async_read(a + 512, get_write_ptr(cb_prev) + 32, 32);
            }
            noc_async_read_barrier();
            cb_push_back(cb_in, 1);
            cb_push_back(cb_prev, 1);
        } else {
            cb_wait_front(cb_rm, 1);
            cb_wait_front(cb_prev, 1);
            const uint32_t rm = get_read_ptr(cb_rm);
            const uint32_t prev = get_read_ptr(cb_prev);
            const uint32_t zeros = MEM_ZEROS_BASE;
            for (uint32_t r = 0; r < 32; ++r) {
                const uint32_t q = t * 32 + r;             // pixel within the batch: (s, y, x)
                if (q >= SUB * MAP) {
                    break;
                }
                const uint32_t s = q / MAP;
                const uint32_t rem = q - s * MAP;
                const uint32_t y = rem / WV;
                const uint32_t x = rem - y * WV;
                const uint32_t m = (bt * HEADS + j) * SUB + s;
                const uint32_t pg = (m * HV + y) * W1 + x;
                const uint32_t cur = rm + r * STK;
                if (x > 0 && r > 0) {
                    noc_async_write(cur - STK, oa.get_noc_addr(pg), PAGE);
                } else {
                    const uint64_t d = oa.get_noc_addr(pg);
                    noc_async_write(x == 0 ? zeros : prev, d, STK);
                    noc_async_write(cur, d + STK, STK);
                }
                if (x == WV - 1) {                         // the end-of-row page (y, WV) = [x = WV - 1 | zeros]
                    const uint64_t d = oa.get_noc_addr(pg + 1);
                    noc_async_write(cur, d, STK);
                    noc_async_write(zeros, d + STK, STK);
                }
            }
            noc_async_writes_flushed();
            cb_pop_front(cb_rm, 1);
            cb_pop_front(cb_prev, 1);
        }
    }
    if constexpr (ROLE == 1) {
        noc_async_write_barrier();
    }
}