更快地规范化下三角矩阵

Normalize lower triangular matrix more quickly

本文关键字:三角 规范化      更新时间:2023-10-16

下面的代码似乎不是瓶颈。

我只是想知道是否有一种更快的方法可以在SSE4.2的cpu上完成这项工作。

该代码在ar_tri:中以以下形式存储为1d数组的矩阵的下三角条目上工作

[ (1,0),
(2,0),(2,1),
(3,0),(3,1),(3,2), 
...,
(n,0)...(n,n-1) ]

其中(x,y)是矩阵在第x行和第yth列的条目。

以及ar_rdia:中以下形式的矩阵的对角线的倒数平方根(rsqrt)

[ rsqrt(0,0), rsqrt(1,1), ... ,rsqrt(n,n) ]

Godbolt编译器资源管理器上的gcc6.1-O3使用SIMD指令(mulps)自动向两个版本矢量化。三角形版本在每行的末尾都有清理代码,因此也有一些标量指令。

使用矩形矩阵作为1d数组存储在连续内存中会提高性能吗?

// Triangular version
#include <iostream>
#include <stdlib.h>
#include <stdint.h>
using namespace std;
int main(void){
size_t n = 10000;
size_t n_tri = n*(n-1)/2;
size_t repeat = 10000;
// test 10000 cycles of the code
float* ar_rdia = (float*)aligned_alloc(16, n*sizeof(float));
//reciprocal square root of diagonal
float* ar_triangular = (float*)aligned_alloc(16, n_tri*sizeof(float));
//lower triangular matrix
size_t i,j,k;
float a,b;
k = 0;
for(i = 0; i < n; ++i){
for(j = 0; j < i; ++j){
ar_triangular[k] *= ar_rdia[i]*ar_rdia[j];
++k;
}
}
cout << k;
free((void*)ar_rdia);
free((void*)ar_triangular);
}

// Square version
#include <iostream>
#include <stdlib.h>
#include <stdint.h>
using namespace std;
int main(void){
size_t n = 10000;
size_t n_sq = n*n;
size_t repeat = 10000;
// test 10000 cycles of the code
float* ar_rdia = (float*)aligned_alloc(16, n*sizeof(float));
//reciprocal square root of diagonal
float* ar_square = (float*)aligned_alloc(16, n_sq*sizeof(float));
//lower triangular matrix
size_t i,j,k;
float a,b;
k = 0;
for(i = 0; i < n; ++i){
for(j = 0; j < n; ++j){
ar_square[k] *= ar_rdia[i]*ar_rdia[j];
++k;
}
}
cout << k;
free((void*)ar_rdia);
free((void*)ar_square);
}

程序集输出:

## Triangular version
main:
...
call    aligned_alloc
movl    $1, %edi
movq    %rax, %rbp
xorl    %esi, %esi
xorl    %eax, %eax
.L2:
testq   %rax, %rax
je      .L3
leaq    -4(%rax), %rcx
leaq    -1(%rax), %r8
movss   (%rbx,%rax,4), %xmm0
shrq    $2, %rcx
addq    $1, %rcx
cmpq    $2, %r8
leaq    0(,%rcx,4), %rdx
jbe     .L9
movaps  %xmm0, %xmm2
leaq    0(%rbp,%rsi,4), %r10
xorl    %r8d, %r8d
xorl    %r9d, %r9d
shufps  $0, %xmm2, %xmm2       # broadcast ar_rdia[i]
.L6:                                # vectorized loop
movaps  (%rbx,%r8), %xmm1
addq    $1, %r9
mulps   %xmm2, %xmm1
movups  (%r10,%r8), %xmm3
mulps   %xmm3, %xmm1
movups  %xmm1, (%r10,%r8)
addq    $16, %r8
cmpq    %rcx, %r9
jb      .L6
cmpq    %rax, %rdx
leaq    (%rsi,%rdx), %rcx
je      .L7
.L4:                                      # scalar cleanup
movss   (%rbx,%rdx,4), %xmm1
leaq    0(%rbp,%rcx,4), %r8
leaq    1(%rdx), %r9
mulss   %xmm0, %xmm1
cmpq    %rax, %r9
mulss   (%r8), %xmm1
movss   %xmm1, (%r8)
leaq    1(%rcx), %r8
jnb     .L7
movss   (%rbx,%r9,4), %xmm1
leaq    0(%rbp,%r8,4), %r8
mulss   %xmm0, %xmm1
addq    $2, %rdx
addq    $2, %rcx
cmpq    %rax, %rdx
mulss   (%r8), %xmm1
movss   %xmm1, (%r8)
jnb     .L7
mulss   (%rbx,%rdx,4), %xmm0
leaq    0(%rbp,%rcx,4), %rcx
mulss   (%rcx), %xmm0
movss   %xmm0, (%rcx)
.L7:
addq    %rax, %rsi
cmpq    $10000, %rdi
je      .L16
.L3:
addq    $1, %rax
addq    $1, %rdi
jmp     .L2
.L9:
movq    %rsi, %rcx
xorl    %edx, %edx
jmp     .L4
.L16:
... print and free
ret

方形案例的程序集的有趣部分:

main:
... allocate both arrays
call    aligned_alloc
leaq    40000(%rbx), %rsi
movq    %rax, %rbp
movq    %rbx, %rcx
movq    %rax, %rdx
.L3:                                       # loop over i
movss   (%rcx), %xmm2
xorl    %eax, %eax
shufps  $0, %xmm2, %xmm2           # broadcast ar_rdia[i]
.L2:                                       # vectorized loop over j
movaps  (%rbx,%rax), %xmm0
mulps   %xmm2, %xmm0
movups  (%rdx,%rax), %xmm1
mulps   %xmm1, %xmm0
movups  %xmm0, (%rdx,%rax)
addq    $16, %rax
cmpq    $40000, %rax
jne     .L2
addq    $4, %rcx             # no scalar cleanup: gcc noticed that the row length is a multiple of 4 elements
addq    $40000, %rdx
cmpq    %rsi, %rcx
jne     .L3
... print and free
ret

存储到三角形数组的循环应该向量化为ok,但每行末尾效率低下。根据您发布的asm,gcc实际上对这两者都进行了自动矢量化。我希望我先看看这个,而不是相信你的话,它需要手动矢量化(

.L6:   # from the first asm dump.
movaps  (%rbx,%r8), %xmm1
addq    $1, %r9
mulps   %xmm2, %xmm1
movups  (%r10,%r8), %xmm3
mulps   %xmm3, %xmm1
movups  %xmm1, (%r10,%r8)
addq    $16, %r8
cmpq    %rcx, %r9
jb      .L6

这看起来就像我的手动矢量化版本将编译到的内部循环。L4是一行最后3个元素的完全展开标量清理。(所以它可能不如我的代码好)。尽管如此,它还是相当不错的,自动矢量化将使您能够在不更改源代码的情况下利用AVX和AVX512

我编辑了你的问题,在godbolt上添加了一个代码链接,两个版本都是单独的函数。我没有花时间将它们转换为将数组作为函数args,因为这样我就必须花时间正确地获取所有__restrict__关键字,并告诉gcc数组在4B*16=64字节的边界上对齐,因此如果愿意,它可以使用对齐的加载。


在一行中,每次都使用相同的ar_rdia[i],因此在该行开始时将其广播到向量中一次。然后您只需在源ar_rdia[j + 0..3]和目标ar_triangular[k + 0..3]之间执行垂直操作。

为了处理行末尾不是矢量大小倍数的最后几个元素,我们有两个选项:

  • 矢量化循环后的标量(或更窄的矢量)回退/清理,处理每行的最后3个元素
  • i上的循环展开4,并使用最佳序列来处理行末尾剩余的奇数0、1、2和3个元素。因此,j上的循环将重复4次,每次循环后都会进行固定的清理这可能是最理想的方法
  • 让最后一个向量迭代超过一行的末尾,而不是在最后一个完整向量之后停止。因此,我们将下一行的开头部分重叠。由于你的运算不是幂等的,所以这个选项不能很好地工作。此外,确保k在下一行开始时正确更新需要一些额外的代码。

    尽管如此,通过使行的最终矢量与乘法器混合,使当前行末尾以外的元素乘以1.0(乘法恒等式),这是可能的。这对于具有1.0的向量的blendvps来替换ar_rdia[i] * ar_rdia[j + 0..3]的一些元素应该是可行的。我们还必须创建一个选择器掩码(可能通过使用j-i作为索引索引到int32_t row_overshoot_blend_window {0, 0, 0, 0, -1, -1, -1}的数组中,以获得4个元素的窗口)。另一个选项是分支,以选择不混合或三个立即混合中的一个(blendps更快,不需要矢量控制掩码,并且分支将具有易于预测的模式)。

    当来自ar_triangular的加载与来自最后一行末尾的存储重叠时,这会在每4行中的3行的开头导致存储转发失败。IDK将表现最佳。

    另一个可能更好的选择是进行超出行末尾的加载,并使用压缩SIMD进行计算,但随后有条件地存储1到4个元素。

    不在分配的内存之外读取可能需要在缓冲区的末尾留下填充,例如,如果最后一行不是4个元素的倍数。

/****** Normalize a triangular matrix using SIMD multiplies,
handling the ends of rows with narrower cleanup code *******/
// size_t i,j,k;   // don't do this in C++ or C99.  Put declarations in the narrowest scope possible.  For types without constructors/destructors, it's still a style / human-readability issue
size_t k = 0;
for(size_t i = 0; i < n; ++i){
// maybe put this inside the for() loop and let the compiler hoist it out, to avoid doing it for small rows where the vector loop doesn't even run once.
__m128 vrdia_i = _mm_set1_ps(ar_rdia[i]);  // broadcast-load: very efficient with AVX, load+shuffle without.  Only done once per row anyway.
size_t j = 0;
for(j = 0; j < (i-3); j+=4){  // vectorize over this loop
__m128 vrdia_j  = _mm_loadu_ps(ar_rdia + j);
__m128 scalefac = _mm_mul_ps(vrdia_j, v_rdia_i);
__m128 vtri       = _mm_loadu_ps(ar_triangular + k);
__m128 normalized = _mm_mul_ps(scalefac , vtri);
_mm_storeu_ps(ar_triangular + k, normalized);
k += 4;
}
// scalar fallback / cleanup for the ends of rows.  Alternative: blend scalefac with 1.0 so it's ok to overlap into the next row.
/*  Fine in theory, but gcc likes to make super-bloated code by auto-vectorizing cleanup loops.  Besides, we can do better than scalar
for ( ; j < i; ++j ){
ar_triangular[k] *= ar_rdia[i]*ar_rdia[j];     ++k;   }
*/
if ((i-j) >= 2) {  // load 2 floats (using movsd to zero the upper 64 bits, so mulps doesn't slow down or raise exceptions on denormals or NaNs
__m128 vrdia_j = _mm_castpd_ps( _mm_load_sd(static_cast<const double*>(ar_rdia+j)) );
__m128 scalefac = _mm_mul_ps(vrdia_j, v_rdia_i);
__m128 vtri       = _mm_castpd_ps( _mm_load_sd(static_cast<const double*>(ar_triangular + k) ));
__m128 normalized = _mm_mul_ps(scalefac , vtri);
_mm_storel_pi(static_cast<__m64*>(ar_triangular + k), normalized);      // movlps.  Agner Fog's table indicates that Nehalem decodes this to 2 uops, instead of 1 for movsd.  Bizarre!
j+=2;
k+=2;
}
if (j<i) {  // last single element
ar_triangular[k] *= ar_rdia[i]*ar_rdia[j];
++k;
//++j;     // end of the row anyway.  A smart compiler would still optimize it away...
}
// another possibility: load 4 elements and do the math, then movss, movsd, movsd + extractps (_mm_extractmem_ps), or movups to store the last 1, 2, 3, or 4 elements of the row.
// don't use maskmovdqu; it bypasses cache
}

movsdmovlps等效为存储器,但不等效为负载。请参阅此评论线程,了解为什么存储表单具有单独的操作码是有意义的。更新:Agner Fog的insn表表明,Nehalem将MOVH/LPS/D解码为2个融合域uop。他们还说,SnB将其解码为1,但IvB将它解码为2个uops。这肯定是错的。对于Haswell,他的表将事物拆分为movlps/d(1个微融合uop)和movhps/d(也是1个微融uop)的单独条目。movlps的存储形式是2个uops并且在任何事情上都需要shuffle端口是没有意义的;它所做的事情与CCD_ 25存储完全相同。

如果您的矩阵真的很大,不要太担心行末处理。如果它们很小,那么更多的总时间将花在行的末尾,因此值得尝试多种方法,并仔细查看asm。


如果源数据是连续的,您可以在这里轻松地实时计算rsqrt。否则是的,只将对角线复制到一个数组中(并在复制时计算rsqrt,而不是像你之前的问题那样再次遍历该数组。要么使用标量rsqrtss,在从矩阵的对角线复制到数组中时没有NR步长,要么手动将元素收集到SIMD向量中(使用_mm_set_ps(a[i][i], a[i+1][i+1], a[i+2][i+2], a[i+3][i+3]),让编译器选择混洗),然后执行rsqrtps+NR步长,然后将4个结果的矢量存储到阵列中。


小问题大小:避免在行末尾不执行完整向量而造成浪费

矩阵的起始是一种特殊情况,因为前6个元素中有三个"结束"是连续的。(第4行有4个元素)。可能值得对此进行特殊封装,并使用两个SSE向量执行前3行。或者可能只是前两行在一起,然后第三行作为一个单独的3人组。实际上,一组4和一组2更为优化,因为SSE可以进行8B和16B的加载/存储,但不能进行12B。

前6个比例因子是ar_rdia的前三个元素的乘积,因此我们可以进行单个向量加载并以几种方式对其进行混洗。

ar_rdia[0]*ar_rdia[0]
ar_rdia[1]*ar_rdia[0], ar_rdia[1]*ar_rdia[1],
ar_rdia[2]*ar_rdia[0], ar_rdia[2]*ar_rdia[1], ar_rdia[2]*ar_rdia[2]
^
end of first vector of 4 elems, start of 2nd.

事实证明,编译器不善于发现和利用这里的模式,所以为了获得前10个元素的最佳代码,我们需要剥离这些迭代,并手动优化洗牌和乘法。我决定执行前4行,因为第4行仍然重用ar_rdia[0..3]的SIMD向量。该向量甚至仍然被第4行(第五行)的第一个向量宽度所使用。

同样值得考虑的是:做2,4,4而不是这个4,2,4。

void triangular_first_4_rows_manual_shuffle(float *tri, const float *ar_rdia)
{
__m128 vr0 = _mm_load_ps(ar_rdia);      // we know ar_rdia is aligned
// elements 0-3    // row 0, row 1, and the first element of row 2
__m128 vi0 = _mm_shuffle_ps(vr0, vr0, _MM_SHUFFLE(2, 1, 1, 0));
__m128 vj0 = _mm_shuffle_ps(vr0, vr0, _MM_SHUFFLE(0, 1, 0, 0));
__m128 sf0  = vi0 * vj0;  // equivalent to _mm_mul_ps(vi0, vj0); // gcc defines __m128 in terms of GNU C vector extensions
__m128 vtri = _mm_load_ps(tri);
vtri *= sf0;
_mm_store_ps(tri, vtri);
tri += 4;
// elements 4 and 5, last two of third row
__m128 vi4 = _mm_shuffle_ps(vr0, vr0, _MM_SHUFFLE(3, 3, 2, 2));   // can compile into unpckhps, saving a byte.  Well spotted by clang
__m128 vj4 = _mm_movehl_ps(vi0, vi0);   // save a mov by reusing a previous shuffle output, instead of a fresh _mm_shuffle_ps(vr0, vr0, _MM_SHUFFLE(2, 1, 2, 1)); // also saves a code byte (no immediate)
// actually, a movsd from ar_ria+1 would get these two elements with no shuffle.  We aren't bottlenecked on load-port uops, so that would be good.
__m128 sf4 = vi4 * vj4;
//sf4 = _mm_movehl_ps(sf4, sf4);        // doesn't save anything compared to shuffling before multiplying
// could use movhps to load and store *tri to/from the high half of an xmm reg, but each of those takes a shuffle uop
// so we shuffle the scale-factor down to the low half of a vector instead.
__m128 vtri4  = _mm_castpd_ps(_mm_load_sd((const double*)tri));  // elements 4 and 5
vtri4 *= sf4;
_mm_storel_pi((__m64*)tri, vtri4);  // 64bit store.  Possibly slower than movsd if Agner's tables are right about movlps, but I doubt it
tri += 2;
// elements 6-9 = row 4, still only needing elements 0-3 of ar_rdia
__m128 vi6 = _mm_shuffle_ps(vr0, vr0, _MM_SHUFFLE(3, 3, 3, 3));  // broadcast.  clang puts this ahead of earlier shuffles.  Maybe we should put this whole block early and load/store this part of tri, too.
//__m128 vi6 = _mm_movehl_ps(vi4, vi4);
__m128 vj6 = vr0; // 3, 2, 1, 0 already in the order we want
__m128 vtri6 = _mm_loadu_ps(tri+6);
vtri6 *= vi6 * vj6;
_mm_storeu_ps(tri+6, vtri6);
tri += 4;
// ... first 4 rows done
}

gcc和clang用-O3 -march=nehalem非常类似地编译它(启用SSE4.2,但不启用AVX)。请参阅Godbolt上的代码,以及其他一些编译不太好的版本:

# gcc 5.3
movaps  xmm0, XMMWORD PTR [rsi]     # D.26921, MEM[(__v4sf *)ar_rdia_2(D)]
movaps  xmm1, xmm0  # tmp108, D.26921
movaps  xmm2, xmm0  # tmp111, D.26921
shufps  xmm1, xmm0, 148     # tmp108, D.26921,
shufps  xmm2, xmm0, 16      # tmp111, D.26921,
mulps   xmm2, xmm1    # sf0, tmp108
movhlps xmm1, xmm1        # tmp119, tmp108
mulps   xmm2, XMMWORD PTR [rdi]       # vtri, MEM[(__v4sf *)tri_5(D)]
movaps  XMMWORD PTR [rdi], xmm2     # MEM[(__v4sf *)tri_5(D)], vtri
movaps  xmm2, xmm0  # tmp116, D.26921
shufps  xmm2, xmm0, 250     # tmp116, D.26921,
mulps   xmm1, xmm2    # sf4, tmp116
movsd   xmm2, QWORD PTR [rdi+16]      # D.26922, MEM[(const double *)tri_5(D) + 16B]
mulps   xmm1, xmm2    # vtri4, D.26922
movaps  xmm2, xmm0  # tmp126, D.26921
shufps  xmm2, xmm0, 255     # tmp126, D.26921,
mulps   xmm0, xmm2    # D.26925, tmp126
movlps  QWORD PTR [rdi+16], xmm1    #, vtri4
movups  xmm1, XMMWORD PTR [rdi+48]  # tmp129,
mulps   xmm0, xmm1    # vtri6, tmp129
movups  XMMWORD PTR [rdi+48], xmm0  #, vtri6
ret

前4行总共只有22条指令,其中4条是movapsreg reg移动。(clang只管理3条指令,总共有21条指令)。我们可能会通过将[ x x 2 1 ]ar_rdia+1中转换为movsd的向量来保存一个,而不是另一个移动+洗牌。并降低混洗端口(以及ALU uops)上的压力。

对于AVX,clang在大多数洗牌中使用vpermips,但这只会浪费一个字节的代码大小。除非它省电(因为它只有1个输入),否则没有理由喜欢它的直接形式而不是shufps,除非你可以将负载折叠到它中


我考虑过使用palignr总是一次4次通过三角矩阵,但这几乎肯定更糟。你会一直需要这些palignr,而不仅仅是在最后。

我认为额外的复杂性/行末较窄的加载/存储只会给无序执行带来麻烦。对于较大的问题规模,你会把大部分时间花在内部循环中,一次处理16B。这可能会成为内存的瓶颈,因此,只要无序执行尽可能快地从内存中提取缓存线,行末占用内存较少的工作基本上是空闲的。

因此三角矩阵对于这个用例来说仍然很好;保持工作集的密集性和连续性似乎很好。根据你下一步要做什么,这可能是理想的,也可能不是理想的。