C++数值计算简单加速技术(五)——矩阵转置优化示例(访存优化和SIMD的使用)

最近在做一个矩阵转置的优化的测试，效果还可以，这里展示一下。

在之前一篇SIMD的文章中已经提过了使用AVX指令集进行转置加速，效率达到了普通转置的三倍左右，效果比较明显，最近恰好看到了一篇讲转置和Cache的文章
http://evol128.is-programmer.com/posts/35453.html
于是又做了一些测试，发现确实还有提升的空间。这里先把SIMD优化的方法再说一下。使用AVX做矩阵转置优化原理上来讲主要就是利用Intel提供的几个数据混洗、重排的函数，将矩阵分成4x4的小块进行转置，然后将转置好的小块再存入完整转置矩阵的相应位置，由于非原址转置代码更好写一些，所以这里只展示非原址变换的代码。

原始操作

首先从最原始的转置开始

void trans(double *x,double *y,size_t n)
{
	double *s=x,*d=y;
    for(int i=0;i<n;i++)
    {
        d=y+i;
        for(int j=0;j<n;j++) 
        {
            *d=*(s++);
            d+=n;
        }
    }
}

这里读取的时候按行读，写入的时候按列写，这样可以在一定程度上利用到cache。
转置一个4096的方阵测试一下速度。

#include <iostream>
#include <ctime>
#define N 4096
#define I 233
#define J 199
using namespace std;
int main()
{
    double* x=(double*)_mm_malloc(sizeof(double)*N*N,32);
    double* y=(double*)_mm_malloc(sizeof(double)*N*N,32);
    for(int i=0;i<N*N;i++) x[i]=i%100;
    clock_t s=clock();
    {
        trans(x,y,N);
            }
    cout<<x[J*N+I]<<'\t'<<y[I*N+J]<<'\n';
    cout<<double(clock()-s)/CLOCKS_PER_SEC<<'\n';
    return 0;
}

I、J是为了检查结果是不是正确，随机取一个位置，同时也避免编译器把循环优化掉。最终测试耗时为：

0.324557

利用AVX优化

接下来使用AVX指令集进行优化，由于AVX指令集可以一次处理256bit的数据，是4个双精度浮点数，因此将矩阵切分为4x4的小块进行变换，具体操作如下：

/*
我们假设4x4的子块是
a0	a1	a2	a3
b0	b1	b2	b3
c0	c1	c2	c3
d0	d1	d2	d3
目标子块即为
a0	b0	c0	d0
a1	b1	c1	d1
a2	b2	c2	d2
a3	b3	c3	d3
*/
void tran_kernel4x4(const __m256d s1,const __m256d s2,const __m256d s3,const __m256d s4,double *d1,double *d2,double *d3,double *d4)//s1,s2,s3,s4是读取到的矢量数据，d1,d2,d3,d4是转置后存放的位置
{
    __m256d t1,t2,t3,t4,t5,t6,t7,t8;
    t1=_mm256_permute4x64_pd(s1,0b01001110);//将第一行数据进行交换得到a2,a3,a0,a1
    t2=_mm256_permute4x64_pd(s2,0b01001110);//将第二行数据进行交换得到b2,b3,b0,b1
    t3=_mm256_permute4x64_pd(s3,0b01001110);//将第三行数据进行交换得到c2,c3,c0,c1
    t4=_mm256_permute4x64_pd(s4,0b01001110);//将第四行数据进行交换得到d2,d3,d0,d1
    t5=_mm256_blend_pd(s1,t3,0b1100);//合并原序列第一行和重排序列第三行得到a0,a1,c0,c1
    t6=_mm256_blend_pd(s2,t4,0b1100);//b0,b1,d0,d1
    t7=_mm256_blend_pd(t1,s3,0b1100);//a2,a3,c2,c3
    t8=_mm256_blend_pd(t2,s4,0b1100);//b2,b3,d2,d3
    _mm256_store_pd(d1,_mm256_unpacklo_pd(t5,t6));//生成转置子块，并写入对应位置a0,b0,c0,d0
    _mm256_store_pd(d2,_mm256_unpackhi_pd(t5,t6));//a1,b1,c1,d1
    _mm256_store_pd(d3,_mm256_unpacklo_pd(t7,t8));//a2,b2,c2,d2
    _mm256_store_pd(d4,_mm256_unpackhi_pd(t7,t8));//a3,b3,c3,d3
}

这时的trans函数就可以写成

void trans(double *x,double *y,size_t n)
{
	double *s1=x,*s2=x+n,*s3=x+2*n,*s4=x+3*n,
           *d1=y,*d2=y+n,*d3=y+2*n,*d4=y+3*n;
    for(int i=0;i<n;i+=4)
    {
        d1=y+i,d2=d1+n,d3=d1+2*n,d4=d1+3*n;
        int j=n;
        while(j>=4)
        {
            __m256d vs1,vs2,vs3,vs4;
            vs1=_mm256_load_pd(s1);
            vs2=_mm256_load_pd(s2);
            vs3=_mm256_load_pd(s3);
            vs4=_mm256_load_pd(s4);
            tran_kernel4x4(vs1,vs2,vs3,vs4,d1,d2,d3,d4);
            s1+=4;s2+=4;s3+=4;s4+=4;
            d1+=4*n;d2+=4*n;d3+=4*n;d4+=4*n;
            j-=4;
        }
        s1+=3*n;s2+=3*n;s3+=3*n;s4+=3*n;
      }
    }

同样转置4096的方阵，耗时为：

0.135692

可以看到在使用AVX指令加速之后，转置的性能有了很大的提升。这一套优化在我之前的文章里也展示过了，基本没有什么差别。

进一步优化Cache

接下来就是更进一步优化Cache了。在做进一步优化之前先说一下之前提到的那篇文章《为什么513的方阵转置比512的方阵更快》。这篇文章写的还是比较模糊的，至少我这种非科班的是看不太懂。这里讲一下我的理解。

文章关键大概是三个问题：
1、Cache从主存中缓存数据的最小单位是一个Cache Line的大小，对于目前大多数CPU(我也没什么高级CPU，只有一个古董Haswell)，Cache Line的大小是64字节，也就是512bit
2、关于多路Cache的内存映射。现在多数CPU(Haswell)是8路的，既8个Cache Line作为一组，以32k的L1 Cache为例，其总共就可以分为32x1024/64/8=64组。这表示任意主存的地址都会被映射到0-63，而这0-63就分别对应L1 Cache中的一个组中。这样的话，Cache中的一个组内(假如叫第0组)就可以保存8个内存地址映射为0的Cache Line，这就叫8路Cache。如果现在需要第9个地址映射为0的数据缓存到L1 Cache中，那么很遗憾这时就必须淘汰掉这一组中的一个Cache Line内的数据，即便此时Cache并没有被完全占用。
3、关于写Cache。现在多数CPU(Haswell)采用一个脏位(dirty bit)来标记该缓存是否被修改，如果被修改则在这块缓存被淘汰时，就将其数据写入到主存，如果没有被修改则直接舍弃。

那么为什么512的方阵转置比513的方阵慢呢？粗略来讲就是由于地址映射的原因导致512的方阵(注意：本文是double数据类型，但是那篇文章是int)，每一列的数据恰好都映射到了Cache中的同一组中去了，尽管转置的时候我们使用的时连续读取，但是写入的时候却是按列写入的，而且前面也提到了写入是会使用Cache的，这时相当于我们在写入某一列的数据时Cache的大小就只用到一个组的大小，即8个Cache Line，512字节。原本32K的L1 Cache，只剩下512B。但是513的方阵则不同，513的方阵每一列的地址映射到Cache中恰好是错开的，这样写入一列的时候就能够把整个Cache利用起来。怎么想都是Cache利用率高的程序效率更高，因此513的方阵转置就会快于512的方阵。

另一方面，一个Cache Line是64字节，读取的时候我们是连续读取的因此可以一下把Cache Line中的数据都读到，但是写入就不一样了，写入一列的时候每一个Cache Line都只用到了一个数据的大小，当后面的数据缓存进来代替前面的数据时，可能这个Cache Line仅仅只修改了4个或者8个字节的数据就要被淘汰掉了。而当下一列的数据再被写入时又要重新刷新缓存，尽管下一列的数据明明和上一列共享的一个Cache Line。这样就增加了Cache和主存的交换次数，尽管乱序执行和流水线也许在一定程度抵消这种开销，但是这个开销终究是存在的。

鉴于以上两点，我们考虑将矩阵划分为8x8的子块，这样子块的每一行都恰好是一个Cache Line的大小，而转置后的子块一行还是一个Cache Line的大小，我们就可以最大限度的利用到Cache了。具体的实现代码如下：

void tran_block8x8(const double *s1,const double *s2,const double* s3,const double *s4,
                          const double *s5,const double *s6,const double *s7,const double *s8,
                          double *d1,double *d2,double *d3,double *d4,
                          double *d5,double *d6,double *d7,double *d8)
{
    __m256d t1,t2,t3,t4,t5,t6,t7,t8,t9,t10,t11,t12,t13,t14,t15,t16;
    t1=_mm256_load_pd(s1);
    t5=_mm256_load_pd(s1+4);
    t2=_mm256_load_pd(s2);
    t6=_mm256_load_pd(s2+4);
    t3=_mm256_load_pd(s3);
    t7=_mm256_load_pd(s3+4);
    t4=_mm256_load_pd(s4);
    t8=_mm256_load_pd(s4+4);
    t9=_mm256_load_pd(s5);
    t13=_mm256_load_pd(s5+4);
    t10=_mm256_load_pd(s6);
    t14=_mm256_load_pd(s6+4);
    t11=_mm256_load_pd(s7);
    t15=_mm256_load_pd(s7+4);
    t12=_mm256_load_pd(s8);
    t16=_mm256_load_pd(s8+4);
    tran_kernel4x4(t1,t2,t3,t4,d1,d2,d3,d4);
    tran_kernel4x4(t9,t10,t11,t12,d1+4,d2+4,d3+4,d4+4);
    tran_kernel4x4(t5,t6,t7,t8,d5,d6,d7,d8);
    tran_kernel4x4(t13,t14,t15,t16,d5+4,d6+4,d7+4,d8+4);
}

这里用到了前面的4x4的kernel函数，即一个8x8的子块调用4个4x4的kernel来完成。这样转置的函数就写成

void trans(double *x,double *y,const size_t n)
{
	double *s1=x,*s2=x+n,*s3=x+2*n,*s4=x+3*n,
           *s5=x+4*n,*s6=x+5*n,*s7=x+6*n,*s8=x+7*n,
           *d1=y,*d2=y+n,*d3=y+2*n,*d4=y+3*n,
           *d5=y+4*n,*d6=y+5*n,*d7=y+6*n,*d8=y+7*n;
    for(int i=0;i<n;i+=8)
    {
        d1=y+i,d2=d1+n,d3=d1+2*n,d4=d1+3*n,
        d5=d1+4*n,d6=d1+5*n,d7=d1+6*n,d8=d1+7*n;
        int j=n;
        while(j>=8)
        {
            tran_block8x8(s1,s2,s3,s4,s5,s6,s7,s8,d1,d2,d3,d4,d5,d6,d7,d8);
            s1+=8;s2+=8;s3+=8;s4+=8;
            s5+=8;s6+=8;s7+=8;s8+=8;
            d1+=8*n;d2+=8*n;d3+=8*n;d4+=8*n;
            d5+=8*n;d6+=8*n;d7+=8*n;d8+=8*n;
            j-=8;
        }
        s1+=7*n;s2+=7*n;s3+=7*n;s4+=7*n;
        s5+=7*n;s6+=7*n;s7+=7*n;s8+=7*n;
    }
}

同样转置4096的方阵进行测试，耗时为：

0.106431

可以看到我们之前的分析是合理的，使用8x8的子块转置时，效率又得到了提高。

再进一步我们可以在AVX指令中找到一个_mm256_stream_pd函数，这个函数和_mm256_store_pd的用法和作用完全相同，唯一的区别是使用_mm256_stream_pd将ymm寄存器中的数据写入主存时可以绕过Cache，直接写入主存，这样写入的时候就不用担心由于写入使用到Cache而淘汰掉Cache中还未读取完的Cache Line了。将4x4的kernel函数中的_mm256_store_pd替换为_mm256_stream_pd，再一次测试，得到4096的方阵转置的耗时为：

0.075976

比较原始的转置从0.324557到0.075976性能达到了之前的4倍，这也和AVX指令的理论加速效果相当，换句话说，尽管有时我们使用了SIMD加速，但是由于访存的问题，我们并不能完全得到预期的加速性能，这时可能就需要更进一步的优化才能达到理想效果。

上一篇：C++数值计算简单加速技术(四)——矩阵乘法优化示例(访存优化和SIMD的使用)
下一篇：没有了