File size: 5,696 Bytes
b9db84c | 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 | // SPDX-License-Identifier: Apache-2.0
// sam3 backbone fc1 GELU (erf form) as a custom SFPU kernel, BLK tiles per DST acquire.
// GELU_VARIANT 0: ttnn's own accurate bf16 path (calculate_gelu<false, false>) -- bit-identical to ttnn.gelu
// GELU_VARIANT 1: ttnn's core polynomial and identity region verbatim (same coefficients / Horner order), negative
// tail (-5.5426, -3.125) by ttnn's exp code (GELU_TAIL 0, bit-identical) or a degree-9 fit (GELU_TAIL 1)
// GELU_VARIANT 3: branch-free clamped odd polynomial (fast, NOT bit-identical; opt-in SAM3_BB_GELU=poly)
// Measured on the real fc1 outputs [5184, 4736] (perf/dev/gelu_kernel_bench.py, 12x10, DRAM in / out, BLK 8):
// ttnn.gelu exact 0.433 ms, variant 0 0.437 ms (SFPU-bound), variant 1/tail 1 0.408 ms, variant 3 0.238 ms
// (DRAM-bound: a plain copy kernel takes 0.229 ms).
#include <cstdint>
#include "api/compute/common.h"
#include "api/compute/tile_move_copy.h"
#include "api/compute/eltwise_unary/eltwise_unary.h"
#include "api/compute/eltwise_unary/gelu.h"
#include "api/compute/compute_kernel_api.h"
#ifndef GELU_TAIL
#define GELU_TAIL 0
#endif
#ifdef TRISC_MATH
#include "sfpi.h"
namespace ckernel::sfpu {
#if GELU_VARIANT == 1
sfpi_inline sfpi::vFloat sam3_gelu_value(sfpi::vFloat x) {
sfpi::vFloat x2 = x * x;
sfpi::vFloat result = 0.0f;
v_if(x > -5.54259443f) {
#if GELU_TAIL == 0
{ // ttnn's exp-region code verbatim (Moroz exp_21f * corr_H, rounded to the 2^-25 grid)
constexpr float NEG_HALF_ONE_LN2 = -0.72134752044f;
sfpi::vFloat xlog2 = x2 * NEG_HALF_ONE_LN2 + 127.0f;
sfpi::vInt z = _float_to_int32_for_exp_21f_(xlog2);
sfpi::vInt exponential_part = sfpi::exexp(sfpi::as<sfpi::vFloat>(z), sfpi::ExponentMode::Biased);
sfpi::vMag fractional_part = sfpi::exman(sfpi::as<sfpi::vFloat>(z));
sfpi::vFloat frac = sfpi::convert<sfpi::vFloat>(fractional_part, sfpi::RoundMode::Nearest);
frac = PolynomialEvaluator::eval(frac, 1.0017248f, 7.839635491371155e-08f, 4.791750143340323e-15f);
sfpi::vFloat exp_val = sfpi::setexp(frac, exponential_part);
sfpi::vFloat H = exp_val * PolynomialEvaluator::eval(x, GELU_HCORR_C0, GELU_HCORR_C1, GELU_HCORR_C2, GELU_HCORR_C3);
constexpr float ROUND_TO_GRID = 0.375f;
sfpi::vFloat Hs = (H + ROUND_TO_GRID) - ROUND_TO_GRID;
result = x * Hs;
}
#elif GELU_TAIL == 1
{ // GELU(x) directly as a degree-9 polynomial in t = (x + c) / h over the tail (rel err ~2e-4)
sfpi::vFloat t = x * GELU_T_SCALE + GELU_T_OFF;
result = PolynomialEvaluator::eval(t, GELU_TP0, GELU_TP1, GELU_TP2, GELU_TP3, GELU_TP4, GELU_TP5, GELU_TP6, GELU_TP7, GELU_TP8, GELU_TP9);
}
#endif
v_and(x >= -3.125f);
sfpi::vFloat odd_poly = PolynomialEvaluator::eval(
x2, GELU_CDF_CORE_C1, GELU_CDF_CORE_C3, GELU_CDF_CORE_C5, GELU_CDF_CORE_C7, GELU_CDF_CORE_C9,
GELU_CDF_CORE_C11, GELU_CDF_CORE_C13);
sfpi::vFloat phi = GELU_CDF_CORE_C0 + x * odd_poly;
result = x * phi;
v_and(x >= 2.78125f);
result = x;
}
v_endif;
return result;
}
inline void sam3_gelu_face() {
#pragma GCC unroll 8
for (int d = 0; d < 8; d++) {
sfpi::vFloat in = sfpi::dst_reg[0];
sfpi::vFloat r = sam3_gelu_value(in);
sfpi::dst_reg[0] = sfpi::convert<sfpi::vFloat16b>(r, sfpi::RoundMode::Nearest);
sfpi::dst_reg++;
}
}
#endif
#if GELU_VARIANT == 3
// cheap branch-free form (moge-2 / rf-detr style): t = clamp(x, -4, 4), Phi = 0.5 + t P(t^2) (degree 7 in t^2,
// Phi(4) = 1 pinned), GELU = t Phi(t) + max(x - 4, 0); max |err| vs erf-GELU 1.35e-4 before the bf16 rounding
inline void sam3_gelu_face() {
#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;
sfpi::vFloat p = PolynomialEvaluator::eval(u, 3.9880555868e-01f, -6.6064164042e-02f, 9.5814364031e-03f, -1.0210596956e-03f,
7.6192176493e-05f, -3.7097588574e-06f, 1.0444770737e-07f, -1.2780509984e-09f);
sfpi::vFloat phi = t * p + 0.5f;
sfpi::vFloat r = t * phi + sfpi::max(x - 4.0f, 0.0f);
sfpi::dst_reg[0] = sfpi::convert<sfpi::vFloat16b>(r, sfpi::RoundMode::Nearest);
sfpi::dst_reg++;
}
}
#endif
} // namespace ckernel::sfpu
#endif
constexpr uint32_t cb_in = get_compile_time_arg_val(0);
constexpr uint32_t cb_out = get_compile_time_arg_val(1);
constexpr uint32_t BLK = get_compile_time_arg_val(2);
void kernel_main() {
const uint32_t n_tiles = get_arg_val<uint32_t>(0);
init_sfpu(cb_in, cb_out);
#if GELU_VARIANT == 0
gelu_tile_init<false>();
#else
MATH((::ckernel::llk_math_eltwise_unary_sfpu_init<::SfpuType::gelu>()));
#endif
for (uint32_t t = 0; t < n_tiles; t += BLK) {
const uint32_t b = (n_tiles - t) < BLK ? (n_tiles - t) : BLK;
cb_wait_front(cb_in, b);
tile_regs_acquire();
for (uint32_t j = 0; j < b; ++j) {
copy_tile(cb_in, j, j);
#if GELU_VARIANT == 0
gelu_tile<false>(j);
#else // variants 1 / 3
MATH(SFPU_UNARY_CALL_NO_TEMPLATE_ARGS(DST_SYNC_MODE, DST_ACCUM_MODE, sam3_gelu_face, j, VectorMode::RC));
#endif
}
tile_regs_commit();
tile_regs_wait();
cb_reserve_back(cb_out, b);
for (uint32_t j = 0; j < b; ++j) {
pack_tile(j, cb_out);
}
cb_push_back(cb_out, b);
cb_pop_front(cb_in, b);
tile_regs_release();
}
}
|