Improving speed of affine transformation of an array using intrinsics - c++

Question

In a performance sensitive code, I have to perform am affine transformation of a vector:
Y=a*X+b
where Y and X are vectors and a and b are scalars.
As a quick-and-dirty way to improve the speed of the computation, I delegated parallelization to openMP
#pragma omp simd directive. Having some spare time, lately I tried to implement it directly using intrinsics, getting more or less the same performance as the omp solution.
Is there a way to beat the OMP vectorization? I can use up AVX2 instructions.
The code below is tested under windows 10, compiled with VS 2019.
#include <iostream>
#include <armadillo>
#include <chrono>
#include <immintrin.h>
///Computes y=alpha*x+beta
inline void SumAndSetOmp(
arma::Col<double>& y /**< Result*/,
const arma::Col<double>& x /**< Input*/,
const double& alpha /**< Coefficient*/,
const double& beta /**< Offset*/)
{
auto* __restrict lhs = y.memptr();
const auto* __restrict add_rhs = x.memptr();
const auto& n = x.n_elem;
#pragma omp simd
for (arma::uword i = 0; i < n; ++i)
{
lhs[i] = add_rhs[i] * alpha + beta;
}
}
inline void SumAndSetSerial(
arma::Col<double>& y /**< Result*/,
const arma::Col<double>& x /**< Input*/,
const double& alpha /**< Coefficient*/,
const double& beta /**< Offset*/)
{
auto* lhs = y.memptr();
const auto* add_rhs = x.memptr();
const auto& n = x.n_elem;
for (arma::uword i = 0; i < n; ++i)
{
lhs[i] = add_rhs[i] * alpha + beta;
}
}
inline void SumAndSetAVX(arma::Col<double>& y /**< Result*/,
const arma::Col<double>& x /**< Input*/,
const double& alpha /**< Coefficient*/,
const double& beta /**< Offset*/)
{
//Allocate coefficients
const auto alphas = _mm256_set1_pd(alpha);
const auto betas = _mm256_set1_pd(beta);
//Extracting memory addresses
auto* __restrict pos_lhs = y.memptr();
const auto* __restrict pos_rhs = x.memptr();
//Computing sizes
const unsigned int length_array = 4;
const unsigned long long n_aligned = x.n_elem / length_array;
const unsigned int remainder = x.n_elem % length_array;
//Performing AVX instruction
for (unsigned long long i = 0; i < n_aligned; i++) {
const __m256d x_avx = _mm256_loadu_pd(pos_rhs);
const __m256d y_avx = _mm256_fmadd_pd(x_avx, alphas, betas);
_mm256_storeu_pd(pos_lhs, y_avx);
pos_rhs += length_array;
pos_lhs += length_array;
}
//Process the rest serially
for (unsigned int i = 0; i < remainder; i++) {
pos_lhs[i] = alpha * pos_rhs[i] + beta;
}
}
enum method
{
serial,
omp,
avx
};
arma::vec perform_test(const arma::vec& x, const method mtd, int trials = 100, const double alpha = 3.0, const double beta = 5.0)
{
arma::Col<double> res(x.n_elem);
const auto beg = std::chrono::steady_clock::now();
switch (mtd) {
case serial:
for (int i = 0; i < trials; i++)
SumAndSetSerial(res, x, alpha, beta);
break;
case omp:
for (int i = 0; i < trials; i++)
SumAndSetOmp(res, x, alpha, beta);
break;
case avx:
for (int i = 0; i < trials; i++)
SumAndSetAVX(res, x, alpha, beta);
break;
}
std::cout << "time:" << std::chrono::duration<double>(std::chrono::steady_clock::now() - beg).count() << "s\n";
return res;
}
//Benchmarking
double test_fun(long long int n,int trials=100, const double alpha = 3.0, const double beta = 5.0)
{
const arma::Col<double> x(n, arma::fill::randn);
const arma::Col<double> reference = alpha*x + beta;
std::cout << "Serial: ";
const auto res_serial = perform_test(x, method::serial, trials, alpha, beta);
std::cout << "OMP: ";
const auto res_omp = perform_test(x, method::omp, trials, alpha, beta);
std::cout << "AVX: ";
const auto res_avx = perform_test(x, method::avx, trials, alpha, beta);
// errors wrt the reference
const double err_serial = arma::max(arma::abs(reference - res_serial));
const double err_avx = arma::max(arma::abs(reference - res_avx));
const double err_omp = arma::max(arma::abs(reference - res_omp));
//Largest error
const double error = std::max(std::max(err_serial, err_avx), err_omp);
if (error> 1e-6)
{
throw std::runtime_error("Something is wrong!");
}
return error;
}
int main()
{
test_fun(10000000);
}

Answer 1

I had a go at improving on your solutions, but I think they are optimal, so you're best sticking with omp.
The following is speculative and goes beyond your question:
If it's an option you could try omp's multithreading, I'm always surprised at how low the number of elements at which you get a reasonable boost is, though 1000 elements with a simple affine transform I expect would be too low. If other parts of your algorithm are parallelisable then this is more likely to be helpful.
If you can afford to change your problem and don't need double precision you could work with floats.