code/tt_frnet/tt/kernels/p2p_writer.cpp

// SPDX-License-Identifier: Apache-2.0
// P2P (OPT round 4, FRNET_P2P): pixel2point gather straight into the KCAT operand, writer (RISCV_1): the Kt tiles of
// block (tile row) r go to X' [.., N, 32 OKt] at tiles (r, off + j) and (r, off + Kt + j): the [x | x] columns of a
// bf16 part (as kcat_writer.cpp mode 1).
// CT args: [0] per-core RT-arg count (P3), [1] Kt, [2] OKt, [3] off, [4] DIAG (0; 2 = timing diagnostic without writes),
// [5] BCOL (tile column of the bias block), [6] BT (bias block tiles; 0 = none; kcat_writer.cpp), then the
// TensorAccessorArgs of X'.
// Common RT args: [out_addr]. Per-core RT args: [b0, nb].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t Kt = get_compile_time_arg_val(1);
constexpr uint32_t OKt = get_compile_time_arg_val(2);
constexpr uint32_t off = get_compile_time_arg_val(3);
constexpr uint32_t DIAG = get_compile_time_arg_val(4);   // timing diagnostic: 2 = no DRAM writes
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_t = 16, 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 b0 = get_arg_val<uint32_t>(0);
    const uint32_t nb = get_arg_val<uint32_t>(1);
    const auto o = TensorAccessor(o_args, get_common_arg_val<uint32_t>(0), TB);
    uint32_t ones = 0;
    if constexpr (BT > 0) {
        ones = fill_bias_tiles(cb_c);
    }
    for (uint32_t r = b0; r < b0 + nb; ++r) {
        if constexpr (BT > 0 && DIAG != 2) {
            write_bias_block(ones, r * OKt + BCOL, BT, o);
        }
        cb_wait_front(cb_t, Kt);
        uint32_t rp = get_read_ptr(cb_t);
        const uint32_t ob = r * OKt + off;
        if constexpr (DIAG != 2) {
            for (uint32_t j = 0; j < Kt; ++j) {
                noc_async_write(rp, o.get_noc_addr(ob + j), TB);
                noc_async_write(rp, o.get_noc_addr(ob + Kt + j), TB);
                rp += TB;
            }
        }
        noc_async_writes_flushed();
        cb_pop_front(cb_t, Kt);
    }
    noc_async_write_barrier();
}