使用ARM NEON进行矩阵乘法运算

1.背景

        最近想要体验一下ARM Helium （MVE）对SIMD指令的支持，于是到网络上想要购买一款搭载ARMv8-M架构CPU的开发板，然而看中的板子需要较长的交货周期，于是暂时放弃。因为NEON和Helium在设计理念上有诸多相似之处，之前笔者对于NEON也不是非常熟悉，所以尝试用NEON进行一个简单的3x3矩阵乘法运算，以对其有个初步的了解，这里顺带做个记录，分享一下。

2. 计算内容

        这里做个最简单的计算。

 求b与a的乘积。

3.代码

        时间有限，用AI先生成一段代码进行计算，结果发现有坑，代码需要调整才能用，调整后的代码如下：

#include <stdio.h>

#include <arm_neon.h>

 

/**

 * matrix_multiply_neon

 */

void matrix_multiply_neon(float32_t *a, float32_t *b, float32_t *c)

{

    float32x4_t a0 = vld1q_f32(a);

    float32x4_t a1 = vld1q_f32(a + 3);

    float32x4_t a2 = vld1q_f32(a + 6);

 

    float32x4_t b0 = vld1q_f32(b);

    float32x4_t b1 = vld1q_f32(b + 3);

    float32x4_t b2 = vld1q_f32(b + 6);

 

    float32x4_t c0 = vmovq_n_f32(0);

    float32x4_t c1 = vmovq_n_f32(0);

    float32x4_t c2 = vmovq_n_f32(0);

 

    c0 = vfmaq_laneq_f32(c0, a0, b0, 0);

    c0 = vfmaq_laneq_f32(c0, a1, b0, 1);

    c0 = vfmaq_laneq_f32(c0, a2, b0, 2);

 

    c1 = vfmaq_laneq_f32(c1, a0, b1, 0);

    c1 = vfmaq_laneq_f32(c1, a1, b1, 1);

    c1 = vfmaq_laneq_f32(c1, a2, b1, 2);

 

    c2 = vfmaq_laneq_f32(c2, a0, b2, 0);

    c2 = vfmaq_laneq_f32(c2, a1, b2, 1);

    c2 = vfmaq_laneq_f32(c2, a2, b2, 2);

 

    vst1q_f32(c, c0);

    vst1q_f32(c + 3, c1);

    vst1q_f32(c + 6, c2);

}

 

/**

 * main

 */

int main()

{

    int i, j;

    float32_t a[10] = {9, 8, 7, 6, 5, 4, 3, 2, 1, 0};

    float32_t b[10] = {1, 2, 3, 4, 5, 6, 7, 8, 9, 0};

    float32_t c[10] = {0};

 

    matrix_multiply_neon(a, b, c);

 

    printf("Result:\n");

 

    for (i = 0; i < 3; i++) {

        for (j = 0; j < 3; j++)

            printf(" %f ", c[i * 3 + j]);

        printf("\n");

    }

 

    return 0;

}

        AI生成的代码会出现调用vst1q_f32向数组c写入越界的问题，这里需要将a,b,c（尤其是c）三个数组声明大一些。原因是vst1q_f32写入的对象是float32x4_t，是4个单精度浮点数。

4.说明

        这里的矩阵乘法运算主要依赖vfmaq_laneq_f32这个API。

        这个API的原型定义为：

        向量b中的每个标量和向量v中lane作为index下标取出的标量数字相乘，计算结果和向量a中的每个标量相加并将计算结果存储在a中。

        需要分解来看。

        上面的矩阵相乘，展开后是：

        上面标红的部分可以体现出对[0][0]这个标量进行的一次计算，对应代码是：

        c0 = vfmaq_laneq_f32(c0, a0, b0, 0);

        之后是第二次计算：

        c0 = vfmaq_laneq_f32(c0, a0, b0, 0);

        第三次：

        c0 = vfmaq_laneq_f32(c0, a2, b0, 2);

        三次计算累加后就得到了第一行的计算结果，30，24和18。

        实机测试结果为：

5.总结

        NEON本质上是利用SIMD的优势加快向量计算，这种计算相较于仅用一般计算指令进行计算一定可以节省更多用于load/store的时钟周期进而加快运算速度（理论上计算规模越大越明显）。但是相较于一些专用加速器在特定领域应该还相对逊色。总之，这种计算方式可以在很多计算场景体现出优势。