循环取矩阵的某行_1.2 震惊! 某大二本科生写的矩阵乘法吊打Mathematica-线性代数库BLAS-矩阵 (上)...

本文是

1. 线性代数库BLASzhuanlan.zhihu.com

系列的第二篇, 将讲述矩阵类的结构和矩阵基础运算的AVX2加速算法.

1. 矩阵类的结构

在讲述矩阵各种算法之前很有必要详解一下普通矩阵和各种常用的特殊矩阵 (包括方阵, 三对角矩阵, 对称矩阵, 上下三角矩阵, 带状矩阵, 上下三角带状矩阵和稀疏矩阵)的基本数据结构.

为了表明某一矩阵是特殊矩阵, 我创建了一个枚举类MatType以指示矩阵的种类. 而为了更快速的区分各种带状矩阵 (储存结构和前面的矩阵不一样) 和稀疏矩阵, 我采用了以下安排方式

enum class MatType
{NormalMat,SquareMat,DiagonalMat,SymmetricMat,LMat,UMat,//for band mat, width == heightBandMat,LBandMat,UBandMat,//sparse matSparseMat,
};

这样只需要用比较matType和BandMat, SparseMat的大小就能判断是否是带状矩阵或者稀疏矩阵. 而对于向量, 还存在这样一种分类:

enum class Type
{Native = 0,Parasitic = 1,Non32Aligened = 2,//must be Parasitic as well...
};

Native即原生的, 构造时是在堆上分配了内存的; Parasitic指的是32字节对齐的寄生向量, Non32Aligened则是没有32字节对齐的寄生向量, 一般是向量的子向量或者矩阵的一行等等. 这些寄生向量在构造和析构时不会进行内存分配和释放, 避免构造中间向量, 提高性能.

而类mat的成员如下:

double* data;
union
{unsigned long long width;unsigned long long halfBandWidth;//if is BandMatunsigned long long* rowIndice;//if is SparseMat
};
union
{unsigned long long height;unsigned long long* colIndice;//if is SparseMat
};
union
{unsigned long long width4d;unsigned long long elementNum;
};
Type type;
MatType matType;

由于一个矩阵不能同时为普通矩阵, 带状矩阵或者稀疏矩阵, 所以可以利用union来节约空间, 并在进行各种矩阵运算的时候根据MatType来使用对应的方式来访问矩阵.

接下来则是详解各种矩阵元素的存储结构.

1. 非带状矩阵

即原始的各种矩阵, 最naive的想法是直接按照行优先 (即同一行的元素是连续的)线性铺开, 但是很快我们就可以发现一个问题: 由于我们分配的内存是4 double对齐的, 所以第一行是一个正常的对齐的向量, 但是若矩阵的列数是一个奇数, 那么4n+1, 4n+2, 4n+3行都不是4 double对齐的向量. 显然这不是我们想要的, 因为在很多情况下, 计算是围绕这一个行向量展开的, 倘若3/4的行向量都不是对齐的, 那对性能的影响是很大的. 所以聪明的你肯定想到了这样的作法: 将行向量的末尾用4 double补齐, 再转到下一行. 这样虽然有部分空间没有利用上, 但是在这个动辄64GB内存的大环境 (自己在做计算物理大作业的时候确实用到过这么多, 即

维电阻网络的两点间电阻计算), 相信大家一定不会在意这一点点的"浪费".

2. 带状矩阵

a. 对于半带宽为halfBandWidth的n阶带状矩阵, 计算物理课上提示的是稠密存储, 即按照

这种模式一个一个的顺序存储 (空位不被储存). 我们发现这样存储对与行向量的4 double对齐非常不友好, 所以同理, 需要在前后加上padding以使每个行向量能够4 double对齐 (注意这里的对齐指的是按照列号来对齐而不是按照该元素是该行第n个非0元的n来对齐, 例如带状矩阵乘向量时, 向量是4 double对齐的, 所以按照矩阵乘法, 矩阵某行的列序号是要对齐的). 按照这种对齐方式, 我们可以计算出每行至少为ceiling4(2*halfBandWidth+4)个元素, 其中ceiling4函数是按4的模向上取整, 例如ceiling4(5)=8等等.

b. 同理, 对于半带宽为halfBandWidth的n阶上/下三角带状矩阵, 即(半带宽为2的6阶下三角矩阵)

每行至少为ceiling4(halfBandWidth+4)个元素.

3. 稀疏矩阵

这里采用行优先的方式顺序存储, 即下面的稀疏矩阵

按照{(1, 1,

), (1, 2,

), (1, 6,

), (2, 2,

), ... , (6, 6,

)}的方式储存, 但是这里的(rowIndex, columnIndex, element)不是直接连续排放在一起的, 而是rowIndex放成一个数组, columnIndex放成一个数组, element放成一个数组.

2. 普通矩阵的基础运算

计算物理中, 很多算法是针对某种特殊矩阵的, 需要与普通矩阵的运算分开. 今天这篇只会讲普通矩阵的运算, 我们有以下常用的矩阵基础运算 (剩下的线性代数相关的运算和算法等会放在下篇):

对应位置上的加减乘除 (包括原址和非原址操作, 原址即结果直接写入原有矩阵而非重新分配一个矩阵)
整体加减乘除常数 (原址和非原址)
矩阵乘向量
矩阵乘矩阵

由于上一篇讲向量的过程中已经讲过了基础的对应位置的各种算法, 这里就不在重复讲述了, 我将把重点放在很多情况下需要深度优化的矩阵乘向量和矩阵乘法, 例如矩阵乘向量在各种迭代法求解线性方程组中是非常重要的操作, 好的优化至关重要.

3. 矩阵乘向量

由于代码过于冗长 (因为需要考虑到各种边界上没对齐问题), 这里只用部分代码来解释具体的并行化方法, 并且只讲述普通矩阵的优化方法 (其他矩阵如带状矩阵和稀疏矩阵并没有进行特殊优化).

矩阵乘向量

, trivial的想法是将矩阵视为多个行向量, 然后和向量b点乘得到结果c的一个元素, 求和的顺序即矩阵行优先 (指的是求和的最内部的循环是同一行).

但是显然这总方式的并行化程度很差, 我们可以从更量化的角度来看这个问题.

在GPU的并行编程中有一个概念叫计算-读写比, 即计算操作与读写操作的比例, 当这个比例为算力/显存带宽时, 计算效率达能到最大化, 而且在很多情况下读写是瓶颈. 以刚才的naive方法为例, 假设我们已经使用了向量点乘的AVX2优化, 那么进行一次元运算 (对于一个AVX2寄存器来说, 一次元操作为对4 double的计算)fmadd需要读写8 double (矩阵a的4 double和向量b的4 double, 由于结果是写在另一个寄存器内, 所以不考虑这部分的写入), 但是只进行了一次对4 double的fmadd操作, 所以计算-读写比为0.5.

这个比例显然没有达到我们的预期. 首先想到读取一次向量b的4 double后可以读取矩阵a的多行的4 double (我选取了4行, 这时候的计算-读写比为0.8, 可以自行验算), 这样就能重复利用向量b的数据, 达到加速的目的. 同时, 由于单核的AVX2寄存器数量很多 (一般至少有16个), 所以为了增大输出的口径 (原理是编译器会将这些相邻的计算分配到不同的寄存器以使得多个寄存器能并行地在流水线上工作), 可以将向量b分成多组, 每一组含有多个__m256d数据 (经过测试8个较为合适, 更大的话一般不会更快, 反而会使得剩余的最后一组变得更大, 耗时更久).

constexpr unsigned long long warp = 8;//将输入向量分组为每组8个__m256d
unsigned long long minWidth4((minDim - 1) / 4 + 1);
__m256d* aData((__m256d*)data);//矩阵
__m256d* bData((__m256d*)source->data);//输入向量
__m256d* rData((__m256d*)b.data);//输出向量
unsigned long long heightFloor4((height >> 2) << 2);
unsigned long long widthWarp((minWidth4 / warp) * warp);
unsigned long long warpLeftFloor((minDim >> 2) - widthWarp);
unsigned long long warpLeftCeiling(minWidth4 - widthWarp);
unsigned long long c0(0);
for (; c0 < heightFloor4; c0 += 4)//每次计算4行
{__m256d ans[4] = { 0 };__m256d tp[warp];unsigned long long c1(0);for (; c1 < widthWarp; c1 += warp){__m256d* s(aData + minWidth4 * c0 + c1);
#pragma unroll(4)for (unsigned long long c2(0); c2 < warp; ++c2)tp[c2] = bData[c1 + c2];for (unsigned long long c2(0); c2 < 4; ++c2, s += minWidth4){#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){__m256d t = s[c3];ans[c2] = _mm256_fmadd_pd(t, tp[c3], ans[c2]);}}}if (c1 < minWidth4)//处理没有分组对齐的部分{__m256d* s(aData + minWidth4 * c0 + c1);
#pragma unroll(4)for (unsigned long long c2(0); c2 < warpLeftCeiling; ++c2)tp[c2] = bData[c1 + c2];unsigned long long finalWidth(minDim - ((minDim >> 2) << 2));for (unsigned long long c2(finalWidth); c2 < 4; ++c2)tp[warpLeftFloor].m256d_f64[c2] = 0;for (unsigned long long c2(0); c2 < 4; ++c2, s += minWidth4){#pragma unroll(4)for (unsigned long long c3(0); c3 < warpLeftCeiling; ++c3){__m256d t = s[c3];ans[c2] = _mm256_fmadd_pd(t, tp[c3], ans[c2]);}}}__m256d s;for (unsigned long long c1(0); c1 < 4; ++c1){s.m256d_f64[c1] = ans[c1].m256d_f64[0];s.m256d_f64[c1] += ans[c1].m256d_f64[1];s.m256d_f64[c1] += ans[c1].m256d_f64[2];s.m256d_f64[c1] += ans[c1].m256d_f64[3];}rData[c0 >> 2] = s;
}
//后面是处理剩余的非4行对齐的部分

4. 矩阵乘矩阵

此处是本文的重头戏, 我才不是标题党呢, 1024阶方阵乘法用Ryzen 3950X实测数据如下 (依次为手写AVX2优化, Mathematica和matlab的计算时间, 由于CPU频率和其他负载等因素导致的时间计算不准确, 计算耗时都已经取了多次的最快的那次, 并且都禁用了打印结果并提前分配好储存答案的内存. 手写版本的正确性已经通过Mathematica计算与正确答案差的范数验证过了, 所有计算方式均已开启多线程优化以利用所有核心和超线程):

可以看出, 手写的矩阵乘法确实比Mathematica快了4倍有余, 和使用mkl的matlab速度几乎一致 (当然仍然会被CUDA的cublas爆锤). 那么, 这么快的矩阵乘法是咋写的呢?

有了上面矩阵乘向量的经验, 我们明白了一件事情: 想要提高速度, 就要让一次读取的数据能尽可能多的被用于计算. 有了这样的想法, 矩阵乘法的AVX2优化也就呼之欲出:

我们需要从另一个角度来看矩阵乘法, 而非和之前的矩阵乘向量一样将右矩阵视为多个列向量. 矩阵乘法

, 这时候从这个乘法求和式可以看出, 矩阵a的第i行的第k列

要遍历地乘以矩阵b的第k行的所有元素, 得到一行向量

(这里i, k是固定的, 所以得到一个以j为列指标的行向量), 对k求和就可以得到最终矩阵c的第i行

(这里i是固定的, 以j为列指标). 到这里并行化的方式已经很明显了, 就是将

利用之前学过的向量乘常数的算法加速. 但是需要注意的是, 由于k的求和范围1~n可以很大 (如1024阶矩阵就是1024), 这样就无法将计算结果

的整个行向量储存在寄存器内, 所以也需要将其分组, 经过测试我选择了16个__m256d为一组, 这样每次得到的也就是矩阵c的第i行

的一组. 而由于矩阵a的两行

等需要乘以b的同一行

, 所以可以利用这一点再同时计算这两行以提高数据利用程度. 经过了以上优化, 大家可以算一下计算-读写比, 以下代码是读了(2+64)个double, 由于是在内循环外, 所以在b的行数较大的时候基本可以忽略写入, 计算了128次fmadd, 所以计算-读写比为1.94, 相比传统的三个循环的算法的读2算1的0.5高了近3倍, 而且这里除了那2个double不是连续读取, 剩下所有的读写都是连续的, 对于内存来说十分友好.

//source是乘法的左矩阵
__m256d* aData((__m256d*)a.data);//乘法的右矩阵
__m256d* bData((__m256d*)b.data);//乘法结果
constexpr unsigned long long warp = 16;
unsigned long long aWidth256d(a.width4d / 4);
unsigned long long aWidthWarpFloor(aWidth256d / warp * warp);
unsigned long long warpLeft(aWidth256d - aWidthWarpFloor);
unsigned long long height2Floor(height & (-2));
unsigned long long c0(0);
for (; c0 < height2Floor; c0 += 2)
{unsigned long long c1(0);for (; c1 < aWidthWarpFloor; c1 += warp){__m256d ans0[warp] = { 0 };__m256d ans1[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source->data[c0 * width4d + c2]);__m256d tp1 = _mm256_set1_pd(source->data[(c0 + 1) * width4d + c2]);
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);ans1[c3] = _mm256_fmadd_pd(tp1, b, ans1[c3]);}}
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){bData[c0 * aWidth256d + c1 + c3] = ans0[c3];bData[(c0 + 1) * aWidth256d + c1 + c3] = ans1[c3];}}if (c1 < aWidth256d)//对行向量分组后剩余的部分{__m256d ans0[warp] = { 0 };__m256d ans1[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source->data[c0 * width4d + c2]);__m256d tp1 = _mm256_set1_pd(source->data[(c0 + 1) * width4d + c2]);for (unsigned long long c3(0); c3 < warpLeft; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);ans1[c3] = _mm256_fmadd_pd(tp1, b, ans1[c3]);}}for (unsigned long long c3(0); c3 < warpLeft; ++c3){bData[c0 * aWidth256d + c1 + c3] = ans0[c3];bData[(c0 + 1) * aWidth256d + c1 + c3] = ans1[c3];}}
}
//后面还需计算因每次计算2行导致的可能的剩余的一行

至于多线程版本的矩阵乘法, 假设矩阵a有n行, 总共有m个线程, 由于是每2行需要同时计算, 而不同的2行之间是独立的, 所以可以将连续的floor2(n/m)行分配给同一线程 (floor2指的是对2取模的向下取整, 例floor2(3) = 2), 并让原线程 (调用矩阵乘法函数的那个线程)来计算最后的那一组和剩下的那行 (如果n是奇数). 由于不想使用全局函数来写每个线程调用的子函数, 而创建线程的函数_beginthread需要的线程函数不能是带this指针的类函数, 所以在矩阵乘法函数内嵌一个lambda表达式是再好不过的了 (不用beginthreadex是因为lambda函数是__cdecl的, 而beginthreadex需要的是__stdcall函数作为线程函数). 最终代码如下:

BLAS::mat& matMultMT(BLAS::mat const& ts, BLAS::mat const& a, BLAS::mat& b)
{using namespace BLAS;SYSTEM_INFO systemInfo;GetSystemInfo(&systemInfo);unsigned long long threadNum(systemInfo.dwNumberOfProcessors);::printf("Number of processors: %llun", threadNum);unsigned long long minDim(ts.width > a.height ? a.height : ts.width);if (minDim){bool overflow(ceiling4(a.width, ts.height) > b.width4d * b.height);if (overflow && b.type != Type::Native)return b;mat const* source(&ts);mat r;if (&b == &ts){source = &r;r = ts;}if (overflow)b.reconstruct(a.width, ts.height, false);__m256d* aData((__m256d*)a.data);__m256d* bData((__m256d*)b.data);constexpr unsigned long long warp = 16;unsigned long long aWidth256d(a.width4d / 4);unsigned long long aWidthWarpFloor(aWidth256d / warp * warp);unsigned long long warpLeft(aWidth256d - aWidthWarpFloor);unsigned long long height2Floor(ts.height & (-2));//ptr: A beginning//     A ending//     B//     C beginning//     width4d//     aWidthWarpFloor//     minDim//     aWidth256d//     warpLeftvoid(*lambda)(void*) = [](void* ptr){void** ptrs((void**)ptr);double* source = ((double*)ptrs[0]);double* sourceEnding((double*)ptrs[1]);__m256d* aData((__m256d*)ptrs[2]);__m256d* bData((__m256d*)ptrs[3]);unsigned long long width4d((unsigned long long)ptrs[4]);unsigned long long aWidthWarpFloor((unsigned long long)ptrs[5]);unsigned long long minDim((unsigned long long)ptrs[6]);unsigned long long aWidth256d((unsigned long long)ptrs[7]);unsigned long long warpLeft((unsigned long long)ptrs[8]);constexpr unsigned long long warp = 16;for (; source < sourceEnding; source += width4d * 2, bData += aWidth256d * 2){unsigned long long c1(0);for (; c1 < aWidthWarpFloor; c1 += warp){__m256d ans0[warp] = { 0 };__m256d ans1[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source[c2]);__m256d tp1 = _mm256_set1_pd(source[width4d + c2]);
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);ans1[c3] = _mm256_fmadd_pd(tp1, b, ans1[c3]);}}
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){bData[c1 + c3] = ans0[c3];bData[aWidth256d + c1 + c3] = ans1[c3];}}if (c1 < aWidth256d){__m256d ans0[warp] = { 0 };__m256d ans1[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source[c2]);__m256d tp1 = _mm256_set1_pd(source[width4d + c2]);for (unsigned long long c3(0); c3 < warpLeft; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);ans1[c3] = _mm256_fmadd_pd(tp1, b, ans1[c3]);}}for (unsigned long long c3(0); c3 < warpLeft; ++c3){bData[c1 + c3] = ans0[c3];bData[aWidth256d + c1 + c3] = ans1[c3];}}}};unsigned long long threadHeight((ts.height / threadNum) & (-2));//floor2unsigned long long rowBeginning;HANDLE* threads(nullptr);void** paras(nullptr);if (threadHeight)//如果每个线程都没法分配到至少2行, 那就直接用单线程算了{rowBeginning = threadHeight * (threadNum - 1);threads = (HANDLE*)::malloc((threadNum - 1) * sizeof(HANDLE));paras = (void**)::malloc((threadNum - 1) * 9 * sizeof(void*));for (unsigned long long c0(0); c0 < threadNum - 1; ++c0){paras[c0 * 9] = (void*)(source->data + threadHeight * ts.width4d * c0);paras[c0 * 9 + 1] = (void*)(source->data + threadHeight * ts.width4d * (c0 + 1));paras[c0 * 9 + 2] = (void*)(a.data);paras[c0 * 9 + 3] = (void*)(b.data + threadHeight * a.width4d * c0);paras[c0 * 9 + 4] = (void*)(ts.width4d);paras[c0 * 9 + 5] = (void*)(aWidthWarpFloor);paras[c0 * 9 + 6] = (void*)(minDim);paras[c0 * 9 + 7] = (void*)(aWidth256d);paras[c0 * 9 + 8] = (void*)(warpLeft);threads[c0] = (HANDLE)_beginthread(lambda, 0, paras + c0 * 9);}}else rowBeginning = 0;//接下来的部分就是原线程求最后剩下的多行unsigned long long c0(rowBeginning);for (; c0 < height2Floor; c0 += 2){unsigned long long c1(0);for (; c1 < aWidthWarpFloor; c1 += warp){__m256d ans0[warp] = { 0 };__m256d ans1[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source->data[c0 * ts.width4d + c2]);__m256d tp1 = _mm256_set1_pd(source->data[(c0 + 1) * ts.width4d + c2]);
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);ans1[c3] = _mm256_fmadd_pd(tp1, b, ans1[c3]);}}
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){bData[c0 * aWidth256d + c1 + c3] = ans0[c3];bData[(c0 + 1) * aWidth256d + c1 + c3] = ans1[c3];}}if (c1 < aWidth256d){__m256d ans0[warp] = { 0 };__m256d ans1[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source->data[c0 * ts.width4d + c2]);__m256d tp1 = _mm256_set1_pd(source->data[(c0 + 1) * ts.width4d + c2]);for (unsigned long long c3(0); c3 < warpLeft; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);ans1[c3] = _mm256_fmadd_pd(tp1, b, ans1[c3]);}}for (unsigned long long c3(0); c3 < warpLeft; ++c3){bData[c0 * aWidth256d + c1 + c3] = ans0[c3];bData[(c0 + 1) * aWidth256d + c1 + c3] = ans1[c3];}}}if (c0 < ts.height){unsigned long long c1(0);for (; c1 < aWidthWarpFloor; c1 += warp){__m256d ans0[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source->data[c0 * ts.width4d + c2]);
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);}}
#pragma unroll(4)for (unsigned long long c3(0); c3 < warp; ++c3)bData[c0 * aWidth256d + c1 + c3] = ans0[c3];}if (c1 < aWidth256d){__m256d ans0[warp] = { 0 };for (unsigned long long c2(0); c2 < minDim; ++c2){__m256d tp0 = _mm256_set1_pd(source->data[c0 * ts.width4d + c2]);for (unsigned long long c3(0); c3 < warpLeft; ++c3){__m256d b = aData[aWidth256d * c2 + c1 + c3];ans0[c3] = _mm256_fmadd_pd(tp0, b, ans0[c3]);}}for (unsigned long long c3(0); c3 < warpLeft; ++c3)bData[c0 * aWidth256d + c1 + c3] = ans0[c3];}}if (threadHeight){WaitForMultipleObjects(threadNum - 1, threads, true, INFINITE);for (unsigned long long c0(0); c0 < threadNum - 1; ++c0)CloseHandle(threads[c0]);::free(threads);::free(paras);}}return b;
}