File size: 3,597 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 67 68 69 70 71 72 73 74 75 76 | // SPDX-License-Identifier: Apache-2.0
// GELU(x) = x * Phi(x) as a branch-free polynomial (moge-2 opt kernel, used through ttnn.generic_op).
// t = clamp(x, -4, 4); u = t^2 / 16; Phi(t) = 0.5 + t * P(u) (P: degree 7, weighted minimax fit of the
// normal CDF on [0, 4] with Phi(4) = 1 pinned); GELU = x * Phi(t) (x > 4: x; x < -4: |y| < 1.3e-7 |x|)
// The four highest-order coefficients are fp16 (one SFPLOADI each, refit after quantization), the three
// lowest fp32 in the programmable constant registers. 27 SFPU instructions per row (v1: 32).
// Max |error| vs the exact (erf) GELU 2.1e-4 (fp32 evaluation, before the bf16 round-to-nearest).
// Reads TILE bf16 pages from c_0, writes c_2. Runtime arg 0: number of tiles for this core.
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/dataflow/dataflow_buffer.h"
#if defined(TRISC_MATH)
#include "sfpi.h"
#include "llk_math_eltwise_unary_sfpu_macros.h"
namespace ckernel::sfpu {
inline void moge_gelu_init() {
sfpi::vConstFloatPrgm0 = 0.3983510434627533f;
sfpi::vConstFloatPrgm1 = -1.0425554513931274f;
sfpi::vConstFloatPrgm2 = 2.3193752765655518f;
}
inline void moge_gelu_poly() {
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat x = sfpi::dst_reg[0];
sfpi::vFloat t = sfpi::clamp(x, -4.0f, 4.0f);
sfpi::vFloat u = t * t;
u = u * 0.0625f; // exact power-of-2 scaling
sfpi::vFloat p = u * sfpi::vFloat(sfpi::sFloat16a(0xafb6u)) + sfpi::vFloat(sfpi::sFloat16a(0x3ae0u)); // -0.1204833984375, 0.859375
p = p * u + sfpi::vFloat(sfpi::sFloat16a(0xc0e2u)); // -2.44140625
p = p * u + sfpi::vFloat(sfpi::sFloat16a(0x438fu)); // 3.779296875
p = p * u + sfpi::vFloat(sfpi::sFloat16a(0xc341u)); // -3.626953125
p = p * u + sfpi::vConstFloatPrgm2;
p = p * u + sfpi::vConstFloatPrgm1;
p = p * u + sfpi::vConstFloatPrgm0;
sfpi::vFloat phi = t * p + 0.5f;
sfpi::vFloat r = x * phi;
sfpi::dst_reg[0] = sfpi::convert<sfpi::vFloat16b>(r, sfpi::RoundMode::Nearest);
sfpi::dst_reg++;
}
}
} // namespace ckernel::sfpu
#endif
void kernel_main() {
uint32_t num_tiles = get_arg_val<uint32_t>(0);
constexpr uint32_t BLOCK = get_compile_time_arg_val(0); // tiles per DST acquire (<= 8, bf16 dest)
constexpr auto cb_input = tt::CBIndex::c_0;
constexpr auto cb_output = tt::CBIndex::c_2;
DataflowBuffer dfb_in(cb_input);
DataflowBuffer dfb_out(cb_output);
init_sfpu(cb_input, cb_output);
MATH((::ckernel::llk_math_eltwise_unary_sfpu_init<::SfpuType::gelu>(::ckernel::sfpu::moge_gelu_init)));
uint32_t done = 0;
while (done < num_tiles) {
const uint32_t n = (num_tiles - done) < BLOCK ? (num_tiles - done) : BLOCK;
dfb_in.wait_front(n);
dfb_out.reserve_back(n);
tile_regs_acquire();
for (uint32_t j = 0; j < n; ++j) {
copy_tile(cb_input, j, j);
MATH((_llk_math_eltwise_unary_sfpu_params_(::ckernel::sfpu::moge_gelu_poly, j, VectorMode::RC)));
}
tile_regs_commit();
tile_regs_wait();
for (uint32_t j = 0; j < n; ++j) {
pack_tile(j, cb_output);
}
tile_regs_release();
dfb_in.pop_front(n);
dfb_out.push_back(n);
done += n;
}
}
|