@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
h 197 lines 12.3 KB
Raw
1 #ifndef CPU_ARCH_MACROS_H
2 #define CPU_ARCH_MACROS_H
3
4 // x86_64
5 #ifdef __x86_64__
6 #define FAST_SPINNING _mm_pause();
7
8 #ifdef __AVX512F__
9 #define DEFINE_FAST_EXP \
10 const __m512 vec_factorial_1 = _mm512_set1_ps(0.999999701f); \
11 const __m512 vec_factorial_2 = _mm512_set1_ps(0.499991506f); \
12 const __m512 vec_factorial_3 = _mm512_set1_ps(0.166676521f); \
13 const __m512 vec_factorial_4 = _mm512_set1_ps(0.0418978221f); \
14 const __m512 vec_factorial_5 = _mm512_set1_ps(0.00828929059f); \
15 const __m512 vec_exp_log2ef = \
16 _mm512_castsi512_ps(_mm512_set1_epi32(0x3fb8aa3b)); \
17 const __m512 vec_half = _mm512_set1_ps(0.5f); \
18 const __m512 vec_one = _mm512_set1_ps(1.f); \
19 const __m512 vec_zero = _mm512_set1_ps(0.f); \
20 const __m512 vec_two = _mm512_set1_ps(2.f); \
21 const __m512 vec_ln2f = \
22 _mm512_castsi512_ps(_mm512_set1_epi32(0x3f317218)); \
23 const __m512 vec_ln_flt_min = \
24 _mm512_castsi512_ps(_mm512_set1_epi32(0xc2aeac50)); \
25 const __m512 vec_ln_flt_max = \
26 _mm512_castsi512_ps(_mm512_set1_epi32(0x42b17218)); \
27 const __m512i vec_127 = _mm512_set1_epi32(0x0000007f); \
28 const int n_mantissa_bits = 23; \
29 auto fast_exp = [&](const vec_op::FP32Vec16& vec) __attribute__(( \
30 always_inline)) { \
31 __m512 values = vec.reg; \
32 auto less_ln_flt_min_mask = \
33 _mm512_cmp_ps_mask(values, vec_ln_flt_min, 1 /*_CMP_LT_OS*/); \
34 auto vec_src = _mm512_min_ps(values, vec_ln_flt_max); \
35 vec_src = _mm512_max_ps(vec_src, vec_ln_flt_min); \
36 auto vec_fx = _mm512_fmadd_ps(vec_src, vec_exp_log2ef, vec_half); \
37 auto vec_fx_i = _mm512_cvt_roundps_epi32( \
38 vec_fx, _MM_FROUND_TO_NEG_INF | _MM_FROUND_NO_EXC); \
39 vec_fx = _mm512_cvtepi32_ps(vec_fx_i); \
40 auto vec_exp_poly = _mm512_fnmadd_ps(vec_fx, vec_ln2f, vec_src); \
41 auto vec_res = \
42 _mm512_fmadd_ps(vec_exp_poly, vec_factorial_5, vec_factorial_4); \
43 vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_3); \
44 vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_2); \
45 vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_factorial_1); \
46 vec_res = _mm512_fmadd_ps(vec_exp_poly, vec_res, vec_one); \
47 auto vec_exp_number = _mm512_sub_ps(vec_fx, vec_one); \
48 auto vec_exp_number_i = _mm512_cvtps_epi32(vec_exp_number); \
49 auto vec_two_pow_n_i = _mm512_add_epi32(vec_exp_number_i, vec_127); \
50 vec_two_pow_n_i = _mm512_slli_epi32(vec_two_pow_n_i, n_mantissa_bits); \
51 auto vec_two_pow_n = _mm512_castsi512_ps(vec_two_pow_n_i); \
52 vec_two_pow_n = _mm512_mask_blend_ps(less_ln_flt_min_mask, \
53 vec_two_pow_n, vec_zero); \
54 vec_res = _mm512_mul_ps(vec_res, vec_two_pow_n); \
55 vec_res = _mm512_mul_ps(vec_res, vec_two); \
56 vec_op::FP32Vec16 res(vec_res); \
57 return res; \
58 };
59 #endif
60
61 #endif
62
63 #ifdef __aarch64__
64 // Implementation of neon_expf copied from Arm Optimized Routines (expf
65 // AdvSIMD)
66 // https://github.com/ARM-software/optimized-routines/blob/master/math/aarch64/advsimd/expf.c
67 //
68 // Additional fast exponential intended for cases where outputs will be
69 // downcasted to FP16 / BF16 (e.g. attention softmax). Accurate within 1 ULP
70 // for FP16 Accurate within 1 ULP for BF16 for inputs in [-87.683, 88.376] &
71 // clamps inputs outside this range to 0 / inf. Implementation is similar to
72 // exp_u20, but:
73 // - uses a third degree polynomial approximation for exp(r) instead of a
74 // fifth degree one, with coefficients re-tuned.
75 // - does not split natural log (ln) into high / low parts
76 // - clamps exp(x) to 0 for x < -87.683113f and inf for x > 88.3762589f
77 // exp(x) = 2^n (exp(r))
78 // r = x - n*ln2, with n = round(x/ln2)
79 // exp(r) ~ poly(r) = 1 + r + r^2 * (c3 + c2 * r)
80 // n = round(x / ln2), r = x - n*ln2
81 #include <limits>
82 #define DEFINE_FAST_EXP \
83 const float32x4_t inv_ln2 = vdupq_n_f32(0x1.715476p+0f); \
84 const float ln2_hi = 0x1.62e4p-1f; \
85 const float ln2_lo = 0x1.7f7d1cp-20f; \
86 const float c0 = 0x1.0e4020p-7f; \
87 const float c2 = 0x1.555e66p-3f; \
88 const float32x4_t ln2_c02 = {ln2_hi, ln2_lo, c0, c2}; \
89 const uint32x4_t exponent_bias = vdupq_n_u32(0x3f800000); \
90 const float32x4_t c1 = vdupq_n_f32(0x1.573e2ep-5f); \
91 const float32x4_t c3 = vdupq_n_f32(0x1.fffdb6p-2f); \
92 const float32x4_t c4 = vdupq_n_f32(0x1.ffffecp-1f); \
93 const float32x4_t pos_special_bound = vdupq_n_f32(0x1.5d5e2ap+6f); \
94 const float32x4_t neg_special_bound = vnegq_f32(pos_special_bound); \
95 const float32x4_t inf = \
96 vdupq_n_f32(std::numeric_limits<float>::infinity()); \
97 const float32x4_t zero = vdupq_n_f32(0.0f); \
98 auto neon_expf = [&](float32x4_t values) __attribute__((always_inline)) { \
99 float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
100 float32x4_t r = vfmsq_laneq_f32(values, n, ln2_c02, 0); \
101 r = vfmsq_laneq_f32(r, n, ln2_c02, 1); \
102 uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
103 float32x4_t scale = vreinterpretq_f32_u32(vaddq_u32(e, exponent_bias)); \
104 float32x4_t r2 = vmulq_f32(r, r); \
105 float32x4_t p = vfmaq_laneq_f32(c1, r, ln2_c02, 2); \
106 float32x4_t q = vfmaq_laneq_f32(c3, r, ln2_c02, 3); \
107 q = vfmaq_f32(q, p, r2); \
108 p = vmulq_f32(c4, r); \
109 float32x4_t poly = vfmaq_f32(p, q, r2); \
110 poly = vfmaq_f32(scale, poly, scale); \
111 const uint32x4_t hi_mask = vcgeq_f32(values, pos_special_bound); \
112 const uint32x4_t lo_mask = vcleq_f32(values, neg_special_bound); \
113 poly = vbslq_f32(hi_mask, inf, poly); \
114 return vbslq_f32(lo_mask, zero, poly); \
115 }; \
116 auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
117 __attribute__((always_inline)) { \
118 float32x4x4_t result; \
119 result.val[0] = neon_expf(vec.reg.val[0]); \
120 result.val[1] = neon_expf(vec.reg.val[1]); \
121 result.val[2] = neon_expf(vec.reg.val[2]); \
122 result.val[3] = neon_expf(vec.reg.val[3]); \
123 return vec_op::FP32Vec16(result); \
124 }; \
125 const float32x4_t lower_bound = vdupq_n_f32(-0x1.5ebb82p+6f); \
126 const float32x4_t upper_bound = vdupq_n_f32(0x1.61814ap+6f); \
127 constexpr float ln2 = 0x1.62e43p-1f; \
128 constexpr float f_c2 = 0x1.5592ecp-3f; \
129 const float32x4_t f_c3 = vdupq_n_f32(0x1.017d34p-1f); \
130 auto neon_expf_f16 = [&](float32x4_t values) __attribute__(( \
131 always_inline)) { \
132 const uint32x4_t lt_lower = vcltq_f32(values, lower_bound); \
133 const uint32x4_t gt_upper = vcgtq_f32(values, upper_bound); \
134 float32x4_t n = vrndaq_f32(vmulq_f32(values, inv_ln2)); \
135 float32x4_t r = vfmsq_n_f32(values, n, ln2); \
136 uint32x4_t e = vshlq_n_u32(vreinterpretq_u32_s32(vcvtq_s32_f32(n)), 23); \
137 float32x4_t r2 = vmulq_f32(r, r); \
138 float32x4_t q = vfmaq_n_f32(f_c3, r, f_c2); \
139 float32x4_t s = vaddq_f32(vdupq_n_f32(1.0f), r); \
140 float32x4_t p = vfmaq_f32(s, q, r2); \
141 float32x4_t y = \
142 vreinterpretq_f32_u32(vaddq_u32(vreinterpretq_u32_f32(p), e)); \
143 y = vbslq_f32(lt_lower, vdupq_n_f32(0.0f), y); \
144 y = vbslq_f32(gt_upper, vdupq_n_f32(INFINITY), y); \
145 return y; \
146 }; \
147 auto fast_exp_f16 = [&](const vec_op::FP32Vec16& vec) \
148 __attribute__((always_inline)) { \
149 float32x4x4_t result; \
150 result.val[0] = neon_expf_f16(vec.reg.val[0]); \
151 result.val[1] = neon_expf_f16(vec.reg.val[1]); \
152 result.val[2] = neon_expf_f16(vec.reg.val[2]); \
153 result.val[3] = neon_expf_f16(vec.reg.val[3]); \
154 return vec_op::FP32Vec16(result); \
155 };
156
157 #endif // __aarch64__
158
159 // RISC-V RVV
160 #ifdef __riscv_v
161 #include <riscv_vector.h>
162
163 #ifdef __riscv_zihintpause
164 #define FAST_SPINNING __riscv_pause();
165 #endif
166
167 // FP32Vec16::exp() in cpu_types_riscv.hpp already implements the full
168 // polynomial approximation for RVV, so we simply delegate to it.
169 #define DEFINE_FAST_EXP \
170 auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
171 __attribute__((always_inline)) { return vec.exp(); };
172
173 #endif // __riscv_v
174
175 // Power VSX
176 #ifdef __powerpc__
177 // FP32Vec16::exp() in cpu_types_vsx.hpp delegates to FP32Vec8::exp(), which
178 // implements a vectorised 5-term minimax polynomial using VSX intrinsics.
179 #define DEFINE_FAST_EXP \
180 auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
181 __attribute__((always_inline)) { return vec.exp(); }; \
182 auto fast_exp_f16 = fast_exp;
183
184 #endif // __powerpc__
185
186 // IBM Z (s390x) VXE
187 #ifdef __s390x__
188 // FP32Vec16::exp() in cpu_types_vxe.hpp delegates to FP32Vec8::exp(), which
189 // implements a vectorised 5-term minimax polynomial using VXE intrinsics.
190 #define DEFINE_FAST_EXP \
191 auto fast_exp = [&](const vec_op::FP32Vec16& vec) \
192 __attribute__((always_inline)) { return vec.exp(); }; \
193 auto fast_exp_f16 = fast_exp;
194
195 #endif // __s390x__
196
197 #endif