// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Tile-granular head split / head concat for the HaMeR ViT and MANO head (hamer/tt_kernels.py).
// One data-movement kernel per core: it owns a contiguous range [start, start+n) of OUTPUT tiles,
// gathers them from the interleaved source with batched async reads into a local scratch CB,
// then scatters them to the interleaved destination(s) with batched async writes.
//
//   MODE_SPLIT  src [R rows x T*H*D tile cols]  ->  T outputs (T <= 12), each [H, R, D] tiles
//               out index i = ((t*H + h)*R + r)*D + c
//               src page   = r*(T*H*D) + t*H*D + h*D + c ; dst page (tensor t) = (h*R + r)*D + c
//   MODE_CONCAT src [H, R, D] tiles  ->  1 output [R x H*D] tiles
//               out index j = r*(H*D) + h*D + c ; src page = (h*R + r)*D + c
//
// compile-time args: T, H, R, D, CHUNK (tiles per batch), then TensorAccessorArgs(src), (dst)
// runtime args: src_addr, start, n, dst0_addr .. dst{T-1}_addr
// Source / destinations may be interleaved or sharded (TensorAccessor resolves page -> bank/core).
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

void kernel_main() {
    const uint32_t src_addr = get_arg_val<uint32_t>(0);
    const uint32_t start = get_arg_val<uint32_t>(1);
    const uint32_t n = get_arg_val<uint32_t>(2);

    constexpr uint32_t T = get_compile_time_arg_val(0);
    constexpr uint32_t H = get_compile_time_arg_val(1);
    constexpr uint32_t R = get_compile_time_arg_val(2);
    constexpr uint32_t D = get_compile_time_arg_val(3);
    constexpr uint32_t CHUNK = get_compile_time_arg_val(4);
    constexpr auto src_args = TensorAccessorArgs<5>();
    constexpr auto dst_args = TensorAccessorArgs<src_args.next_compile_time_args_offset()>();

    constexpr uint32_t cb = 0;
    const uint32_t page = get_local_cb_interface(cb).fifo_page_size;
    const uint32_t scratch = get_write_ptr(cb);

    const auto src = TensorAccessor(src_args, src_addr, page);

    for (uint32_t base = 0; base < n; base += CHUNK) {
        const uint32_t m = (n - base) < CHUNK ? (n - base) : CHUNK;
        for (uint32_t j = 0; j < m; ++j) {
            const uint32_t i = start + base + j;
            const uint32_t c = i % D;
#ifdef MODE_SPLIT
            const uint32_t r = (i / D) % R;
            const uint32_t th = i / (D * R);  // t*H + h
            const uint32_t src_page = r * (T * H * D) + th * D + c;
#else
            const uint32_t hc = i % (H * D);  // h*D + c
            const uint32_t r = i / (H * D);
            const uint32_t h = hc / D;
            const uint32_t src_page = (h * R + r) * D + c;
#endif
            noc_async_read(src.get_noc_addr(src_page), scratch + j * page, page);
        }
        noc_async_read_barrier();
        for (uint32_t j = 0; j < m; ++j) {
            const uint32_t i = start + base + j;
#ifdef MODE_SPLIT
            const uint32_t t = i / (H * R * D);
            const uint32_t dst_page = i % (H * R * D);
#else
            const uint32_t t = 0;
            const uint32_t dst_page = i;
#endif
            const auto dst = TensorAccessor(dst_args, get_arg_val<uint32_t>(3 + t), page);
            const uint64_t a = dst.get_noc_addr(dst_page);
            noc_async_write(scratch + j * page, a, page);
        }
        noc_async_write_barrier();
    }
}