使用 SSE Intrinsics 在浮点 x、y、z 数组上向量化循环计算长度和差异

Vectorizing a loop over float x,y,z arrays calculating length and differences using SSE Intrinsics

我正在尝试将我拥有的循环转换为 SSE 内在函数。我似乎取得了相当大的进步,我的意思是它的方向是正确的,但是我似乎在某处做了一些翻译错误,因为我没有得到相同的 "correct" 答案,这是非-代码。

我以 4 倍展开的原始循环如下所示:

int unroll_n = (N/4)*4;

for (int j = 0; j < unroll_n; j++) {
        for (int i = 0; i < unroll_n; i+=4) {
            float rx = x[j] - x[i];
            float ry = y[j] - y[i];
            float rz = z[j] - z[i];
            float r2 = rx*rx + ry*ry + rz*rz + eps;
            float r2inv = 1.0f / sqrt(r2);
            float r6inv = r2inv * r2inv * r2inv;
            float s = m[j] * r6inv;
            ax[i] += s * rx;
            ay[i] += s * ry;
            az[i] += s * rz;
            //u
            rx = x[j] - x[i+1];
             ry = y[j] - y[i+1];
             rz = z[j] - z[i+1];
             r2 = rx*rx + ry*ry + rz*rz + eps;
             r2inv = 1.0f / sqrt(r2);
             r6inv = r2inv * r2inv * r2inv;
             s = m[j] * r6inv;
            ax[i+1] += s * rx;
            ay[i+1] += s * ry;
            az[i+1] += s * rz;
            //unroll i 3
             rx = x[j] - x[i+2];
             ry = y[j] - y[i+2];
             rz = z[j] - z[i+2];
             r2 = rx*rx + ry*ry + rz*rz + eps;
             r2inv = 1.0f / sqrt(r2);
             r6inv = r2inv * r2inv * r2inv;
             s = m[j] * r6inv;
            ax[i+2] += s * rx;
            ay[i+2] += s * ry;
            az[i+2] += s * rz;
            //unroll i 4
             rx = x[j] - x[i+3];
             ry = y[j] - y[i+3];
             rz = z[j] - z[i+3];
             r2 = rx*rx + ry*ry + rz*rz + eps;
             r2inv = 1.0f / sqrt(r2);
             r6inv = r2inv * r2inv * r2inv;
             s = m[j] * r6inv;
            ax[i+3] += s * rx;
            ay[i+3] += s * ry;
            az[i+3] += s * rz;
    }
}

然后我基本上逐行处理顶部部分并将其转换为 SSE 内在函数。代码如下。我不完全确定是否需要前三行,但我知道我的数据需要 16 位对齐才能正确和最佳地工作。

float *x = malloc(sizeof(float) * N);
float *y = malloc(sizeof(float) * N);
float *z = malloc(sizeof(float) * N); 

for (int j = 0; j < N; j++) {
    for (int i = 0; i < N; i+=4) {
        __m128 xj_v = _mm_set1_ps(x[j]);
        __m128 xi_v = _mm_load_ps(&x[i]);
        __m128 rx_v = _mm_sub_ps(xj_v, xi_v);

        __m128 yj_v = _mm_set1_ps(y[j]);
        __m128 yi_v = _mm_load_ps(&y[i]);
        __m128 ry_v = _mm_sub_ps(yj_v, yi_v);

        __m128 zj_v = _mm_set1_ps(z[j]);
        __m128 zi_v = _mm_load_ps(&z[i]);
        __m128 rz_v = _mm_sub_ps(zj_v, zi_v);

    __m128 r2_v = _mm_mul_ps(rx_v, rx_v) + _mm_mul_ps(ry_v, ry_v) + _mm_mul_ps(rz_v, rz_v) + _mm_set1_ps(eps);

    __m128 r2inv_v = _mm_div_ps(_mm_set1_ps(1.0f),_mm_sqrt_ps(r2_v));

    __m128 r6inv_1v = _mm_mul_ps(r2inv_v, r2inv_v);
    __m128 r6inv_v = _mm_mul_ps(r6inv_1v, r2inv_v);

    __m128 mj_v = _mm_set1_ps(m[j]);
    __m128 s_v = _mm_mul_ps(mj_v, r6inv_v);

    __m128 axi_v = _mm_load_ps(&ax[i]);
    __m128 ayi_v = _mm_load_ps(&ay[i]);
    __m128 azi_v = _mm_load_ps(&az[i]);

    __m128 srx_v = _mm_mul_ps(s_v, rx_v);
    __m128 sry_v = _mm_mul_ps(s_v, ry_v);
    __m128 srz_v = _mm_mul_ps(s_v, rz_v);

    axi_v = _mm_add_ps(axi_v, srx_v);
    ayi_v = _mm_add_ps(ayi_v, srx_v);
    azi_v = _mm_add_ps(azi_v, srx_v);

    _mm_store_ps(ax, axi_v);
    _mm_store_ps(ay, ayi_v);
    _mm_store_ps(az, azi_v);
    }
}

我觉得主要观点是正确的,但是由于结果答案不正确,某处有 an/some 个错误。

我认为你唯一的错误是简单的拼写错误,而不是逻辑错误,见下文。

你不能只使用 clang 的自动矢量化吗?或者您是否需要为此代码使用 gcc?自动矢量化可让您从同一来源制作 SSE、AVX 和(将来)AVX512 版本,无需修改。不幸的是,内部函数不能扩展到不同的向量大小。

根据你对矢量化的开始,我做了一个优化版本。您应该尝试一下,我很想知道它是否比您修复错误的版本或 clang 的自动矢量化版本更快。 :)


这看起来不对:

_mm_store_ps(ax, axi_v);
_mm_store_ps(ay, ayi_v);
_mm_store_ps(az, azi_v);

您从 ax[i] 加载,但现在正在存储到 ax[0]


此外,clang's unused-variable warning 发现了这个错误:

axi_v = _mm_add_ps(axi_v, srx_v);
ayi_v = _mm_add_ps(ayi_v, srx_v);  // should be sry_v
azi_v = _mm_add_ps(azi_v, srx_v);  // should be srz_v

就像我在回答您之前的问题时所说的那样,您可能应该互换循环,所以相同的 ax[i+0..3]、ay[i+0..3] 和 az[i +0..3] 被使用,避免 load/store.

此外,如果您不打算使用 ,您应该使用我在上一个回答中指出的转换:divide m[j] sqrt(k2) ^3。 div使用 divps 乘以 1.0f,然后再乘以一次是没有意义的。

rsqrt 可能 实际上不是胜利,因为总 uop 吞吐量可能比 div / sqrt 吞吐量或延迟更成为瓶颈。 比 sqrtps + divps 多得多。 rsqrt 对 AVX 256b 向量更有帮助,因为 divide / sqrt 单位在 SnB 到 Broadwell 上不是全角。 Skylake 有 12c 延迟 sqrtps ymm,与 xmm 相同,但 xmm 的吞吐量仍然更好(每 3c 一个而不是每 6c 一个)。

。 (当然,只有 clang 使用打包版本。)


如果不交换循环,则应手动将仅依赖于 j 的所有内容提升到内部循环之外。编译器通常擅长于此,但让源代码反映您期望编译器能够执行的操作似乎仍然是一个很好的 idea。这有助于“看到”循环实际在做什么。


这是一个对您的原始版本进行了一些优化的版本:

  • 为了让 gcc/clang 将 mul/adds 融合到 FMA 中,我使用了 -ffp-contract=fast。这会在不使用 -ffast-math 的情况下获取高吞吐量的 FMA 指令。 (三个独立的累加器有很多并行性,因此与 addps 相比,FMA 的延迟增加根本不会造成伤害。我预计 port0/1 吞吐量是这里的限制因素。)我 ,但似乎没有 -ffast-math

  • 注意 v3/2 = sqrt(v)3 = sqrt(v)*v。这具有更低的延迟和更少的指令。

  • 交换了循环,并在内部循环中使用广播负载来改善局部性(使用 AVX 将带宽要求减少 4 或 8)。内部循环的每次迭代仅从四个源数组中的每一个读取 4B 的新数据。 (x、y、z 和 m)。所以它在热的时候大量使用每个缓存行。

    在内循环中使用广播负载也意味着我们并行累积 ax[i + 0..3],避免需要 horizontal sum, which takes extra code. (See a previous version of this answer for code with the loops interchanged, but that used vector loads in the inner loop, with stride = 16B。)

它编译得很好 for Haswell with gcc, using FMA。 (不过仍然只有 128b 向量大小,这与 clang 的自动向量化 256b 版本不同)。内部循环只有 20 条指令,其中只有 13 条是 FPU ALU 指令,需要英特尔 SnB 系列上的端口 0/1。即使使用基线 SSE2,它也能生成不错的代码:没有 FMA,并且需要广播负载的 shufps,但它们不会与 add/mul.

竞争执行单元
#include <immintrin.h>

void ffunc_sse128(float *restrict ax, float *restrict ay, float *restrict az,
           const float *x, const float *y, const float *z,
           int N, float eps, const float *restrict m)
{
  for (int i = 0; i < N; i+=4) {
    __m128 xi_v = _mm_load_ps(&x[i]);
    __m128 yi_v = _mm_load_ps(&y[i]);
    __m128 zi_v = _mm_load_ps(&z[i]);

    // vector accumulators for ax[i + 0..3] etc.
    __m128 axi_v = _mm_setzero_ps();
    __m128 ayi_v = _mm_setzero_ps();
    __m128 azi_v = _mm_setzero_ps();
    
    // AVX broadcast-loads are as cheap as normal loads,
    // and data-reuse meant that stand-alone load instructions were used anyway.
    // so we're not even losing out on folding loads into other insns
    // An inner-loop stride of only 4B is a huge win if memory / cache bandwidth is a bottleneck
    // even without AVX, the shufps instructions are cheap,
    // and don't compete with add/mul for execution units on Intel
    
    for (int j = 0; j < N; j++) {
      __m128 xj_v = _mm_set1_ps(x[j]);
      __m128 rx_v = _mm_sub_ps(xj_v, xi_v);

      __m128 yj_v = _mm_set1_ps(y[j]);
      __m128 ry_v = _mm_sub_ps(yj_v, yi_v);

      __m128 zj_v = _mm_set1_ps(z[j]);
      __m128 rz_v = _mm_sub_ps(zj_v, zi_v);

      __m128 mj_v = _mm_set1_ps(m[j]);   // do the load early

      // sum of squared differences
      __m128 r2_v = _mm_set1_ps(eps) + rx_v*rx_v + ry_v*ry_v + rz_v*rz_v;   // GNU extension
      /* __m128 r2_v = _mm_add_ps(_mm_set1_ps(eps), _mm_mul_ps(rx_v, rx_v));
      r2_v = _mm_add_ps(r2_v, _mm_mul_ps(ry_v, ry_v));
      r2_v = _mm_add_ps(r2_v, _mm_mul_ps(rz_v, rz_v));
      */

      // rsqrt and a Newton-Raphson iteration might have lower latency
      // but there's enough surrounding instructions and cross-iteration parallelism
      // that the single-uop sqrtps and divps instructions prob. aren't be a bottleneck
      __m128 r2sqrt = _mm_sqrt_ps(r2_v);
      __m128 r6sqrt = _mm_mul_ps(r2_v, r2sqrt);  // v^(3/2) = sqrt(v)^3 = sqrt(v)*v
      __m128 s_v = _mm_div_ps(mj_v, r6sqrt);

      __m128 srx_v = _mm_mul_ps(s_v, rx_v);
      __m128 sry_v = _mm_mul_ps(s_v, ry_v);
      __m128 srz_v = _mm_mul_ps(s_v, rz_v);

      axi_v = _mm_add_ps(axi_v, srx_v);
      ayi_v = _mm_add_ps(ayi_v, sry_v);
      azi_v = _mm_add_ps(azi_v, srz_v);
    }
    _mm_store_ps(&ax[i], axi_v);
    _mm_store_ps(&ay[i], ayi_v);
    _mm_store_ps(&az[i], azi_v);
  }
}

我也试过a version with rcpps,但如果它会更快,我不知道。请注意,对于 -ffast-math,gcc 和 clang 会将 division 转换为 rcpps + 牛顿迭代。 (出于某种原因,他们不会将 1.0f/sqrtf(x) 转换为 rsqrt + Newton,即使在独立函数中也是如此)。 clang 做得更好,在迭代步骤中使用 FMA。

#define USE_RSQRT
#ifndef USE_RSQRT
      // even with -mrecip=vec-sqrt after -ffast-math, this still does sqrt(v)*v, then rcpps
      __m128 r2sqrt = _mm_sqrt_ps(r2_v);
      __m128 r6sqrt = _mm_mul_ps(r2_v, r2sqrt);  // v^(3/2) = sqrt(v)^3 = sqrt(v)*v
      __m128 s_v = _mm_div_ps(mj_v, r6sqrt);
#else
      __m128 r2isqrt = rsqrt_float4_single(r2_v);
      // can't use the sqrt(v)*v trick, unless we either do normal sqrt first then rcpps
      // or rsqrtps and rcpps.  Maybe it's possible to do a Netwon Raphson iteration on that product
      // instead of refining them both separately?
      __m128 r6isqrt = r2isqrt * r2isqrt * r2isqrt;
      __m128 s_v = _mm_mul_ps(mj_v, r6isqrt);
#endif