// SPDX-FileCopyrightText: © 2026 Tenstorrent Inc.
// SPDX-License-Identifier: Apache-2.0
// Two-hop activation broadcast micro-benchmark (round 7).  NSH shards of SB bytes must reach every core.
// Hop A: a head core multicasts the shard it holds into rect A (its row or column), then sets flag s there.
// Hop B: inside rect B (the other dimension) the NB holders send one after the other (chain: sender k waits until the
// shard of sender k-1 has arrived at itself), each multicast + flag into the whole rect B.
// Every multicast includes the sender (loopback); the head of a rect first waits until every member has started
// (semaphores are reset at launch, so nobody may signal a core that has not started yet).
// runtime args:
//  0 slot (shard held after hop A, 0xff: none)   1 isHeadA   2..5 rectA (virtual min x,y max x,y)   6 nA
//  7..10 rectB   11 nB   12 kB (position in the B chain, 0xff: not a sender)   13 prev slot (chain)
//  14 inA (member of a rect A, not head)  15,16 headA x,y   17 inB  18,19 headB x,y   20 isHeadB
//  21 NSH   22 SB (shard bytes)
// semaphores: 0 START_A, 1 START_B, 2+s flag of shard s.  CB 0: NSH * SB bytes (shard s at s * SB).  CB 1: scratch.
#include <stdint.h>
#include "api/dataflow/dataflow_api.h"

static inline void wait_ge(uint32_t sem, uint32_t v) {
    volatile tt_l1_ptr uint32_t* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(get_semaphore(sem));
    while (true) {
        invalidate_l1_cache();
        if (*p >= v) {
            break;
        }
    }
}

void kernel_main() {
    const uint32_t slot = get_arg_val<uint32_t>(0);
    const uint32_t isHeadA = get_arg_val<uint32_t>(1);
    const uint32_t ax0 = get_arg_val<uint32_t>(2), ay0 = get_arg_val<uint32_t>(3);
    const uint32_t ax1 = get_arg_val<uint32_t>(4), ay1 = get_arg_val<uint32_t>(5);
    const uint32_t nA = get_arg_val<uint32_t>(6);
    const uint32_t bx0 = get_arg_val<uint32_t>(7), by0 = get_arg_val<uint32_t>(8);
    const uint32_t bx1 = get_arg_val<uint32_t>(9), by1 = get_arg_val<uint32_t>(10);
    const uint32_t nB = get_arg_val<uint32_t>(11);
    const uint32_t kB = get_arg_val<uint32_t>(12);
    const uint32_t prev = get_arg_val<uint32_t>(13);
    const uint32_t inA = get_arg_val<uint32_t>(14);
    const uint32_t hax = get_arg_val<uint32_t>(15), hay = get_arg_val<uint32_t>(16);
    const uint32_t inB = get_arg_val<uint32_t>(17);
    const uint32_t hbx = get_arg_val<uint32_t>(18), hby = get_arg_val<uint32_t>(19);
    const uint32_t isHeadB = get_arg_val<uint32_t>(20);
    const uint32_t NSH = get_arg_val<uint32_t>(21);
    const uint32_t SB = get_arg_val<uint32_t>(22);
    const uint32_t base = get_write_ptr(0);
    const uint32_t scr = get_write_ptr(1);
    auto mc = [&](uint32_t x0, uint32_t y0, uint32_t x1, uint32_t y1, uint32_t addr) -> uint64_t {
        if (noc_index == 1) {
            return get_noc_multicast_addr(x1, y1, x0, y0, addr);
        }
        return get_noc_multicast_addr(x0, y0, x1, y1, addr);
    };
    *reinterpret_cast<volatile tt_l1_ptr uint32_t*>(scr) = 1;

    if (inA) {
        noc_semaphore_inc(get_noc_addr(hax, hay, get_semaphore(0)), 1);
    }
    if (inB) {
        noc_semaphore_inc(get_noc_addr(hbx, hby, get_semaphore(1)), 1);
    }
    if (isHeadA) {
#ifdef BC_VERIFY
        {
            volatile tt_l1_ptr uint32_t* d = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base + slot * SB);
            for (uint32_t i = 0; i < SB / 4; ++i) {
                d[i] = (slot << 20) ^ i ^ 0x5a5a0000u;
            }
        }
#endif
        wait_ge(0, nA - 1);
#ifndef BC_NO_DATA
        noc_async_write_multicast_loopback_src(base + slot * SB, mc(ax0, ay0, ax1, ay1, base + slot * SB), SB, nA);
#endif
        noc_async_writes_flushed();
        noc_semaphore_set_multicast_loopback_src(scr, mc(ax0, ay0, ax1, ay1, get_semaphore(2 + slot)), nA);
    }
    if (kB != 0xff) {
#ifdef BC_CHAIN2
        // two chains: even positions on NOC_0, odd on NOC_1, each waits for the previous sender of its parity
        const uint8_t nn = kB & 1;
        auto mcn = [&](uint32_t addr) -> uint64_t {
            return nn == 1 ? get_noc_multicast_addr(bx1, by1, bx0, by0, addr, 1) : get_noc_multicast_addr(bx0, by0, bx1, by1, addr, 0);
        };
        // NOT VALIDATED (round 7): the first version waited START_B on chain positions 0 AND 1, but only the
        // chain head (position 0) receives START_B -> position 1 spun forever (chip 28 marked FAULT, software
        // deadlock).  Now: position 0 waits START_B, position 1 waits the shard of position 0 (prev), positions
        // k >= 2 wait the shard of position k - 2 (rc layout: slot == position).
        wait_ge(2 + slot, 1);
        if (kB == 0) {
            wait_ge(1, nB - 1);
        } else if (kB == 1) {
            wait_ge(2 + prev, 1);
        } else {
            wait_ge(2 + slot - 2, 1);
        }
#ifndef BC_NO_DATA
        noc_async_write_multicast_loopback_src(base + slot * SB, mcn(base + slot * SB), SB, nB, false, nn);
#endif
        noc_async_writes_flushed(nn);
        noc_semaphore_set_multicast_loopback_src(scr, mcn(get_semaphore(2 + slot)), nB, false, nn);
#else
        wait_ge(2 + slot, 1);
#ifndef BC_NOCHAIN
        if (isHeadB) {
            wait_ge(1, nB - 1);
        } else {
            wait_ge(2 + prev, 1);
        }
#else
        wait_ge(1, nB - 1);
#endif
#ifndef BC_NO_DATA
        noc_async_write_multicast_loopback_src(base + slot * SB, mc(bx0, by0, bx1, by1, base + slot * SB), SB, nB);
#endif
        noc_async_writes_flushed();
        noc_semaphore_set_multicast_loopback_src(scr, mc(bx0, by0, bx1, by1, get_semaphore(2 + slot)), nB);
#endif
    }
    for (uint32_t s = 0; s < NSH; ++s) {
        wait_ge(2 + s, 1);
    }
    noc_async_write_barrier(0);
    noc_async_write_barrier(1);
    noc_async_atomic_barrier(0);
    noc_async_atomic_barrier(1);
#ifdef BC_VERIFY
    {
        uint32_t bad = 0;
        for (uint32_t s = 0; s < NSH; ++s) {
            volatile tt_l1_ptr uint32_t* d = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base + s * SB);
            for (uint32_t i = 0; i < SB / 4; ++i) {
                bad += d[i] != ((s << 20) ^ i ^ 0x5a5a0000u);
            }
        }
        volatile tt_l1_ptr uint32_t* r = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(scr + 64);
        r[0] = bad;
        r[1] = 0xC0DE0000u | get_arg_val<uint32_t>(24);
        constexpr auto o_args = TensorAccessorArgs<0>();
        const auto o = TensorAccessor(o_args, get_arg_val<uint32_t>(23), 32);
        noc_async_write(scr + 64, o.get_noc_addr(get_arg_val<uint32_t>(24)), 32);
        noc_async_write_barrier();
    }
#endif
}