// SPDX-License-Identifier: Apache-2.0
// SE-K reader (sceneseg-p150 OPT_PLAN SE-K, tt/sek.py): the operands of one squeeze-excitation MLP, in the order the
// compute kernel consumes them: the fp32 row sum x (KT1 tiles, CB 0), the fc1 weight tiles (2 terms x KT1 x NT1
// bf16, CB 1), the fc1 bias tiles (NT1 fp32, CB 2), then this core's fc2 weight tiles (2 terms x NT1 x its output
// tiles, CB 3) and fc2 bias tiles (CB 4). Every core computes fc1 in full (no inter-core sync); core i of the first
// NC (row-major) owns fc2 output tiles [i * B + min(i, E), + B + (i < E)), B = NT2 / NC, E = NT2 % NC.
// CT args: [0] KT1, [1] NT1, [2] NT2, [3] NC, [4] GX, then the TensorAccessorArgs of x, W1, b1, W2, b2.
// Common RT args: [x, w1, b1, w2, b2].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t KT1 = get_compile_time_arg_val(0);
constexpr uint32_t NT1 = get_compile_time_arg_val(1);
constexpr uint32_t NT2 = get_compile_time_arg_val(2);
constexpr uint32_t NC = get_compile_time_arg_val(3);
constexpr uint32_t GX = get_compile_time_arg_val(4);
constexpr auto x_args = TensorAccessorArgs<5>();
constexpr auto w1_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto b1_args = TensorAccessorArgs<w1_args.next_compile_time_args_offset()>();
constexpr auto w2_args = TensorAccessorArgs<b1_args.next_compile_time_args_offset()>();
constexpr auto b2_args = TensorAccessorArgs<w2_args.next_compile_time_args_offset()>();
constexpr uint32_t TB32 = 4096, TB16 = 2048;
constexpr uint32_t cb_x = 0, cb_w1 = 1, cb_b1 = 2, cb_w2 = 3, cb_b2 = 4;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NT2 / NC, E = NT2 % NC;
    const uint32_t n0 = core * B + (core < E ? core : E);
    const uint32_t nn = B + (core < E ? 1 : 0);
    const auto ax = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB32);
    const auto aw1 = TensorAccessor(w1_args, get_common_arg_val<uint32_t>(1), TB16);
    const auto ab1 = TensorAccessor(b1_args, get_common_arg_val<uint32_t>(2), TB32);
    const auto aw2 = TensorAccessor(w2_args, get_common_arg_val<uint32_t>(3), TB16);
    const auto ab2 = TensorAccessor(b2_args, get_common_arg_val<uint32_t>(4), TB32);

    cb_reserve_back(cb_x, KT1);
    uint32_t l1 = get_write_ptr(cb_x);
    for (uint32_t k = 0; k < KT1; ++k, l1 += TB32) {
        noc_async_read(ax.get_noc_addr(k), l1, TB32);
    }
    noc_async_read_barrier();
    cb_push_back(cb_x, KT1);

    cb_reserve_back(cb_w1, 2 * KT1 * NT1);
    l1 = get_write_ptr(cb_w1);
    for (uint32_t i = 0; i < 2 * KT1 * NT1; ++i, l1 += TB16) {
        noc_async_read(aw1.get_noc_addr(i), l1, TB16);
    }
    cb_reserve_back(cb_b1, NT1);
    l1 = get_write_ptr(cb_b1);
    for (uint32_t n = 0; n < NT1; ++n, l1 += TB32) {
        noc_async_read(ab1.get_noc_addr(n), l1, TB32);
    }
    noc_async_read_barrier();
    cb_push_back(cb_w1, 2 * KT1 * NT1);
    cb_push_back(cb_b1, NT1);

    // fc2: tile (t * NT1 + k2) * NT2 + n of W2, pushed term-major, then k2, then this core's n
    cb_reserve_back(cb_w2, 2 * NT1 * nn);
    l1 = get_write_ptr(cb_w2);
    for (uint32_t t = 0; t < 2; ++t) {
        for (uint32_t k2 = 0; k2 < NT1; ++k2) {
            for (uint32_t j = 0; j < nn; ++j, l1 += TB16) {
                noc_async_read(aw2.get_noc_addr((t * NT1 + k2) * NT2 + n0 + j), l1, TB16);
            }
        }
    }
    cb_reserve_back(cb_b2, nn);
    l1 = get_write_ptr(cb_b2);
    for (uint32_t j = 0; j < nn; ++j, l1 += TB32) {
        noc_async_read(ab2.get_noc_addr(n0 + j), l1, TB32);
    }
    noc_async_read_barrier();
    cb_push_back(cb_w2, 2 * NT1 * nn);
    cb_push_back(cb_b2, nn);
}