File size: 3,220 Bytes
7a24154 | 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 67 68 69 70 71 72 | // SPDX-License-Identifier: Apache-2.0
// KCAT operand build (tt/kcat.py), compute: per fp32 x tile, x_hi = bf16(x) (typecast_tile Float32 -> Float16_b,
// the LLK of the stock ttnn.typecast), x_lo = bf16(x - x_hi) (sub_binary_tile in fp32: exact; then the same
// typecast), both packed as bf16 (exact: the DEST values are bf16-representable). The same values as the stock
// typecast -> typecast -> subtract -> typecast chain of PointLinear._split (bit-identity checked on the device).
// Mode 2 (OPT round 3, HEAD-SKIP): x = x + skip first (add_binary_tile<NearestEven> on the fp32 tiles of CBs 0 and
// 1 in DEST: binary_ng's fp32 SFPU add, the LLK of the stock ttnn.add(y, skip) it replaces; as tt/epi.py), then the
// same split. CT args: [0] per-core RT-arg count (P3), [1] mode (0 = fp32 split; 1 = bf16 input: nothing to do here,
// the writer reads CB 0 itself; 2 = fp32 x + fp32 skip, then the split). Per-core RT args: [n, 0].
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/compute_kernel_api.h"
#include "api/compute/copy_dest_values.h"
#include "api/compute/eltwise_binary_sfpu.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/tile_move_copy.h"
using namespace ckernel;
constexpr uint32_t mode = get_compile_time_arg_val(1);
constexpr uint32_t cb_x = 0, cb_s = 1, cb_hi = 16, cb_lo = 17;
void kernel_main() {
if constexpr (mode == 1) {
return;
} else {
const uint32_t n = get_arg_val<uint32_t>(0);
if (n == 0) {
return;
}
compute_kernel_hw_startup(cb_x, cb_hi);
for (uint32_t t = 0; t < n; ++t) {
cb_wait_front(cb_x, 1);
tile_regs_acquire();
copy_init(cb_x);
copy_tile(cb_x, 0, 0);
if constexpr (mode == 2) {
cb_wait_front(cb_s, 1);
copy_init(cb_s);
copy_tile(cb_s, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
copy_dest_values_init();
copy_dest_values<DataFormat::Float32>(0, 1);
} else {
copy_tile(cb_x, 0, 1);
}
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(0);
sub_binary_tile_init();
sub_binary_tile<ckernel::DstRoundingMode::NearestEven>(1, 0, 1);
typecast_tile_init<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>();
typecast_tile<(uint32_t)DataFormat::Float32, (uint32_t)DataFormat::Float16_b>(1);
cb_reserve_back(cb_hi, 1);
cb_reserve_back(cb_lo, 1);
tile_regs_commit();
tile_regs_wait();
pack_tile(0, cb_hi);
pack_tile(1, cb_lo);
tile_regs_release();
cb_push_back(cb_hi, 1);
cb_push_back(cb_lo, 1);
cb_pop_front(cb_x, 1);
if constexpr (mode == 2) {
cb_pop_front(cb_s, 1);
}
}
}
}
|