code/tt_diffusion_planner/tt/kernels/ln32s_writer.cpp

// SPDX-License-Identifier: Apache-2.0
// Split-row fused fp32 LayerNorm (tt/ln_kernel.py, LN_SPLIT), writer (RISCV_1). Per row: (write_h) the new stream
// tile; sends this member's tile for the mean fold (c_11, from the compute) to page j of the root's gather CB c_10
// and raises the root's semaphore 0; on the root, broadcasts the packed mean (c_4, column 0) over the 32 columns and
// sends it to page 0 of every member's c_5 (semaphore 1 raised); the same for the squared tile and rstd (page 1);
// then writes the output tile.
// With kcat_out (KCAT_EMIT) out is the split operand [y_hi | y_hi | y_lo | 1 | 0..] of the next K-concatenated
// linear (row tiles Ktp, the kcat_writer layout): y_hi to (r, j) and (r, Wt + j), y_lo to (r, 2 Wt + j); the root
// also writes the ones tile (r, 3 Wt) and the zero pad tiles.
// CT args: [0] Wt, [1] per-core RT-arg count (P3), [2] write_h, [3] sfpu_bcast, [4] kcat_out, [5] Ktp, then the
//          TensorAccessorArgs of out and h.
// Common RT args: [out_addr, h_addr]. Per-core RT args: as the reader's.
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t Wt = get_compile_time_arg_val(0);
constexpr uint32_t write_h = get_compile_time_arg_val(2);
constexpr uint32_t sfpu_bcast = get_compile_time_arg_val(3);   // the root's compute packs a broadcast tile
constexpr uint32_t kcat_out = get_compile_time_arg_val(4);
constexpr uint32_t Ktp = get_compile_time_arg_val(5);
constexpr auto out_args = TensorAccessorArgs<6>();
constexpr auto h_args = TensorAccessorArgs<out_args.next_compile_time_args_offset()>();

constexpr uint32_t cb_stat = 4, cb_bc = 5, cb_gat = 10, cb_snd = 11, cb_scr = 12, cb_out = 16, cb_hout = 17,
                   cb_lo = 19, cb_one = 20, cb_zero = 21;
constexpr uint32_t TB = 4096;

FORCE_INLINE void bcast_col(uint32_t src, uint32_t dst) {
    auto* s = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(src);
    auto* d = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(dst);
    for (uint32_t fo = 0; fo < 1024; fo += 512) {
        for (uint32_t ro = 0; ro < 256; ro += 16) {
            const uint32_t v = s[fo + ro];
            volatile tt_l1_ptr uint32_t* l = d + fo + ro;
            volatile tt_l1_ptr uint32_t* r = l + 256;
            for (uint32_t c = 0; c < 16; ++c) {
                l[c] = v;
                r[c] = v;
            }
        }
    }
}

void kernel_main() {
    const uint32_t out_addr = get_common_arg_val<uint32_t>(0);
    const uint32_t h_addr = get_common_arg_val<uint32_t>(1);
    const uint32_t row0 = get_arg_val<uint32_t>(0);
    const uint32_t n_rows = get_arg_val<uint32_t>(1);
    const uint32_t stride = get_arg_val<uint32_t>(2);
    const uint32_t j = get_arg_val<uint32_t>(3);
    const uint32_t root_x = get_arg_val<uint32_t>(4);
    const uint32_t root_y = get_arg_val<uint32_t>(5);
    if (n_rows == 0) {
        return;
    }
    const bool root = j == 0;
    const auto out = TensorAccessor(out_args, out_addr, TB);
    const auto hacc = TensorAccessor(h_args, h_addr, TB);
    const uint32_t gat_base = get_write_ptr(cb_gat);   // the same L1 address on every core of the program
    const uint32_t bc_base = get_write_ptr(cb_bc);
    const uint64_t root_sem = get_noc_addr(root_x, root_y, get_semaphore(0));
    const uint32_t sem_b_addr = get_semaphore(1);
    cb_reserve_back(cb_scr, 1);
    const uint32_t scr = get_write_ptr(cb_scr);

    auto send_to_root = [&]() {
        cb_wait_front(cb_snd, 1);
        noc_async_write(get_read_ptr(cb_snd), get_noc_addr(root_x, root_y, gat_base + j * TB), TB);
        noc_async_write_barrier();
        noc_semaphore_inc(root_sem, 1);
        cb_pop_front(cb_snd, 1);
    };
    auto bcast_out = [&](uint32_t page) {
        cb_wait_front(cb_stat, 1);
        uint32_t src = get_read_ptr(cb_stat);
        if constexpr (!sfpu_bcast) {
            bcast_col(src, scr);
            src = scr;
        }
        for (uint32_t m = 0; m < Wt; ++m) {
            const uint32_t mx = get_arg_val<uint32_t>(6 + 2 * m), my = get_arg_val<uint32_t>(7 + 2 * m);
            noc_async_write(src, get_noc_addr(mx, my, bc_base + page * TB), TB);
        }
        noc_async_write_barrier();
        cb_pop_front(cb_stat, 1);
        for (uint32_t m = 0; m < Wt; ++m) {
            const uint32_t mx = get_arg_val<uint32_t>(6 + 2 * m), my = get_arg_val<uint32_t>(7 + 2 * m);
            noc_semaphore_inc(get_noc_addr(mx, my, sem_b_addr), 1);
        }
    };

    uint32_t one_ptr = 0, zero_ptr = 0;
    if constexpr (kcat_out) {
        if (root) {
            // ones tile: element (i, c) at face (i / 16) * 2 + (c / 16), offset (i % 16) * 16 + c % 16
            cb_reserve_back(cb_one, 1);
            one_ptr = get_write_ptr(cb_one);
            auto* p = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(one_ptr);
            for (uint32_t i = 0; i < 1024; ++i) {
                p[i] = 0;
            }
            for (uint32_t i = 0; i < 32; ++i) {
                const uint32_t base = (i / 16) * 512 + (i % 16) * 16;
                p[base] = 0x3F800000u;
                p[base + 1] = 0x3F800000u;
            }
            if constexpr (Ktp > 3 * Wt + 1) {
                cb_reserve_back(cb_zero, 1);
                zero_ptr = get_write_ptr(cb_zero);
                auto* z = reinterpret_cast<volatile tt_l1_ptr uint32_t*>(zero_ptr);
                for (uint32_t i = 0; i < 1024; ++i) {
                    z[i] = 0;
                }
            }
        }
    }

    for (uint32_t i = 0; i < n_rows; ++i) {
        const uint32_t r = row0 + i * stride;
        if constexpr (write_h) {
            cb_wait_front(cb_hout, 1);
            noc_async_write(get_read_ptr(cb_hout), hacc.get_noc_addr(r * Wt + j), TB);
            noc_async_writes_flushed();
            cb_pop_front(cb_hout, 1);
        }
        send_to_root();
        if (root) {
            bcast_out(0);
        }
        send_to_root();
        if (root) {
            bcast_out(1);
        }
        if constexpr (kcat_out) {
            const uint32_t ob = r * Ktp;
            cb_wait_front(cb_out, 1);
            const uint32_t hp = get_read_ptr(cb_out);
            noc_async_write(hp, out.get_noc_addr(ob + j), TB);
            noc_async_write(hp, out.get_noc_addr(ob + Wt + j), TB);
            cb_wait_front(cb_lo, 1);
            noc_async_write(get_read_ptr(cb_lo), out.get_noc_addr(ob + 2 * Wt + j), TB);
            if (root) {
                noc_async_write(one_ptr, out.get_noc_addr(ob + 3 * Wt), TB);
                for (uint32_t z = 3 * Wt + 1; z < Ktp; ++z) {
                    noc_async_write(zero_ptr, out.get_noc_addr(ob + z), TB);
                }
            }
            noc_async_writes_flushed();
            cb_pop_front(cb_out, 1);
            cb_pop_front(cb_lo, 1);
        } else {
            cb_wait_front(cb_out, 1);
            noc_async_write(get_read_ptr(cb_out), out.get_noc_addr(r * Wt + j), TB);
            noc_async_writes_flushed();
            cb_pop_front(cb_out, 1);
        }
    }
    noc_async_write_barrier();
    noc_async_atomic_barrier();
}