Vectorized function to count numbers in an array when a number is a specified power

user_16558006 0
Viewed 100

I am attempting to vectorize this fairly expensive function (Scaler Now working!):

template<typename N, typename POW>
inline constexpr bool isPower(const N n, const POW p) noexcept
{
    double x = std::log(static_cast<double>(n)) / std::log(static_cast<double>(p));

    return (x - std::trunc(x)) < 0.000001;
}//End of isPower

Here's what I have so far (for 32-bit int only):

 template<typename RETURN_T>
 inline RETURN_T count_powers_of(const std::vector<int32_t>& arr, const int32_t power)
 {
      RETURN_T cnt = 0;
      
      const __m256 _MAGIC = _mm256_set1_ps(0.000001f);
      const __m256 _POWER_D = _mm256_set1_ps(static_cast<float>(para));

      const __m256 LOG_OF_POWER = _mm256_log_ps(_POWER_D);

      __m256i _count = _mm256_setzero_si256();
      __m256i _N_INT = _mm256_setzero_si256();
      __m256 _N_DBL = _mm256_setzero_ps();
      __m256 LOG_OF_N = _mm256_setzero_ps();
      __m256 DIVIDE_LOG = _mm256_setzero_ps();
      __m256 TRUNCATED = _mm256_setzero_ps();

      __m256 CMP_MASK = _mm256_setzero_ps();
                                                  
      for (size_t i = 0uz; (i + 8uz) < end; i += 8uz)
      {
           //Set Values
           _N_INT = _mm256_load_si256((__m256i*) &arr[i]);
           _N_DBL = _mm256_cvtepi32_ps(_N_INT);
                            
           LOG_OF_N = _mm256_log_ps(_N_DBL);

           DIVIDE_LOG = _mm256_div_ps(LOG_OF_N, LOG_OF_POWER);

           TRUNCATED = _mm256_sub_ps(DIVIDE_LOG, _mm256_trunc_ps(DIVIDE_LOG));

           CMP_MASK = _mm256_cmp_ps(TRUNCATED, _MAGIC, _CMP_LT_OQ);

           _count = _mm256_sub_epi32(_count, _mm256_castps_si256(CMP_MASK));
      }//End for

     cnt = static_cast<RETURN_T>(util::_mm256_sum_epi32(_count));
 }//End of count_powers_of

The scaler version runs in about 14.1 seconds. The scaler version called from std::count_if with par_unseq runs in 4.5 seconds.

The vectorized version runs in just 155 milliseconds but produces the wrong result. Albeit vastly closer now.

Testing:

 int64_t count = 0;
 for (size_t i = 0; i < vec.size(); ++i)
 {              
    if (isPower(vec[i], 4))
    {
        ++count;
    }//End if
}//End for

std::cout << "Counted " << count << " powers of 4.\n";//produces 4,996,215 powers of 4 in a vector of 1 billion 32-bit ints consisting of a uniform distribution of 0 to 1000

std::cout << "Counted " << count_powers_of<int32_t>(vec, 4) << " powers of 4.\n";//produces 4,996,865 powers of 4 on the same array

This new vastly simplified code often produces results that are either slightly off the correct number of powers found (usually higher). I think the problem is my reinterpret cast from __m256 to _m256i but when I try use a conversation (with floor) instead I get a number that's way off (in the billions again).

It could also be this sum function (based off of code by @PeterCordes ):

 inline uint32_t _mm_sum_epi32(__m128i& x)
{
    __m128i hi64 = _mm_unpackhi_epi64(x, x);           
    __m128i sum64 = _mm_add_epi32(hi64, x);
    __m128i hi32 = _mm_shuffle_epi32(sum64, _MM_SHUFFLE(2, 3, 0, 1));
    __m128i sum32 = _mm_add_epi32(sum64, hi32);
    return _mm_cvtsi128_si32(sum32);
}

 inline uint32_t _mm256_sum_epi32(__m256i& v)
{
    __m128i sum128 = _mm_add_epi32(
        _mm256_castsi256_si128(v),
        _mm256_extracti128_si256(v, 1));
    return _mm_sum_epi32(sum128);
}

I know this has got to be a floating-point precision/comparison issue; Is there a better way to approach this?

Thanks for all your insights and suggestions thus far.

avx2 counting simd math c++
1 Answers

A more sensible unit-test would be to non-random: Check all powers in a loop to make sure they're all true, like x *= base;, and count how many powers there are <= n. Then check all numbers from 0..n in a loop, once each to verify the right total. If both those checks succeed, that means it returned false in all the cases it should have, otherwise the count would be wrong.

Re: the original version:

This seems to depend on there being no floating-point rounding error. You do d == (N)d which (if N is an integral type) checks that the ratio of two logs is an exact integer; even 1 bit in the mantissa will make it unequal. Hardly surprising that a different log implementation would give different results, if one has different rounding error.

Except your scalar code at least is even more broken because it takes d = floor(log ratio) so it's already always an exact integer.

I just tried your scalar version for a testcase like return isPower(5, 4) to ask if 5 is a power of 4. It returns true: https://godbolt.org/z/aMT94ro6o . So yeah, your code is super broken, and is in fact only checking that n>0 or something. That would explain why 999 of 1000 of your "random" inputs from 0..999 were counted as powers of 4, which is obviously super broken.

I think it's impossible to achieve correctness with your FP log ratio idea: FP rounding error means you can't expect exact equality, but allowing a range would probably let in non-exact powers.

You might want to special-case integral N, power-of-2 pow. That can go vastly vaster by checking that n has a single bit set (n & (n-1) == 0) and that it's at a valid position. (e.g. for pow=4, n & 0b...10101010 != 0). You can construct the constant by multiplying and adding until overflow or something. Or 32/pow times? Anyway, one psubd/pand/pcmpeqd, pand/pcmpeqd, and pand/psubd per 8 elements, with maybe some room to optimize that further.

Otherwise, in the general case, you can brute-force check 32-bit integers one at a time against the 32 or fewer possible powers that fit in an int32_t. e.g. broadcast-load, 4x vpcmpeqd / vpsubd into multiple accumulators. (The smallest possible base, 2, can have exponents up to 2^31` and still fit in an unsigned int). log_3(2^31) is 19, so you'd only need three AVX2 vectors of powers. Or log_4(2^31) is 15.5 so you'd only need 2 vectors to hold every non-overflowing power.

That only handles 1 input element per vector instead of 4 doubles, but it's probably faster than your current FP attempt, as well as fixing the correctness problems. I could see that running more than 4x the throughput per iteration of what you're doing now, or even 8x, so it should be good for speed. And of course has the advantage that correctness is possible!!

Speed gets even better for bases of 4 or greater, only 2x compare/sub per input element, or 1x for bases of 16 or greater. (<= 8 elements to compare against can fit in one vector).

Implementation mistakes in the attempt to vectorize this probably-unfixable algorithm:

  • _mm256_rem_epi32 is slow library function, but you're using it with a constant divisor of 2! Integer mod 2 is just n & 1 for non-negative. Or if you need to handle negative remainders, you can use the tricks compilers use to implement int % 2: https://godbolt.org/z/b89eWqEzK where it shifts down the sign bit as a correction to do signed division.

Updated version using (x - std::trunc(x)) < 0.000001;

This might work, especially if you limit it to small n. I'd worry that with large n, the difference between an exact power and off-by-1 would be a small ratio. (I haven't really looked at the details, though.)

Your vectorization with __m256 vectors of single-precision float is doomed for large n, but could be ok for small n: float32 can't represent every int32_t, so large odd integers (above 2^24) get rounded to multiples of 2, or multiples of 4 above 2^25, etc.

float has less relative precision in general, so it might not have enough to spare for this algorithm. Or maybe there's something that could be fixed, IDK, I haven't looked closely since the update.

I'd still recommend trying a simple compare-for-equality against all possible powers in the range, broadcast-loading each element. That will definitely work exactly, and if it's as fast then there's no need to try to fix this version using FP logs.

__m256 _N_DBL = _mm256_setzero_ps(); is a confusing name; it's a vector of float, not double. (And it's not part of a standard library header so it shouldn't be using a leading underscore.)

Also, there's zero point initializing it with zero there, since it gets written unconditionally inside the loop. In fact it's only ever used inside the loop, so it could just be declared at that scope, when you're ready to give it a value. Only declare variables in outer scopes if you need them after a loop.

user_2241320
Related