File size: 2,949 Bytes
13b4736 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 | // 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();
}
|