最佳的SIMD算法用于旋转或转置一个数组

Question

我正在处理一个数据结构，其中我有一个包含16个uint64的数组。它们在内存中的布局如下（每个表示单个int64）：

A0 A1 A2 A3 B0 B1 B2 B3 C0 C1 C2 C3 D0 D1 D2 D3

期望的结果是将数组转置为以下形式：

A0 B0 C0 D0 A1 B1 C1 D1 A2 B2 C2 D2 A3 B3 C3 D3

将数组旋转90度也是未来循环的可接受解决方案：

D0 C0 B0 A0 D1 C1 B1 A1 D2 C2 B2 A2 D3 C3 B3 A3

我需要这个操作，以便在稍后的时间中快速处理箭头（使用另一个SIMD trip逐个遍历，每次4个）。

到目前为止，我已经尝试通过加载A的4 x 64位向量、按位掩码和重新排列元素并与B等进行或运算等来“混合”数据，然后为C等重复这个过程…不幸的是，这是数组中4个元素的每个段的5 x 4 SIMD指令（一次加载、一次掩码、一次洗牌、一次与下一个元素的或运算和最后一次存储）。看起来我应该能做得更好。

我拥有AVX2，并且正在使用clang进行编译。

- Thomas Kejser

---

“C1 C1”是打字错误吗？请展示正确的输出。 - 2501

抱歉，打错字了...是的，我想要转置矩阵（将其旋转90度）。 - Thomas Kejser

让我看看我是否理解了你的问题。你想转置一个4x4的uint64矩阵？ - Z boson

@Zboson：是的，那是另一种看待它的方式。 - Thomas Kejser

转置并不等同于将内容旋转90度 - 根据你的例子，看起来你想要的是前者而不是后者。 - Paul R

@PaulR：你是对的。实际上，我可以使用90度旋转或转置。我会更新问题。 - Thomas Kejser

Answer 1

uint64_t A[16] = {0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15};
__m256i row0 = _mm256_loadu_si256((__m256i*)&A[ 0]); //0 1 2 3
__m256i row1 = _mm256_loadu_si256((__m256i*)&A[ 4]); //4 5 6 7
__m256i row2 = _mm256_loadu_si256((__m256i*)&A[ 8]); //8 9 a b
__m256i row3 = _mm256_loadu_si256((__m256i*)&A[12]); //c d e f

我现在没有硬件来测试这个，但类似以下的代码应该可以满足您的需求。

__m256i tmp3, tmp2, tmp1, tmp0;
tmp0 = _mm256_unpacklo_epi64(row0, row1);            //0 4 2 6
tmp1 = _mm256_unpackhi_epi64(row0, row1);            //1 5 3 7
tmp2 = _mm256_unpacklo_epi64(row2, row3);            //8 c a e
tmp3 = _mm256_unpackhi_epi64(row2, row3);            //9 d b f
//now select the appropriate 128-bit lanes
row0 = _mm256_permute2x128_si256(tmp0, tmp2, 0x20);  //0 4 8 c
row1 = _mm256_permute2x128_si256(tmp1, tmp3, 0x20);  //1 5 9 d
row2 = _mm256_permute2x128_si256(tmp0, tmp2, 0x31);  //2 6 a e
row3 = _mm256_permute2x128_si256(tmp1, tmp3, 0x31);  //3 7 b f

这

__m256i _mm256_permute2x128_si256 (__m256i a, __m256i b, const int imm)

intrinsic是从两个来源中选择128位的通道。您可以在Intel Intrinsic Guide中了解它。有一个版本_mm256_permute2f128_si256，它只需要AVX并作用于浮点数域。我使用它来检查我是否使用了正确的控制字。

- Z boson

---

+1：不错的转置操作 - 我在代码和注释中修复了一些小错误，现在已经在 Haswell CPU 上进行了测试验证。 - Paul R

@Zboson：这是一个很棒的解决方案。只需要8个指令！我想知道是否可以通过90度旋转（也是目标数组的可接受布局）来减少指令数。 - Thomas Kejser

Answer 2

另一种方法是使用gather指令，您可以直接加载转置后的矩阵。以下五行代码在i7-Haswell上使用gcc没有问题。

  int32_t stride = 4 * sizeof(A[0]);
  __m128i r128_gather_idx = _mm_set_epi32(3 * stride, 2 * stride, 1 * stride, 0 * stride);
  __m256i row0 = _mm256_i32gather_epi64(reinterpret_cast<long long const *>(&A[ 0]), r128_gather_idx, 1);
  __m256i row1 = _mm256_i32gather_epi64(reinterpret_cast<long long const *>(&A[ 1]), r128_gather_idx, 1);
  __m256i row2 = _mm256_i32gather_epi64(reinterpret_cast<long long const *>(&A[ 2]), r128_gather_idx, 1);

- user3636086

---

在Haswell架构上，gather提供了功能，但性能并不是很好（当然，这可能会在未来的µarches上改变）。基本上，任何时候你可以使用固定置换来执行相同的操作，你都应该这样做。 - Stephen Canon

我发现聚集操作比固定排列操作慢了大约两倍。所以@Zbosons的答案是最快的。虽然这个完整性很好。 - Thomas Kejser

我没有想到使用gather。这是一个有趣的想法。 - Z boson