SIMD Finding the maximum Euclidean distance to origin (c++)

Question

I am trying to find the longest Euclidean distance to the origin for each point. Their positions are stored in two float arrays, denoting each point's position on the Cartesian plane. I tried to speed up the function using SIMD (WITHOUT SVML, so no _mm256_hypot_ps), however it still seemed a little bit slow. The code loop through the array by a chunk size of 8, on each iteration updates the maximum squared Euclidean distance and finally returns the sqrt of the maximum element.

#include <immintrin.h>
#include <math.h>
#include <stdio.h>
float dist(const float* x, const float* y, unsigned int len) {
  if (len < 16) {
    float d = 0.0f, a;
    unsigned i = 0;
    for (i = 0; i < len; i++) {
      a = x[i] * x[i] + y[i] * y[i];
      if (a > d) d = a;
    }
    return sqrt(d);
  }
  __m256 a = _mm256_loadu_ps(x), b = _mm256_loadu_ps(y);
  a = _mm256_mul_ps(a, a);
  b = _mm256_mul_ps(b, b);
  __m256 max = _mm256_add_ps(a, b);  // x[i]*x[i] + y[i]*y[i]
  unsigned i;
  for (i = 8; i + 8 < len; i += 8) {
    a = _mm256_loadu_ps(x + i);
    a = _mm256_mul_ps(a, a);
    b = _mm256_loadu_ps(y + i);
    b = _mm256_mul_ps(b, b);
    max = _mm256_max_ps(
        max, _mm256_add_ps(a, b));  // max(max[i], x[i]*x[i] + y[i]*y[i])
  }
  float d = 0.0f, t;
  for (; i < len; ++i) {
    t = x[i] * x[i] + y[i] * y[i];
    if (t > d) d = t;
  }
  const __m128 hiQuad = _mm256_extractf128_ps(max, 1); // high 128 bits in max
  const __m128 loQuad = _mm256_castps256_ps128(max); // low 128 bits in max
  const __m128 maxQuad = _mm_max_ps(loQuad, hiQuad);
  const __m128 hiDual = _mm_movehl_ps(maxQuad, maxQuad); 
  const __m128 loDual = maxQuad;
  const __m128 maxDual = _mm_max_ps(loDual, hiDual);
  const __m128 lo = maxDual;
  const __m128 hi = _mm_shuffle_ps(maxDual, maxDual, 0x1);
  const __m128 res = _mm_max_ss(lo, hi);
  t = _mm_cvtss_f32(res);
  if (t > d) d = t;
  return sqrt(d);
}

The function will be called many times. I would like to speed it up further.

Answer 1

The most important part of that code, the vector loop, is held up by the loop-carried dependency through vmaxps. vmaxps on eg Skylake takes 4 cycles (other modern Intel processors are similar, but AMD Ryzen is substantially different), limiting the loop to run one iteration per 4 cycles (in steady state), which leaves a lot of available throughput unused. Optimistically, in terms of throughput the expectation is 1 iteration every cycle, 4 times as much. Therefore, unroll the loop by 4, using multiple accumulators:

  // initialized max0, max1, max2, max3 as appropriate
  for (i = 32; i + 32 < len; i += 32) {
    a = _mm256_loadu_ps(x + i);
    a = _mm256_mul_ps(a, a);
    b = _mm256_loadu_ps(y + i);
    b = _mm256_mul_ps(b, b);
    max0 = _mm256_max_ps(max0, _mm256_add_ps(a, b));
    a = _mm256_loadu_ps(x + i + 8);
    a = _mm256_mul_ps(a, a);
    b = _mm256_loadu_ps(y + i + 8);
    b = _mm256_mul_ps(b, b);
    max1 = _mm256_max_ps(max1, _mm256_add_ps(a, b));
    a = _mm256_loadu_ps(x + i + 16);
    a = _mm256_mul_ps(a, a);
    b = _mm256_loadu_ps(y + i + 16);
    b = _mm256_mul_ps(b, b);
    max2 = _mm256_max_ps(max2, _mm256_add_ps(a, b));
    a = _mm256_loadu_ps(x + i + 24);
    a = _mm256_mul_ps(a, a);
    b = _mm256_loadu_ps(y + i + 24);
    b = _mm256_mul_ps(b, b);
    max3 = _mm256_max_ps(max3, _mm256_add_ps(a, b));
  }
  max0 = _mm256_max_ps(_mm256_max_ps(max0, max1), _mm256_max_ps(max2, max3));

Since the total unroll factor is so high, it becomes more interesting to try to eliminate the "scalar epilog" (the scalar loop after the main loop), for example by using the "step back from the end"-technique in which when i + 32 < len becomes false, i or the pointers x and y are adjusted so that an iteration of the main loop (or a duplicate of it) exactly handles the remaining part of the data, without reading beyond the end of the arrays. That last iteration will overlap somewhat with the second-to-last iteration, but that is OK in this case because the maximum of all results will be taken anyway, and having some duplicates won't affect that maximum.