Optimization: SIMD polynomial exp inside the j loop

Author: Linus
Predecessor: variant_inlining_7_SoA_AVX2.cpp
File(s): variants/variant_inlining_7_SoA_AVX2_polyexp.cpp
Status: TBD, positive locally (2 times)

What this changes

The predecessor still called libm exp() 4 times into a scratch array inside the j loop. A profile pointed at those calls as ~54% of inner-loop cycles. This variant replaces them with one SIMD polynomial approximation, exp256_pd, that operates on a __m256d directly.

The approximation is the textbook 2^n bit trick combined with a degree-7 Taylor polynomial. We write exp(x) = 2^n * exp(r) with n = round(x / ln2) and r = x - n*ln2, so after range reduction r lives in roughly [-0.347, 0.347]. On that range a degree-7 Taylor evaluated with Horner is accurate to ~1e-9 absolute, way under the 1e-3 tolerance verify.py uses. The 2^n part is free in IEEE 754 doubles, you just pack (n + 1023) into the exponent slot via a left shift by 52.

Code (relevant snippet)

See variants/variant_inlining_7_SoA_AVX2_polyexp.cpp for the full exp256_pd helper. Quick walkthrough: a round-to-nearest of x / ln2 gives n, then an fnmadd strips n*ln2 out of x to get the reduced argument r, which lives in roughly [-0.347, 0.347]. Seven fmadds evaluate a degree-7 Horner polynomial for exp(r) on that range. In parallel we build 2^n straight into the IEEE 754 exponent field by converting n to int32, adding the bias 1023, left shifting by 52, then reinterpreting as double. No actual pow call needed. Final multiply of polynomial value with 2^n gives exp(x).

In the j loop the old scratch-array path (4 libm calls into tmp_arg / tmp_val) becomes a single __m256d e_val = exp256_pd(e_arg);.

Benchmark

Conditions: min of N runs (TBD), N_ITER=1000, -O3 -march=native -flto. Diff is vs variant_inlining_7_SoA_AVX2.

local runs show a sizeable speedup over the AVX2-with-scalar-exp predecessor (approx 2times)

x86

Machine: TBD

Frequency: TBD

N	W(n) [Flops]	T(n) [Cycles]	Performance [Flops/Cycle]	Diff vs Predecessor
50
100
200
500
1000

ARM (Apple)

Not applicable, AVX2-only. A NEON port would need its own vectorized exp.

Analysis

exp was the single biggest cost in the predecessor's inner loop, so the expectation was that killing it would lift the whole loop by a large fraction. The SIMD polynomial is roughly 9 fmadds + a round + a couple of converts + a shift, no library call, no scratch array, no scalar fallback. Should pipeline well.

Accuracy: degree-7 Taylor on the reduced range is good to ~1e-9 absolute. verify.py's tolerance is 1e-3, so there is a lot of headroom. If we ever need tighter accuracy we could bump the degree or move to an Estrin schedule, but right now there is no point.

Takeaway

Biggest single step on top of the first AVX2 pass. The inner loop is now essentially fma / sqrt / div bound with no libm calls left in it.

References

Schraudolph, N. N. (1999). "A Fast, Compact Approximation of the Exponential Function." Neural Computation 11(4), 853-862.
IEEE 754 double-precision format. Wikipedia. https://en.wikipedia.org/wiki/Double-precision_floating-point_format
Horner's method. Wikipedia. https://en.wikipedia.org/wiki/Horner%27s_method

3.8 KiB Raw Blame History