@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
hpp 128 lines 3.98 KB
Raw
1 // SPDX-License-Identifier: Apache-2.0
2 // SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 #ifndef CPU_TANHF_NEON_HPP
5 #define CPU_TANHF_NEON_HPP
6
7 #include <cstdint>
8 #include <arm_neon.h>
9
10 namespace vec_op {
11
12 namespace {
13
14 struct TanhfConstants {
15 float32x4_t special_bound;
16 float32x4_t two;
17 float32x4_t c0;
18 float32x4_t c2;
19 int32x4_t exponent_bias;
20 float c1;
21 float c3;
22 float two_over_ln2;
23 float c4;
24 float ln2_hi;
25 float ln2_lo;
26 };
27
28 const TanhfConstants kTanhfConstants = {
29 // 9.01, above which tanhf rounds to 1 (or -1 for negative).
30 .special_bound = vdupq_n_f32(0x1.205966p+3f),
31 .two = vdupq_n_f32(0x1.0p+1f),
32 .c0 = vdupq_n_f32(0x1.fffffep-2f),
33 .c2 = vdupq_n_f32(0x1.555736p-5f),
34 .exponent_bias = vdupq_n_s32(0x3f800000),
35 .c1 = 0x1.5554aep-3f,
36 .c3 = 0x1.12287cp-7f,
37 .two_over_ln2 = 0x1.715476p+1f,
38 .c4 = 0x1.6b55a2p-10f,
39 .ln2_hi = 0x1.62e4p-1f,
40 .ln2_lo = 0x1.7f7d1cp-20f,
41 };
42
43 // Return the ptr but hide it's value from the compiler so accesses
44 // through it can't be optimised based on contents.
45 template <typename T>
46 inline const T* ptr_barrier(const T* ptr) {
47 const T* opaque_ptr = ptr;
48 __asm__("" : "+r"(opaque_ptr));
49 return opaque_ptr;
50 }
51
52 // Check whether any lanes in the mask are set
53 inline bool any_u32(uint32x4_t x) { return vmaxvq_u32(x) != 0; }
54
55 // e^2x - 1 inline helper
56 inline float32x4_t e2xm1f_inline(float32x4_t x, const TanhfConstants* d) {
57 float32x2_t ln2 = vld1_f32(&d->ln2_hi);
58 float32x4_t lane_consts = vld1q_f32(&d->c1);
59
60 // Reduce argument: f in [-ln2/2, ln2/2], i is exact.
61 float32x4_t j = vrndaq_f32(vmulq_laneq_f32(x, lane_consts, 2));
62 int32x4_t i = vcvtq_s32_f32(j);
63 float32x4_t f = vaddq_f32(x, x);
64 f = vfmsq_lane_f32(f, j, ln2, 0);
65 f = vfmsq_lane_f32(f, j, ln2, 1);
66
67 // Approximate expm1(f) with polynomial P, expm1(f) ~= f + f^2 * P(f)
68 float32x4_t f2 = vmulq_f32(f, f);
69 float32x4_t f4 = vmulq_f32(f2, f2);
70 float32x4_t p01 = vfmaq_laneq_f32(d->c0, f, lane_consts, 0);
71 float32x4_t p23 = vfmaq_laneq_f32(d->c2, f, lane_consts, 1);
72 float32x4_t poly = vfmaq_f32(p01, f2, p23);
73 poly = vfmaq_laneq_f32(poly, f4, lane_consts, 3);
74 poly = vfmaq_f32(f, f2, poly);
75
76 // scale = 2^i
77 int32x4_t u = vaddq_s32(vshlq_n_s32(i, 23), d->exponent_bias);
78 float32x4_t scale = vreinterpretq_f32_s32(u);
79 return vfmaq_f32(vsubq_f32(scale, vdupq_n_f32(1.0f)), poly, scale);
80 }
81
82 // Calculate the result tanh(x) = q / (q+2) and set special lanes to ±1
83 inline float32x4_t special_case(float32x4_t x, float32x4_t q,
84 uint32x4_t special) {
85 const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
86
87 float32x4_t y = vdivq_f32(q, vaddq_f32(q, d->two));
88 uint32x4_t ix = vreinterpretq_u32_f32(x);
89 uint32x4_t one_bits = vreinterpretq_u32_s32(d->exponent_bias);
90 uint32x4_t sign_mask = vdupq_n_u32(0x80000000u);
91 uint32x4_t special_bits = vbslq_u32(sign_mask, ix, one_bits);
92 float32x4_t special_y = vreinterpretq_f32_u32(special_bits);
93 return vbslq_f32(special, special_y, y);
94 }
95
96 } // namespace
97
98 // Implementation of tanhf adapted from Arm Optimized Routines (tanhf
99 // AdvSIMD)
100 // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/tanhf.c
101 //
102 // Approximation for single-precision vector tanh(x), using a simplified
103 // version of expm1f. The maximum error is 2.08 + 0.5 ULP:
104 // _ZGVnN4v_tanhf (0x1.fa5eep-5) got 0x1.f9ba02p-5 want 0x1.f9ba08p-5.
105 inline float32x4_t fast_tanhf_f32x4(float32x4_t x) {
106 const TanhfConstants* d = ptr_barrier(&kTanhfConstants);
107
108 // tanh(x) = (e^2x - 1) / (e^2x + 1)
109 // q = e^2x -1
110 float32x4_t q = e2xm1f_inline(x, d);
111
112 // Check for special cases
113 uint32x4_t special = vcagtq_f32(x, d->special_bound);
114
115 // Fall back to vectorised special case for any lanes which would cause
116 // expm1 to overflow
117 if (any_u32(special)) {
118 return special_case(x, q, special);
119 }
120
121 // Complete fast path if no special lanes
122 // tanh(x) = q / (q+2)
123 return vdivq_f32(q, vaddq_f32(q, d->two));
124 }
125
126 } // namespace vec_op
127
128 #endif // CPU_TANHF_NEON_HPP