c++ - 如何优化 SIMD 转置函数(8x4 => 4x8)？

Question

我需要使用 AVX 优化 8x4 和 4x8 浮点矩阵的转置。我使用 Agner Fog 的 vector class library .

青色任务 - 构建 BVH 并求和最小值-最大值。转置用于每个循环的最后阶段(它们也通过多线程进行了优化，但任务确实很多)。

代码现在看起来像:

void transpose(register Vec4f (&fin)[8], register Vec8f (&mat)[4]) {
    for (int i = 0;i < 8;i++) {
        fin[i] = lookup<28>(Vec4i(0, 8, 16, 24) + i, (float *)mat);
    }
}

需要优化的变体。如何为 SIMD 优化此功能？

---

我最近用 vector 类编写了自己的转置变体(4x8 和 8x4)。版本 1.0。

void transpose(register Vec4f(&fin)[8], register Vec8f(&mat)[4]) {
    register Vec8f a00 = blend8f<0, 8, 1, 9, 2, 10, 3, 11>(mat[0], mat[1]);
    register Vec8f a10 = blend8f<0, 8, 1, 9, 2, 10, 3, 11>(mat[2], mat[3]);
    register Vec8f a01 = blend8f<4, 12, 5, 13, 6, 14, 7, 15>(mat[0], mat[1]);
    register Vec8f a11 = blend8f<4, 12, 5, 13, 6, 14, 7, 15>(mat[2], mat[3]);

    register Vec8f v0_1 = blend8f<0, 1, 8, 9, 2, 3, 10, 11>(a00, a10);
    register Vec8f v2_3 = blend8f<4, 5, 12, 13, 6, 7, 14, 15>(a00, a10);
    register Vec8f v4_5 = blend8f<0, 1, 8, 9, 2, 3, 10, 11>(a01, a11);
    register Vec8f v6_7 = blend8f<4, 5, 12, 13, 6, 7, 14, 15>(a01, a11);

    fin[0] = v0_1.get_low();
    fin[1] = v0_1.get_high();
    fin[2] = v2_3.get_low();
    fin[3] = v2_3.get_high();
    fin[4] = v4_5.get_low();
    fin[5] = v4_5.get_high();
    fin[6] = v6_7.get_low();
    fin[7] = v6_7.get_high();
}

void transpose(register Vec8f(&fin)[4], register Vec4f(&mat)[8]) {
    register Vec8f a0_1 = Vec8f(mat[0], mat[1]);
    register Vec8f a2_3 = Vec8f(mat[2], mat[3]);
    register Vec8f a4_5 = Vec8f(mat[4], mat[5]);
    register Vec8f a6_7 = Vec8f(mat[6], mat[7]);

    register Vec8f a00 = blend8f<0, 4, 8 , 12, 1, 5, 9 , 13>(a0_1, a2_3);
    register Vec8f a10 = blend8f<0, 4, 8 , 12, 1, 5, 9 , 13>(a4_5, a6_7);
    register Vec8f a01 = blend8f<2, 6, 10, 14, 3, 7, 11, 15>(a0_1, a2_3);
    register Vec8f a11 = blend8f<2, 6, 10, 14, 3, 7, 11, 15>(a4_5, a6_7);

    fin[0] = blend8f<0, 1, 2, 3, 8, 9, 10, 11>(a00, a10);
    fin[1] = blend8f<4, 5, 6, 7, 12, 13, 14, 15>(a00, a10);
    fin[2] = blend8f<0, 1, 2, 3, 8, 9, 10, 11>(a01, a11);
    fin[3] = blend8f<4, 5, 6, 7, 12, 13, 14, 15>(a01, a11);
}

需要 2.0 版。

Answer 1

我没有使用 vectorclass 库的经验，但是通过简要浏览 lookup 模板函数的源代码，您似乎在做一些非常低效的事情。

我在下面提出了一个使用 SSE/AVX 内在函数的简单高效的解决方案。我不知道如何根据 vectorclass 库对其进行完整编码。但是，您可以使用转换运算符从类 Vec4f 和 Vec8f 中提取原始数据作为 __m128 和 __m256。适当的构造函数允许您将原始结果转换回 vector 类。

---

在具有内部函数的纯 SSE 中， header xmmintrin.h 中有一个宏 _MM_TRANSPOSE4_PS。它将 4x4 浮点矩阵与单独的 128 位寄存器中的每一行进行转置。如果您只有 SSE(即没有 AVX)，那么您只需调用该宏两次即可。这是代码:

#define _MM_TRANSPOSE4_PS(row0, row1, row2, row3) {    \
  __m128 tmp3, tmp2, tmp1, tmp0;                      \
  tmp0 = _mm_shuffle_ps(row0, row1, 0x44);            \
  tmp2 = _mm_shuffle_ps(row0, row1, 0xEE);            \
  tmp1 = _mm_shuffle_ps(row2, row3, 0x44);            \
  tmp3 = _mm_shuffle_ps(row2, row3, 0xEE);            \
  row0 = _mm_shuffle_ps(tmp0, tmp1, 0x88);            \
  row1 = _mm_shuffle_ps(tmp0, tmp1, 0xDD);            \
  row2 = _mm_shuffle_ps(tmp2, tmp3, 0x88);            \
  row3 = _mm_shuffle_ps(tmp2, tmp3, 0xDD);            \
}

在 AVX 中，具有 256 位操作数的指令通常只对操作数的两半(称为 channel )执行 SSE 等效操作。内在的 _mm256_shuffle_ps 也不异常(exception):它只是将两个 128 位 channel 洗牌，就像它的 _mm 等效项一样。如果意味着如果我们在宏中将_mm前缀更改为_mm256前缀，它将转置两个4x4矩阵:位于四个256位的较低 channel 的那个寄存器，以及位于四个 256 位寄存器的上 channel 的那个。我们只需将生成的 256 位寄存器分成两半并正确排序即可。

生成的代码如下所示。我已经检查过它是否正常工作。它似乎只有 12 条指令，所以我想它会很快。

void Transpose4x8(__m128 dst[8], __m256 src[4]) {
  __m256 row0 = src[0], row1 = src[1], row2 = src[2], row3 = src[3];
  __m256 tmp3, tmp2, tmp1, tmp0;
  tmp0 = _mm256_shuffle_ps(row0, row1, 0x44);
  tmp2 = _mm256_shuffle_ps(row0, row1, 0xEE);
  tmp1 = _mm256_shuffle_ps(row2, row3, 0x44);
  tmp3 = _mm256_shuffle_ps(row2, row3, 0xEE);
  row0 = _mm256_shuffle_ps(tmp0, tmp1, 0x88);
  row1 = _mm256_shuffle_ps(tmp0, tmp1, 0xDD);
  row2 = _mm256_shuffle_ps(tmp2, tmp3, 0x88);
  row3 = _mm256_shuffle_ps(tmp2, tmp3, 0xDD);
  dst[0] = _mm256_castps256_ps128(row0);
  dst[1] = _mm256_castps256_ps128(row1);
  dst[2] = _mm256_castps256_ps128(row2);
  dst[3] = _mm256_castps256_ps128(row3);
  dst[4] = _mm256_extractf128_ps(row0, 1);
  dst[5] = _mm256_extractf128_ps(row1, 1);
  dst[6] = _mm256_extractf128_ps(row2, 1);
  dst[7] = _mm256_extractf128_ps(row3, 1);
}

---

UPDATE 逆向换位以完全相同的方式完成，只是有些事情以相反的顺序进行。这是代码:

void Transpose8x4(__m256 dst[4], __m128 src[8]) {
  __m256 row0 = _mm256_setr_m128(src[0], src[4]);
  __m256 row1 = _mm256_setr_m128(src[1], src[5]);
  __m256 row2 = _mm256_setr_m128(src[2], src[6]);
  __m256 row3 = _mm256_setr_m128(src[3], src[7]);
  __m256 tmp3, tmp2, tmp1, tmp0;
  tmp0 = _mm256_shuffle_ps(row0, row1, 0x44);
  tmp2 = _mm256_shuffle_ps(row0, row1, 0xEE);
  tmp1 = _mm256_shuffle_ps(row2, row3, 0x44);
  tmp3 = _mm256_shuffle_ps(row2, row3, 0xEE);
  row0 = _mm256_shuffle_ps(tmp0, tmp1, 0x88);
  row1 = _mm256_shuffle_ps(tmp0, tmp1, 0xDD);
  row2 = _mm256_shuffle_ps(tmp2, tmp3, 0x88);
  row3 = _mm256_shuffle_ps(tmp2, tmp3, 0xDD);
  dst[0] = row0; dst[1] = row1; dst[2] = row2; dst[3] = row3;
}