File size: 1,731 Bytes
e91e989 | 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 | // SPDX-License-Identifier: Apache-2.0
// FOUT compute (meteor-p150 OPT2 MK-FO, tt/fout.py): per unit fused_k = other_k + t_k for the 3 tile columns, one FPU
// eltwise add per tile (SrcA = the other operand (raw / r4), SrcB = the tilized t), bf16 in / out, as ttnn.add's
// binary_ng FPU kernel (no fp32 DEST accumulation); work split as the reader.
// CT args: [0] NU, [1] NC, [2] GX.
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/compute_kernel_hw_startup.h"
#include "api/compute/eltwise_binary.h"
using namespace ckernel;
constexpr uint32_t NU = get_compile_time_arg_val(0);
constexpr uint32_t NC = get_compile_time_arg_val(1);
constexpr uint32_t GX = get_compile_time_arg_val(2);
constexpr uint32_t cb_a = 0, cb_b = 2, cb_out = 16;
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = NU / NC, E = NU % NC;
const uint32_t n_units = B + (core < E ? 1 : 0);
if (n_units == 0) {
return;
}
compute_kernel_hw_startup(cb_b, cb_a, cb_out);
add_init(cb_b, cb_a);
for (uint32_t u = 0; u < n_units; ++u) {
cb_wait_front(cb_b, 3);
cb_wait_front(cb_a, 3);
cb_reserve_back(cb_out, 3);
tile_regs_acquire();
for (uint32_t k = 0; k < 3; ++k) {
add_tiles(cb_b, cb_a, k, k, k);
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t k = 0; k < 3; ++k) {
pack_tile(k, cb_out);
}
tile_regs_release();
cb_push_back(cb_out, 3);
cb_pop_front(cb_a, 3);
cb_pop_front(cb_b, 3);
}
}
|