code/tt_moge/kernels/addln_writer.cpp

// SPDX-License-Identifier: Apache-2.0
// Fused residual add + LayerNorm (moge-2 opt kernel, round 7): BRISC. Generates the reduce scaler (c_2), reads this core's half row of y into
// c_1 (RB tiles per barrier), writes z = x + y (c_17; z may alias x: in place) and the even BLK-blocks of
// LN(z) (c_26; the reader writes the odd ones).
// Runtime args: y_addr, z_addr, out_addr, row, col0. Compile args: HALF, WT, BLK, TensorAccessorArgs y, z, out.
#include "api/dataflow/dataflow_api.h"
#include "api/dataflow/noc.h"
#include "api/dataflow/dataflow_buffer.h"
#include "api/tensor/noc_traits.h"
#include "ttnn/cpp/ttnn/kernel_lib/reduce_helpers_dataflow.hpp"
#ifndef RB
#define RB 2
#endif

void kernel_main() {
    const uint32_t y_addr = get_arg_val<uint32_t>(0);
    const uint32_t z_addr = get_arg_val<uint32_t>(1);
    const uint32_t o_addr = get_arg_val<uint32_t>(2);
    const uint32_t row = get_arg_val<uint32_t>(3);
    const uint32_t col0 = get_arg_val<uint32_t>(4);
    constexpr uint32_t HALF = get_compile_time_arg_val(0);
    constexpr uint32_t WT = get_compile_time_arg_val(1);
    constexpr uint32_t BLK = get_compile_time_arg_val(2);
    constexpr auto y_args = TensorAccessorArgs<3>();
    constexpr auto z_args = TensorAccessorArgs<y_args.next_compile_time_args_offset()>();
    constexpr auto o_args = TensorAccessorArgs<z_args.next_compile_time_args_offset()>();
    constexpr uint32_t PAGE = 2048;
    const auto sy = TensorAccessor(y_args, y_addr);
    const auto sz = TensorAccessor(z_args, z_addr);
    const auto so = TensorAccessor(o_args, o_addr);
    DataflowBuffer dy(1);
    DataflowBuffer dz(17);
    DataflowBuffer dout(26);
    const uint32_t base = row * WT + col0;
    dataflow_kernel_lib::calculate_and_prepare_reduce_scaler<2, ckernel::PoolType::SUM, ckernel::ReduceDim::REDUCE_ROW,
                                                             dataflow_kernel_lib::SUM_AND_MAX_REDUCE_FACTOR>();
    for (uint32_t b = 0; b < HALF; b += RB) {
        dy.reserve_back(RB);
        const uint32_t l = dy.get_write_ptr();
        for (uint32_t j = 0; j < RB; ++j) {
            noc_async_read(sy.get_noc_addr(base + b + j), l + j * PAGE, PAGE);
        }
        noc_async_read_barrier();
        dy.push_back(RB);
    }
    for (uint32_t b = 0; b < HALF; b += BLK) {
        dz.wait_front(BLK);
        const uint32_t l = dz.get_read_ptr();
        for (uint32_t j = 0; j < BLK; ++j) {
            noc_async_write(l + j * PAGE, sz.get_noc_addr(base + b + j), PAGE);
        }
        noc_async_writes_flushed();
        dz.pop_front(BLK);
    }
    for (uint32_t b = 0; b < HALF; b += 2 * BLK) {
        dout.wait_front(BLK);
        const uint32_t l = dout.get_read_ptr();
        for (uint32_t j = 0; j < BLK; ++j) {
            noc_async_write(l + j * PAGE, so.get_noc_addr(base + b + j), PAGE);
        }
        noc_async_writes_flushed();
        dout.pop_front(BLK);
    }
    noc_async_write_barrier();
}