打包和解交错两个__m256寄存器

Question

我有一个行的行数组(~20 cols x~1M行),我需要一次从两个__m256寄存器中提取两列.

...a0.........b0......
...a1.........b1......
// ...
...a7.........b7......
// end first __m256

一个天真的方法是这样做

__m256i vindex = _mm256_setr_epi32(
    0,
    1 * stride,
    2 * stride,
    // ...
    7 * stride);
__m256 colA = _mm256_i32gather_ps(baseAddrColA, vindex, sizeof(float));
__m256 colB = _mm256_i32gather_ps(baseAddrColB, vindex, sizeof(float));

但是,我想知道我是否会通过a0, b0, a1, b1, a2, b2, a3, b3在一个中检索gather而a4, b4, ... a7, b7在另一个中获得更好的性能,因为它们在内存中更接近,然后对它们进行反交错.那是:

// __m256   lo = a0 b0 a1 b1 a2 b2 a3 b3 // load proximal elements
// __m256   hi = a4 b4 a5 b5 a6 b6 a7 b7
// __m256 colA = a0 a1 a2 a3 a4 a5 a6 a7 // goal
// __m256 colB = b0 b1 b2 b3 b4 b5 b6 b7

我无法弄清楚如何很好地交错lo和hi.我基本上需要相反的_mm256_unpacklo_ps.我提出的最好的是:

__m256i idxA = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
__m256i idxB = _mm256_setr_epi32(1, 3, 5, 7, 0, 2, 4, 6);

__m256 permLA = _mm256_permutevar8x32_ps(lo, idxA);        // a0 a1 a2 a3 b0 b1 b2 b3
__m256 permHB = _mm256_permutevar8x32_ps(hi, idxB);        // b4 b5 b6 b7 a4 a5 a6 a7
__m256 colA = _mm256_blend_ps(permLA, permHB, 0b11110000); // a0 a1 a2 a3 a4 a5 a6 a7
__m256 colB = _mm256_setr_m128(
                          _mm256_extractf128_ps(permLA, 1), 
                          _mm256_castps256_ps128(permHB)); // b0 b1 b2 b3 b4 b5 b6 b7

这是13个周期.有没有更好的办法？

(据我所知,prefetch已经尽可能地优化了天真的方法,但是缺乏这些知识,我希望能够对第二种方法进行基准测试.如果有人已经知道这将是什么结果,请分享.在去交错方法之上,它比天真的方法慢了约8%.)

编辑即使没有去隔行扫描,"近端"聚集方法比天真的恒定步幅聚集方法慢约6%.我认为这意味着这种访问模式过多地混淆了硬件预取,无法进行有价值的优化.

Answer 1

// __m256   lo = a0 b0 a1 b1 a2 b2 a3 b3 // load proximal elements
// __m256   hi = a4 b4 a5 b5 a6 b6 a7 b7
// __m256 colA = a0 a1 a2 a3 a4 a5 a6 a7 // goal
// __m256 colB = b0 b1 b2 b3 b4 b5 b6 b7

看来我们可以比我原来的答案更快地进行洗牌：

void unpack_cols(__m256i lo, __m256i hi, __m256i& colA, __m256i& colB) {
    const __m256i mask = _mm256_setr_epi32(0, 2, 4, 6, 1, 3, 5, 7);
    // group cols crossing lanes: 
    // a0 a1 a2 a3 b0 b1 b2 b3
    // a4 a5 a6 a7 b4 b5 b6 b7
    auto lo_grouped = _mm256_permutevar8x32_epi32(lo, mask);
    auto hi_grouped = _mm256_permutevar8x32_epi32(hi, mask);

    // swap lanes: 
    // a0 a1 a2 a3 a4 a5 a6 a7
    // b0 b1 b2 b3 b4 b5 b6 b7
    colA = _mm256_permute2x128_si256(lo_grouped, hi_grouped, 0 | (2 << 4));
    colB = _mm256_permute2x128_si256(lo_grouped, hi_grouped, 1 | (3 << 4));
}

虽然这两条指令在 Haswell 上都有 3 个周期的延迟（参见Agner Fog），但它们具有单周期的吞吐量。这意味着它具有4 个周期的吞吐量和8 个周期的延迟。如果您有一个可以保存掩码的备用寄存器，这应该会更好。仅并行执行其中两项即可让您完全隐藏其延迟。请参阅godbolt和reextester。

老答案，留作参考：

进行此随机播放的最快方法如下：

void unpack_cols(__m256i lo, __m256i hi, __m256i& colA, __m256i& colB) {
    // group cols within lanes: 
    // a0 a1 b0 b1 a2 a3 b2 b3
    // a4 a5 b4 b5 a6 a7 b6 b7
    auto lo_shuffled = _mm256_shuffle_epi32(lo, _MM_SHUFFLE(3, 1, 2, 0));
    auto hi_shuffled = _mm256_shuffle_epi32(hi, _MM_SHUFFLE(3, 1, 2, 0));

    // unpack lo + hi a 64 bit
    // a0 a1 a4 a5 a2 a3 a6 a7
    // b0 b1 b4 b5 b2 b3 b6 b7
    auto colA_shuffled = _mm256_unpacklo_epi64(lo_shuffled, hi_shuffled);
    auto colB_shuffled = _mm256_unpackhi_epi64(lo_shuffled, hi_shuffled);

    // swap crossing lanes: 
    // a0 a1 a2 a3 a4 a5 a6 a7
    // b0 b1 b2 b3 b4 b5 b6 b7
    colA = _mm256_permute4x64_epi64(colA_shuffled, _MM_SHUFFLE(3, 1, 2, 0));
    colB = _mm256_permute4x64_epi64(colB_shuffled, _MM_SHUFFLE(3, 1, 2, 0));
}

从 Haswell 开始，吞吐量为 6 个周期（遗憾的是端口 5 上有 6 个指令）。据Agner 称，Fog 的 _mm256_permute4x64_epi64延迟为 3 个周期。这意味着unpack_cols延迟为11 8 个周期。

您可以在godbolt.org上检查代码或在rextester上测试它，它支持 AVX2，但遗憾的是没有像 godbolt 这样的永久链接。

请注意，这也非常接近我收集 64 位整数并需要将高 32 位和低 32 位分开的问题。

请注意，Haswell 中的收集性能确实很差，但根据 Agner Fog Skylake 的说法，Skylake 在这方面做得更好（约 12 个周期吞吐量降至约 5）。仍然围绕这些简单的模式进行洗牌应该仍然比聚集快得多。