// SPDX-License-Identifier: Apache-2.0
// K-IN rows (bevformer-p150 OPT_PLAN R2-2): one camera of the uploaded bf16 image rows -> the height-sharded L1
// input of the 7x7 / stride-2 stem conv (ttnn's L1 conv path), replacing the device reshape of the image rows to
// 6-byte pixel pages and the stem's 10 DRAM PaddedSlices.
//
//   in : [1, 1, CAMS * H, W * C_IN] bf16 ROW_MAJOR, DRAM interleaved (one page per image row, ROW_BYTES = W * C_IN * 2)
//   out: [1, 1, H * W, C_IN] bf16 ROW_MAJOR, L1 HEIGHT_SHARDED, shard [S, C_PAD] (C_PAD = 8: 16-byte rows), camera
//        CAM's image rows ROW0 = CAM * H .. ROW0 + H - 1
//
// Core k (row-major over a GX-wide grid at the origin; ROW_MAJOR shard orientation) owns pixels
// [k * S, min((k + 1) * S, N)), split between SLOTS data-movement RISC-Vs (BRISC / NCRISC, each its own scratch CB).
// Each reads the image rows that hold its pixels into its scratch CB (whole pages), then copies each pixel's C_IN
// bf16 values and zeros in channels C_IN .. C_PAD - 1 straight into its own shard (local L1 stores; no NoC write).
// The copy is exact (bits moved, no conversion). Data-movement only, no CB push / pop, no semaphores, no multicast;
// every read waited on by noc_async_read_barrier before use. Runtime args: COMMON only [in_addr, out_addr] (P3).
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t N = get_compile_time_arg_val(0);          // pixels of one camera (H * W)
constexpr uint32_t S = get_compile_time_arg_val(1);          // shard rows (pixels per core)
constexpr uint32_t W = get_compile_time_arg_val(2);          // pixels per image row
constexpr uint32_t C_IN = get_compile_time_arg_val(3);       // channels (3)
constexpr uint32_t C_PAD = get_compile_time_arg_val(4);      // shard width in elements (8: 16 B rows)
constexpr uint32_t NUM_CORES = get_compile_time_arg_val(5);  // cores holding shards
constexpr uint32_t GX = get_compile_time_arg_val(6);         // grid width (cores)
constexpr uint32_t CB_SCRATCH = get_compile_time_arg_val(7); // scratch CB: MAX_ROWS image rows
constexpr uint32_t MAX_ROWS = get_compile_time_arg_val(8);   // rows a slot's pixels can touch
constexpr uint32_t SLOTS = get_compile_time_arg_val(9);      // RISC-Vs sharing a core's shard (1 or 2)
constexpr uint32_t SLOT = get_compile_time_arg_val(10);      // this kernel's slot: pixels [SLOT * S / SLOTS, ...)
constexpr uint32_t ROW0 = get_compile_time_arg_val(11);      // first image row of the camera in the input
constexpr auto in_args = TensorAccessorArgs<12>();

constexpr uint32_t ROW_BYTES = W * C_IN * 2;
static_assert(C_IN == 3 && C_PAD == 8, "the copy loop handles 3 channels into 16-byte rows");
static_assert(ROW_BYTES % 64 == 0, "image rows must be 64-byte multiples (DRAM read alignment)");

void kernel_main() {
    const uint32_t in_addr = get_common_arg_val<uint32_t>(0);
    const uint32_t out_addr = get_common_arg_val<uint32_t>(1);
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NUM_CORES) {
        return;
    }
    constexpr uint32_t PART = (S + SLOTS - 1) / SLOTS;
    const uint32_t shard0 = core * S;
    const uint32_t p0 = shard0 + SLOT * PART;
    const uint32_t end = shard0 + ((SLOT + 1) * PART < S ? (SLOT + 1) * PART : S);
    const uint32_t p1 = end < N ? end : N;
    if (p0 >= p1) {
        return;
    }
    const uint32_t row0 = p0 / W;
    const uint32_t row1 = (p1 - 1) / W + 1;                   // exclusive
    const auto ia = TensorAccessor(in_args, in_addr, ROW_BYTES);
    const uint32_t scratch = (get_write_ptr(CB_SCRATCH) + 63u) & ~63u;
    for (uint32_t r = row0; r < row1; ++r) {
        noc_async_read(ia.get_noc_addr(ROW0 + r), scratch + (r - row0) * ROW_BYTES, ROW_BYTES);
    }
    noc_async_read_barrier();
    asm volatile("" ::: "memory");                            // no load of the rows may move above the barrier

    // pixel p's 3 bf16 = 6 bytes at scratch + 6 (p - row0 W); an even pixel starts on a 4-byte boundary
    const uint32_t src0 = scratch + (p0 - row0 * W) * C_IN * 2;
    uint32_t* dst = reinterpret_cast<uint32_t*>(out_addr + (p0 - shard0) * C_PAD * 2);
    uint32_t n = p1 - p0;
    const uint16_t* sh = reinterpret_cast<const uint16_t*>(src0);
    if (n && (reinterpret_cast<uint32_t>(sh) & 3u)) {          // odd start pixel: one single
        dst[0] = static_cast<uint32_t>(sh[0]) | (static_cast<uint32_t>(sh[1]) << 16);
        dst[1] = sh[2];
        dst[2] = 0;
        dst[3] = 0;
        sh += 3;
        dst += 4;
        --n;
    }
    const uint32_t* sw = reinterpret_cast<const uint32_t*>(sh);
    for (; n >= 2; n -= 2) {                                  // 2 pixels = 3 words: a0 b0 | c0 a1 | b1 c1
        const uint32_t w0 = sw[0], w1 = sw[1], w2 = sw[2];
        dst[0] = w0;
        dst[1] = w1 & 0xFFFFu;
        dst[2] = 0;
        dst[3] = 0;
        dst[4] = (w1 >> 16) | (w2 << 16);
        dst[5] = w2 >> 16;
        dst[6] = 0;
        dst[7] = 0;
        sw += 3;
        dst += 8;
    }
    if (n) {
        sh = reinterpret_cast<const uint16_t*>(sw);
        dst[0] = static_cast<uint32_t>(sh[0]) | (static_cast<uint32_t>(sh[1]) << 16);
        dst[1] = sh[2];
        dst[2] = 0;
        dst[3] = 0;
    }
    asm volatile("" ::: "memory");
}