矩陣乘法示例
矩陣乘法是在許多數據密集型應用程序中執行的操作。 它由以簡單方式重複的算術運算組組成:
矩陣乘法過程如下:
A-在第一個矩陣中進行一行
B-執行該行的點積與第二個矩陣中的一列
C-將結果存儲在新矩陣的相應行和列中
對於32位浮點矩陣,乘法可以寫爲:
void matrix_multiply_c(float32_t *A, float32_t *B, float32_t *C, uint32_t n, uint32_t m, uint32_t k) {
for (int i_idx=0; i_idx < n; i_idx++) {
for (int j_idx=0; j_idx < m; j_idx++) {
C[n*j_idx + i_idx] = 0;
for (int k_idx=0; k_idx < k; k_idx++) {
C[n*j_idx + i_idx] += A[n*k_idx + i_idx]*B[k*j_idx + k_idx];
}
}
}
}
我們假設內存中矩陣的主要列布局。 即,將n×m矩陣M表示爲陣列M_array,其中Mij = M_array [n * j + i]。
該代碼不是最佳代碼,因爲它沒有充分利用Neon。 我們可以開始使用內在函數對其進行改進,但是讓我們先解決一個簡單的問題,即先查看固定大小的小型矩陣,然後再處理較大的矩陣。
以下代碼使用內部函數將兩個4x4矩陣相乘。 由於我們要處理的值很小且固定,所有這些值都可以一次放入處理器的Neon寄存器中,因此我們可以完全展開循環。
void matrix_multiply_4x4_neon(float32_t *A, float32_t *B, float32_t *C) {
// these are the columns A
float32x4_t A0;
float32x4_t A1;
float32x4_t A2;
float32x4_t A3;
// these are the columns B
float32x4_t B0;
float32x4_t B1;
float32x4_t B2;
float32x4_t B3;
// these are the columns C
float32x4_t C0;
float32x4_t C1;
float32x4_t C2;
float32x4_t C3;
A0 = vld1q_f32(A);
A1 = vld1q_f32(A+4);
A2 = vld1q_f32(A+8);
A3 = vld1q_f32(A+12);
// Zero accumulators for C values
C0 = vmovq_n_f32(0);
C1 = vmovq_n_f32(0);
C2 = vmovq_n_f32(0);
C3 = vmovq_n_f32(0);
// Multiply accumulate in 4x1 blocks, i.e. each column in C
B0 = vld1q_f32(B);
C0 = vfmaq_laneq_f32(C0, A0, B0, 0);
C0 = vfmaq_laneq_f32(C0, A1, B0, 1);
C0 = vfmaq_laneq_f32(C0, A2, B0, 2);
C0 = vfmaq_laneq_f32(C0, A3, B0, 3);
vst1q_f32(C, C0);
B1 = vld1q_f32(B+4);
C1 = vfmaq_laneq_f32(C1, A0, B1, 0);
C1 = vfmaq_laneq_f32(C1, A1, B1, 1);
C1 = vfmaq_laneq_f32(C1, A2, B1, 2);
C1 = vfmaq_laneq_f32(C1, A3, B1, 3);
vst1q_f32(C+4, C1);
B2 = vld1q_f32(B+8);
C2 = vfmaq_laneq_f32(C2, A0, B2, 0);
C2 = vfmaq_laneq_f32(C2, A1, B2, 1);
C2 = vfmaq_laneq_f32(C2, A2, B2, 2);
C2 = vfmaq_laneq_f32(C2, A3, B2, 3);
vst1q_f32(C+8, C2);
B3 = vld1q_f32(B+12);
C3 = vfmaq_laneq_f32(C3, A0, B3, 0);
C3 = vfmaq_laneq_f32(C3, A1, B3, 1);
C3 = vfmaq_laneq_f32(C3, A2, B3, 2);
C3 = vfmaq_laneq_f32(C3, A3, B3, 3);
vst1q_f32(C+12, C3);
}
由於某些原因,我們選擇將固定大小的4x4矩陣相乘:
一些應用程序特別需要4x4矩陣,例如圖形或相對論物理學。
Neon向量寄存器具有四個32位值,因此將程序與體系結構進行匹配將使其更易於優化。
我們可以採用這種4x4內核,並在更通用的內核中使用它。
讓我們總結一下這裏使用的內在函數:
現在我們可以將一個4x4矩陣相乘,可以將較大的矩陣視爲4x4矩陣的塊來相乘。 這種方法的一個缺點是,它僅適用於二維尺寸均爲四的倍數的矩陣大小,但是通過將任何矩陣都填充爲零,您可以使用此方法而無需對其進行更改。
下面列出了用於更通用的矩陣乘法的代碼。 內核的結構變化很小,主要的變化是增加了循環和地址計算。 與在4x4內核中一樣,即使我們可以使用一個變量並重新加載,我們也爲B列使用了唯一的變量名。 這提示編譯器將不同的寄存器分配給這些變量,這將使處理器能夠在等待另一列加載的同時完成一列的算術指令。
void matrix_multiply_neon(float32_t *A, float32_t *B, float32_t *C, uint32_t n, uint32_t m, uint32_t k) {
/*
* Multiply matrices A and B, store the result in C.
* It is the user's responsibility to make sure the matrices are compatible.
*/
int A_idx;
int B_idx;
int C_idx;
// these are the columns of a 4x4 sub matrix of A
float32x4_t A0;
float32x4_t A1;
float32x4_t A2;
float32x4_t A3;
// these are the columns of a 4x4 sub matrix of B
float32x4_t B0;
float32x4_t B1;
float32x4_t B2;
float32x4_t B3;
// these are the columns of a 4x4 sub matrix of C
float32x4_t C0;
float32x4_t C1;
float32x4_t C2;
float32x4_t C3;
for (int i_idx=0; i_idx<n; i_idx+=4 {
for (int j_idx=0; j_idx<m; j_idx+=4){
// zero accumulators before matrix op
c0=vmovq_n_f32(0);
c1=vmovq_n_f32(0);
c2=vmovq_n_f32(0);
c3=vmovq_n_f32(0);
for (int k_idx=0; k_idx<k; k_idx+=4){
// compute base index to 4x4 block
a_idx = i_idx + n*k_idx;
b_idx = k*j_idx k_idx;
// load most current a values in row
A0=vld1q_f32(A+A_idx);
A1=vld1q_f32(A+A_idx+n);
A2=vld1q_f32(A+A_idx+2*n);
A3=vld1q_f32(A+A_idx+3*n);
// multiply accumulate 4x1 blocks, i.e. each column C
B0=vld1q_f32(B+B_idx);
C0=vfmaq_laneq_f32(C0,A0,B0,0);
C0=vfmaq_laneq_f32(C0,A1,B0,1);
C0=vfmaq_laneq_f32(C0,A2,B0,2);
C0=vfmaq_laneq_f32(C0,A3,B0,3);
B1=v1d1q_f32(B+B_idx+k);
C1=vfmaq_laneq_f32(C1,A0,B1,0);
C1=vfmaq_laneq_f32(C1,A1,B1,1);
C1=vfmaq_laneq_f32(C1,A2,B1,2);
C1=vfmaq_laneq_f32(C1,A3,B1,3);
B2=vld1q_f32(B+B_idx+2*k);
C2=vfmaq_laneq_f32(C2,A0,B2,0);
C2=vfmaq_laneq_f32(C2,A1,B2,1);
C2=vfmaq_laneq_f32(C2,A2,B2,2);
C2=vfmaq_laneq_f32(C2,A3,B3,3);
B3=vld1q_f32(B+B_idx+3*k);
C3=vfmaq_laneq_f32(C3,A0,B3,0);
C3=vfmaq_laneq_f32(C3,A1,B3,1);
C3=vfmaq_laneq_f32(C3,A2,B3,2);
C3=vfmaq_laneq_f32(C3,A3,B3,3);
}
//Compute base index for stores
C_idx = n*j_idx + i_idx;
vstlq_f32(C+C_idx, C0);
vstlq_f32(C+C_idx+n,Cl);
vstlq_f32(C+C_idx+2*n,C2);
vstlq_f32(C+C_idx+3*n,C3);
}
}
}
編譯和反彙編此函數,並將其與我們的C函數進行比較,結果顯示:
給定矩陣乘法的算術指令較少,因爲我們利用具有完整寄存器打包功能的Advanced SIMD技術。 純C代碼通常不這樣做。
FMLA代替FMUL指令。 如內在函數所指定。
更少的循環迭代。 如果正確使用內在函數,則可以輕鬆展開循環。
但是,由於內存分配和數據類型(例如,float32x4_t)的初始化而導致不必要的加載和存儲,而純C代碼中未使用這些數據類型。
可以使用以下命令在Arm機器上編譯和反彙編上面的完整源代碼:
gcc -g -o3 matrix.c -o exe_matrix_o3
objdump -d exe_ matrix _o3 > disasm_matrix_o3
如果您無權使用基於Arm的硬件,則可以使用Arm DS-5 Community Edition和Armv8-A Foundation Platform。