没有任何技巧或任何东西,矩阵向量乘法只是向量和矩阵行之间的一堆点积。您的代码实际上并不具有该结构。实际上将其写为点积(未测试):
for (int row = 0; row < nrows; ++row) {
__m128 acc = _mm_setzero_ps();
// I'm just going to assume the number of columns is a multiple of 4
for (int col = 0; col < ncols; col += 4) {
__m128 vec = _mm_load_ps(&v[col]);
// don't forget it's a matrix, do 2d addressing
__m128 mat = _mm_load_ps(&m[col + ncols * row]);
acc = _mm_add_ps(acc, _mm_mul_ps(mat, vec));
}
// now we have 4 floats in acc and they have to be summed
// can use two horizontal adds for this, they kind of suck but this
// isn't the inner loop anyway.
acc = _mm_hadd_ps(acc, acc);
acc = _mm_hadd_ps(acc, acc);
// store result, which is a single float
_mm_store_ss(&result[row], acc);
}
有一些明显的技巧,例如一次处理多行、重用向量中的负载以及创建多个独立的依赖链,以便您可以更好地利用吞吐量(见下文)。另外一个非常简单的技巧是使用 FMA 进行 mul/add 组合,但支持还没有那么广泛(2015 年还没有,但现在到 2020 年已经相当广泛了)。
您可以由此构建矩阵-矩阵乘法(如果您更改结果所在的位置),但这不是最佳的(请参阅下文)。
一次取四行(未测试):
for (int row = 0; row < nrows; row += 4) {
__m128 acc0 = _mm_setzero_ps();
__m128 acc1 = _mm_setzero_ps();
__m128 acc2 = _mm_setzero_ps();
__m128 acc3 = _mm_setzero_ps();
for (int col = 0; col < ncols; col += 4) {
__m128 vec = _mm_load_ps(&v[col]);
__m128 mat0 = _mm_load_ps(&m[col + ncols * row]);
__m128 mat1 = _mm_load_ps(&m[col + ncols * (row + 1)]);
__m128 mat2 = _mm_load_ps(&m[col + ncols * (row + 2)]);
__m128 mat3 = _mm_load_ps(&m[col + ncols * (row + 3)]);
acc0 = _mm_add_ps(acc0, _mm_mul_ps(mat0, vec));
acc1 = _mm_add_ps(acc1, _mm_mul_ps(mat1, vec));
acc2 = _mm_add_ps(acc2, _mm_mul_ps(mat2, vec));
acc3 = _mm_add_ps(acc3, _mm_mul_ps(mat3, vec));
}
acc0 = _mm_hadd_ps(acc0, acc1);
acc2 = _mm_hadd_ps(acc2, acc3);
acc0 = _mm_hadd_ps(acc0, acc2);
_mm_store_ps(&result[row], acc0);
}
现在,每 4 个 FMA 只有 5 个负载,而在未行展开的版本中,每 1 个 FMA 有 2 个负载。此外,还有 4 个独立的 FMA,或无 FMA 收缩的加法/乘法对,无论哪种方式,它都增加了流水线/同时执行的潜力。实际上,您可能想要展开更多,例如 Skylake 每个周期可以启动 2 个独立的 FMA,它们需要 4 个周期才能完成,因此要完全占用两个 FMA 单元,您需要 8 个独立的 FMA。作为奖励,这 3 个水平相加最终对于水平求和来说效果相对较好。
不同的数据布局最初似乎是一个缺点,不再可能简单地从矩阵和向量进行向量加载并将它们相乘(这会将第一个矩阵的微小行向量乘以一个微小的行向量)row再次是第二个矩阵的向量,这是错误的)。但是,完整的矩阵-矩阵乘法可以利用这样一个事实:它本质上是将一个矩阵乘以许多独立的向量,它充满了需要完成的独立工作。水平总和也可以轻松避免。所以实际上它比矩阵向量乘法更方便。
关键是从矩阵 A 中取出一点列向量,从矩阵 B 中取出一点行向量,然后将它们相乘成一个小矩阵。与您习惯的方式相比,这听起来可能是相反的,但使用 SIMD 这样做效果更好,因为计算始终保持独立且无需水平操作。
例如(未经测试,假设矩阵的维度可被展开因子整除,需要 x64,否则会耗尽寄存器)
for (size_t i = 0; i < mat1rows; i += 4) {
for (size_t j = 0; j < mat2cols; j += 8) {
float* mat1ptr = &mat1[i * mat1cols];
__m256 sumA_1, sumB_1, sumA_2, sumB_2, sumA_3, sumB_3, sumA_4, sumB_4;
sumA_1 = _mm_setzero_ps();
sumB_1 = _mm_setzero_ps();
sumA_2 = _mm_setzero_ps();
sumB_2 = _mm_setzero_ps();
sumA_3 = _mm_setzero_ps();
sumB_3 = _mm_setzero_ps();
sumA_4 = _mm_setzero_ps();
sumB_4 = _mm_setzero_ps();
for (size_t k = 0; k < mat2rows; ++k) {
auto bc_mat1_1 = _mm_set1_ps(mat1ptr[0]);
auto vecA_mat2 = _mm_load_ps(mat2 + m2idx);
auto vecB_mat2 = _mm_load_ps(mat2 + m2idx + 4);
sumA_1 = _mm_add_ps(_mm_mul_ps(bc_mat1_1, vecA_mat2), sumA_1);
sumB_1 = _mm_add_ps(_mm_mul_ps(bc_mat1_1, vecB_mat2), sumB_1);
auto bc_mat1_2 = _mm_set1_ps(mat1ptr[N]);
sumA_2 = _mm_add_ps(_mm_mul_ps(bc_mat1_2, vecA_mat2), sumA_2);
sumB_2 = _mm_add_ps(_mm_mul_ps(bc_mat1_2, vecB_mat2), sumB_2);
auto bc_mat1_3 = _mm_set1_ps(mat1ptr[N * 2]);
sumA_3 = _mm_add_ps(_mm_mul_ps(bc_mat1_3, vecA_mat2), sumA_3);
sumB_3 = _mm_add_ps(_mm_mul_ps(bc_mat1_3, vecB_mat2), sumB_3);
auto bc_mat1_4 = _mm_set1_ps(mat1ptr[N * 3]);
sumA_4 = _mm_add_ps(_mm_mul_ps(bc_mat1_4, vecA_mat2), sumA_4);
sumB_4 = _mm_add_ps(_mm_mul_ps(bc_mat1_4, vecB_mat2), sumB_4);
m2idx += 8;
mat1ptr++;
}
_mm_store_ps(&result[i * mat2cols + j], sumA_1);
_mm_store_ps(&result[i * mat2cols + j + 4], sumB_1);
_mm_store_ps(&result[(i + 1) * mat2cols + j], sumA_2);
_mm_store_ps(&result[(i + 1) * mat2cols + j + 4], sumB_2);
_mm_store_ps(&result[(i + 2) * mat2cols + j], sumA_3);
_mm_store_ps(&result[(i + 2) * mat2cols + j + 4], sumB_3);
_mm_store_ps(&result[(i + 3) * mat2cols + j], sumA_4);
_mm_store_ps(&result[(i + 3) * mat2cols + j + 4], sumB_4);
}
}
该代码的要点在于,很容易将计算安排得非常 SIMD 友好,使用大量独立算术来使浮点单元饱和,同时使用相对较少的负载(否则可能会成为瓶颈) ,即使不考虑它们可能会错过 L1 缓存,只是因为它们太多了)。
您甚至可以使用此代码,但它与英特尔 MKL 没有竞争力。特别是对于中型或大型矩阵,平铺极其重要。升级到 AVX 很容易。它根本不适合微小矩阵,例如将两个 4x4 矩阵相乘,请参见高效的 4x4 矩阵乘法.