转置 8x8 64 位矩阵

Transpose 8x8 64-bits matrix

针对 AVX2,转置包含 64 位整数(或双精度)的 8x8 矩阵的最快方法是什么?

我搜索了这个站点,发现了几种进行 8x8 转置的方法,但主要用于 32 位浮点数。所以我问的主要是因为我不确定使这些算法快速转换的原理是否很容易转化为 64 位,其次,显然 AVX2 只有 16 个寄存器,所以只加载所有值会占用所有寄存器。

一种方法是调用 2x2 ,但我想知道这是否是最优的:

  #define _MM_TRANSPOSE4_PD(row0,row1,row2,row3)                \
        {                                                       \
            __m256d tmp3, tmp2, tmp1, tmp0;                     \
                                                                \
            tmp0 = _mm256_shuffle_pd((row0),(row1), 0x0);       \
            tmp2 = _mm256_shuffle_pd((row0),(row1), 0xF);       \
            tmp1 = _mm256_shuffle_pd((row2),(row3), 0x0);       \
            tmp3 = _mm256_shuffle_pd((row2),(row3), 0xF);       \
                                                                \
            (row0) = _mm256_permute2f128_pd(tmp0, tmp1, 0x20);  \
            (row1) = _mm256_permute2f128_pd(tmp2, tmp3, 0x20);  \
            (row2) = _mm256_permute2f128_pd(tmp0, tmp1, 0x31);  \
            (row3) = _mm256_permute2f128_pd(tmp2, tmp3, 0x31);  \
        }

仍然假设 AVX2,转置 double[8][8]int64_t[8][8] 原则上大致相同?

PS:只是好奇,拥有 AVX512 会大大改变事情,对吗?

经过一些思考和评论中的讨论,我认为这是最有效的版本,至少当源数据和目标数据位于 RAM 中时。不需要AVX2,AVX1就够了

主要思想,与存储相比,现代 CPU 可以执行两倍的加载微操作,并且在许多 CPU 上,使用 vinsertf128 将内容加载到向量的较高一半与常规 16 字节加载的成本相同.与您的宏相比,此版本不再需要这些相对昂贵的(大多数 CPU 上有 3 个延迟周期)vperm2f128 随机播放。

struct Matrix4x4
{
    __m256d r0, r1, r2, r3;
};

inline void loadTransposed( Matrix4x4& mat, const double* rsi, size_t stride = 8 )
{
    // Load top half of the matrix into low half of 4 registers
    __m256d t0 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi ) );     // 00, 01
    __m256d t1 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi + 2 ) ); // 02, 03
    rsi += stride;
    __m256d t2 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi ) );     // 10, 11
    __m256d t3 = _mm256_castpd128_pd256( _mm_loadu_pd( rsi + 2 ) ); // 12, 13
    rsi += stride;
    // Load bottom half of the matrix into high half of these registers
    t0 = _mm256_insertf128_pd( t0, _mm_loadu_pd( rsi ), 1 );    // 00, 01, 20, 21
    t1 = _mm256_insertf128_pd( t1, _mm_loadu_pd( rsi + 2 ), 1 );// 02, 03, 22, 23
    rsi += stride;
    t2 = _mm256_insertf128_pd( t2, _mm_loadu_pd( rsi ), 1 );    // 10, 11, 30, 31
    t3 = _mm256_insertf128_pd( t3, _mm_loadu_pd( rsi + 2 ), 1 );// 12, 13, 32, 33

    // Transpose 2x2 blocks in registers.
    // Due to the tricky way we loaded stuff, that's enough to transpose the complete 4x4 matrix.
    mat.r0 = _mm256_unpacklo_pd( t0, t2 ); // 00, 10, 20, 30
    mat.r1 = _mm256_unpackhi_pd( t0, t2 ); // 01, 11, 21, 31
    mat.r2 = _mm256_unpacklo_pd( t1, t3 ); // 02, 12, 22, 32
    mat.r3 = _mm256_unpackhi_pd( t1, t3 ); // 03, 13, 23, 33
}

inline void store( const Matrix4x4& mat, double* rdi, size_t stride = 8 )
{
    _mm256_storeu_pd( rdi, mat.r0 );
    _mm256_storeu_pd( rdi + stride, mat.r1 );
    _mm256_storeu_pd( rdi + stride * 2, mat.r2 );
    _mm256_storeu_pd( rdi + stride * 3, mat.r3 );
}

// Transpose 8x8 matrix of double values
void transpose8x8( double* rdi, const double* rsi )
{
    Matrix4x4 block;
    // Top-left corner
    loadTransposed( block, rsi );
    store( block, rdi );

#if 1
    // Using another instance of the block to support in-place transpose
    Matrix4x4 block2;
    loadTransposed( block, rsi + 4 );       // top right block
    loadTransposed( block2, rsi + 8 * 4 ); // bottom left block

    store( block2, rdi + 4 );
    store( block, rdi + 8 * 4 );
#else
    // Flip the #if if you can guarantee ( rsi != rdi )
    // Performance is about the same, but this version uses 4 less vector registers,
    // slightly more efficient when some registers need to be backed up / restored.
    assert( rsi != rdi );
    loadTransposed( block, rsi + 4 );
    store( block, rdi + 8 * 4 );

    loadTransposed( block, rsi + 8 * 4 );
    store( block, rdi + 4 );
#endif
    // Bottom-right corner
    loadTransposed( block, rsi + 8 * 4 + 4 );
    store( block, rdi + 8 * 4 + 4 );
}

为了完整起见,这里有一个使用与您的宏非常相似的代码的版本,加载次数是原来的两倍,存储次数相同,随机播放次数更多。尚未进行基准测试,但我预计它会稍微慢一些。

struct Matrix4x4
{
    __m256d r0, r1, r2, r3;
};

inline void load( Matrix4x4& mat, const double* rsi, size_t stride = 8 )
{
    mat.r0 = _mm256_loadu_pd( rsi );
    mat.r1 = _mm256_loadu_pd( rsi + stride );
    mat.r2 = _mm256_loadu_pd( rsi + stride * 2 );
    mat.r3 = _mm256_loadu_pd( rsi + stride * 3 );
}

inline void store( const Matrix4x4& mat, double* rdi, size_t stride = 8 )
{
    _mm256_storeu_pd( rdi, mat.r0 );
    _mm256_storeu_pd( rdi + stride, mat.r1 );
    _mm256_storeu_pd( rdi + stride * 2, mat.r2 );
    _mm256_storeu_pd( rdi + stride * 3, mat.r3 );
}

inline void transpose( Matrix4x4& m4 )
{
    // These unpack instructions transpose lanes within 2x2 blocks of the matrix
    const __m256d t0 = _mm256_unpacklo_pd( m4.r0, m4.r1 );
    const __m256d t1 = _mm256_unpacklo_pd( m4.r2, m4.r3 );
    const __m256d t2 = _mm256_unpackhi_pd( m4.r0, m4.r1 );
    const __m256d t3 = _mm256_unpackhi_pd( m4.r2, m4.r3 );
    // Produce the transposed matrix by combining these blocks
    m4.r0 = _mm256_permute2f128_pd( t0, t1, 0x20 );
    m4.r1 = _mm256_permute2f128_pd( t2, t3, 0x20 );
    m4.r2 = _mm256_permute2f128_pd( t0, t1, 0x31 );
    m4.r3 = _mm256_permute2f128_pd( t2, t3, 0x31 );
}

// Transpose 8x8 matrix with double values
void transpose8x8( double* rdi, const double* rsi )
{
    Matrix4x4 block;
    // Top-left corner
    load( block, rsi );
    transpose( block );
    store( block, rdi );

    // Using another instance of the block to support in-place transpose, with very small overhead
    Matrix4x4 block2;
    load( block, rsi + 4 );     // top right block
    load( block2, rsi + 8 * 4 ); // bottom left block

    transpose( block2 );
    store( block2, rdi + 4 );
    transpose( block );
    store( block, rdi + 8 * 4 );

    // Bottom-right corner
    load( block, rsi + 8 * 4 + 4 );
    transpose( block );
    store( block, rdi + 8 * 4 + 4 );
}

对于单个 SIMD 向量中可以容纳 1 行以上的小矩阵,AVX-512 具有非常好的 2 输入通道交叉洗牌,具有 32 位或 64 位粒度,具有向量控制。 (与 _mm512_unpacklo_pd 不同,后者基本上是 4 个独立的 128 位随机播放。)

一个 4x4 double 矩阵“只有”128 字节,两个 ZMM __m512d 向量,所以你只需要两个 vpermt2ps (_mm512_permutex2var_pd ) 以产生两个输出向量:每个输出向量一个洗牌,加载和存储都是全宽的。不过,您确实需要控制向量常量。

使用 512 位向量指令有一些缺点(时钟速度和执行端口吞吐量),但如果您的程序可以在使用 512 位向量的代码中花费大量时间,则抛出可能会显着提高吞吐量每条指令围绕更多数据,并具有更强大的随机播放。

对于 256 位向量,vpermt2pd ymm 可能对 4x4 没有用,因为对于每个 __m256d 输出行,您想要的 4 个元素中的每一个都来自不同的输入行。所以一个 2-input shuffle 不能产生你想要的输出。

我认为小于 128 位粒度的跨车道洗牌没有用,除非你的矩阵小到足以在一个 SIMD 向量中容纳多行。 对于一些关于 32 位元素的算法复杂性推理 - AVX1 的 32 位元素的 8x8 xpose 与 AVX-512 的 64 位元素的 8x8 大致相同,其中每个 SIMD 向量正好包含一整行.

因此不需要向量常量,只需立即随机播放 128 位块,并且 unpacklo/hi


用 512 位向量(8 个双精度)转置一个 8x8 会遇到同样的问题:8 个双精度的每个输出行需要 8 个输入向量中的每一个的 1 个双精度。 所以最终我认为你想要一个与 Soonts 的 AVX 答案类似的策略,从 _mm512_insertf64x4(v, load, 1) 开始作为将 2 个输入行的前半部分放入一个向量的第一步。

(如果您关心 KNL / Xeon Phi,@ZBoson 在 上的其他回答展示了一些有趣的想法,使用合并屏蔽和 1-input 洗牌,如 vpermpdvpermq ,而不是像 vunpcklpdvpermt2pd)

这样的 2 输入随机播放

使用更宽的向量意味着更少的加载和存储,甚至可能更少的总混洗,因为每个混洗组合了更多的数据。但是您还需要做更多的洗牌工作,将一行的所有 8 个元素放入一个向量中,而不是仅仅加载和存储到行大小一半的块中的不同位置。不是显而易见的更好;如果我开始实际编写代码,我会更新这个答案。

请注意,Ice Lake(第一个使用 AVX-512 的消费者 CPU)每个时钟可以执行 2 次加载和 2 次存储。对于 some 洗牌,它比 Skylake-X 具有更好的洗牌吞吐量,但对于任何对此或 Soonts 的答案有用的东西都没有。 (对于 ymm 和 zmm 版本,所有 vperm2f128vunpcklpdvpermt2pd 仅 运行 在端口 5 上。https://uops.info/. vinsertf64x4 zmm, mem, 1 is 2 uops for the front-end, and needs a load port and a uop for p0/p5. (Not p1 because it's a 512-bit uop, and see also SIMD instructions lowering CPU frequency)。)