Is it possible to vectorize non-trivial loop in C with SIMD? (multiple length 5 double-precision dot products reusing one input)

Viewed 200

I have a performance critical C code where > 90% of the time is spent doing one basic operation:

Operation

The C code I am using is:

static void function(double *X1, double *Y1, double *X2, double *Y2, double *output) {
    double Z1, Z2;
    int i, j, k;
    for (i = 0, j = 0; i < 25; j++) { // sweep Y
        Z1 = 0;
        Z2 = 0;
        for (k = 0; k < 5; k++, i++) { // sweep X
            Z1 += X1[k] * Y1[i];
            Z2 += X2[k] * Y2[i];
        }
        output[j] = Z1*Z2;
    }
}

The lengths are fixed (X is 5; Y is 25; the output is 5). I have tried everything I know to make this faster. When I compile this code using clang with -O3 -march=native -Rpass-analysis=loop-vectorize -Rpass=loop-vectorize -Rpass-missed=loop-vectorize, I get this message:

remark: the cost-model indicates that vectorization is not beneficial [-Rpass-missed=loop-vectorize]

But I assume the way to make this faster is with SIMD somehow. Any suggestions would be appreciated.

3 Answers

Try the following version, it requires SSE2 and FMA3. Untested.

void function_fma( const double* X1, const double* Y1, const double* X2, const double* Y2, double* output )
{
    // Load X1 and X2 vectors into 6 registers; the instruction set has 16 of them available, BTW.
    const __m128d x1_0 = _mm_loadu_pd( X1 );
    const __m128d x1_1 = _mm_loadu_pd( X1 + 2 );
    const __m128d x1_2 = _mm_load_sd( X1 + 4 );

    const __m128d x2_0 = _mm_loadu_pd( X2 );
    const __m128d x2_1 = _mm_loadu_pd( X2 + 2 );
    const __m128d x2_2 = _mm_load_sd( X2 + 4 );

    // 5 iterations of the outer loop
    const double* const y1End = Y1 + 25;
    while( Y1 < y1End )
    {
        // Multiply first 2 values
        __m128d z1 = _mm_mul_pd( x1_0, _mm_loadu_pd( Y1 ) );
        __m128d z2 = _mm_mul_pd( x2_0, _mm_loadu_pd( Y2 ) );

        // Multiply + accumulate next 2 values
        z1 = _mm_fmadd_pd( x1_1, _mm_loadu_pd( Y1 + 2 ), z1 );
        z2 = _mm_fmadd_pd( x2_1, _mm_loadu_pd( Y2 + 2 ), z2 );

        // Horizontal sum both vectors
        z1 = _mm_add_sd( z1, _mm_unpackhi_pd( z1, z1 ) );
        z2 = _mm_add_sd( z2, _mm_unpackhi_pd( z2, z2 ) );

        // Multiply + accumulate the last 5-th value
        z1 = _mm_fmadd_sd( x1_2, _mm_load_sd( Y1 + 4 ), z1 );
        z2 = _mm_fmadd_sd( x2_2, _mm_load_sd( Y2 + 4 ), z2 );

        // Advance Y pointers
        Y1 += 5;
        Y2 += 5;

        // Compute and store z1 * z2
        z1 = _mm_mul_sd( z1, z2 );
        _mm_store_sd( output, z1 );

        // Advance output pointer
        output++;
    }
}

It’s possible to micro-optimize further by using AVX, but I’m not sure it’s going to help much because the inner loop is too short. I think that these two extra FMA instructions are cheaper than the overhead of computing horizontal sum of the 32-byte AVX vectors.

Update: here's another version, it takes less instructions overall at the cost of a few shuffles. May of may not be faster for your use case. Requires SSE 4.1 but I think all CPUs which have FMA3 have SSE 4.1 as well.

void function_fma_v2( const double* X1, const double* Y1, const double* X2, const double* Y2, double* output )
{
    // Load X1 and X2 vectors into 5 registers
    const __m128d x1_0 = _mm_loadu_pd( X1 );
    const __m128d x1_1 = _mm_loadu_pd( X1 + 2 );
    __m128d xLast = _mm_load_sd( X1 + 4 );

    const __m128d x2_0 = _mm_loadu_pd( X2 );
    const __m128d x2_1 = _mm_loadu_pd( X2 + 2 );
    xLast = _mm_loadh_pd( xLast, X2 + 4 );

    // 5 iterations of the outer loop
    const double* const y1End = Y1 + 25;
    while( Y1 < y1End )
    {
        // Multiply first 2 values
        __m128d z1 = _mm_mul_pd( x1_0, _mm_loadu_pd( Y1 ) );
        __m128d z2 = _mm_mul_pd( x2_0, _mm_loadu_pd( Y2 ) );

        // Multiply + accumulate next 2 values
        z1 = _mm_fmadd_pd( x1_1, _mm_loadu_pd( Y1 + 2 ), z1 );
        z2 = _mm_fmadd_pd( x2_1, _mm_loadu_pd( Y2 + 2 ), z2 );

        // Horizontal sum both vectors while transposing
        __m128d res = _mm_shuffle_pd( z1, z2, _MM_SHUFFLE2( 0, 1 ) );   // [ z1.y, z2.x ]
        // On Intel CPUs that blend SSE4 instruction doesn't use shuffle port,
        // throughput is 3x better than shuffle or unpack. On AMD they're equal.
        res = _mm_add_pd( res, _mm_blend_pd( z1, z2, 0b10 ) );  // [ z1.x + z1.y, z2.x + z2.y ]

        // Load the last 5-th Y values into a single vector
        __m128d yLast = _mm_load_sd( Y1 + 4 );
        yLast = _mm_loadh_pd( yLast, Y2 + 4 );

        // Advance Y pointers
        Y1 += 5;
        Y2 += 5;

        // Multiply + accumulate the last 5-th value
        res = _mm_fmadd_pd( xLast, yLast, res );

        // Compute and store z1 * z2
        res = _mm_mul_sd( res, _mm_unpackhi_pd( res, res ) );
        _mm_store_sd( output, res );
        // Advance output pointer
        output++;
    }
}

From the extended discussions in the comments it looks like mostly you are interested in reducing the latency between reading X1, X2 and writing output. What you are calculating is the element-wise product of two matrix-vector products. The two MV products can happen quasi-parallel (with OOO-excecution), but both MV products require the sum of five products, which you can do either in a sequence (as you are doing at the moment) or in a tree-like reduction:

Z = ((X[0]*Y[0] + X[1]*Y[1]) + X[2]*Y[2])  +  ([X[3]*Y[3] + [X[4]*Y[4]);

This results in a critical path of mulsd - fmaddsd - fmaddsd - addsd, which is followed by a multiplication of Z1*Z2. That means, assuming 4 cycles latency per FLOP, you'll have a latency of 20 cycles plus latency of reading and writing memory (unless you are able to keep everything in registers -- this requires you to show the surounding code). If you accumulate the values linearily, you'll have a critical path of mulsd - fmaddsd - fmaddsd - fmaddsd - fmaddsd - mulsd (i.e. 24 cycles + read/write)

Now if you are able to change the memory order of Y, it would be beneficial to transpose these matrices, because then you can easily calculate output[0 ~ 3] in parallel (assuming you have AVX), by first broadcast-loading each entry of X and doing packed accumulation.

void function_fma( const double* X1, const double* Y1, const double* X2, const double* Y2, double* output )
{
    // Load X1 and X2 vectors into 10 registers.
    const __m256d x1_0 = _mm256_broadcast_sd( X1 );
    const __m256d x1_1 = _mm256_broadcast_sd( X1 + 1 );
    const __m256d x1_2 = _mm256_broadcast_sd( X1 + 2 );
    const __m256d x1_3 = _mm256_broadcast_sd( X1 + 3 );
    const __m256d x1_4 = _mm256_broadcast_sd( X1 + 4 );

    const __m256d x2_0 = _mm256_broadcast_sd( X2 );
    const __m256d x2_1 = _mm256_broadcast_sd( X2 + 1 );
    const __m256d x2_2 = _mm256_broadcast_sd( X2 + 2 );
    const __m256d x2_3 = _mm256_broadcast_sd( X2 + 3 );
    const __m256d x2_4 = _mm256_broadcast_sd( X2 + 4 );

    // first four values:
    {
        // Multiply column 0
        __m256d z1 = _mm256_mul_pd( x1_0, _mm256_loadu_pd( Y1 ) );
        __m256d z2 = _mm256_mul_pd( x2_0, _mm256_loadu_pd( Y2 ) );

        // Multiply + accumulate column 1 and column 2
        z1 = _mm256_fmadd_pd( x1_1, _mm256_loadu_pd( Y1 + 5 ), z1 );
        z2 = _mm256_fmadd_pd( x2_1, _mm256_loadu_pd( Y2 + 5 ), z2 );
        z1 = _mm256_fmadd_pd( x1_2, _mm256_loadu_pd( Y1 + 10 ), z1 );
        z2 = _mm256_fmadd_pd( x2_2, _mm256_loadu_pd( Y2 + 10 ), z2 );

        // Multiply column 3
        __m256d z1_ = _mm256_mul_pd( x1_3, _mm256_loadu_pd( Y1 + 15 ) );
        __m256d z2_ = _mm256_mul_pd( x2_3, _mm256_loadu_pd( Y2 + 15 ) );

        // Multiply + accumulate column 4
        z1_ = _mm256_fmadd_pd( x1_4, _mm256_loadu_pd( Y1 + 20 ), z1_ );
        z2_ = _mm256_fmadd_pd( x2_4, _mm256_loadu_pd( Y2 + 20 ), z2_ );

        // Add both partial sum
        z1 = _mm256_add_pd( z1, z1_ );
        z2 = _mm256_add_pd( z2, z2_ );

        // Multiply and store result
        _mm256_store_pd(output, _mm256_mul_pd(z1, z2));
    }
    // last value:
    {
        // Multiply column 0
        __m128d z1 = _mm_mul_sd( _mm256_castpd256_pd128(x1_0), _mm_load_sd( Y1 + 4) );
        __m128d z2 = _mm_mul_sd( _mm256_castpd256_pd128(x2_0), _mm_load_sd( Y2 + 4) );

        // Multiply + accumulate column 1 and column 2
        z1 = _mm_fmadd_sd( _mm256_castpd256_pd128(x1_1), _mm_load_sd( Y1 + 9 ), z1 );
        z2 = _mm_fmadd_sd( _mm256_castpd256_pd128(x2_1), _mm_load_sd( Y2 + 9 ), z2 );
        z1 = _mm_fmadd_sd( _mm256_castpd256_pd128(x1_2), _mm_load_sd( Y1 + 14 ), z1 );
        z2 = _mm_fmadd_sd( _mm256_castpd256_pd128(x2_2), _mm_load_sd( Y2 + 14 ), z2 );

        // Multiply column 3
        __m128d z1_ = _mm_mul_sd( _mm256_castpd256_pd128(x1_3), _mm_load_sd( Y1 + 19 ) );
        __m128d z2_ = _mm_mul_sd( _mm256_castpd256_pd128(x2_3), _mm_load_sd( Y2 + 19 ) );

        // Multiply + accumulate column 4
        z1_ = _mm_fmadd_sd( _mm256_castpd256_pd128(x1_4), _mm_load_sd( Y1 + 24 ), z1_ );
        z2_ = _mm_fmadd_sd( _mm256_castpd256_pd128(x2_4), _mm_load_sd( Y2 + 24 ), z2_ );

        // Add both partial sum
        z1 = _mm_add_sd( z1, z1_ );
        z2 = _mm_add_sd( z2, z2_ );

        // Multiply and store result
        _mm_store_sd(output+4, _mm_mul_sd(z1, z2));
    }
}

If you don't have FMA, you can replace these by multiplication and addition (this will not change the latency a lot, since only the additions are in the critical path -- throughput could be about 50% worse, of course). Also, if you don't have AVX, the first four values can be calculated by doing two times two values.

You can at least process 2 elements at a time by loading the lower and upper half registers separately. Unrolling i by two may give a small edge...

The __restrict keyword, if applicable, allows the five constant coefficients X1[0..4], X2[0..4] to be preloaded. If X1 or X2 partially alias output, it's better to let the compiler know it (by using the same array). In this way, as the complete function is unrolled, the compiler will not reload any element unnecessarily.

typedef double __attribute__((vector_size(16))) f2;

void function2(double *X1, double *Y1, double *X2, double *Y2, double *__restrict output) {
    double Z1, Z2;
    int i = 0, j, k;
    for (j = 0; j < 5; j++) { // sweep Y
        f2 Z12 = {0.0, 0.0};
        for (k = 0; k < 5; k++, i++) { 
            f2 Y12 = {Y1[i], Y2[i]};
            f2 X12 = {X1[k], X2[k]};
            Z12 += X12 * Y12;
        }
        output[j] = Z12[0]*Z12[1];
    }
}

If possible, please consider interleaving Y1Y2, X1X3:

void function2(f2 const *X12, f2 const *Y12, double *output) {
    int i = 0, j, k;
    for (j = 0; j < 5; j++) { // sweep Y
        f2 Z12 = X12[0] * Y12[0]; 
        for (k = 1; k < 5; k++, i++) { 
            Z12 += X12[k] * Y12[k];
        }
        output[j] = Z12[0]*Z12[1]; // possibly [j * 2]?
    }
}

Slightly better performance is probably available through intrinsics, however, this answer emphasises auto-vectorisation.

Related