// SPDX-License-Identifier: Apache-2.0
// KCAT operand build (tt/kcat.py), writer (RISCV_1): input tile (r, j) of part [.., N, K] (Kt tiles per row) goes to
// the bf16 operand X' [.., N, 32 OKt] at tile (r, off + j) and (r, off + Kt + j) as x_hi and, for an fp32 input, at
// (r, off + 2 Kt + j) as x_lo. A bf16 input (mode 1) is its own x_hi (x_lo = 0: two weight terms only) and is
// copied from CB 0 as it is. Mode 2 (fp32 x + skip, OPT round 3) writes like mode 0.
// CT args: [0] Kt, [1] per-core RT-arg count (P3), [2] OKt, [3] off (tile column of the part in X'), [4] mode,
// [5] BCOL (tile column of the bias block), [6] BT (bias block tiles; 0 = none), then the TensorAccessorArgs of X'.
// Common RT args: [out_addr]. Per-core RT args: [t0, n].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t Kt = get_compile_time_arg_val(0);
constexpr uint32_t OKt = get_compile_time_arg_val(2);
constexpr uint32_t off = get_compile_time_arg_val(3);
constexpr uint32_t mode = get_compile_time_arg_val(4);
constexpr uint32_t BCOL = get_compile_time_arg_val(5);
constexpr uint32_t BT = get_compile_time_arg_val(6);
constexpr auto o_args = TensorAccessorArgs<7>();
constexpr uint32_t cb_x = 0, cb_hi = 16, cb_lo = 17, cb_c = 18;
constexpr uint32_t TB = 2048;

// The bias block (FRNET_KCAT_BIAS=col; BT > 0 only in the writer of part 0): tile 0 of the block is 1.0 in columns
// 0..2 of every row (x the W' rows b_hi, b_lo, b_lo2), tiles 1..BT-1 are zero. Built once in CB 18 (2 pages; never
// pushed) with 32-bit L1 stores: faces of 16 x 16 bf16, columns 0..3 of row i at face (i / 16) * 2.
FORCE_INLINE uint32_t fill_bias_tiles(uint32_t cb_c) {
    cb_reserve_back(cb_c, 2);
    const uint32_t base = get_write_ptr(cb_c);
    volatile tt_l1_ptr uint32_t* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(base);
    for (uint32_t i = 0; i < 2 * 2048 / 4; ++i) {
        p[i] = 0;
    }
    for (uint32_t i = 0; i < 32; ++i) {
        const uint32_t e = (i >> 4) * 512 + (i & 15) * 16;   // bf16 element index of (row i, column 0)
        p[e / 2] = 0x3F803F80u;                               // columns 0, 1 = 1.0
        p[e / 2 + 1] = 0x00003F80u;                           // column 2 = 1.0, column 3 = 0
    }
    return base;
}

template <typename A>
FORCE_INLINE void write_bias_block(uint32_t ones, uint32_t row_tile0, uint32_t bt, const A& o) {
    noc_async_write(ones, o.get_noc_addr(row_tile0), 2048);
    for (uint32_t b = 1; b < bt; ++b) {
        noc_async_write(ones + 2048, o.get_noc_addr(row_tile0 + b), 2048);
    }
}

void kernel_main() {
    const uint32_t o_addr = get_common_arg_val<uint32_t>(0);
    const uint32_t t0 = get_arg_val<uint32_t>(0);
    const uint32_t n = get_arg_val<uint32_t>(1);
    const auto o = TensorAccessor(o_args, o_addr, TB);
    constexpr uint32_t cb_a = mode != 1 ? cb_hi : cb_x;
    uint32_t ones = 0;
    if constexpr (BT > 0) {
        ones = fill_bias_tiles(cb_c);
    }
    uint32_t r = t0 / Kt, j = t0 - (t0 / Kt) * Kt;
    for (uint32_t t = 0; t < n; ++t) {
        const uint32_t ob = r * OKt + off + j;
        if constexpr (BT > 0) {
            if (j == 0) {   // the core that writes the row's first part tile also writes its bias block
                write_bias_block(ones, r * OKt + BCOL, BT, o);
            }
        }
        cb_wait_front(cb_a, 1);
        const uint32_t hp = get_read_ptr(cb_a);
        noc_async_write(hp, o.get_noc_addr(ob), TB);
        noc_async_write(hp, o.get_noc_addr(ob + Kt), TB);
        if constexpr (mode != 1) {
            cb_wait_front(cb_lo, 1);
            noc_async_write(get_read_ptr(cb_lo), o.get_noc_addr(ob + 2 * Kt), TB);
        }
        noc_async_writes_flushed();
        cb_pop_front(cb_a, 1);
        if constexpr (mode != 1) {
            cb_pop_front(cb_lo, 1);
        }
        if (++j == Kt) {
            j = 0;
            ++r;
        }
    }
    noc_async_write_barrier();
}