// SPDX-License-Identifier: Apache-2.0
// DPOOL (meteor-p150 OPT2 round 4 R4.4, tt/dpool.py): the four 2x2-pool quadrant maps of an fp32 TILE map
// [N * H * W, C] (NHWC rows, H even, W % 64 == 0) in ONE data-movement program: map m = (dy, dx) = (m / 2, m % 2) holds
// at output row (n * H / 2 + i) * W / 2 + j the input pixel (n, 2 i + dy, 2 j + dx), all C channels. The depth head's
// avgpool2 then is ((q00 + q10) + (q01 + q11)) * 0.25 with stock fp32 adds: the same operands and order as the port's
// slice / tilize / reshape chain (bit-identical), without its RM round trips.
// Work unit = one output tile row u (32 output pixels of one output image row: W / 2 % 32 == 0): its input pixels are
// the 64 consecutive pixels from (2 q) * W + 2 j0 (top) and (2 q + 1) * W + 2 j0 (bottom), q = u * 32 / (W / 2),
// j0 = u * 32 % (W / 2): input tile rows T, T + 1 (top) and T + W / 32, T + W / 32 + 1 (bottom), CT tiles each.
// Output row k of map (dy, dx) <- input pixel 2 k + dx of the top / bottom span: tile (2 k + dx) / 32, row
// (2 k + dx) % 32; each (row, face column) is a 16-word face row copied by the RISC.
// Element (r, c) of a 32 x 32 fp32 tile: face (r / 16) * 2 + (c % 32) / 16, word (r % 16) * 16 + c % 16.
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns units
// [i * B + min(i, E), + B + (i < E)), B = NU / NC, E = NU % NC. Reads and writes on one RISC / NOC (SYS-1419).
// CT args: [0] NU, [1] NC, [2] GX, [3] CT (channel tiles), [4] WT (= W / 32 input tile rows per image row), then the
// TensorAccessorArgs of the input and of the outputs (all four alike). Common RT args: [in, out0, out1, out2, out3].
// CB 0: 2 x 4 * CT tiles (input slots); CB 1: 2 x 4 * CT tiles (output slots); plain L1 scratch, never pushed.
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NU = 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 CT = get_compile_time_arg_val(3);
constexpr uint32_t WT = get_compile_time_arg_val(4);
constexpr auto in_args = TensorAccessorArgs<5>();
constexpr auto out_args = TensorAccessorArgs<in_args.next_compile_time_args_offset()>();
constexpr uint32_t TB = 4096;              // fp32 tile bytes
constexpr uint32_t TW = 1024;              // fp32 tile words
constexpr uint32_t SLOT = 4 * CT * TB;     // 4 input tile rows (top 2, bottom 2) / 4 output maps

FORCE_INLINE void cbar() { asm volatile("" ::: "memory"); }

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NU / NC, E = NU % NC;
    const uint32_t u0 = core * B + (core < E ? core : E);
    const uint32_t nu = B + (core < E ? 1 : 0);
    if (nu == 0) {
        return;
    }
    const auto in = TensorAccessor(in_args, get_common_arg_val<uint32_t>(0), TB);
    const uint32_t ibuf = get_write_ptr(0);
    const uint32_t obuf = get_write_ptr(1);
    constexpr uint32_t HALF = WT / 2;      // output tile rows per output image row (W / 2 / 32)

    // input slot layout: [s][t][ct] tiles, s = 0 top / 1 bottom, t = 0 / 1 the tile row of the 64-pixel span
    auto issue_reads = [&](uint32_t u, uint32_t buf) {
        const uint32_t q = u / HALF;
        const uint32_t jb = u - q * HALF;
        const uint32_t top = 2 * q * WT + 2 * jb;
        for (uint32_t s = 0; s < 2; ++s) {
            const uint32_t tr = top + s * WT;
            for (uint32_t t = 0; t < 2; ++t) {
                for (uint32_t ct = 0; ct < CT; ++ct) {
                    noc_async_read(in.get_noc_addr((tr + t) * CT + ct), buf + ((s * 2 + t) * CT + ct) * TB, TB);
                }
            }
        }
    };
    issue_reads(u0, ibuf);
    for (uint32_t i = 0; i < nu; ++i) {
        const uint32_t ib = ibuf + (i & 1) * SLOT;
        const uint32_t ob = obuf + (i & 1) * SLOT;
        noc_async_read_barrier();
        if (i + 1 < nu) {
            issue_reads(u0 + i + 1, ibuf + ((i + 1) & 1) * SLOT);
        }
        if (i >= 2) {
            noc_async_writes_flushed();        // the writes of unit i - 2 (output slot i & 1) have left L1
        }
        cbar();
        const uint32_t* src = reinterpret_cast<const uint32_t*>(ib);
        uint32_t* dst = reinterpret_cast<uint32_t*>(ob);
        // output slot layout: [m][ct] tiles, m = dy * 2 + dx
        for (uint32_t m = 0; m < 4; ++m) {
            const uint32_t dy = m >> 1, dx = m & 1;
            for (uint32_t ct = 0; ct < CT; ++ct) {
                uint32_t* o = dst + (m * CT + ct) * TW;
                for (uint32_t k = 0; k < 32; ++k) {
                    const uint32_t p = 2 * k + dx;                  // pixel of the 64-pixel span
                    const uint32_t t = p >> 5, r = p & 31;
                    const uint32_t* it = src + ((dy * 2 + t) * CT + ct) * TW + (r >> 4) * 512 + (r & 15) * 16;
                    uint32_t* ot = o + (k >> 4) * 512 + (k & 15) * 16;
#pragma GCC unroll 16
                    for (uint32_t w = 0; w < 16; ++w) {
                        ot[w] = it[w];                              // face column 0
                    }
#pragma GCC unroll 16
                    for (uint32_t w = 0; w < 16; ++w) {
                        ot[256 + w] = it[256 + w];                  // face column 1
                    }
                }
            }
        }
        cbar();
        const uint32_t up = u0 + i;
        for (uint32_t m = 0; m < 4; ++m) {
            const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(1 + m), TB);
            for (uint32_t ct = 0; ct < CT; ++ct) {
                noc_async_write(ob + (m * CT + ct) * TB, out.get_noc_addr(up * CT + ct), TB);
            }
        }
    }
    noc_async_write_barrier();
}