5
\$\begingroup\$

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.

\$\endgroup\$

1 Answer 1

3
\$\begingroup\$

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.

\$\endgroup\$
4
  • \$\begingroup\$ Could I optimise it even further? \$\endgroup\$
    – PunnyBunny
    Jan 30, 2021 at 7:34
  • \$\begingroup\$ @FunnyBunnyleaksmemory yes but not much, perhaps more could be done in larger context \$\endgroup\$
    – harold
    Jan 30, 2021 at 7:43
  • \$\begingroup\$ Is there a faster algorithm? \$\endgroup\$
    – PunnyBunny
    Jan 30, 2021 at 8:41
  • \$\begingroup\$ @FunnyBunnyleaksmemory not as a drop-in replacement (if fewer than len elements are looked at, one of the skipped elements might actually be the point that's furthest from the origin) but maybe in a broader context \$\endgroup\$
    – harold
    Jan 30, 2021 at 12:55

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.