SIMD Поиск максимального евклидова расстояния до начала координат (c ++)

Я пытаюсь найти самое длинное евклидово расстояние до начала координат для каждой точки. Их позиции хранятся в двух float массивы, обозначающие положение каждой точки на декартовой плоскости. Пытался ускорить работу с помощью SIMD (БЕЗ SVML, так что нет _mm256_hypot_ps), однако это все еще казалось немного медленным. Цикл кода через массив размером блока 8, на каждой итерации обновляет максимальный квадрат евклидова расстояния и, наконец, возвращает sqrt максимального элемента.

#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);
}

Функция будет вызываться много раз. Хотелось бы еще больше ускорить.

Самая важная часть этого кода, векторный цикл, поддерживается зависимостью, переносимой циклом через vmaxps. vmaxps Например, Skylake занимает 4 цикла (другие современные процессоры Intel аналогичны, но AMD Ryzen существенно отличается), ограничивая цикл выполнением одной итерации за 4 цикла (в устойчивом состоянии), что оставляет большую доступную пропускную способность неиспользованной. Оптимистично с точки зрения пропускной способности ожидается 1 итерация за цикл, что в 4 раза больше. Поэтому разверните цикл на 4, используя несколько аккумуляторов:

  // 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));

Поскольку общий коэффициент развертывания настолько высок, становится более интересным попытаться устранить «скалярный эпилог» (скалярный цикл после основного цикла), например, используя технику «шаг назад от конца», в которой когда i + 32 < len становится false, i или указатели x и y настраиваются таким образом, чтобы итерация основного цикла (или его дубликата) точно обрабатывала оставшуюся часть данных, не считывая за пределами конца массивов. Эта последняя итерация будет частично перекрываться с предпоследней итерацией, но в данном случае это нормально, потому что в любом случае будет взят максимум всех результатов, а наличие некоторых дубликатов не повлияет на этот максимум.