更好的 8x8 字节矩阵转置与 SSE?

A better 8x8 bytes matrix transpose with SSE?

我发现 this post that explains how to transpose an 8x8 bytes matrix with 24 operations, and a few scrolls later there's the code that implements the transpose. However, this method does not exploit the fact that we can block the 8x8 transpose into four 4x4 transposes, and each one can be done in one shuffle instruction only (this post 是参考)。所以我想出了这个解决方案:

__m128i transpose4x4mask = _mm_set_epi8(15, 11, 7, 3, 14, 10, 6, 2, 13,  9, 5, 1, 12,  8, 4, 0);
__m128i shuffle8x8Mask = _mm_setr_epi8(0, 1, 2, 3, 8, 9, 10, 11, 4,  5, 6, 7, 12,  13, 14, 15);

void TransposeBlock8x8(uint8_t *src, uint8_t *dst, int srcStride, int dstStride) {
    __m128i load0 = _mm_set_epi64x(*(uint64_t*)(src + 1 * srcStride), *(uint64_t*)(src + 0 * srcStride));
    __m128i load1 = _mm_set_epi64x(*(uint64_t*)(src + 3 * srcStride), *(uint64_t*)(src + 2 * srcStride));
    __m128i load2 = _mm_set_epi64x(*(uint64_t*)(src + 5 * srcStride), *(uint64_t*)(src + 4 * srcStride));
    __m128i load3 = _mm_set_epi64x(*(uint64_t*)(src + 7 * srcStride), *(uint64_t*)(src + 6 * srcStride));

    __m128i shuffle0 = _mm_shuffle_epi8(load0, shuffle8x8Mask);
    __m128i shuffle1 = _mm_shuffle_epi8(load1, shuffle8x8Mask);
    __m128i shuffle2 = _mm_shuffle_epi8(load2, shuffle8x8Mask);
    __m128i shuffle3 = _mm_shuffle_epi8(load3, shuffle8x8Mask);

    __m128i block0 = _mm_unpacklo_epi64(shuffle0, shuffle1);
    __m128i block1 = _mm_unpackhi_epi64(shuffle0, shuffle1);
    __m128i block2 = _mm_unpacklo_epi64(shuffle2, shuffle3);
    __m128i block3 = _mm_unpackhi_epi64(shuffle2, shuffle3);

    __m128i transposed0 = _mm_shuffle_epi8(block0, transpose4x4mask);   
    __m128i transposed1 = _mm_shuffle_epi8(block1, transpose4x4mask);   
    __m128i transposed2 = _mm_shuffle_epi8(block2, transpose4x4mask);   
    __m128i transposed3 = _mm_shuffle_epi8(block3, transpose4x4mask);   

    __m128i store0 = _mm_unpacklo_epi32(transposed0, transposed2);
    __m128i store1 = _mm_unpackhi_epi32(transposed0, transposed2);
    __m128i store2 = _mm_unpacklo_epi32(transposed1, transposed3);
    __m128i store3 = _mm_unpackhi_epi32(transposed1, transposed3);

    *((uint64_t*)(dst + 0 * dstStride)) = _mm_extract_epi64(store0, 0);
    *((uint64_t*)(dst + 1 * dstStride)) = _mm_extract_epi64(store0, 1);
    *((uint64_t*)(dst + 2 * dstStride)) = _mm_extract_epi64(store1, 0);
    *((uint64_t*)(dst + 3 * dstStride)) = _mm_extract_epi64(store1, 1);
    *((uint64_t*)(dst + 4 * dstStride)) = _mm_extract_epi64(store2, 0);
    *((uint64_t*)(dst + 5 * dstStride)) = _mm_extract_epi64(store2, 1);
    *((uint64_t*)(dst + 6 * dstStride)) = _mm_extract_epi64(store3, 0);
    *((uint64_t*)(dst + 7 * dstStride)) = _mm_extract_epi64(store3, 1);
}

不包括 load/store 操作,此过程仅包含 16 条指令而不是 24 条指令。

我错过了什么?

通常当加载和存储指令不被计算在内时,这是因为代码正在使用寄存器中的矩阵,例如除了循环中的转置之外,还执行多项操作。本例中的加载和存储不计算在内,因为它们不是主循环的一部分。

但是在您的代码中,加载和存储(或者说集合和提取)正在执行部分转置。

GCC 在您的代码中使用 _mm_insert_epi64_mm_loadl_epi64 为 SSE4.1 实现 _mm_set_epi64x。插入指令正在执行部分转置,即转置从 load0,1,2,3 开始,而不是 shuffle0,1,2,3。然后你的最终 store0,1,2,3 值也不包含转置。您必须使用八个 _mm_extract_epi64 指令才能在内存中完成转置。所以不计算集合和提取内在函数真的没有意义。

无论如何,事实证明你可以只使用 SSSE3 仅使用 16 条指令从寄存器进行转置,如下所示:

//__m128i B0, __m128i B1, __m128i B2, __m128i B3
__m128i mask = _mm_setr_epi8(0x0,0x04,0x01,0x05, 0x02,0x06,0x03,0x07, 0x08,0x0c,0x09,0x0d, 0x0a,0x0e,0x0b,0x0f);

__m128i T0, T1, T2, T3;
T0 = _mm_unpacklo_epi8(B0,B1);
T1 = _mm_unpackhi_epi8(B0,B1);
T2 = _mm_unpacklo_epi8(B2,B3);
T3 = _mm_unpackhi_epi8(B2,B3);

B0 = _mm_unpacklo_epi16(T0,T2);
B1 = _mm_unpackhi_epi16(T0,T2);
B2 = _mm_unpacklo_epi16(T1,T3);
B3 = _mm_unpackhi_epi16(T1,T3);

T0 = _mm_unpacklo_epi32(B0,B2);
T1 = _mm_unpackhi_epi32(B0,B2);
T2 = _mm_unpacklo_epi32(B1,B3);
T3 = _mm_unpackhi_epi32(B1,B3);

B0 = _mm_shuffle_epi8(T0,mask);
B1 = _mm_shuffle_epi8(T1,mask);
B2 = _mm_shuffle_epi8(T2,mask);
B3 = _mm_shuffle_epi8(T3,mask);

我不确定在这里排除加载和存储是否有意义,因为我不确定在四个 128 位寄存器中使用 8x8 字节矩阵有多方便。

这是测试这个的代码:

#include <stdio.h>
#include <x86intrin.h>

void print8x8b(char *A) {
  for(int i=0; i<8; i++) {
    for(int j=0; j<8; j++) {
      printf("%2d ", A[i*8+j]);
    } puts("");
  } puts("");
}

void tran8x8b(char *A, char *B) {
  for(int i=0; i<8; i++) {
    for(int j=0; j<8; j++) {
      B[j*8+i] = A[i*8+j];
    }
  }
}

void tran8x8b_SSE(char *A, char *B) {
  __m128i mask = _mm_setr_epi8(0x0,0x04,0x01,0x05, 0x02,0x06,0x03,0x07, 0x08,0x0c,0x09,0x0d, 0x0a,0x0e,0x0b,0x0f);

  __m128i B0, B1, B2, B3, T0, T1, T2, T3;
  B0 = _mm_loadu_si128((__m128i*)&A[ 0]);
  B1 = _mm_loadu_si128((__m128i*)&A[16]);
  B2 = _mm_loadu_si128((__m128i*)&A[32]);
  B3 = _mm_loadu_si128((__m128i*)&A[48]);

  T0 = _mm_unpacklo_epi8(B0,B1);
  T1 = _mm_unpackhi_epi8(B0,B1);
  T2 = _mm_unpacklo_epi8(B2,B3);
  T3 = _mm_unpackhi_epi8(B2,B3);

  B0 = _mm_unpacklo_epi16(T0,T2);
  B1 = _mm_unpackhi_epi16(T0,T2);
  B2 = _mm_unpacklo_epi16(T1,T3);
  B3 = _mm_unpackhi_epi16(T1,T3);

  T0 = _mm_unpacklo_epi32(B0,B2);
  T1 = _mm_unpackhi_epi32(B0,B2);
  T2 = _mm_unpacklo_epi32(B1,B3);
  T3 = _mm_unpackhi_epi32(B1,B3);

  B0 = _mm_shuffle_epi8(T0,mask);
  B1 = _mm_shuffle_epi8(T1,mask);
  B2 = _mm_shuffle_epi8(T2,mask);
  B3 = _mm_shuffle_epi8(T3,mask);

  _mm_storeu_si128((__m128i*)&B[ 0], B0);
  _mm_storeu_si128((__m128i*)&B[16], B1);
  _mm_storeu_si128((__m128i*)&B[32], B2);
  _mm_storeu_si128((__m128i*)&B[48], B3);
}

int main(void) {
  char A[64], B[64], C[64];
  for(int i=0; i<64; i++) A[i] = i;
  print8x8b(A);
  tran8x8b(A,B);
  print8x8b(B);
  tran8x8b_SSE(A,C);
  print8x8b(C);
}

除了读取和写入内存的加载、存储和 pinsrq-s 之外,步幅可能不等于 8 个字节, 你可以只用 12 条指令进行转置(这段代码可以很容易地与 Z 玻色子的测试代码结合使用):

void tran8x8b_SSE_v2(char *A, char *B) {
  __m128i pshufbcnst = _mm_set_epi8(15,11,7,3, 14,10,6,2, 13,9,5,1, 12,8,4,0);

  __m128i B0, B1, B2, B3, T0, T1, T2, T3;
  B0 = _mm_loadu_si128((__m128i*)&A[ 0]);
  B1 = _mm_loadu_si128((__m128i*)&A[16]);
  B2 = _mm_loadu_si128((__m128i*)&A[32]);
  B3 = _mm_loadu_si128((__m128i*)&A[48]);


  T0 = _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(B0),_mm_castsi128_ps(B1),0b10001000));
  T1 = _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(B2),_mm_castsi128_ps(B3),0b10001000));
  T2 = _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(B0),_mm_castsi128_ps(B1),0b11011101));
  T3 = _mm_castps_si128(_mm_shuffle_ps(_mm_castsi128_ps(B2),_mm_castsi128_ps(B3),0b11011101));

  B0 = _mm_shuffle_epi8(T0,pshufbcnst);
  B1 = _mm_shuffle_epi8(T1,pshufbcnst);
  B2 = _mm_shuffle_epi8(T2,pshufbcnst);
  B3 = _mm_shuffle_epi8(T3,pshufbcnst);

  T0 = _mm_unpacklo_epi32(B0,B1);
  T1 = _mm_unpackhi_epi32(B0,B1);
  T2 = _mm_unpacklo_epi32(B2,B3);
  T3 = _mm_unpackhi_epi32(B2,B3);

  _mm_storeu_si128((__m128i*)&B[ 0], T0);
  _mm_storeu_si128((__m128i*)&B[16], T1);
  _mm_storeu_si128((__m128i*)&B[32], T2);
  _mm_storeu_si128((__m128i*)&B[48], T3);
}


这里我们使用32位浮点数shuffle,比epi32 shuffle更灵活。 转换不会生成额外的指令(使用 gcc 5.4 生成的代码):

tran8x8b_SSE_v2:
.LFB4885:
    .cfi_startproc
    vmovdqu 48(%rdi), %xmm5
    vmovdqu 32(%rdi), %xmm2
    vmovdqu 16(%rdi), %xmm0
    vmovdqu (%rdi), %xmm1
    vshufps 6, %xmm5, %xmm2, %xmm4
    vshufps 1, %xmm5, %xmm2, %xmm2
    vmovdqa .LC6(%rip), %xmm5
    vshufps 6, %xmm0, %xmm1, %xmm3
    vshufps 1, %xmm0, %xmm1, %xmm1
    vpshufb %xmm5, %xmm3, %xmm3
    vpshufb %xmm5, %xmm1, %xmm0
    vpshufb %xmm5, %xmm4, %xmm4
    vpshufb %xmm5, %xmm2, %xmm1
    vpunpckldq  %xmm4, %xmm3, %xmm5
    vpunpckldq  %xmm1, %xmm0, %xmm2
    vpunpckhdq  %xmm4, %xmm3, %xmm3
    vpunpckhdq  %xmm1, %xmm0, %xmm0
    vmovups %xmm5, (%rsi)
    vmovups %xmm3, 16(%rsi)
    vmovups %xmm2, 32(%rsi)
    vmovups %xmm0, 48(%rsi)
    ret
    .cfi_endproc



在某些(但不是全部)较旧的 cpus 上,在 整数和浮点数单位。这会增加函数的延迟,但不一定会影响 代码的吞吐量。

具有 1e9 换位的简单延迟测试:

  for (int i=0;i<500000000;i++){
     tran8x8b_SSE(A,C);
     tran8x8b_SSE(C,A);
  }
  print8x8b(A);

tran8x8b_SSE 大约需要 5.5 秒(19.7e9 个周期),tran8x8b_SSE_v2(英特尔酷睿 i5-6500)大约需要 4.5 秒(16.0e9 个周期)。注意 加载和存储并没有被编译器消除,尽管这些函数是内联在 for 循环中的。


更新:AVX2-128 / SSE 4.1 混合解决方案。

'shuffles'(解包、随机播放)由端口 5 处理,在现代 cpus 上每个 cpu 周期有 1 条指令。 有时用两种混合物代替一种 'shuffle' 是值得的。在 Skylake 上,32 位混合指令可以 运行 在端口 0、1 或 5 上。

不幸的是,_mm_blend_epi32 只是 AVX2-128。一个有效的 SSE 4.1 替代方案是 _mm_blend_ps 组合 有一些演员表(通常是免费的)。 12 'shuffles' 替换为 8 次随机播放与 8 次混合。

简单的延迟测试现在 运行 秒大约需要 3.6 秒(13e9 cpu 周期),比 tran8x8b_SSE_v2.

的结果快 18%

代码:

/* AVX2-128 version, sse 4.1 version see ---------------->       SSE 4.1 version of tran8x8b_AVX2_128()                                                              */
void tran8x8b_AVX2_128(char *A, char *B) {                   /*  void tran8x8b_SSE4_1(char *A, char *B) {                                                            */                                    
  __m128i pshufbcnst_0 = _mm_set_epi8(15, 7,11, 3,  
               13, 5, 9, 1,  14, 6,10, 2,  12, 4, 8, 0);     /*    __m128i pshufbcnst_0 = _mm_set_epi8(15, 7,11, 3,  13, 5, 9, 1,  14, 6,10, 2,  12, 4, 8, 0);       */                                    
  __m128i pshufbcnst_1 = _mm_set_epi8(13, 5, 9, 1,  
               15, 7,11, 3,  12, 4, 8, 0,  14, 6,10, 2);     /*    __m128i pshufbcnst_1 = _mm_set_epi8(13, 5, 9, 1,  15, 7,11, 3,  12, 4, 8, 0,  14, 6,10, 2);       */                                    
  __m128i pshufbcnst_2 = _mm_set_epi8(11, 3,15, 7,  
                9, 1,13, 5,  10, 2,14, 6,   8, 0,12, 4);     /*    __m128i pshufbcnst_2 = _mm_set_epi8(11, 3,15, 7,   9, 1,13, 5,  10, 2,14, 6,   8, 0,12, 4);       */                                    
  __m128i pshufbcnst_3 = _mm_set_epi8( 9, 1,13, 5,  
               11, 3,15, 7,   8, 0,12, 4,  10, 2,14, 6);     /*    __m128i pshufbcnst_3 = _mm_set_epi8( 9, 1,13, 5,  11, 3,15, 7,   8, 0,12, 4,  10, 2,14, 6);       */                                    
  __m128i B0, B1, B2, B3, T0, T1, T2, T3;                    /*    __m128 B0, B1, B2, B3, T0, T1, T2, T3;                                                            */                                    
                                                             /*                                                                                                      */                                    
  B0 = _mm_loadu_si128((__m128i*)&A[ 0]);                    /*    B0 = _mm_loadu_ps((float*)&A[ 0]);                                                                */                                    
  B1 = _mm_loadu_si128((__m128i*)&A[16]);                    /*    B1 = _mm_loadu_ps((float*)&A[16]);                                                                */                                    
  B2 = _mm_loadu_si128((__m128i*)&A[32]);                    /*    B2 = _mm_loadu_ps((float*)&A[32]);                                                                */                                    
  B3 = _mm_loadu_si128((__m128i*)&A[48]);                    /*    B3 = _mm_loadu_ps((float*)&A[48]);                                                                */                                    
                                                             /*                                                                                                      */                                    
  B1 = _mm_shuffle_epi32(B1,0b10110001);                     /*    B1 = _mm_shuffle_ps(B1,B1,0b10110001);                                                            */                                    
  B3 = _mm_shuffle_epi32(B3,0b10110001);                     /*    B3 = _mm_shuffle_ps(B3,B3,0b10110001);                                                            */                                    
  T0 = _mm_blend_epi32(B0,B1,0b1010);                        /*    T0 = _mm_blend_ps(B0,B1,0b1010);                                                                  */                                    
  T1 = _mm_blend_epi32(B2,B3,0b1010);                        /*    T1 = _mm_blend_ps(B2,B3,0b1010);                                                                  */                                    
  T2 = _mm_blend_epi32(B0,B1,0b0101);                        /*    T2 = _mm_blend_ps(B0,B1,0b0101);                                                                  */                                    
  T3 = _mm_blend_epi32(B2,B3,0b0101);                        /*    T3 = _mm_blend_ps(B2,B3,0b0101);                                                                  */                                    
                                                             /*                                                                                                      */                                    
  B0 = _mm_shuffle_epi8(T0,pshufbcnst_0);                    /*    B0 = _mm_castsi128_ps(_mm_shuffle_epi8(_mm_castps_si128(T0),pshufbcnst_0));                       */                                    
  B1 = _mm_shuffle_epi8(T1,pshufbcnst_1);                    /*    B1 = _mm_castsi128_ps(_mm_shuffle_epi8(_mm_castps_si128(T1),pshufbcnst_1));                       */                                    
  B2 = _mm_shuffle_epi8(T2,pshufbcnst_2);                    /*    B2 = _mm_castsi128_ps(_mm_shuffle_epi8(_mm_castps_si128(T2),pshufbcnst_2));                       */                                    
  B3 = _mm_shuffle_epi8(T3,pshufbcnst_3);                    /*    B3 = _mm_castsi128_ps(_mm_shuffle_epi8(_mm_castps_si128(T3),pshufbcnst_3));                       */                                    
                                                             /*                                                                                                      */                                    
  T0 = _mm_blend_epi32(B0,B1,0b1010);                        /*    T0 = _mm_blend_ps(B0,B1,0b1010);                                                                  */                                    
  T1 = _mm_blend_epi32(B0,B1,0b0101);                        /*    T1 = _mm_blend_ps(B0,B1,0b0101);                                                                  */                                    
  T2 = _mm_blend_epi32(B2,B3,0b1010);                        /*    T2 = _mm_blend_ps(B2,B3,0b1010);                                                                  */                                    
  T3 = _mm_blend_epi32(B2,B3,0b0101);                        /*    T3 = _mm_blend_ps(B2,B3,0b0101);                                                                  */                                    
  T1 = _mm_shuffle_epi32(T1,0b10110001);                     /*    T1 = _mm_shuffle_ps(T1,T1,0b10110001);                                                            */                                    
  T3 = _mm_shuffle_epi32(T3,0b10110001);                     /*    T3 = _mm_shuffle_ps(T3,T3,0b10110001);                                                            */                                    
                                                             /*                                                                                                      */                                    
  _mm_storeu_si128((__m128i*)&B[ 0], T0);                    /*    _mm_storeu_ps((float*)&B[ 0], T0);                                                                */                                    
  _mm_storeu_si128((__m128i*)&B[16], T1);                    /*    _mm_storeu_ps((float*)&B[16], T1);                                                                */                                    
  _mm_storeu_si128((__m128i*)&B[32], T2);                    /*    _mm_storeu_ps((float*)&B[32], T2);                                                                */                                    
  _mm_storeu_si128((__m128i*)&B[48], T3);                    /*    _mm_storeu_ps((float*)&B[48], T3);                                                                */                                    
}                                                            /*  }                                                                                                   */                                    

将此作为答案发布。由于到目前为止收到的一些答案和评论,我还将问题的标题从“...使用 SSE”更改为“...使用 SIMD”。

我仅用 8 条指令就成功地使用 AVX2 转置了矩阵,其中 10 条包括 load/store(不包括掩码加载)。 编辑:我找到了一个较短的版本。见下文。这是矩阵在内存中都是连续的情况,所以直接load/store即可。

这是 C 代码:

void tran8x8b_AVX2(char *src, char *dst) {
    __m256i perm = _mm256_set_epi8(
        0, 0, 0, 7,
        0, 0, 0, 5,
        0, 0, 0, 3,
        0, 0, 0, 1,

        0, 0, 0, 6,
        0, 0, 0, 4,
        0, 0, 0, 2,
        0, 0, 0, 0
    );

    __m256i tm = _mm256_set_epi8(
        15, 11, 7, 3,
        14, 10, 6, 2,
        13,  9, 5, 1,
        12,  8, 4, 0,

        15, 11, 7, 3,
        14, 10, 6, 2,
        13,  9, 5, 1,
        12,  8, 4, 0
    );

    __m256i load0 = _mm256_loadu_si256((__m256i*)&src[ 0]);
    __m256i load1 = _mm256_loadu_si256((__m256i*)&src[32]);  

    __m256i perm0 = _mm256_permutevar8x32_epi32(load0, perm);   
    __m256i perm1 = _mm256_permutevar8x32_epi32(load1, perm);   

    __m256i transpose0 = _mm256_shuffle_epi8(perm0, tm);    
    __m256i transpose1 = _mm256_shuffle_epi8(perm1, tm);    

    __m256i unpack0 = _mm256_unpacklo_epi32(transpose0, transpose1);    
    __m256i unpack1 = _mm256_unpackhi_epi32(transpose0, transpose1);

    perm0 = _mm256_castps_si256(_mm256_permute2f128_ps(_mm256_castsi256_ps(unpack0), _mm256_castsi256_ps(unpack1), 32));    
    perm1 = _mm256_castps_si256(_mm256_permute2f128_ps(_mm256_castsi256_ps(unpack0), _mm256_castsi256_ps(unpack1), 49));    

    _mm256_storeu_si256((__m256i*)&dst[ 0], perm0);
    _mm256_storeu_si256((__m256i*)&dst[32], perm1);
}

GCC 足够聪明,可以在 AVX 加载期间执行排列,从而节省了两条指令。这是编译器输出:

tran8x8b_AVX2(char*, char*):
        vmovdqa ymm1, YMMWORD PTR .LC0[rip]
        vmovdqa ymm2, YMMWORD PTR .LC1[rip]
        vpermd  ymm0, ymm1, YMMWORD PTR [rdi]
        vpermd  ymm1, ymm1, YMMWORD PTR [rdi+32]
        vpshufb ymm0, ymm0, ymm2
        vpshufb ymm1, ymm1, ymm2
        vpunpckldq      ymm2, ymm0, ymm1
        vpunpckhdq      ymm0, ymm0, ymm1
        vinsertf128     ymm1, ymm2, xmm0, 1
        vperm2f128      ymm0, ymm2, ymm0, 49
        vmovdqu YMMWORD PTR [rsi], ymm1
        vmovdqu YMMWORD PTR [rsi+32], ymm0
        vzeroupper
        ret

它发出带有 -O3 的 vzerupper 指令,但是向下到 -O1 会删除它。

对于我原来的问题(一个大矩阵,我正在放大它的 8x8 部分),处理步幅会以非常糟糕的方式破坏输出:

void tran8x8b_AVX2(char *src, char *dst, int srcStride, int dstStride) {
    __m256i load0 = _mm256_set_epi64x(*(uint64_t*)(src + 3 * srcStride), *(uint64_t*)(src + 2 * srcStride), *(uint64_t*)(src + 1 * srcStride), *(uint64_t*)(src + 0 * srcStride));
    __m256i load1 = _mm256_set_epi64x(*(uint64_t*)(src + 7 * srcStride), *(uint64_t*)(src + 6 * srcStride), *(uint64_t*)(src + 5 * srcStride), *(uint64_t*)(src + 4 * srcStride));

    // ... the same as before, however we can skip the final permutations because we need to handle the destination stride...

    *((uint64_t*)(dst + 0 * dstStride)) = _mm256_extract_epi64(unpack0, 0);
    *((uint64_t*)(dst + 1 * dstStride)) = _mm256_extract_epi64(unpack0, 1);
    *((uint64_t*)(dst + 2 * dstStride)) = _mm256_extract_epi64(unpack1, 0);
    *((uint64_t*)(dst + 3 * dstStride)) = _mm256_extract_epi64(unpack1, 1);
    *((uint64_t*)(dst + 4 * dstStride)) = _mm256_extract_epi64(unpack0, 2);
    *((uint64_t*)(dst + 5 * dstStride)) = _mm256_extract_epi64(unpack0, 3);
    *((uint64_t*)(dst + 6 * dstStride)) = _mm256_extract_epi64(unpack1, 2);
    *((uint64_t*)(dst + 7 * dstStride)) = _mm256_extract_epi64(unpack1, 3);
}

编译器输出如下:

tran8x8b_AVX2(char*, char*, int, int):
        movsx   rdx, edx
        vmovq   xmm5, QWORD PTR [rdi]
        lea     r9, [rdi+rdx]
        vmovdqa ymm3, YMMWORD PTR .LC0[rip]
        movsx   rcx, ecx
        lea     r11, [r9+rdx]
        vpinsrq xmm0, xmm5, QWORD PTR [r9], 1
        lea     r10, [r11+rdx]
        vmovq   xmm4, QWORD PTR [r11]
        vpinsrq xmm1, xmm4, QWORD PTR [r10], 1
        lea     r8, [r10+rdx]
        lea     rax, [r8+rdx]
        vmovq   xmm7, QWORD PTR [r8]
        vmovq   xmm6, QWORD PTR [rax+rdx]
        vpinsrq xmm2, xmm7, QWORD PTR [rax], 1
        vinserti128     ymm1, ymm0, xmm1, 0x1
        vpinsrq xmm0, xmm6, QWORD PTR [rax+rdx*2], 1
        lea     rax, [rsi+rcx]
        vpermd  ymm1, ymm3, ymm1
        vinserti128     ymm0, ymm2, xmm0, 0x1
        vmovdqa ymm2, YMMWORD PTR .LC1[rip]
        vpshufb ymm1, ymm1, ymm2
        vpermd  ymm0, ymm3, ymm0
        vpshufb ymm0, ymm0, ymm2
        vpunpckldq      ymm2, ymm1, ymm0
        vpunpckhdq      ymm0, ymm1, ymm0
        vmovdqa xmm1, xmm2
        vmovq   QWORD PTR [rsi], xmm1
        vpextrq QWORD PTR [rax], xmm1, 1
        vmovdqa xmm1, xmm0
        add     rax, rcx
        vextracti128    xmm0, ymm0, 0x1
        vmovq   QWORD PTR [rax], xmm1
        add     rax, rcx
        vpextrq QWORD PTR [rax], xmm1, 1
        add     rax, rcx
        vextracti128    xmm1, ymm2, 0x1
        vmovq   QWORD PTR [rax], xmm1
        add     rax, rcx
        vpextrq QWORD PTR [rax], xmm1, 1
        vmovq   QWORD PTR [rax+rcx], xmm0
        vpextrq QWORD PTR [rax+rcx*2], xmm0, 1
        vzeroupper
        ret

但是,如果与我的原始代码的输出相比,这似乎没什么大不了的。


编辑:我找到了一个较短的版本。总共 4 条指令,8 条同时计算 load/stores。 这是可能的,因为我以不同的方式读取矩阵,在加载期间在 "gather" 指令中隐藏了一些 "shuffles" .另外请注意,执行存储需要最终排列,因为 AVX2 没有 "scatter" 指令。使用分散指令会将所有内容减少到仅 2 条指令。另外,请注意,我可以通过更改 vindex 向量的内容来轻松处理 src 步幅。

不幸的是,这个 AVX_v2 似乎比上一个慢。这是代码:

void tran8x8b_AVX2_v2(char *src1, char *dst1) {
    __m256i tm = _mm256_set_epi8(
        15, 11, 7, 3,
        14, 10, 6, 2,
        13,  9, 5, 1,
        12,  8, 4, 0,

        15, 11, 7, 3,
        14, 10, 6, 2,
        13,  9, 5, 1,
        12,  8, 4, 0
    );

    __m256i vindex = _mm256_setr_epi32(0, 8, 16, 24, 32, 40, 48, 56);
    __m256i perm = _mm256_setr_epi32(0, 4, 1, 5, 2, 6, 3, 7);

     __m256i load0 = _mm256_i32gather_epi32((int*)src1, vindex, 1);
    __m256i load1 = _mm256_i32gather_epi32((int*)(src1 + 4), vindex, 1); 

    __m256i transpose0 = _mm256_shuffle_epi8(load0, tm);    
    __m256i transpose1 = _mm256_shuffle_epi8(load1, tm);    

    __m256i final0 = _mm256_permutevar8x32_epi32(transpose0, perm);    
    __m256i final1 = _mm256_permutevar8x32_epi32(transpose1, perm);    

    _mm256_storeu_si256((__m256i*)&dst1[ 0], final0);
    _mm256_storeu_si256((__m256i*)&dst1[32], final1);
}

这是编译器的输出:

tran8x8b_AVX2_v2(char*, char*):
        vpcmpeqd        ymm3, ymm3, ymm3
        vmovdqa ymm2, YMMWORD PTR .LC0[rip]
        vmovdqa ymm4, ymm3
        vpgatherdd      ymm0, DWORD PTR [rdi+4+ymm2*8], ymm3
        vpgatherdd      ymm1, DWORD PTR [rdi+ymm2*8], ymm4
        vmovdqa ymm2, YMMWORD PTR .LC1[rip]
        vpshufb ymm1, ymm1, ymm2
        vpshufb ymm0, ymm0, ymm2
        vmovdqa ymm2, YMMWORD PTR .LC2[rip]
        vpermd  ymm1, ymm2, ymm1
        vpermd  ymm0, ymm2, ymm0
        vmovdqu YMMWORD PTR [rsi], ymm1
        vmovdqu YMMWORD PTR [rsi+32], ymm0
        vzeroupper
        ret

简化版

void tp128_8x8(char *A, char *B) {
  __m128i sv = _mm_set_epi8(15, 7, 14, 6, 13, 5, 12, 4, 11, 3, 10, 2, 9, 1, 8,  0);
  __m128i iv[4], ov[4];

  ov[0] = _mm_shuffle_epi8(_mm_loadu_si128((__m128i*)A), sv);
  ov[1] = _mm_shuffle_epi8(_mm_loadu_si128((__m128i*)(A+16)), sv);
  ov[2] = _mm_shuffle_epi8(_mm_loadu_si128((__m128i*)(A+32)), sv);
  ov[3] = _mm_shuffle_epi8(_mm_loadu_si128((__m128i*)(A+48)), sv);

  iv[0] = _mm_unpacklo_epi16(ov[0], ov[1]); 
  iv[1] = _mm_unpackhi_epi16(ov[0], ov[1]); 
  iv[2] = _mm_unpacklo_epi16(ov[2], ov[3]); 
  iv[3] = _mm_unpackhi_epi16(ov[2], ov[3]); 

  _mm_storeu_si128((__m128i*)B,      _mm_unpacklo_epi32(iv[0], iv[2]));
  _mm_storeu_si128((__m128i*)(B+16), _mm_unpackhi_epi32(iv[0], iv[2]));
  _mm_storeu_si128((__m128i*)(B+32), _mm_unpacklo_epi32(iv[1], iv[3]));
  _mm_storeu_si128((__m128i*)(B+48), _mm_unpackhi_epi32(iv[1], iv[3]));
}



Benchmark:i5-5300U 2.3GHz (cycles per byte)
tran8x8b           : 2.140
tran8x8b_SSE       : 1.602
tran8x8b_SSE_v2    : 1.551
tp128_8x8          : 1.535
tran8x8b_AVX2      : 1.563
tran8x8b_AVX2_v2   : 1.731

这对我来说真的很有趣,我正想做这个,但由于各种原因,我最终需要用 Go 而不是 C 来做,而且我没有向量内在函数,所以我以为 "well, I'll just write something and see how it does".

我报告的时间,在 ~3.6GHz CPU 上,对于一个简单的实现,每个 64 字节块转置大约需要 28ns,而使用位移完成的每个块大约需要 19ns。我使用 perf 来确认数字,这对我来说似乎有点不太可能,而且它们似乎相加。奇特的移位实现有点超过 250 条指令,每个周期大约有 3.6 条指令,因此每次操作大约需要 69-70 个周期。

这是 Go,但老实说,实现起来应该很简单;它只是将 64 字节的输入数组视为 8 uint64_t.

通过将其中一些东西声明为新变量来提示寄存器分配器,您可以获得另一个纳秒左右的时间。


import (
"unsafe"
)

const (
    hi16 = uint64(0xFFFF0000FFFF0000)
    lo16 = uint64(0x0000FFFF0000FFFF)
    hi8  = uint64(0xFF00FF00FF00FF00)
    lo8  = uint64(0x00FF00FF00FF00FF)
)

// Okay, this might take some explaining. We are working on a logical
// 8x8 matrix of bytes, which we get as a 64-byte array. We want to transpose
// it (row/column).
//
// start:
// [[00 08 16 24 32 40 48 56]
//  [01 09 17 25 33 41 49 57]
//  [02 10 18 26 34 42 50 58]
//  [03 11 19 27 35 43 51 59]
//  [04 12 20 28 36 44 52 60]
//  [05 13 21 29 37 45 53 61]
//  [06 14 22 30 38 46 54 62]
//  [07 15 23 31 39 47 55 63]]
//
// First, let's make sure everything under 32 is in the top four rows,
// and everything over 32 is in the bottom four rows. We do this by
// swapping pairs of 32-bit words.
// swap32:
// [[00 08 16 24 04 12 20 28]
//  [01 09 17 25 05 13 21 29]
//  [02 10 18 26 06 14 22 30]
//  [03 11 19 27 07 15 23 31]
//  [32 40 48 56 36 44 52 60]
//  [33 41 49 57 37 45 53 61]
//  [34 42 50 58 38 46 54 62]
//  [35 43 51 59 39 47 55 63]]
//
// Next, let's make sure everything over 16 or 48 is in the bottom two
// rows of the two four-row sections, and everything under 16 or 48 is
// in the top two rows of the section. We do this by swapping masked
// pairs in much the same way:
// swap16:
// [[00 08 02 10 04 12 06 14]
//  [01 09 03 11 05 13 07 15]
//  [16 24 18 26 20 28 22 30]
//  [17 25 19 27 21 29 23 31]
//  [32 40 34 42 36 44 38 46]
//  [33 41 35 43 37 45 39 47]
//  [48 56 50 58 52 60 54 62]
//  [49 57 51 59 53 61 55 63]]
//
// Now, we will do the same thing to each pair -- but because of
// clever choices in the specific arrange ment leading up to this, that's
// just one more byte swap, where each 2x2 block has its upper right
// and lower left corners swapped, and that turns out to be an easy
// shift and mask.
func UnswizzleLazy(m *[64]uint8) {
    // m32 treats the 8x8 array as a 2x8 array, because
    // it turns out we only need to swap a handful of the
    // bits...
    m32 := (*[16]uint32)(unsafe.Pointer(&m[0]))
    m32[1], m32[8] = m32[8], m32[1]
    m32[3], m32[10] = m32[10], m32[3]
    m32[5], m32[12] = m32[12], m32[5]
    m32[7], m32[14] = m32[14], m32[7]
    m64 := (*[8]uint64)(unsafe.Pointer(&m[0]))
    // we're now at the state described above as "swap32"
    tmp0, tmp1, tmp2, tmp3 :=
        (m64[0]&lo16)|(m64[2]&lo16)<<16,
        (m64[1]&lo16)|(m64[3]&lo16)<<16,
        (m64[0]&hi16)>>16|(m64[2]&hi16),
        (m64[1]&hi16)>>16|(m64[3]&hi16)
    tmp4, tmp5, tmp6, tmp7 :=
        (m64[4]&lo16)|(m64[6]&lo16)<<16,
        (m64[5]&lo16)|(m64[7]&lo16)<<16,
        (m64[4]&hi16)>>16|(m64[6]&hi16),
        (m64[5]&hi16)>>16|(m64[7]&hi16)
    // now we're at "swap16".
    lo8 := lo8
    hi8 := hi8
    m64[0], m64[1] = (tmp0&lo8)|(tmp1&lo8)<<8, (tmp0&hi8)>>8|tmp1&hi8
    m64[2], m64[3] = (tmp2&lo8)|(tmp3&lo8)<<8, (tmp2&hi8)>>8|tmp3&hi8
    m64[4], m64[5] = (tmp4&lo8)|(tmp5&lo8)<<8, (tmp4&hi8)>>8|tmp5&hi8
    m64[6], m64[7] = (tmp6&lo8)|(tmp7&lo8)<<8, (tmp6&hi8)>>8|tmp7&hi8
}

我希望这样做是相当明显的:打乱半字,所以前四个字具有属于它们的所有值,而后四个具有属于它们的所有值.然后对每组四个单词做类似的事情,所以你最终得到属于前两个单词的东西,等等。

直到我意识到,如果上面的 cycles/byte 数字是正确的,这实际上优于 shuffle/unpack 解决方案,我才打算发表评论。

(请注意,这是就地转置,但中间步骤使用临时文件并在其他地方进行最终存储很容易。实际上可能更快。)

更新:我最初描述我的算法有点不正确,然后我意识到我实际上可以做我描述的事情。这个 运行 每 64 位大约 65.7 个周期。

编辑 #2:在此机器上尝试了上述 AVX 版本之一。在我的硬件(Xeon E3-1505M,标称 3GHz)上,我每 64 字节块得到 10 个多一点的周期,因此,每个周期大约 6 个字节。这对我来说似乎比每字节 1.5 个周期更合理。

编辑#3:进一步下降,每 64 位大约 45 个周期,只需将第一部分写为 uint64 上的移位和掩码,而不是尝试 "smart" 并移动 32 位我很关心。

AVX512VBMI (Cascade Lake / Ice Lake)

AVX512VBMI 引入了 vpermb,一种具有字节粒度的 64 字节通道交叉随机播放。
_mm512_permutexvar_epi8( __m512i idx, __m512i a);

现有 CPU 支持它 运行 它作为单个 uop,具有 1/clock 吞吐量。 (https://www.uops.info/html-tp/CNL/VPERMB_ZMM_ZMM_M512-Measurements.html)

这使问题变得微不足道,使 1 条指令成为可能(至少对于整个 8x8 块是连续的 stride=8 的情况)。否则,您应该查看 vpermt2b 以将来自 2 个源的字节混在一起。但那是 CannonLake 上的 3 微码。

// TODO: strided loads / stores somehow for stride != 8
// AVX512VBMI
void TransposeBlock8x8_contiguous(uint8_t *src, uint8_t *dst)
{
    const __m512i trans8x8shuf = _mm512_set_epi8(
        63, 63-8*1, 63-8*2, 63-8*3, 63-8*4, ...
        ...
        57, 49, 41, 33, 25, 17, 9, 1,
        56, 48, 40, 32, 24, 16, 8, 0
   );

    __m512i vsrc = _mm512_loadu_si512(src);
    __m512i shuffled = _mm512_permutexvar_epi8(trans8x8shuf, vsrc);
    _mm512_storeu_si512(dst, shuffled);
}

https://godbolt.org/z/wrfyy3

显然 _mm512_setr_epi8 对于 gcc/clang 不存在(仅 256 和 128 版本)所以你必须按照从后到前的顺序定义常量,与 C 数组初始化顺序相反.


vpermb 甚至可以将数据用作内存源操作数,因此它可以在一条指令中加载+随机播放。但是根据解码为 1 个融合域 uop 的 https://uops.info/, it doesn't micro-fuse on CannonLake: unlike vpermd zmm, zmm, [r14](注意 "retire_slots: 1.0")
vpermd zmm, zmm, [r14] 为前端/融合域解码为 2 个独立的 uops:"retire_slots: 2.0")。这是在真实的 CannonLake CPU 上使用性能计数器进行的实验测试。 uops.info 尚无可用的 Cascade Lake 或 Ice Lake,因此它可能会更有效率。

uops.info 表无用地计算未融合域 uops 的总数,因此您必须单击指令以查看它是否微融合。


源或 dst 跨越,不连续,数据

我猜你想要将 qword(8 字节)加载到 XMM 寄存器中并将输入对混洗在一起,或者将它们与 movhpspinsrq 连接起来。使用跨步索引的 qword-gather 加载可能是值得的,但这通常不值得。

我不确定就 YMM 寄存器而言是否值得合并,更不用说 ZMM 了,或者最好只与 XMM 寄存器一样宽,这样我们就可以使用 [=20 手动有效地将 qwords 分散回内存=] 和 vmovhps(在 Intel CPUs 上不需要 shuffle uop,只需要一个存储)。如果 dst 是连续的,合并一个非连续的 strided src 更有意义。

AVX512VBMI vpermt2b ymm 看起来它对洗牌+合并很有用,就像车道交叉 punpcklbw,从另外两个 32 字节 YMM 寄存器的串联中选择任意 32 字节。 (或 ZMM 版本的 2x 64 字节 regs 中的 64)。但不幸的是,在 CannonLake 上它需要 3 微码,就像 vpermt2w 在 Skylake-X 和 Cannon Lake 上一样。

如果我们以后可以担心字节,vpermt2d 在支持它的 CPU 上是有效的(单 uop)! Skylake-X 及更高版本。

Ice Lake 对于 vpermt2binstlat)每 2 个周期有一个吞吐量,可能是因为它有一个额外的洗牌单元,可以 运行 一些(但不是全部)洗牌哎呀。请注意,例如 vpshufb xmmymm 是 0.5c 的吞吐量,但 vpshufb zmm 是 1c 的吞吐量。但是 vpermb 总是只有 1c 的吞吐量。

我想知道我们是否可以利用合并屏蔽?就像 vpmovzxbq 将输入字节零扩展到 qwords。 (一个 8 字节行 -> 64 字节 ZMM 寄存器)。那么也许双字左移合并屏蔽到另一个寄存器?不,这没有帮助,有用的数据在两个输入的相同 dword 元素中,除非你先对一个寄存器做一些事情,破坏了目的。


重叠字节掩码存储 (vmovdqu8 [rdi + 0..7]{k1}, zmm0..7) of vpmovzxbq 加载结果也是可能的,但可能效率不高。充其量,除了其中一个之外,所有其他人都将错位。不过,存储缓冲区 and/or 缓存硬件可能能够有效地提交 8 倍屏蔽存储。

在寄存器和一些掩码存储中进行一些移动的混合策略可能很有趣,可以平衡 shuffle/blend 与连续 dst 的存储工作。特别是如果所有商店都可以对齐,这将需要在每个向量中四处移动数据以使其位于正确的位置。

Ice Lake 有 2 个商店执行单元。 (IDK 如果 L1d 缓存提交可以跟上它,或者如果在存储缓冲区中合并通常有帮助,或者如果这只是有助于突发工作。)

这里的大多数答案使用 _mm_shuffle_epi8 的不同大小的随机排列和排列的组合,这仅适用于 SSSE3 及更高版本。

可以通过连续三次将前 32 个元素与后 32 个元素交错来形成具有 12* 指令内核的纯 SSE2 实现:

void systolic_kernel(__m128i a[4]) {
    __m128i a0 = _mm_unpacklo_epi8(a[0], a[2]);
    __m128i a1 = _mm_unpackhi_epi8(a[0], a[2]);
    __m128i a2 = _mm_unpacklo_epi8(a[1], a[3]);
    __m128i a3 = _mm_unpackhi_epi8(a[1], a[3]);
    a[0] = a0;
    a[1] = a1;
    a[2] = a2;
    a[3] = a3;
}

void transpose(__m128i a[4]) {
    systolic_kernel(a);
    systolic_kernel(a);
    systolic_kernel(a);
}

*没有 VEX 编码(对于三个操作数指令),将添加 6 个潜在的零成本 movdqa 指令。

相同的策略可以更容易地应用于 4x4、16x16 转置等,因为要置换的索引的计算和块大小已从等式中分解出来。