Threshold on absolute values on large float array in C

Question

I would like to start by pointing out my poor C skills - and abysmal knowledge of SIMD instructions - so I apologize for any stupid mistake I may have made.

This question is related to another one I posted in the past: Better way to find nonzero indices and values in a 2D array, but with a little twist.

I have somewhat large 1D arrays of float values (with some million elements), and they are very sparse: sparsity varies between 0.1% (or even lower) and 1-2%. They will contain subnormal numbers (i.e., -0.0), and also smallish numbers (i.e., 1.4e-12). The numbers can be positive or negative.

I am trying to find all the indices of the elements that have their absolute value greater than some threshold (that I set at 1e-10).

I have implemented this in C using 3 different approaches:

1. A "naïve" implementation, in which I loop through all the elements and check them one by one using fabsf. If they pass the test, their index is recorded.
2. Similar to (1), but avoids branching with an if condition.
3. Using SIMD instruction (I can assume at least AVX2/Skylake)

This is my attempt - with a benchmark in the main function:

#include <stdio.h>
#include <stdlib.h>
#include <stddef.h>
#include <math.h>
#include <inttypes.h>
#include <immintrin.h>
#include <omp.h>

#define SMALL 1e-10f


int naive(const float *M, int *T, int N) {   
    int count = 0;
    for(int i = 0; i < N; i++) {
        if (fabsf(M[i]) > SMALL) {
            T[count] = i;
            count += 1;
        }
    }
    return count;
}


int naive_no_branch(const float *M, int *T, int N) {   
    int count = 0;
    for(int i = 0; i < N; i++) {
        T[count] = i;
        count += fabsf(M[i]) > SMALL;
    }
    return count;
}


int avx2(const float *M, int *T, int N) {
    
    int i, k;
    
    int count = 0;
    const int reminder = (N / 8) * 8;

    const __m256 zeros = _mm256_setzero_ps();

    for(i = 0; i < reminder; i += 8) {
        
        const __m256 vec = _mm256_load_ps(M + i);
        
        const __m256 comparison_out = _mm256_cmp_ps(zeros, vec, _CMP_NEQ_UQ); // 3 cycles latency, throughput 1
        // Does it ignore subnormals?
        uint64_t mask = _mm256_movemask_ps(comparison_out);                   // 2-3 cycles latency

        if (!mask)
            // Mask is zero, everything is zero
            continue;

        for(k = 0; k < 8; k++) {
            if (fabsf(M[k+i]) > SMALL) {
                T[count] = k+i;
                count += 1;
            }
        }
    }

    // Handle reminder if N is not divisible by 8
    for(i = reminder; i < N; i++) {
        if (fabsf(M[i]) > SMALL) {
            T[count] = i;
            count += 1;
        }
    }

    return count;
}


int main() {
            
    int array_size, N;
    double dtime1, dtime2, dtime3;
    
    printf("\n");
    
    for(array_size = 19; array_size < 26; array_size++) {
        // Make an array with a few million elements
        N = (2 << array_size) + 5;
        printf("Array size: %d\n", N);

        float* M  = (float*) malloc(sizeof(float)*N);
        int* T1  = (int*) malloc(sizeof(int)*N);
        int* T2  = (int*) malloc(sizeof(int)*N);
        int* T3  = (int*) malloc(sizeof(int)*N);

        for(int i=0; i<N; i++) {
            if ((rand()%500)==0) {
                M[i] = 3.14159265f;
            }
            else if ((rand()%1000)==0) {
                M[i] = -0.0f;
            }
            else  {
                M[i] = 0.0f;
            }       
        }

        int repeat = 10;
        int count1, count2, count3;

        dtime1 = omp_get_wtime();
        for(int i = 0; i < repeat; i++) 
            count1 = naive(M, T1, N);

        dtime1 = omp_get_wtime() - dtime1;
        printf("Time naive          : %f  (x %0.2f)\n", dtime1, dtime1/dtime1);

        dtime2 = omp_get_wtime();
        for(int i = 0; i < repeat; i++) 
            count2 = naive_no_branch(M, T2, N);

        dtime2 = omp_get_wtime() - dtime2;
        printf("Time naive_no_branch: %f  (x %0.2f)\n", dtime2, dtime1/dtime2);

        dtime3 = omp_get_wtime();
        for(int i = 0; i < repeat; i++) 
            count3 = avx2(M, T3, N);

        dtime3 = omp_get_wtime() - dtime3;
        printf("Time avx2           : %f  (x %0.2f)\n\n", dtime3, dtime1/dtime3);
                
        printf("NNZ  naive          : %d\n", count1);
        printf("NNZ  naive_no_branch: %d\n", count2);
        printf("NNZ  avx2           : %d\n\n", count3);
        
        free(M);
        free(T1);
        free(T2);
        free(T3);
    }
    
}

On my machine:

• Windows 10 64 bit
• Cascadelake
• Intel(R) Xeon(R) Gold 6248R CPU @ 3.00GHz 2.99 GHz
• GCC 11.3
• 128 GB RAM

Compiled as:

gcc -O3 -ffast-math -fopenmp -march=skylake filter_small.c -o filter_small.exe

I get these timings:

Array size: 1048581
Time naive          : 0.009000  (x 1.00)
Time naive_no_branch: 0.015000  (x 0.60)
Time avx2           : 0.002000  (x 4.50)

NNZ  naive          : 2135
NNZ  naive_no_branch: 2135
NNZ  avx2           : 2135

Array size: 2097157
Time naive          : 0.018000  (x 1.00)
Time naive_no_branch: 0.029000  (x 0.62)
Time avx2           : 0.004000  (x 4.50)

NNZ  naive          : 4168
NNZ  naive_no_branch: 4168
NNZ  avx2           : 4168

Array size: 4194309
Time naive          : 0.040000  (x 1.00)
Time naive_no_branch: 0.059000  (x 0.68)
Time avx2           : 0.010000  (x 4.00)

NNZ  naive          : 8289
NNZ  naive_no_branch: 8289
NNZ  avx2           : 8289

Array size: 8388613
Time naive          : 0.078000  (x 1.00)
Time naive_no_branch: 0.121000  (x 0.64)
Time avx2           : 0.031000  (x 2.52)

NNZ  naive          : 16847
NNZ  naive_no_branch: 16847
NNZ  avx2           : 16847

Array size: 16777221
Time naive          : 0.168000  (x 1.00)
Time naive_no_branch: 0.245000  (x 0.69)
Time avx2           : 0.081000  (x 2.07)

NNZ  naive          : 33620
NNZ  naive_no_branch: 33620
NNZ  avx2           : 33620

Array size: 33554437
Time naive          : 0.341000  (x 1.00)
Time naive_no_branch: 0.499000  (x 0.68)
Time avx2           : 0.176000  (x 1.94)

NNZ  naive          : 67261
NNZ  naive_no_branch: 67261
NNZ  avx2           : 67261

Array size: 67108869
Time naive          : 0.671000  (x 1.00)
Time naive_no_branch: 0.992000  (x 0.68)
Time avx2           : 0.347000  (x 1.93)

NNZ  naive          : 134652
NNZ  naive_no_branch: 134652
NNZ  avx2           : 134652

Seems like avoiding branching is not very helpful, but I am very pleased with the results of the SIMD approach, which is between 2 and 4.5 times faster.

Now, for the question: is there anything more (better?) I could do to the SIMD implementation (beside using omp) to make it marginally/significantly faster? Happy to try any suggestion from more experienced people, it's always a chance for me to learn something new.

Note: I can assume at least AVX2/Skylake, but unfortunately not AVX512 as we have some Alderlake machines and as far as I know Intel has removed AVX512 support for those...

Answer 1

// Does it ignore subnormals?

No, unless you enable DAZ. (A program linked with -ffast-math will include CRT startup code that sets DAZ and FTZ before main.)

Now, for the question: is there anything more (better?) I could do to the SIMD implementation (beside using omp) to make it marginally/significantly faster?

Instead of only testing for unequal to zero, you can do the entire "check if absolute value is large enough" test, and then end up with an all-zero mask more often. It costs more operations, so whether that's worth it depends on the data.

I also changed the load to unaligned, in this function we don't really know that the data is aligned, although practically speaking you'll probably get code that accepts an unaligned address anyway (eg a memory operand in an AVX instruction, as opposed to legacy-SSE).

__m256 data = _mm256_loadu_ps(M + i);
__m256 data_abs = _mm256_and_ps(data, _mm256_castsi256_ps(_mm256_set1_epi32(0x7fffffff)));
__m256 compout = _mm256_cmp_ps(data_abs, _mm256_set1_ps(SMALL), _CMP_GT_OQ);
int mask = _mm256_movemask_ps(compout);

As mentioned by Peter, we can use integer arithmetic instead of floating point arithmetic so we can handle subnormals even if DAZ is enabled:

__m256i data = _mm256_loadu_si256((__m256i*)(M + i));
__m256i data_abs = _mm256_and_si256(data, _mm256_set1_epi32(0x7fffffff));
__m256i compout = _mm256_cmpgt_epi32(data_abs, _mm256_castps_si256(_mm256_set1_ps(SMALL)));
int mask = _mm256_movemask_ps(_mm256_castsi256_ps(compout));

(FTZ and DAZ are normally enabled together and FTZ stops FP math ops from producing subnormals in the first place. But you might change that, or another thread with different FP settings could have produced subnormals, or you could be reading binary floats from outside the program.)

The small loop that records the indices can be written like this (similar idea as naive_no_branch, but here we reuse the mask that was just computed with SIMD):

for(k = 0; k < 8; k++) {
    T[count] = k+i;
    count += (mask >> k) & 1;
}

The trade-off for doing it here is different than in naive_no_branch, in naive_no_branch we're replacing a branch that usually goes the same way (ie more predictable, the code in the if only runs about 1/1000th of the times) and changing it into code that usually does some redundant work, but here it's different: the only way we're even in this code is if there is at least one index to store. So less of the work is redundant, and the branch is less predictable (the if has to be taken at least 1/8th of the times). So we should not a-priori assume or even expect that since it wasn't profitable to use this technique in naive_no_branch that it also won't be profitable in this context.

If a non-zero mask usually has only 1 set bit (but that has to be very often the case to make this worth doing), you could consider adding something like this:

if ((mask & (mask - 1)) == 0) {
    T[count++] = i + std::countr_zero((unsigned)mask);
}
else {
    // do the little loop

That the little loop that records the indices can also be replaced with left pack-ing the indices based on the mask. With AVX2 that is relatively annoying, with AVX512 (or AVX10 when it comes out) it's easy and fairly cheap (cheaper than in AVX2) with vcompressps aka _mm256_maskz_compress_ps (or _mm512_maskz_compress_ps if you also switch to 512-bit vectors).

Something that all of the above has in common is that the performance characteristics and ranking between the techniques depends on the data, so this takes extra care to evaluate. Benchmarking this on non-representative test data can easily mislead you into choosing a technique that is not the best for your actual data.

Answer 2

Try the following version. In addition to vectorized absolute value test mentioned in another answer, this uses more efficient method to enumerate set bits in the bitmap.

size_t avx2_opt( const float* rsi, int* rdi, size_t length )
{
    const __m256 signBits = _mm256_set1_ps( -0.0f );
    const __m256 smallVec = _mm256_set1_ps( SMALL );

    const size_t lengthAligned = length & ( ~(size_t)7 );
    size_t count = 0;
    size_t i;
    for( i = 0; i < lengthAligned; i += 8 )
    {
        // Load 8 numbers
        __m256 vec = _mm256_loadu_ps( rsi + i );
        // Use bitwise tricks to compute absolute value
        // The above load should be merged into `vandnps` instruction
        vec = _mm256_andnot_ps( signBits, vec );
        // Compare for e > SMALL
        vec = _mm256_cmp_ps( vec, smallVec, _CMP_GT_OQ );
        // Move results to bitmap in scalar registers
        uint32_t mask = (uint32_t)_mm256_movemask_ps( vec );
        // Enumerate set bits in the bitmap using BMI instructions
        while( 0 != mask )
        {
            uint32_t idx = _tzcnt_u32( mask );
            mask = _blsr_u32( mask );
            rdi[ count ] = (int)( (uint32_t)i + idx );
            count++;
        }
    }

    // Handle reminder if N is not divisible by 8
    for( ; i < length; i++ )
    {
        if( fabsf( rsi[ i ] ) > SMALL )
        {
            rdi[ count ] = (int)i;
            count++;
        }
    }

    return count;
}

The above code is untested.