// SPDX-License-Identifier: Apache-2.0
// pairmap2 (tt/pairmap_k.py, BEVFORMER_PAIRMAP=v2), compute: each item's tile untilized (pack_untilize, one tile
// wide) into cb_rm: 32 pixel sticks of 32 channels (64 B), consecutive, for the writer's 128-byte pair pages.
#include <cstdint>

#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/pack_untilize.h"

constexpr uint32_t NITEMS = 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 cb_in = 0, 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;
    if (w0 >= w1) {
        return;
    }
    compute_kernel_hw_startup(cb_in, cb_rm);
    pack_untilize_init<1, 1>(cb_in, cb_rm);
    for (uint32_t w = w0; w < w1; ++w) {
        cb_wait_front(cb_in, 1);
        cb_reserve_back(cb_rm, 1);
        pack_untilize_block<1, 1>(cb_in, 1, cb_rm);
        cb_push_back(cb_rm, 1);
        cb_pop_front(cb_in, 1);
    }
    pack_untilize_uninit(cb_rm);
}