code/tt_diffusion_planner/tt/kernels/fattn_writer.cpp

// SPDX-License-Identifier: Apache-2.0
// Fused fp32 matmul attention (tt/fattn_kernel.py), writer (RISCV_1): per unit u (head h = u / Mt, query tile row
// r = u % Mt) the output tile as tile (r, h) of the merged-heads [1, 1, Sq, H * 32] tensor (what nlp_concat_heads
// makes of the [1, H, Sq, 32] P V output).
// With kcat_out (KCAT_EMIT) the output is instead the split operand [o_hi | o_hi | o_lo | 1 | 0..] of the next
// K-concatenated linear (row tiles Ktp, kernels/kcat_writer.cpp's layout): o_hi to (r, h) and (r, H + h), o_lo to
// (r, 2 H + h); the unit of head 0 also writes the ones tile (r, 3 H) and the zero pad tiles.
// CT args: [0] Mt, [1] H, [2] per-core RT-arg count (P3), [3] kcat_out, [4] Ktp, then the TensorAccessorArgs of out.
// Common RT args: [out_addr]. Per-core RT args: [u0, n].
#include <cstdint>

#include "api/dataflow/dataflow_api.h"

constexpr uint32_t Mt = get_compile_time_arg_val(0);
constexpr uint32_t H = get_compile_time_arg_val(1);
constexpr uint32_t kcat_out = get_compile_time_arg_val(3);
constexpr uint32_t Ktp = get_compile_time_arg_val(4);
constexpr auto o_args = TensorAccessorArgs<5>();
constexpr uint32_t cb_o = 16, cb_lo = 17, cb_one = 18, cb_zero = 19;
constexpr uint32_t TB = 4096;

void kernel_main() {
    const uint32_t o_addr = get_common_arg_val<uint32_t>(0);
    const uint32_t u0 = 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);
    if constexpr (!kcat_out) {
        for (uint32_t u = u0; u < u0 + n; ++u) {
            const uint32_t h = u / Mt;
            const uint32_t r = u % Mt;
            cb_wait_front(cb_o, 1);
            noc_async_write(get_read_ptr(cb_o), o.get_noc_addr(r * H + h), TB);
            noc_async_writes_flushed();
            cb_pop_front(cb_o, 1);
        }
    } else {
        // ones tile: element (i, c) at face (i / 16) * 2 + (c / 16), offset (i % 16) * 16 + c % 16
        uint32_t one_ptr = 0, zero_ptr = 0;
        if (u0 < Mt) {
            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 * H + 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 u = u0; u < u0 + n; ++u) {
            const uint32_t h = u / Mt;
            const uint32_t r = u % Mt;
            const uint32_t ob = r * Ktp;
            cb_wait_front(cb_o, 1);
            const uint32_t hp = get_read_ptr(cb_o);
            noc_async_write(hp, o.get_noc_addr(ob + h), TB);
            noc_async_write(hp, o.get_noc_addr(ob + H + h), TB);
            cb_wait_front(cb_lo, 1);
            noc_async_write(get_read_ptr(cb_lo), o.get_noc_addr(ob + 2 * H + h), TB);
            if (h == 0) {
                noc_async_write(one_ptr, o.get_noc_addr(ob + 3 * H), TB);
                for (uint32_t z = 3 * H + 1; z < Ktp; ++z) {
                    noc_async_write(zero_ptr, o.get_noc_addr(ob + z), TB);
                }
            }
            noc_async_writes_flushed();
            cb_pop_front(cb_o, 1);
            cb_pop_front(cb_lo, 1);
        }
    }
    noc_async_write_barrier();
}