行大小大于向量宽度时的 SIMD 转置答案

【问题标题】：SIMD transpose when row size is greater than vector width行大小大于向量宽度时的 SIMD 转置
【发布时间】：2020-02-15 16:19:47
【问题描述】：

你可以找到很多 good answers 用于转置一个符合 SIMD 指令集的自然大小的矩阵，特别是在一行大小不再比向量宽度。例如，SSE 中的 4x4 float 转置，或 AVX/AVX2 中的 4x4 double 或 8x8 float 转置（对于 AVX-512，再次加倍）。

但是，当矩阵大于那个时，有什么选择？例如，使用 AVX2 的 16x16 float 矩阵？是否可以使用 SIMD shuffle 来加快速度，或者是聚集 + 顺序写入的唯一方法？

【问题讨论】：

标签： matrix transpose simd avx avx2

【解决方案1】：

如果您的所有矩阵维度都是数据包大小的倍数，您可以按块执行操作并根据需要交换块。使用 SSE2 的 4x4 双矩阵示例：

// transpose vectors i0 and i1 and store the result to addresses r0 and r1
void transpose2x2(double *r0, double* r1, __m128d i0, __m128d i1)
{
    __m128d t0 = _mm_unpacklo_pd(i0,i1);
    __m128d t1 = _mm_unpackhi_pd(i0,i1);
    _mm_storeu_pd(r0, t0);
    _mm_storeu_pd(r1, t1);
}


void transpose(double mat[4][4])
{
    // transpose [00]-block in-place
    transpose2x2(mat[0]+0, mat[1]+0,_mm_loadu_pd(mat[0]+0),_mm_loadu_pd(mat[1]+0));

    // load [20]-block
    __m128d t20 = _mm_loadu_pd(mat[2]+0), t30 = _mm_loadu_pd(mat[3]+0);
    // transpose [02]-block and store it to [20] position
    transpose2x2(mat[2]+0,mat[3]+0, _mm_loadu_pd(mat[0]+2),_mm_loadu_pd(mat[1]+2));
    // transpose temp-block and store it to [02] position
    transpose2x2(mat[0]+2,mat[1]+2, t20, t30);

    // transpose [22]-block in-place
    transpose2x2(mat[2]+2, mat[3]+2,_mm_loadu_pd(mat[2]+2),_mm_loadu_pd(mat[3]+2));
}

这应该相对容易扩展到其他方阵、其他标量类型和其他架构。不是数据包大小倍数的矩阵可能更复杂（如果它们足够大，可能值得使用矢量化来完成大部分工作并手动完成最后的行/列）。

对于某些尺寸，例如3x4 或 3x8 矩阵有特殊算法 [1] - 如果您有一个 1003x1003 矩阵，您可以将其用于最后的行/列（并且可能还有其他奇数大小的算法）。

通过一些努力，您也可以为矩形矩阵编写此代码（必须考虑如何避免一次缓存多个块，但这是可能的）。

Godbolt 演示：https://godbolt.org/z/tVk_Bc

[1]https://software.intel.com/en-us/articles/3d-vector-normalization-using-256-bit-intel-advanced-vector-extensions-intel-avx

【讨论】：

谢谢！公平地说，对于大型矩阵（例如 100x10,000），这仍然会比基于标量或基于集合的加载存储显着加速吗？
我猜相对加速应该与大小无关（只要没有像缓存这样的效果生效——我猜一个具有良好缓存局部性的标量算法将击败一个糟糕实现的块转置） .但可以肯定的是，您需要对此进行基准测试（矩形情况可能比我最初想象的要复杂一些，尤其是如果您在编译时不知道大小）。

【解决方案2】：

也许可以将 fortran TRANSPOSE intrisic 与 ISO_C_BINDING 一起使用，并将其与 C 链接作为子例程或函数调用。

TRANSPOSE 在 fortran 中进行了相当优化。

混合语言技能有时对一般知识很有用。我什至将 F90 与 GO 关联起来。

【讨论】：

好的，但是在任何给定的 Fortran 实现上，它最终运行什么 asm 使它运行得很快？如果我们知道这一点，我们也可以在 C 中使用内部函数或手写 asm 来实现它。
gfortran 似乎正在将TRANSPOSE 优化为高效的 SIMD 代码，如果矩阵是正方形并且大小是四/八的倍数（在编译时已知）。对于我查看的所有其他情况，我只发现它生成了一个微不足道的嵌套循环：godbolt.org/z/Uiv1iR——我也没有设法让它在不分配临时的情况下生成就地转置（但这可能是由于我的有限fortran 技能...）
以易于使用的语言使用功能并不少见。如果混合语言可能不是每个人都喜欢，但它是相当普遍的。
对于一个小矩阵，比如 4x4 doubles，调用另一个函数而不是仅仅内联几个 SIMD 指令的开销很大。