NEON优化3：矩阵转置的指令优化案例

背景

---

矩阵运算中经常用到转置操作，这里将原子矩阵4x4的转置NEON优化案例总结一下。

原始C函数负责将M[4][4]转置为MT[4][4]，效果如下：

• M[4][4]
  • a1, b1, c1, d1
  • a2, b2, c2, d2
  • a3, b3, c3, d3
  • a4, b4, c4, d4
• M[4][4]^T
  • a1, a2, a3, a4
  • b1, b2, b3, b4
  • c1, c2, c3, c4
  • d1, d2, d3, d4

优化思路

---

有两种并行计算的方法将其转置。

• 方法1
  • 用到4条指令：vld4q_f32/vtrnq_f32/vuzpq_f32/vst4q_f32
  • ld4q先从内存一批交叉读入到寄存器中
  • trnq利用内部两两转置功能，将部分行列进行转置
  • uzpq利用解交织的读写功能，将部分行列进行转置
  • 然后利用寄存器单独赋值给不同行
  • st4q将寄存器结果写入到内存中
• 方法2
  • 利用内存和寄存器的交叉读写关系，两条指令实现转置，不用在寄存器中倒来倒去
  • ld1q实现按行读取数据到寄存器
  • st4q实现交叉读取寄存器值然后写入到内存中

样例代码

---

#include <stdio.h>
#include <stdint.h>
#include <arm_neon.h>

#define ROW_NUM   4
#define COL_NUM   4

int main(void)
{
    // initial
    float M[ROW_NUM][COL_NUM] = {
        {0, 1, 2, 3},
        {4, 5, 6, 7},
        {8, 9, 10, 11},
        {12, 13, 14, 15},
    };
    float MT[ROW_NUM][COL_NUM] = {0};
	
    // to do this:
    // a1 b1 c1 d1 => a1 a2 a3 a4
    // a2 b2 c2 d2 => b1 b2 b3 b4
    // ...
    // a4 b4 c4 d4 => d1 d2 d3 d4
 
    // origin
    int32_t i, j;
    for (i = 0; i < ROW_NUM; i++) {
        for (j = 0; j < COL_NUM; j++) {
            MT[j][i] = M[i][j];
        }
    }

    printf("ver1:\n");
    for (i = 0; i < ROW_NUM; i++) {
        for (j = 0; j < COL_NUM; j++) {
            printf("%f ", MT[i][j]);
            MT[i][j] = 0.;
        }
        printf("\n");
    }

    // method1
    float32x4x4_t vf32x4x4fTmpABCD = vld4q_f32(&M[0][0]); // vf32x4x4fTmpABCD中val[0]： a1 b1 c1 d1, val[1]: a2 b2 c2 d2
    float32x4x2_t vf32x4x2fTmpABCD01 = vtrnq_f32(vf32x4x4fTmpABCD.val[0], vf32x4x4fTmpABCD.val[1]); // vf32x4x2fTmpABCD01中val[0]： a1 a2 c1 c2, val[1]: b1 b2 d1 d2
    float32x4x2_t vf32x4x2fTmpABCD23 = vtrnq_f32(vf32x4x4fTmpABCD.val[2], vf32x4x4fTmpABCD.val[3]);
    float32x4x2_t vf32x4x2fTmpABCD02 = vuzpq_f32(vf32x4x2fTmpABCD01.val[0], vf32x4x2fTmpABCD23.val[0]); // row02, 按行组合
    float32x4x2_t vf32x4x2fTmpABCD13 = vuzpq_f32(vf32x4x2fTmpABCD01.val[1], vf32x4x2fTmpABCD23.val[1]); // row13, 按行组合
    vf32x4x2fTmpABCD02 = vtrnq_f32(vf32x4x2fTmpABCD02.val[0], vf32x4x2fTmpABCD02.val[1]);
    vf32x4x2fTmpABCD13 = vtrnq_f32(vf32x4x2fTmpABCD13.val[0], vf32x4x2fTmpABCD13.val[1]);
    vf32x4x4fTmpABCD.val[0] = vf32x4x2fTmpABCD02.val[0]; // a0 a1 a2 a3
    vf32x4x4fTmpABCD.val[2] = vf32x4x2fTmpABCD02.val[1];
    vf32x4x4fTmpABCD.val[1] = vf32x4x2fTmpABCD13.val[0];
    vf32x4x4fTmpABCD.val[3] = vf32x4x2fTmpABCD13.val[1]; // d0 d1 d2 d3
    vst4q_f32(&MT[0][0], vf32x4x4fTmpABCD);

    printf("ver2:\n");
    for (i = 0; i < ROW_NUM; i++) {
        for (j = 0; j < COL_NUM; j++) {
            printf("%f ", MT[i][j]);
            MT[i][j] = 0.;
        }
        printf("\n");
    }


    // method2
    float32x4x4_t vf32x4x4fTmp1ABCD;
    vf32x4x4fTmp1ABCD.val[0] = vld1q_f32(&M[0][0]); // a1 b1 c1 d1
    vf32x4x4fTmp1ABCD.val[1] = vld1q_f32(&M[1][0]);
    vf32x4x4fTmp1ABCD.val[2] = vld1q_f32(&M[2][0]);
    vf32x4x4fTmp1ABCD.val[3] = vld1q_f32(&M[3][0]); // a4 b4 c4 d4
    vst4q_f32(&MT[0][0], vf32x4x4fTmp1ABCD); // 利用交叉读写特性，放入到MT数组

    printf("ver3:\n");
    for (i = 0; i < ROW_NUM; i++) {
        for (j = 0; j < COL_NUM; j++) {
            printf("%f ", MT[i][j]);
            MT[i][j] = 0.;
        }
        printf("\n");
    }

    // 仅在寄存器内转置不输出到内存
    // float fTmpABCD4x4[4][4]; // 临时中转数组
    // vst4q_f32(&fTmpABCD4x4[0][0], vf32x4x4fTmpABCD);  // 假设待转置的数据为vf32x4x4fTmpABCD
    // vf32x4x4fTmpABCD.val[0] = vld1q_f32(&fTmpABCD4x4[0][0]);
    // vf32x4x4fTmpABCD.val[1] = vld1q_f32(&fTmpABCD4x4[1][0]);
    // vf32x4x4fTmpABCD.val[2] = vld1q_f32(&fTmpABCD4x4[2][0]);
    // vf32x4x4fTmpABCD.val[3] = vld1q_f32(&fTmpABCD4x4[3][0]); // 转置结果放到寄存器中
    
    return 0;
}
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92

以上代码末尾，附有仅在寄存器中实现转置的demo，可根据具体场景使用。

小结

---

方法1

• 在寄存器内转置后输出到内存

• 只在寄存器内操作，6条指令，加4个赋值

方法2

• 直接通过内存与寄存器的交叉读取实现转置功能，命令降为3条。
• 寄存器和内存之间交叉读取实现，5条指令

总之，实践得知，只在寄存器内操作场景，法1更优，比内存和寄存器之间读写交互更高效，哪怕多了一两条指令。而需要将结果输出到内存的时候，法2更优。