// SPDX-License-Identifier: Apache-2.0
// ARGMAX v2 (meteor-p150 OPT2 round 4 R4.2, tt/argmax.py, METEOR_ARGMAX=fast): argmax_rows.cpp with the same data
// movement and the same compare (ArgMaxNC semantics: fp32 values, strict `>` in channel order, first maximum on ties,
// -0.0 == +0.0), only the per-row reduction rewritten. v1 spent ~32 RISC cycles per element (the per-element face index
// arithmetic and a branchy key transform); the reads are not the limit (the kernel is RISC-bound: 42 MB in 1.9 ms).
//   - the key transform is branchless: z = (b << 1) ? b : 0 (-0.0 -> +0.0), key = z ^ ((int32(z) >> 31) | 0x80000000)
//     (sign set: ~z; sign clear: z | 0x80000000): the monotone unsigned order of the floats, as v1's key_of;
//   - the channels of one row are walked face by face (16 contiguous words), the loop bounds are compile-time (C);
//   - two rows (r and r + 16: the same word offsets in face rows 0 and 1) are reduced together, so their loads and
//     compares interleave.
// Work unit, work split, CT args, common RT args and CBs: as argmax_rows.cpp (no per-core RT args, probe P3; reads
// and writes on one RISC / NOC, SYS-1419).
#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 C = get_compile_time_arg_val(3);
constexpr uint32_t CT = 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 SLOT = CT * TB;
constexpr uint32_t OB = 128;           // one output row: 32 uint32
constexpr uint32_t NF = (C + 15) / 16; // faces (16-channel groups) per row

FORCE_INLINE uint32_t key_of(uint32_t b) {
    const uint32_t z = (b << 1) ? b : 0u;                                         // -0.0 == +0.0
    return z ^ (static_cast<uint32_t>(static_cast<int32_t>(z) >> 31) | 0x80000000u);
}

// channels [16 g, 16 g + n) of rows r (word offset `a` of face row 0) and r + 16 (a + 512)
template <uint32_t G>
FORCE_INLINE void face(const uint32_t* src, uint32_t a, uint32_t& k0, uint32_t& b0, uint32_t& k1, uint32_t& b1) {
    constexpr uint32_t n = (C - 16 * G < 16) ? C - 16 * G : 16;
    constexpr uint32_t off = (G >> 1) * 1024 + (G & 1) * 256;                     // tile G / 2, face column G % 2
    const uint32_t* p = src + off + a;
#pragma GCC unroll 16
    for (uint32_t j = (G == 0 ? 1 : 0); j < n; ++j) {
        const uint32_t x0 = key_of(p[j]);
        const uint32_t x1 = key_of(p[j + 512]);
        if (x0 > k0) {
            k0 = x0;
            b0 = 16 * G + j;
        }
        if (x1 > k1) {
            k1 = x1;
            b1 = 16 * G + j;
        }
    }
}

template <uint32_t G>
FORCE_INLINE void faces(const uint32_t* src, uint32_t a, uint32_t& k0, uint32_t& b0, uint32_t& k1, uint32_t& b1) {
    if constexpr (G < NF) {
        face<G>(src, a, k0, b0, k1, b1);
        faces<G + 1>(src, a, k0, b0, k1, b1);
    }
}

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 auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(1), OB);
    const uint32_t ibuf = get_write_ptr(0);
    const uint32_t obuf = get_write_ptr(1);

    auto issue_reads = [&](uint32_t u, uint32_t buf) {
        for (uint32_t k = 0; k < CT; ++k) {
            noc_async_read(in.get_noc_addr(u * CT + k), buf + k * 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) * OB;
        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 write of unit i - 2 (same slot) is out
        }
        const uint32_t* src = reinterpret_cast<const uint32_t*>(ib);
        uint32_t* dst = reinterpret_cast<uint32_t*>(ob);
        for (uint32_t r = 0; r < 16; ++r) {
            const uint32_t a = r * 16;                                // row r of face row 0; row r + 16 at a + 512
            uint32_t k0 = key_of(src[a]), b0 = 0;
            uint32_t k1 = key_of(src[a + 512]), b1 = 0;
            faces<0>(src, a, k0, b0, k1, b1);
            dst[r] = b0;
            dst[r + 16] = b1;
        }
        noc_async_write(ob, out.get_noc_addr(u0 + i), OB);
    }
    noc_async_write_barrier();
}