How to vectorize the pow function (with negative bases)?

Viewed 812

I'm trying to vectorize (SSE/AVX) the pow function. In all implementations that I have found, it simply vectorized using log and exp:

pow(x, y) = exp(y * log(x))

It works well for positive x, but this won't work for negative x since the log of a negative number is a complex number. Is it possible to efficiently vectorize the pow while keeping the ability to handle negative x numbers ?

2 Answers

This is a generic answer that doesn't take advantage of any specifics of how you might actually vectorize pow().


You could check if any of the elements of the base-vector are negative, and branch on that to select between a fast path and a slow path.

Return two vectors, of real parts and imaginary parts, so the fast path can return _mm_setzero_ps() for the imaginary part. Callers that don't want the imaginary part can ignore it (instead of having to shuffle to extract the real part for vectors of alternating real/imaginary).

So callers that pass only non-negative bases get behaviour that's nearly as fast as a vectorized real-only version.

But callers that pass a mix of negative and non-negative will get the slow version. If you can vectorize the slow version, that's perfect.

If it doesn't work for positive bases, when there's a mix you could run both and blend (based on the same compare-mask that you checked to see if you needed the slow version).

For an AVX version, type an extra 256 into the intrinsic names. (And change the check to == 0xff, because you have 4 more bits in the movemask result).

// SSE4.1 for BLENDVPS
__m128  pow_complexresult(__m128 base, __m128 exp, __m128 &imag_result)
{
    __m128 negbase_vec = _mm_cmplt_ps(base, _mm_setzero_ps());
    unsigned negbase_mask = _mm_movemask_ps(negbase_vec);

    if (negbase_mask == 0) {               // all elements false
        imag_result = _mm_setzero_ps();
        return pow_nonegative(base, exp);   // fast path
    } else if (negbase_mask == 0xf) {      // all elements true
        return pow_negative(base, exp, imag_result);
    } else {
        // Only needed if pow_negative doesn't work for non-negative inputs.
        __m128 negpow = pow_negative(base, exp, imag_result);
        __m128 pospow = pow_simple(base, exp);
        imag_result = _mm_andn_ps(negbase_mask, imag_result);  // blend imaginary part
        return _mm_blendv_ps(pospow, negpow, negbase_vec);  // blend real part
    }
}

Make sure the helper functions inline so you aren't really passing a vector by reference through memory.

And/or inline this wrapper into the callers, which may let the check optimize away for constant vectors.

I don't think either Windows or System V ABIs will return a struct of two __m256 vectors in two ymm registers, so a 2nd by-reference arg is probably the best you're going to get.

Notice that imag_result is the last arg, so even in the Windows x64 ABI, this function can still forward its args in the same registers to pow_nonegative(base, exp);. Although you want it to inline anyway.

Well, if you do some arithmetic with negative numbers, you can use the principal branch of log() function, that introduces an imaginary +i*b*pi number to the result, which translates then in a exp(+i*b*pi) factor..... that maps to real only in the cases your exponent b is an exact integer number..... this is not the general case you have, as the function pow(3) is declared for double type.... you can solve this issue, writing a wrapper function that checks for negative first parameter, and gives a negative also only in case the second parameter is an odd integer number. The implementation details are hidden as an exercise for you to complete... is quite easy to implement a double generalized_pow(double b, double e); that does this checking and gives the appropiate result.

But if the exponent b parameter is going to be integer... then you have better to implement the fast power algorithm, that will give you the answer probably faster than going through the exp(3) and log(3) functions. Look in Google for an implementation of the fast exponential algorithm.

Related