// SPDX-License-Identifier: Apache-2.0
// TFUSE reader (meteor-p150 OPT2 MK-TF, tt/tgate.py fuse_mm): the temporal fuse's merged 1x1 conv AND its gate glue
// in one program. Per 32-row block b of the bf16 raw map x (flat [R, 96] or the 4-D [1, H, W, 96] TILE form; both
// are 3 tiles per 32-row block): tiles b * 3 + 0..2 into CB 8. Once: the 12 fp32 weight tiles of W^T [96, 128]
// (tile k * 4 + j: input channels 32k.., output columns 32j..; columns 0..95 tfuse.0, 96..98 the gate logit
// differences, the rest zero) into CB 6, the ones tile (CB 2), the 3 tfuse.0 bias tiles (CB 3) and the two gate
// bias tiles (hi / lo bf16 terms of b_j - b_0, every row; CB 7).
// Work split (no per-core RT args, probe P3): core i = y * GX + x of the first NC cores owns the contiguous block
// range [i * B + min(i, E), + B + (i < E)) of the NB blocks, B = NB / NC, E = NB % NC.
// CT args: [0] NB, [1] NC, [2] GX, then the TensorAccessorArgs of x, w, ones, bias, gbias.
// Common RT args: [x_addr, w_addr, ones_addr, bias_addr, gbias_addr].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t NB = 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 auto x_args = TensorAccessorArgs<3>();
constexpr auto w_args = TensorAccessorArgs<x_args.next_compile_time_args_offset()>();
constexpr auto o_args = TensorAccessorArgs<w_args.next_compile_time_args_offset()>();
constexpr auto b_args = TensorAccessorArgs<o_args.next_compile_time_args_offset()>();
constexpr auto gb_args = TensorAccessorArgs<b_args.next_compile_time_args_offset()>();
constexpr uint32_t TB32 = 4096, TB16 = 2048;
constexpr uint32_t cb_one = 2, cb_bias = 3, cb_w = 6, cb_gb = 7, cb_x = 8;

void kernel_main() {
    const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
    if (core >= NC) {
        return;
    }
    constexpr uint32_t B = NB / NC, E = NB % NC;
    const uint32_t b0 = core * B + (core < E ? core : E);
    const uint32_t n = B + (core < E ? 1 : 0);
    if (n == 0) {
        return;
    }
    const auto x = TensorAccessor(x_args, get_common_arg_val<uint32_t>(0), TB16);
    const auto w = TensorAccessor(w_args, get_common_arg_val<uint32_t>(1), TB32);
    const auto ones = TensorAccessor(o_args, get_common_arg_val<uint32_t>(2), TB32);
    const auto bias = TensorAccessor(b_args, get_common_arg_val<uint32_t>(3), TB32);
    const auto gbias = TensorAccessor(gb_args, get_common_arg_val<uint32_t>(4), TB32);
    cb_reserve_back(cb_w, 12);
    for (uint32_t t = 0; t < 12; ++t) {
        noc_async_read(w.get_noc_addr(t), get_write_ptr(cb_w) + t * TB32, TB32);
    }
    cb_reserve_back(cb_one, 1);
    noc_async_read(ones.get_noc_addr(0), get_write_ptr(cb_one), TB32);
    cb_reserve_back(cb_bias, 3);
    for (uint32_t k = 0; k < 3; ++k) {
        noc_async_read(bias.get_noc_addr(k), get_write_ptr(cb_bias) + k * TB32, TB32);
    }
    cb_reserve_back(cb_gb, 2);
    for (uint32_t k = 0; k < 2; ++k) {
        noc_async_read(gbias.get_noc_addr(k), get_write_ptr(cb_gb) + k * TB32, TB32);
    }
    noc_async_read_barrier();
    cb_push_back(cb_w, 12);
    cb_push_back(cb_one, 1);
    cb_push_back(cb_bias, 3);
    cb_push_back(cb_gb, 2);
    for (uint32_t b = b0; b < b0 + n; ++b) {
        cb_reserve_back(cb_x, 3);
        const uint32_t wx = get_write_ptr(cb_x);
        for (uint32_t k = 0; k < 3; ++k) {
            noc_async_read(x.get_noc_addr(b * 3 + k), wx + k * TB16, TB16);
        }
        noc_async_read_barrier();
        cb_push_back(cb_x, 3);
    }
}