C语言 使用 SSE 内联函数对浮点 x,y,z 数组上的循环进行矢量化,计算长度和差异



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

我以 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);
    }
}

我觉得主要思想是正确的,但是由于结果答案不正确,因此某处存在/一些错误。

我认为您唯一的错误是简单的错别字,而不是逻辑错误,请参见下文。

你不能只使用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 的未使用变量警告发现了这个错误:

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],避免加载/存储。

另外,如果你不打算使用rsqrtps + Newton-Raphson,你应该使用我在上一个答案中指出的变换:将m[j]除以sqrt(k2) ^3。 用divps除以1.0f是没有意义的,然后只乘以一次。

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

Clang和GCC在使用-ffast-math编译代码时都使用rsqrtps/rsqrtss。 (当然,只使用打包版本。

<小时 />

如果你不互换循环,你应该手动将所有只依赖于j的东西从内循环中提升出来。 编译器通常擅长于此,但让源代码反映您期望编译器能够执行的操作似乎仍然是一个好主意。 这有助于"看到"循环实际在做什么。

<小时 />

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

  • 为了让 gcc/clang 将 mul/add 融合到 FMA 中,我使用了 -ffp-contract=fast . 这获得了 FMA 指令以实现高吞吐量,而无需使用-ffast-math。 (三个独立的累加器有很多并行性,因此与addps相比,FMA增加的延迟根本不会受到伤害。 我希望端口 0/1 吞吐量是这里的限制因素。 我以为 gcc 会自动执行此操作,但似乎没有-ffast-math在这里它不会.

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

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

    在内部循环中使用广播负载也意味着我们并行累积ax[i + 0..3],避免了需要额外代码的水平总和。 (有关循环互换的代码,请参阅此答案的先前版本,但在内部循环中使用向量加载,步幅 = 16B。

它使用 FMA 为 gcc 的 Haswell 编译得很好。 (虽然仍然只有 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);
  }
}

我也尝试了一个带有rcpps的版本,但是IDK是否会更快。 注意,使用 -ffast-math 时,gcc 和 clang 会将除法转换为 rcpps + 牛顿迭代。 (由于某种原因,它们不会将1.0f/sqrtf(x)转换为 rsqrt + 牛顿,即使在独立函数中也是如此)。 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

最新更新