// SPDX-License-Identifier: Apache-2.0
// MK-T writer (sceneseg-p150 OPT_PLAN 3.4, the D19 seg-tail module megakernel, tt/mkt.py), BRISC.
//
// D store (CB 6, the same L1 address on every core): slot s (0 .. n + 1) = image row ya - 1 + s of the bf16
// 64-channel decode_layer_9 map, 4 channel groups of RSB = (W + 2) * 32 bytes each: stick (pixel p, group g) at
// slot + g * RSB + (p + 1) * 32 (pixels -1 and W are the zero padding sticks).
//   0. zero the padding sticks of every slot, and the halo slot outside the image (slot 0 of core 0, slot n + 1 of
//      the last core).
//   1. phase 1, per tile column x: the n x 2 bf16 tiles of compute (CB 16; tile (i, c) = pixels x0 .. x0 + 31 of row
//      ya + i, channels 32 c ..) -> slot i + 1 as four 512-byte face copies (face f = group 2c + (f & 1), pixels
//      x0 + 16 (f >> 1) ..); MKT_DBG >= 1 also writes them to the debug map.
//   2. exchange (no multicast, no grid barrier): after its own rows have landed, each core increments semaphore 0 of
//      its neighbours c - 1 and c + 1, waits until its own semaphore 0 counts its neighbours (and resets it), copies
//      slot n(c - 1) of core c - 1 into its slot 0 and slot 1 of core c + 1 into its slot n + 1 (82 KB reads from the
//      neighbour's L1), then increments semaphore 1 of both neighbours (the "done reading" acknowledgement, waited on
//      before exit so no core leaves while a neighbour still reads its D store) and hands the reader a token (CB 7).
//   3. phase 2, per tile column: the n fp32 margin tiles (CB 17: [0, m1, m2, pad] per pixel) -> the class of each
//      pixel exactly as K8 / ttnn.argmax(dim=1) (strict '>' in the SFPU's sign-magnitude order, lowest index on ties,
//      both-zero / subnormal compares equal) into n row buffers; after the last column the n uint8 rows (one W-byte
//      page each) are written. MKT_DBG 2 also writes the margin tiles.
// Semaphores are self-cleaning (reset after each wait) for trace replays.
// CT args: [0] H, [1] W, [2] NC, [3] GX, [4] RB, [5] RE, then the TensorAccessorArgs of the class map, the debug
// decode_layer_9 map and the debug margins. Common RT args: [class map, dbg9, dbgm, NoC coords of cores 0 .. NC-1].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t H = get_compile_time_arg_val(0);
constexpr uint32_t W = get_compile_time_arg_val(1);
constexpr uint32_t NC = get_compile_time_arg_val(2);
constexpr uint32_t GX = get_compile_time_arg_val(3);
constexpr uint32_t RB = get_compile_time_arg_val(4);
constexpr uint32_t RE = get_compile_time_arg_val(5);
constexpr auto out_args = TensorAccessorArgs<6>();
constexpr auto d9_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();
constexpr auto dm_args = TensorAccessorArgs<d9_args.next_compile_time_args_offset()>();

constexpr uint32_t XT = W / 32;
constexpr uint32_t TB16 = 2048, TB32 = 4096, FACE = 512, STICK = 32;
constexpr uint32_t RSB = (W + 2) * STICK;
constexpr uint32_t SLOTB = 4 * RSB;
constexpr uint32_t cb_d = 6, cb_sync = 7, cb_scr = 9, cb_d9 = 16, cb_m = 17;

#ifndef MKT_DBG
#define MKT_DBG 0
#endif

FORCE_INLINE uint64_t core_noc_addr(uint32_t c, uint32_t l1_addr) {
    const uint32_t packed = get_common_arg_val<uint32_t>(3 + c);
    return get_noc_addr(packed & 0xFFFFu, packed >> 16, l1_addr);
}

// SFPU gt(a, b) on fp32 bit patterns (no NaN in the margins): as K8 (epi_writer.cpp)
FORCE_INLINE bool sfpu_gt(uint32_t a, uint32_t b) {
    if ((a & 0x7F800000u) == 0 && (b & 0x7F800000u) == 0) {
        return false;
    }
    const int32_t ka = (a & 0x80000000u) ? -(int32_t)(a & 0x7FFFFFFFu) : (int32_t)a;
    const int32_t kb = (b & 0x80000000u) ? -(int32_t)(b & 0x7FFFFFFFu) : (int32_t)b;
    return ka > kb;
}

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    const uint32_t n = RB + (core < RE ? 1 : 0);
    const uint32_t ya = core * RB + (core < RE ? core : RE);
    const uint32_t d = get_write_ptr(cb_d);
    const uint32_t zero = (get_write_ptr(cb_scr) + 63u) & ~63u;
    const uint32_t rows = zero + FACE + 64;              // n class-map rows of W bytes (64-aligned)
    {
        volatile tt_l1_ptr uint32_t* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero);
        for (uint32_t i = 0; i < FACE / 4; ++i) {
            z[i] = 0;
        }
        asm volatile("" ::: "memory");
        invalidate_l1_cache();
    }
    const uint64_t zero_noc = get_noc_addr(zero);
    const uint32_t nr = n + 2;
    for (uint32_t s = 0; s < nr; ++s) {
        for (uint32_t g = 0; g < 4; ++g) {
            const uint32_t reg = d + s * SLOTB + g * RSB;
            noc_async_read(zero_noc, reg, STICK);
            noc_async_read(zero_noc, reg + (W + 1) * STICK, STICK);
        }
    }
    const bool first = core == 0, last = core == NC - 1;
    if (first || last) {
        const uint32_t reg = d + (first ? 0 : (n + 1)) * SLOTB;
        for (uint32_t o = 0; o < SLOTB; o += FACE) {
            noc_async_read(zero_noc, reg + o, (SLOTB - o) < FACE ? (SLOTB - o) : FACE);
        }
        if (first && last) {
            const uint32_t reg2 = d + (n + 1) * SLOTB;
            for (uint32_t o = 0; o < SLOTB; o += FACE) {
                noc_async_read(zero_noc, reg2 + o, (SLOTB - o) < FACE ? (SLOTB - o) : FACE);
            }
        }
    }
    noc_async_read_barrier();

    // ---------------- phase 1: decode_layer_9 tiles into the D store
    const uint32_t no = 2 * n;
#if MKT_DBG >= 1
    const auto d9 = TensorAccessor(d9_args, get_common_arg_val<uint32_t>(1), TB16);
#endif
    for (uint32_t x = 0; x < XT; ++x) {
        cb_wait_front(cb_d9, no);
        const uint32_t src = get_read_ptr(cb_d9);
        for (uint32_t i = 0; i < n; ++i) {
            for (uint32_t c = 0; c < 2; ++c) {
                const uint32_t tile = src + (i * 2 + c) * TB16;
                const uint32_t slot = d + (i + 1) * SLOTB;
                for (uint32_t f = 0; f < 4; ++f) {
                    const uint32_t dst = slot + (2 * c + (f & 1)) * RSB + (x * 32 + (f >> 1) * 16 + 1) * STICK;
                    noc_async_read(get_noc_addr(tile + f * FACE), dst, FACE);
                }
#if MKT_DBG >= 1
                noc_async_write(tile, d9.get_noc_addr(((ya + i) * XT + x) * 2 + c), TB16);
#endif
            }
        }
        noc_async_read_barrier();
#if MKT_DBG >= 1
        noc_async_write_barrier();
#endif
        cb_pop_front(cb_d9, no);
    }

#if MKT_DBG != 1
    // ---------------- exchange with the neighbour cores
    volatile tt_l1_ptr uint32_t* ready = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(0));
    volatile tt_l1_ptr uint32_t* done = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(1));
    const uint32_t nn = (first ? 0 : 1) + (last ? 0 : 1);
    if (!first) {
        noc_semaphore_inc(core_noc_addr(core - 1, get_semaphore(0)), 1);
    }
    if (!last) {
        noc_semaphore_inc(core_noc_addr(core + 1, get_semaphore(0)), 1);
    }
    if (nn) {
        noc_semaphore_wait(ready, nn);
        noc_semaphore_set(ready, 0);
    }
    if (!first) {
        const uint32_t nprev = RB + (core - 1 < RE ? 1 : 0);
        noc_async_read(core_noc_addr(core - 1, d + nprev * SLOTB), d, SLOTB);
    }
    if (!last) {
        noc_async_read(core_noc_addr(core + 1, d + SLOTB), d + (n + 1) * SLOTB, SLOTB);
    }
    noc_async_read_barrier();
    if (!first) {
        noc_semaphore_inc(core_noc_addr(core - 1, get_semaphore(1)), 1);
    }
    if (!last) {
        noc_semaphore_inc(core_noc_addr(core + 1, get_semaphore(1)), 1);
    }
    cb_reserve_back(cb_sync, 1);
    cb_push_back(cb_sync, 1);

    // ---------------- phase 2: class decision
#if MKT_DBG >= 2
    const auto dm = TensorAccessor(dm_args, get_common_arg_val<uint32_t>(2), TB32);
#endif
    volatile tt_l1_ptr uint8_t* cls = reinterpret_cast<volatile tt_l1_ptr uint8_t*>(rows);
    for (uint32_t x = 0; x < XT; ++x) {
        cb_wait_front(cb_m, n);
        const uint32_t src = get_read_ptr(cb_m);
        for (uint32_t i = 0; i < n; ++i) {
            volatile tt_l1_ptr uint32_t* m = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(src + i * TB32);
            for (uint32_t r = 0; r < 32; ++r) {
                const uint32_t w = (r >> 4) * 512 + (r & 15) * 16;
                uint32_t best = m[w];
                uint8_t c = 0;
                const uint32_t m1 = m[w + 1], m2 = m[w + 2];
                if (sfpu_gt(m1, best)) {
                    best = m1;
                    c = 1;
                }
                if (sfpu_gt(m2, best)) {
                    c = 2;
                }
                cls[i * W + x * 32 + r] = c;
            }
#if MKT_DBG >= 2
            noc_async_write(src + i * TB32, dm.get_noc_addr((ya + i) * XT + x), TB32);
#endif
        }
#if MKT_DBG >= 2
        noc_async_write_barrier();
#endif
        cb_pop_front(cb_m, n);
    }
    const auto out = TensorAccessor(out_args, get_common_arg_val<uint32_t>(0), W);
    for (uint32_t i = 0; i < n; ++i) {
        noc_async_write(rows + i * W, out.get_noc_addr(ya + i), W);
    }
    noc_async_write_barrier();
    if (nn) {
        noc_semaphore_wait(done, nn);
        noc_semaphore_set(done, 0);
    }
    noc_async_atomic_barrier();
#endif
}