matrix - 行大小大于向量宽度时的 SIMD 转置

Question

你可以找到很多good answers用于转置符合 SIMD 指令集的自然大小的矩阵，特别是其中一行的大小不超过向量宽度的情况。示例是 SSE 中的 4x4 float 转置，或 AVX/AVX2 中的 4x4 double 或 8x8 float 转置(对于 AVX- 512).

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

Answer 1

如果您所有的矩阵维度都是数据包大小的倍数，您可以逐 block 执行操作并根据需要交换 block 。使用 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 矩阵，你可以利用它来处理最后的行/列(并且可能还有其他奇数大小的算法)。

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

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