// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
//
// Face-granular (16x16) head split / concat for the ViT with UNPADDED heads (hamer/tt_kernels.py).
// HaMeR ViT-H heads are 80 wide = 5 faces (2.5 tiles).  The qkv matmul produces the unpadded
// [192, 3*16*80] activation (120 tile columns = one per core on a 12x10 grid); SDPA needs head
// dims that are tile multiples, so every head is re-laid out to 96 = 3 tiles (6 faces, the 6th
// face zero).  A bf16 TILE is 4 contiguous 512-byte faces: f0 rows 0-15 x cols 0-15, f1 rows 0-15
// x cols 16-31, f2 rows 16-31 x cols 0-15, f3 rows 16-31 x cols 16-31, so moving a 16-column
// slice of a tile row means moving the 2 faces (fr = 0, 1) of that face column.
//
//   MODE_SPLIT  src [R tile rows x W tile cols], W*2 = T*H*F face cols (F = faces per head)
//               ->  T outputs [H, R, DP] tiles (DP*2 >= F face cols per head, extra faces zero)
//               out index i = ((t*H + h)*R + r)*DP + c ; out face (fr, fc): j = 2c + fc
//               j < F: src face col g = (t*H + h)*F + j -> src page r*W + g/2, face fr*2 + g%2
//   MODE_CONCAT src [H, R, DP] tiles  ->  1 output [R x W] tiles, W*2 = H*F
//               out index i = r*W + oc ; out face (fr, fc): g = 2*oc + fc, h = g / F, j = g % F
//               -> src page (h*R + r)*DP + j/2, face fr*2 + j%2
//
// compile-time args: T, H, R, W, DP, F, CHUNK, then TensorAccessorArgs(src), (dst)
// runtime args: src_addr, start, n, dst0_addr .. dst{T-1}_addr
#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 W = get_compile_time_arg_val(3);
    constexpr uint32_t DP = get_compile_time_arg_val(4);
    constexpr uint32_t F = get_compile_time_arg_val(5);
    constexpr uint32_t CHUNK = get_compile_time_arg_val(6);
    constexpr auto src_args = TensorAccessorArgs<7>();
    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 face = page / 4;
    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 jt = 0; jt < m; ++jt) {
            const uint32_t i = start + base + jt;
            const uint32_t out_tile = scratch + jt * page;
            for (uint32_t fc = 0; fc < 2; ++fc) {
#ifdef MODE_SPLIT
                const uint32_t c = i % DP;
                const uint32_t r = (i / DP) % R;
                const uint32_t th = i / (DP * R);  // t*H + h
                const uint32_t j = 2 * c + fc;
                if (j >= F) {
                    // zero pad faces (rows 0-15 and 16-31 of this face column)
                    for (uint32_t fr = 0; fr < 2; ++fr) {
                        volatile tt_l1_ptr uint32_t* p =
                            reinterpret_cast<volatile tt_l1_ptr uint32_t*>(out_tile + (fr * 2 + fc) * face);
                        for (uint32_t w = 0; w < face / 4; ++w) {
                            p[w] = 0;
                        }
                    }
                    continue;
                }
                const uint32_t g = th * F + j;
                const uint32_t src_page = r * W + g / 2;
                const uint32_t sfc = g % 2;
#else
                const uint32_t oc = i % W;
                const uint32_t r = i / W;
                const uint32_t g = 2 * oc + fc;
                const uint32_t h = g / F;
                const uint32_t j = g % F;
                const uint32_t src_page = (h * R + r) * DP + j / 2;
                const uint32_t sfc = j % 2;
#endif
                for (uint32_t fr = 0; fr < 2; ++fr) {
                    noc_async_read(
                        src.get_noc_addr(src_page, (fr * 2 + sfc) * face), out_tile + (fr * 2 + fc) * face, face);
                }
            }
        }
        noc_async_read_barrier();
        for (uint32_t jt = 0; jt < m; ++jt) {
            const uint32_t i = start + base + jt;
#ifdef MODE_SPLIT
            const uint32_t t = i / (H * R * DP);
            const uint32_t dst_page = i % (H * R * DP);
#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);
            noc_async_write(scratch + jt * page, dst.get_noc_addr(dst_page), page);
        }
        noc_async_write_barrier();
    }
}