Writing one AVX-5121 operation is difficult enough; building just a standalone function requires an intimate understanding of the ISA and its registers and instructions. For example, the inverse square root function2 must be built on an approximate reciprocal square root instruction, a fixture of x86 since SSE arrived with the Pentium III in 1999. The AVX-512 form, VRSQRT14PS, is good to 14 bits — its SSE ancestor guaranteed only 12 — so a Newton-Raphson refinement is needed to bring it within a few ULP:

// 1/sqrt(v) to within 2 ULP. rsqrt14 supplies 14 bits; one Newton-Raphson
// step roughly doubles that, and the FMA holds the correction to a single
// rounding. Requires positive input.
__m512 rsqrt_ps( __m512 v )
{
    const __m512 half  = _mm512_set1_ps( 0.5f );
    const __m512 three = _mm512_set1_ps( 3.0f );

    __m512 y = _mm512_rsqrt14_ps( v );                     // y ~ 1/sqrt(v)
    __m512 t = _mm512_fnmadd_ps( _mm512_mul_ps( v, y ),
                                 y, three );               // 3 - v*y*y
    return _mm512_mul_ps( _mm512_mul_ps( y, half ), t );
}

When programming SIMD instruction sets that definitionally perform k of the same operation, in parallel, in one instruction, it’s extremely common to process arrays. We’ll cover a more realistic real-world example in a subsequent article, but for now let’s wrap the standalone rsqrt_ps in a loop that applies it to every element of an array:

// Replace every element of p[0..N-1] with its inverse square root.
// p must be 16-byte aligned; N must be a multiple of 16.
void rsqrt_array( float *p, size_t N )
{
    for ( size_t i = 0; i < N; i += 16 ) {
        _mm512_storeu_ps( p+i, rsqrt_ps( _mm512_loadu_ps( p+i ) ) );
    }
}

Of course, problem sizes that aren’t multiples of 16 require an accommodation, such as a masked epilogue:

// Same, for any N. The remainder is a single masked iteration through the
// same code path, rather than a scalar loop with different arithmetic in it.
void rsqrt_array( float *p, size_t N )
{
    size_t i = 0;
    for ( ; i + 16 <= N; i += 16 ) {
        _mm512_storeu_ps( p+i, rsqrt_ps( _mm512_loadu_ps( p+i ) ) );
    }
    if ( i < N ) {
        __mmask16 m = (__mmask16) ( ( 1u << ( N - i ) ) - 1 );
        _mm512_mask_storeu_ps( p+i, m,
                               rsqrt_ps( _mm512_maskz_loadu_ps( m, p+i ) ) );
    }
}

By now, it should be clear why developers may want a separation of concerns between the loop and the operation(s) being performed on the array elements: the epilogue code invokes rsqrt_ps again, so even our modest improvement of the first loop implementation introduces some code duplication. Other variations of these loops may include ones that use aligned loads and stores, perhaps requiring prologue as well as epilogue code to deal with misalignments and odd problem sizes; and differing degrees of loop unrolling.

Let’s see how we can enlist variadic templates to enable that separation.

The First Parameter Pack: Operator Arity

Applying a chain of elementwise operations to a large buffer is a common pattern in numerical code. But invoking a series of functions akin to rsqrt_array to process the array in turn, leaves a lot of performance on the table. By decoupling the loop from evaluation of the function, we can deliver higher performance by keeping intermediate values in registers and minimizing the amount of external memory traffic. Here’s a variadic template function that takes an arbitrary list of AVX-512 operators to apply elementwise to an array:

template <typename... Ops>
void fused_avx512_transform(const float* in, float* out, std::size_t n, Ops... ops) {
    std::size_t i = 0;
    for (; i + 16 <= n; i += 16) {
        __m512 v = _mm512_loadu_ps(in + i);
        ((v = ops(v)), ...);            // the chain, in a register
        _mm512_storeu_ps(out + i, v);
    }
    avx512_masked_tail(in, out, i, n, ops...); // n need not be a multiple of 16
}

Ops... is a parameter pack, the C++11 feature that lets a template accept any number of arguments of any types.

All the caller writes is the chain; and, as shown below, the developer can even extemporaneously create functions to specify in the list:

float  gain = read_gain();              // not known until runtime
__m512 v_gain = _mm512_set1_ps( gain );

fused_avx512_transform( in, out, n,
    [v_gain]( __m512 v ) { return _mm512_mul_ps( v, v_gain ); }, // scale
    FmaBias( 1.25f, 0.125f ),                            // v*a + b
    Clamp( -8.0f, 8.0f ),                                // min(max(v,lo),hi)
    Poly<3>( coeffs ),                                   // Horner, degree 3
    Relu( 0.0f ) );                                      // max(v, 0)

The first operator is a lambda that captures a broadcasted constant, and the rest are small classes that do the same in a constructor. Both formulations work because a pack accepts any set of types and a lambda’s closure type serves just as well as a named class.

Note that the caller did not have to write a loop! fused_avx512_transform takes care of that: how many vectors are processed per pass, what happens when the length is not a multiple of 16, and propagating intermediate values between operators without a round trip to memory. If we wished to refactor how the loop were managed (say, optimize with the knowledge that the problem size was evenly divisible by 16), we could invoke a different, related version of fused_avx512_transform.

Measured on a Ryzen 7 7700X with g++ 9.4 at -O2, chaining five operators and sweeping the working set:

working set    fused speedup
 16 KB (L1)        1.95×
256 KB (L2)        2.85×
  4 MB (L3)        3.30×
 64 MB (DRAM)      6.33×

As we can see from the timing data, the performance benefits of fusion are most pronounced for larger problem sizes, because fusion reduces the number of round trips to DRAM. The abstraction is free: against a hand-written chain of the same five operations, the pack generates identical code.

The Second Parameter Pack: Loop Unroll Factor

Processing one SIMD vector per iteration may leave performance on the table. Every operator in the chain depends on its predecessor, introducing serial dependencies that limit the CPU’s ability to identify and exploit ILP. We can increase the amount of independent work by unrolling the loop, and parameterizing the unroll factor as U is an excellent application for templates:

template <std::size_t U, typename... Ops>
void fused_avx512_transform( const float *in, float *out, std::size_t n, Ops... ops );

U and Ops... now sit in the same parameter list, and they are there for unrelated reasons. Ops... must be a pack: the caller supplies a different number of operators on each call, and their types are unrelated. U is a single number the caller already knows at compile time, and the signature above expresses it without a pack at all.

But to deliver the benefits of a loop unroll, we enlist variadic templates: we introduce a second parameter pack, one level further down, in how the repetition over those U accumulators is written. The non-parameter-pack implementation would loop:

__m512 v[U];
for (std::size_t u = 0; u < U; ++u) v[u] = _mm512_loadu_ps(in + i + 16 * u);
for (std::size_t u = 0; u < U; ++u) ((v[u] = ops(v[u])), ...);
for (std::size_t u = 0; u < U; ++u) _mm512_storeu_ps(out + i + 16 * u, v[u]);

An index_sequence pack unrolls those loops explicitly:

template <std::size_t... I, typename... Ops>
inline void avx512_block_seq(const float* in, float* out, std::size_t i,
                      std::index_sequence<I...>, const Ops&... ops) {
    __m512 v[sizeof...(I)];
    ((v[I] = _mm512_loadu_ps(in + i + 16 * I)), ...);
    ((v[I] = avx512_chain(v[I], ops...)), ...);
    ((_mm512_storeu_ps(out + i + 16 * I, v[I])), ...);
}

Timing the two at 4 MB, sweeping the polynomial degree to vary the chain length, in milliseconds per call:

degree  `-O2` loop   `-O2` pack  `-O3` loop  `-O3` pack
  3       0.132      **0.070**     0.081        0.081
  5       0.121        0.113       0.113        0.113
  7       0.216        0.179       0.143        0.143
  9       0.253        0.203       0.178        0.177
 11       0.272        0.234       0.210        0.210

At -O3, the two deliver near-identical performance. At -O2, the pack wins everywhere, by 47% at degree 3.

The mechanism is worth spelling out: In the loop version, v[u] is indexed by a loop variable, so the accumulator array becomes register-allocatable only if the compiler unrolls the loop, which GCC does not do at -O2. The array therefore stays in memory and every operator round-trips through the stack. In the pack version, v[I] is a compile-time constant, so scalar replacement promotes the array to registers whether or not an unrolling pass runs.

I was surprised by one result: I expected that applying each operator in turn, rather than running each accumulator through the whole chain, would result in better instruction scheduling. It does not: it is slower in nearly every configuration, because Zen 4’s reorder window is wide enough to find the independence on its own, while the other formulation forces all U accumulators to stay live across the entire chain, increasing register pressure.

The operator pack is the ordinary C++11 idiom, and the unroll pack is std::index_sequence used the way the standard library uses it internally.

Conclusion

The operator pack enables the caller to enumerate any series of operations to be performed on array elements, such that every operator is inlined and the intermediate values remain in registers. Fusing five operations that way is almost 2x faster even for small arrays that fit in L1, and >6x faster when the memory traffic goes to DRAM.

The unroll pack does something narrower. Writing the repetition over accumulators as a pack rather than a loop makes every accumulator index a compile-time constant, so good code generation does not require a higher optimization level. The same effect benefited the unrolled binary search: a probe sequence generated by a pack beats a loop over the same powers of two by 22% at -O2, and the two converge when the compiler unrolls the loop.

The header and its benchmark are on GitHub, in variadic-avx512-fusion, under a BSD license — fusion.hpp for the two packs, ops.hpp for the operators, and bench.cpp for the numbers in the tables above.

What the two have in common is the separation they enforce. An operator is a function of __m512 that knows nothing about iteration; fused_avx512_transform knows all about iteration and nothing about the operators it is applying. Neither has to be updated when the other changes, and the compiler still emits what a competent programmer would have written by hand. Writing AVX-512 is as difficult as ever, as the inverse square root at the top of this article attests. But the task can be made a bit easier with the aid of modern C++.