如何用SSE更有效地乘A*B^T或A^T*B^T (T为转置)矩阵
How do I more efficiently multiply A*B^T or A^T*B^T (T for transpose) matrices using SSE?
我一直在用这件事打自己的头。我有一个基于sse的算法用于矩阵A乘以矩阵B。我还需要实现A, B，或者两者都转置的运算。我做了一个天真的实现，4x4矩阵代码如下所示(这是相当标准的SSE操作，我认为)，但A*B^T操作大约需要A*B的两倍。ATLAS实现对于A*B返回类似的值，对于乘以转置返回几乎相同的结果，这表明有一种有效的方法可以做到这一点。
MM-Multiplication:
m1 = (mat1.m_>>2)<<2;
n2 = (mat2.n_>>2)<<2;
n = (mat1.n_>>2)<<2;
for (k=0; k<n; k+=4) {
for (i=0; i<m1; i+=4) {
// fetch: get 4x4 matrix from mat1
// row-major storage, so get 4 rows
Float* a0 = mat1.el_[i]+k;
Float* a1 = mat1.el_[i+1]+k;
Float* a2 = mat1.el_[i+2]+k;
Float* a3 = mat1.el_[i+3]+k;
for (j=0; j<n2; j+=4) {
// fetch: get 4x4 matrix from mat2
// row-major storage, so get 4 rows
Float* b0 = mat2.el_[k]+j;
Float* b1 = mat2.el_[k+1]+j;
Float* b2 = mat2.el_[k+2]+j;
Float* b3 = mat2.el_[k+3]+j;
__m128 b0r = _mm_loadu_ps(b0);
__m128 b1r = _mm_loadu_ps(b1);
__m128 b2r = _mm_loadu_ps(b2);
__m128 b3r = _mm_loadu_ps(b3);
{ // first row of result += first row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a0+0), b0r), _mm_mul_ps(_mm_load_ps1(a0+1), b1r));
__m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a0+2), b2r), _mm_mul_ps(_mm_load_ps1(a0+3), b3r));
Float* c0 = this->el_[i]+j;
_mm_storeu_ps(c0, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c0)));
}
{ // second row of result += second row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a1+0), b0r), _mm_mul_ps(_mm_load_ps1(a1+1), b1r));
__m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a1+2), b2r), _mm_mul_ps(_mm_load_ps1(a1+3), b3r));
Float* c1 = this->el_[i+1]+j;
_mm_storeu_ps(c1, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c1)));
}
{ // third row of result += third row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a2+0), b0r), _mm_mul_ps(_mm_load_ps1(a2+1), b1r));
__m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a2+2), b2r), _mm_mul_ps(_mm_load_ps1(a2+3), b3r));
Float* c2 = this->el_[i+2]+j;
_mm_storeu_ps(c2, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c2)));
}
{ // fourth row of result += fourth row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a3+0), b0r), _mm_mul_ps(_mm_load_ps1(a3+1), b1r));
__m128 cX2 = _mm_add_ps(_mm_mul_ps(_mm_load_ps1(a3+2), b2r), _mm_mul_ps(_mm_load_ps1(a3+3), b3r));
Float* c3 = this->el_[i+3]+j;
_mm_storeu_ps(c3, _mm_add_ps(_mm_add_ps(cX1, cX2), _mm_loadu_ps(c3)));
}
}
// Code omitted to handle remaining rows and columns
}
对于MT乘法(矩阵乘以转置矩阵)，我使用以下命令将b0r存储到b3r，并适当地更改了循环变量:
__m128 b0r = _mm_set_ps(b3[0], b2[0], b1[0], b0[0]);
__m128 b1r = _mm_set_ps(b3[1], b2[1], b1[1], b0[1]);
__m128 b2r = _mm_set_ps(b3[2], b2[2], b1[2], b0[2]);
__m128 b3r = _mm_set_ps(b3[3], b2[3], b1[3], b0[3]);
我怀疑减速部分是由于一次拉入一行和每次必须存储4个值以获得列之间的差异，但我觉得另一种方式是，拉入B行，然后乘以a列，只会将成本转移到存储4列的结果。
我也试过拉B的行作为行，然后使用_MM_TRANSPOSE4_PS(b0r, b1r, b2r, b3r);做换位(我认为可能有一些额外的优化在该宏)，但没有真正的改进。
表面上，我觉得这应该更快…所涉及的点积将是一行接一行，这似乎天生更有效率，但试图直接做点积只会导致必须做同样的事情来存储结果。
我在这里错过了什么?
添加:只是为了澄清，我试图不转置矩阵。我更喜欢沿着它们进行迭代。据我所知，问题在于_mm_set_ps命令比_mm_load_ps命令慢得多。
我还尝试了一种变体，其中我存储了a矩阵的4行，然后用4个乘法指令和3个hadds替换了包含1个负载，4个乘法和2个加法的4个花括号段，但收效甚微。时间保持不变(是的，我尝试用调试语句来验证代码在我的测试编译中是否发生了变化)。上述调试语句在分析之前被删除，当然):
{ // first row of result += first row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a0r, b0r), _mm_mul_ps(a0r, b1r));
__m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a0r, b2r), _mm_mul_ps(a0r, b3r));
Float* c0 = this->el_[i]+j;
_mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
}
{ // second row of result += second row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a1r, b0r), _mm_mul_ps(a1r, b1r));
__m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a1r, b2r), _mm_mul_ps(a1r, b3r));
Float* c0 = this->el_[i+1]+j;
_mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
}
{ // third row of result += third row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a2r, b0r), _mm_mul_ps(a2r, b1r));
__m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a2r, b2r), _mm_mul_ps(a2r, b3r));
Float* c0 = this->el_[i+2]+j;
_mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
}
{ // fourth row of result += fourth row of mat1 * 4x4 of mat2
__m128 cX1 = _mm_hadd_ps(_mm_mul_ps(a3r, b0r), _mm_mul_ps(a3r, b1r));
__m128 cX2 = _mm_hadd_ps(_mm_mul_ps(a3r, b2r), _mm_mul_ps(a3r, b3r));
Float* c0 = this->el_[i+3]+j;
_mm_storeu_ps(c0, _mm_add_ps(_mm_hadd_ps(cX1, cX2), _mm_loadu_ps(c0)));
}
更新:对，将a0r到a3r的行加载到花括号中以避免寄存器抖动也失败了。
我认为这是水平添加有用的少数情况之一。你想要C = AB^T但是B不是作为转置存储在内存中。这就是问题所在。它就像一个aop而不是SoA。在这种情况下，取B的转置然后做垂直加法比用水平加法要慢。这至少对于矩阵向量是正确的，使用SSE进行高效的4x4矩阵向量乘法:水平加法和点积-重点是什么?在下面的代码中，函数m4x4是非SSE 4x4矩阵积，m4x4_vec使用SSE，m4x4T在不使用SSE的情况下实现C=AB^T, m4x4T_vec在使用SSE的情况下实现C=AB^T。我想最后一个才是你想要的。
注意:对于较大的矩阵，我不会使用这种方法。在这种情况下，首先进行转置并使用垂直相加会更快(对于SSE/AVX，你要做一些更复杂的事情，你要用SSE/AVX宽度转置条带)。这是因为转置等于O(n^2)矩阵乘积等于O(n^3)所以对于大矩阵，转置是不重要的。然而，对于4x4，转置是重要的，所以水平添加。
编辑:我误解了你想要什么。C = (AB)^T。这应该和(AB)一样快，代码几乎是一样的，你基本上只是交换了A和B的角色。
我们可以这样写数学:
C = A*B in Einstein notation is C_i,j = A_i,k * B_k,j.
Since (A*B)^T = B^T*A^T we can write
C = (A*B)^T in Einstein notation is C_i,j = B^T_i,k * A^T_k,j = A_j,k * B_k,i
如果你比较这两个，唯一改变的是我们交换了j和i的角色。我在这个答案的末尾放了一些代码来做这个。
#include "stdio.h"
#include <nmmintrin.h>
void m4x4(const float *A, const float *B, float *C) {
for(int i=0; i<4; i++) {
for(int j=0; j<4; j++) {
float sum = 0.0f;
for(int k=0; k<4; k++) {
sum += A[i*4+k]*B[k*4+j];
}
C[i*4 + j] = sum;
}
}
}
void m4x4T(const float *A, const float *B, float *C) {
for(int i=0; i<4; i++) {
for(int j=0; j<4; j++) {
float sum = 0.0f;
for(int k=0; k<4; k++) {
sum += A[i*4+k]*B[j*4+k];
}
C[i*4 + j] = sum;
}
}
}
void m4x4_vec(const float *A, const float *B, float *C) {
__m128 Brow[4], Mrow[4];
for(int i=0; i<4; i++) {
Brow[i] = _mm_load_ps(&B[4*i]);
}
for(int i=0; i<4; i++) {
Mrow[i] = _mm_set1_ps(0.0f);
for(int j=0; j<4; j++) {
__m128 a = _mm_set1_ps(A[4*i +j]);
Mrow[i] = _mm_add_ps(Mrow[i], _mm_mul_ps(a, Brow[j]));
}
}
for(int i=0; i<4; i++) {
_mm_store_ps(&C[4*i], Mrow[i]);
}
}
void m4x4T_vec(const float *A, const float *B, float *C) {
__m128 Arow[4], Brow[4], Mrow[4];
for(int i=0; i<4; i++) {
Arow[i] = _mm_load_ps(&A[4*i]);
Brow[i] = _mm_load_ps(&B[4*i]);
}
for(int i=0; i<4; i++) {
__m128 prod[4];
for(int j=0; j<4; j++) {
prod[j] = _mm_mul_ps(Arow[i], Brow[j]);
}
Mrow[i] = _mm_hadd_ps(_mm_hadd_ps(prod[0], prod[1]), _mm_hadd_ps(prod[2], prod[3]));
}
for(int i=0; i<4; i++) {
_mm_store_ps(&C[4*i], Mrow[i]);
}
}
float compare_4x4(const float* A, const float*B) {
float diff = 0.0f;
for(int i=0; i<4; i++) {
for(int j=0; j<4; j++) {
diff += A[i*4 +j] - B[i*4+j];
printf("A %f, B %fn", A[i*4 +j], B[i*4 +j]);
}
}
return diff;
}
int main() {
float *A = (float*)_mm_malloc(sizeof(float)*16,16);
float *B = (float*)_mm_malloc(sizeof(float)*16,16);
float *C1 = (float*)_mm_malloc(sizeof(float)*16,16);
float *C2 = (float*)_mm_malloc(sizeof(float)*16,16);
for(int i=0; i<4; i++) {
for(int j=0; j<4; j++) {
A[i*4 +j] = i*4+j;
B[i*4 +j] = i*4+j;
C1[i*4 +j] = 0.0f;
C2[i*4 +j] = 0.0f;
}
}
m4x4T(A, B, C1);
m4x4T_vec(A, B, C2);
printf("compare %fn", compare_4x4(C1,C2));
}
这里是C = (AB)^T的标量函数和SSE函数。它们应该和AB版本一样快。
void m4x4TT(const float *A, const float *B, float *C) {
for(int i=0; i<4; i++) {
for(int j=0; j<4; j++) {
float sum = 0.0f;
for(int k=0; k<4; k++) {
sum += A[j*4+k]*B[k*4+i];
}
C[i*4 + j] = sum;
}
}
}
void m4x4TT_vec(const float *A, const float *B, float *C) {
__m128 Arow[4], Crow[4];
for(int i=0; i<4; i++) {
Arow[i] = _mm_load_ps(&A[4*i]);
}
for(int i=0; i<4; i++) {
Crow[i] = _mm_set1_ps(0.0f);
for(int j=0; j<4; j++) {
__m128 a = _mm_set1_ps(B[4*i +j]);
Crow[i] = _mm_add_ps(Crow[i], _mm_mul_ps(a, Arow[j]));
}
}
for(int i=0; i<4; i++) {
_mm_store_ps(&C[4*i], Crow[i]);
}
}
一些可能有所帮助的建议:
- 不要使用未对齐的内存(那些_mm_loadu*很慢)。
- 您没有顺序访问内存，这会杀死缓存。在实际访问该内存之前尝试转置矩阵，这将使CPU尽可能多地获取和使用缓存。这样就不需要下面的
__m128 b0r = _mm_set_ps(b3[0], b2[0], b1[0], b0[0]); // and b1r, etc..了。我们的想法是按顺序抓取全部4个组件。如果您需要在SSE代码调用之前重新组织内存，请这样做。 - 你在内部循环中加载:
_mm_load_ps1(a0+0)(与a1, a2和a3相同)，但对内循环中的所有迭代都是恒定的。您可以在外部加载这些值并节省一些周期。关注一下你可以从之前的迭代中重用什么。 - 概要文件。使用英特尔VTune或类似的软件，它会告诉你瓶颈在哪里。
相关文章：
- 转置矩阵：交换元素不会更改值
- 我的转置矩阵代码有什么问题?
- 在C++中使用矢量转置 2D 矩阵
- MKL矩形矩阵Inplace转置:不使用多个核心
- 并行转置不同的矩阵
- 关于次级对角线的转置(翻转）矩阵
- 输出是从您输入的矩阵中打印出矩阵的转置，但我的代码只是打印出您输入的第一个矩阵
- CUDA矩阵与共享内存转置
- 在CUDA中具有共享MEM的非方面矩阵转置
- 通过矩阵转置优化矩阵乘法
- 矩阵矩形部分转置Cuda
- 矩阵转置模板
- 为什么在乘法之前转置矩阵会导致很大的加速
- 使用 c++ 转置矩阵
- 基于模板C++创建转置矩阵函数
- 内存黑客转置矩阵破坏堆栈，C++
- 用c++最快的转置矩阵的方法是什么?
- 如何用SSE更有效地乘A*B^T或A^T*B^T (T为转置)矩阵
- 为什么Eigen(版本3.2.5)不能得到正确的转置矩阵
- 转置矩阵会导致错误的输出
