// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
//
// On-device im2col of HaMeR's patch embedding Conv2d(3->1280, k=16, s=16, pad=2) on the 256x192
// crop, written straight into fp32 TILE layout (no host F.unfold, no tilize op).
//
// src : fp32 ROW_MAJOR [1, 3, 256, 192] in L1 (the normalized crop as the caller hands it),
//       page = one image row (192 floats = 768 B).
// dst : fp32 TILE [1, 1, 192, 768]: token t = gh*12 + gw (16 x 12 patch grid), feature
//       f = c*256 + kh*16 + kw (F.unfold / patch_embed_weight order), value = img[c, gh*16+kh-2, gw*16+kw-2]
//       (zero outside the image = the conv's zero padding).
// A tile (tr, tc) covers tokens tr*32..+31 and features tc*32..+31, i.e. one channel c = tc/8 and the
// kernel-row pair kh = 2*(tc%8) + {0,1}; each (token, kh) is ONE 16-float face row (kw = 0..15) =
// image row y = gh*16+kh-2, columns gw*16-2 .. gw*16+13.  The source column offset is 8 B off the
// 16 B NOC alignment, so the core reads the aligned 80 B window [gw*16-4, gw*16+16) and shifts locally.
// Work unit = HALF a tile (tokens r0..r0+15 = faces 0,1 or 2,3 = 2 KB contiguous); the program runs
// this kernel on both data-movement RISCs of every core (defines SCR_TILE / SCR_WIN pick the CBs).
//
// compile-time args: TensorAccessorArgs(src), TensorAccessorArgs(dst)
// runtime args: src_addr, dst_addr, first_unit, num_units
#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 dst_addr = get_arg_val<uint32_t>(1);
    const uint32_t first = get_arg_val<uint32_t>(2);
    const uint32_t nunits = get_arg_val<uint32_t>(3);

    constexpr auto src_args = TensorAccessorArgs<0>();
    constexpr auto dst_args = TensorAccessorArgs<src_args.next_compile_time_args_offset()>();
#ifdef IM2COL_BF16
    // bf16 source (the raw 0..255 pixel values, exact in bf16; the normalization is folded into the patch weights
    // and a per-token bias): 384 B rows, the 16 B-aligned 48 B window [gw*16-8, gw*16+16) px, bf16 -> fp32 by <<16
    constexpr uint32_t ROW_BYTES = 192 * 2;
    constexpr uint32_t WIN = 48;
#else
    constexpr uint32_t ROW_BYTES = 192 * 4;
    constexpr uint32_t WIN = 80;  // bytes per read window (20 floats)
#endif
    constexpr uint32_t TILE_BYTES = 32 * 32 * 4;
    constexpr uint32_t HALF_BYTES = TILE_BYTES / 2;
    constexpr uint32_t TILES_W = 768 / 32;

    const auto src = TensorAccessor(src_args, src_addr, ROW_BYTES);
    const auto dst = TensorAccessor(dst_args, dst_addr, TILE_BYTES);
    const uint32_t half_buf = get_write_ptr(SCR_TILE);
    const uint32_t win_buf = get_write_ptr(SCR_WIN);
    uint32_t* half = reinterpret_cast<uint32_t*>(half_buf);
    uint32_t* win = reinterpret_cast<uint32_t*>(win_buf);

    for (uint32_t k = 0; k < nunits; ++k) {
        const uint32_t unit = first + k;
        const uint32_t tid = unit >> 1, hsel = unit & 1;
        const uint32_t tr = tid / TILES_W, tc = tid - tr * TILES_W;
        const uint32_t c = tc >> 3, kh0 = 2 * (tc & 7);
        const uint32_t t0 = tr * 32 + hsel * 16;
        uint32_t gh = t0 / 12, gw = t0 - gh * 12;
        // 1) issue the 32 window reads of this half tile (16 tokens x 2 kernel rows)
        uint32_t g_h = gh, g_w = gw;
        for (uint32_t r = 0; r < 16; ++r) {
#ifdef IM2COL_BF16
            const uint32_t off = g_w == 0 ? 0 : (g_w * 16 - 8) * 2;
#else
            const uint32_t off = g_w == 0 ? 0 : (g_w * 16 - 4) * 4;
#endif
            const int32_t ybase = int32_t(g_h * 16 + kh0) - 2;
            for (uint32_t j = 0; j < 2; ++j) {
                const int32_t y = ybase + int32_t(j);
                if (y >= 0) {
                    noc_async_read(src.get_noc_addr(c * 256 + uint32_t(y), off), win_buf + (r * 2 + j) * WIN, WIN);
                }
            }
            if (++g_w == 12) {
                g_w = 0;
                ++g_h;
            }
        }
        noc_async_read_barrier();
        // 2) shift each window into its face row (half tile = faces {0,1} of rows r: face j, row r)
        g_h = gh;
        g_w = gw;
        for (uint32_t r = 0; r < 16; ++r) {
            const int32_t ybase = int32_t(g_h * 16 + kh0) - 2;
            for (uint32_t j = 0; j < 2; ++j) {
                uint32_t* d = half + j * 256 + r * 16;
#ifdef IM2COL_BF16
                const uint16_t* w = reinterpret_cast<const uint16_t*>(win_buf + (r * 2 + j) * WIN);
                if (ybase + int32_t(j) < 0) {
#pragma GCC unroll 16
                    for (uint32_t e = 0; e < 16; ++e) d[e] = 0;
                } else if (g_w == 0) {
                    d[0] = 0;
                    d[1] = 0;
#pragma GCC unroll 14
                    for (uint32_t e = 2; e < 16; ++e) d[e] = uint32_t(w[e - 2]) << 16;
                } else {
#pragma GCC unroll 16
                    for (uint32_t e = 0; e < 16; ++e) d[e] = uint32_t(w[e + 6]) << 16;
                }
#else
                const uint32_t* w = win + (r * 2 + j) * (WIN / 4);
                if (ybase + int32_t(j) < 0) {
#pragma GCC unroll 16
                    for (uint32_t e = 0; e < 16; ++e) d[e] = 0;
                } else if (g_w == 0) {
                    d[0] = 0;
                    d[1] = 0;
#pragma GCC unroll 14
                    for (uint32_t e = 2; e < 16; ++e) d[e] = w[e - 2];
                } else {
#pragma GCC unroll 16
                    for (uint32_t e = 0; e < 16; ++e) d[e] = w[e + 2];
                }
#endif
            }
            if (++g_w == 12) {
                g_w = 0;
                ++g_h;
            }
        }
        // 3) one 2 KB write: faces 2*hsel, 2*hsel+1 of the tile
        noc_async_write(half_buf, dst.get_noc_addr(tid, hsel * HALF_BYTES), HALF_BYTES);
        noc_async_write_barrier();
    }
}