File size: 4,791 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 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 | // SPDX-License-Identifier: Apache-2.0
// EPI compute (frnet-p150 OPT round 2, tt/epi.py): one image-branch epilogue per tile with the LLKs and the order of
// the stock programs it replaces, so both outputs are bit-identical to them:
// y = ((p0 + p1) + p2) + ... add_binary_tile<NearestEven>: binary_ng's fp32 SFPU add (KSplitConv's partial sums
// in order; the BasicBlock's ttnn.add(y, identity)), operands unpacked to DEST as fp32
// y = hardswish(y) (ACT 1) ttnn.hardswish's fp32 chain (hardswish_kernel.cpp, kIsFloat32):
// D1 = y; hardsigmoid_tile(D0); mul_binary_tile(D0, D1, D0)
// out32 = y fp32 (OUT32)
// out16 = bf16(y) typecast_tile<Float32, Float16_b> (ttnn.typecast; OUT16): the next conv's operand
// DEST tiles: 0 = y, 1 = operand / copy. CT args: [0] T, [1] NC, [2] GX, [3] NP, [4] ACT, [5] OUT32, [6] OUT16.
#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/activations.h"
#include "api/compute/eltwise_unary/typecast.h"
#include "api/compute/reconfig_data_format.h"
#include "api/compute/tile_move_copy.h"
using namespace ckernel;
constexpr uint32_t T = 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 NP = get_compile_time_arg_val(3);
constexpr uint32_t ACT = get_compile_time_arg_val(4);
constexpr uint32_t OUT32 = get_compile_time_arg_val(5);
constexpr uint32_t OUT16 = get_compile_time_arg_val(6);
constexpr uint32_t cb32 = 16, cb16 = 17;
FORCE_INLINE void add_part(uint32_t cb) {
copy_init(cb);
copy_tile(cb, 0, 1);
add_binary_tile_init();
add_binary_tile<ckernel::DstRoundingMode::NearestEven>(0, 1, 0);
}
void kernel_main() {
const uint32_t core = get_absolute_logical_y() * GX + get_absolute_logical_x();
if (core >= NC) {
return;
}
constexpr uint32_t B = T / NC, E = T % NC;
const uint32_t t0 = core * B + (core < E ? core : E);
const uint32_t n = B + (core < E ? 1 : 0);
if (n == 0) {
return;
}
compute_kernel_hw_startup(0, OUT32 ? cb32 : cb16);
for (uint32_t t = t0; t < t0 + n; ++t) {
cb_wait_front(0, 1);
if constexpr (NP > 1) {
cb_wait_front(1, 1);
}
if constexpr (NP > 2) {
cb_wait_front(2, 1);
}
if constexpr (NP > 3) {
cb_wait_front(3, 1);
}
if constexpr (NP > 4) {
cb_wait_front(4, 1);
}
tile_regs_acquire();
copy_init(0);
copy_tile(0, 0, 0);
if constexpr (NP > 1) {
add_part(1);
}
if constexpr (NP > 2) {
add_part(2);
}
if constexpr (NP > 3) {
add_part(3);
}
if constexpr (NP > 4) {
add_part(4);
}
if constexpr (ACT == 1) {
copy_dest_values_init();
copy_dest_values<DataFormat::Float32>(0, 1);
hardsigmoid_tile_init();
hardsigmoid_tile(0);
mul_binary_tile_init();
mul_binary_tile(0, 1, 0);
}
if constexpr (OUT16) {
copy_dest_values_init();
copy_dest_values<DataFormat::Float32>(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);
}
tile_regs_commit();
if constexpr (OUT32) {
cb_reserve_back(cb32, 1);
}
if constexpr (OUT16) {
cb_reserve_back(cb16, 1);
}
tile_regs_wait();
if constexpr (OUT32) {
if constexpr (OUT16) {
pack_reconfig_data_format(cb32);
}
pack_tile(0, cb32);
}
if constexpr (OUT16) {
if constexpr (OUT32) {
pack_reconfig_data_format(cb16);
}
pack_tile(1, cb16);
}
tile_regs_release();
if constexpr (OUT32) {
cb_push_back(cb32, 1);
}
if constexpr (OUT16) {
cb_push_back(cb16, 1);
}
cb_pop_front(0, 1);
if constexpr (NP > 1) {
cb_pop_front(1, 1);
}
if constexpr (NP > 2) {
cb_pop_front(2, 1);
}
if constexpr (NP > 3) {
cb_pop_front(3, 1);
}
if constexpr (NP > 4) {
cb_pop_front(4, 1);
}
}
}
|