ncnn avx2/armv8.2 基础架构

作者：nihui
转自：知乎

问题的出现

2020年开始，新手机 CPU 几乎都是 armv8.2 架构，这个架构引入了新的 fp16 运算和 int8 dot 指令，优化得当就能大幅加速深度学习框架的推理效率。类似于 x86 CPU 的 AVX，不是全部 CPU 都支持，得考虑兼容性，做到老CPU上用老指令，新CPU用新指令。

之前社区大佬

@圈圈虫

为 ncnn 实现了初版 avx2 优化加速，做了个 cmake 选项 NCNN\_AVX2，如果开启，编译时会全局增加 -mavx2 参数，同时编译器会自动定义 __AVX__ 宏。-mavx2 等于告诉编译器，cpu 一定支持 avx2，请利用 avx2 指令尽情优化（例如将 sse intrinsic 编译为 avx 特有的 VEX encoding）。

#if __AVX__
    // use avx code
#else
    // use good-old sse code
#endif

这样编译出来的 libncnn.a 只包含了 __AVX__ 内的代码，即使不在 __AVX__ 内的代码也会可能被编译器优化为 avx 的指令提高效率，如果在不支持 avx 的 cpu 跑就会直接 Illegal instruction 崩溃。如果硬件是自己能控制的，比如针对某个开发板或者 cpu 定制，这样是可以的。但是开发客户端是不行的，良好的兼容性是首要条件，这么编译，不行。

通常有这样的做法

打包两种架构的 so 库，运行时判断 cpu 指令集支持情况，动态加载 so
针对核心优化的函数，用新编译参数单独编译这些源码，库的主体依然使用基本指令，运行时判断 cpu 指令集支持情况，动态调用对应的优化函数

方法1，编译打包太麻烦，用起来也麻烦，两个 so 有很多重复的部分（如上层推理逻辑，gpu优化等），增大了安装包体积，我很不满意！

方法2，这是很多底层优化库采用的方法，对于应用开发者完全透明，一个库包含所有功能，对于新指令集支持的扩展性也不错

我还是不满意！

优化的核心代码得写两份，而且要写在两个文件里，这也就算了。在每次调用前，要判断 cpu 指令集支持情况，这就麻烦了，怎么个麻烦呢？来看看这些 if，太长太丑了，这只是大概的片段，实际有各式各样的 convolution 参数组合，还有 pack4 pack1to4 pack4to1 pack8 pack4to8 pack8to4 pack1to8 pack8to1 fp32 fp16 bf16s 各式各样的参数组合，再来个 cpu sse avx 组合一下？打算放在个 function table 里？convolution 的各种参数太多了，有的优化了有的没优化，table 也脏脏的。

if (conv1x1s1)
{
    if (cpu_avx)
    {
        conv1x1s1_avx();
    }
    else
    {
        conv1x1s1_sse();
    }
}
if (conv1x1s2)
{
    if (cpu_avx)
    {
        conv1x1s2_avx();
    }
    else
    {
        conv1x1s2_sse();
    }
}
if (conv3x3s1)
{
    if (cpu_avx)
    {
        conv3x3s1_avx();
    }
    else
    {
        conv3x3s1_sse();
    }
}

ncnn 的基础架构

很多时候，avx 和 sse 的区别只是各种 for loop unroll 的区别，原先就是直接可以写在一个源文件里头，我也挺喜欢这样，相当于共享了全部控制逻辑。

编两个库又不想编译，写成两个函数又不想写，那么就只好让 cmake 帮忙了。cmake 脏点就脏点吧，cpp 写着舒服就行了。

# must define SRC DST CLASS

file(READ ${SRC} source_data)

# replace
string(TOUPPER ${CLASS} CLASS_UPPER)
string(TOLOWER ${CLASS} CLASS_LOWER)

string(REGEX REPLACE "LAYER_${CLASS_UPPER}_ARM_H" "LAYER_${CLASS_UPPER}_ARM_ARM82_H" source_data "${source_data}")
string(REGEX REPLACE "${CLASS}_arm" "${CLASS}_arm_arm82" source_data "${source_data}")
string(REGEX REPLACE "#include \"${CLASS_LOWER}_arm.h\"" "#include \"${CLASS_LOWER}_arm_arm82.h\"" source_data "${source_data}")

file(WRITE ${DST} "${source_data}")

这段 cmake 是简单的全局查找替换，XXX\_arm 换成 XXX\_arm\_arm82，存到新的文件里。比如我有 src/layer/arm/clip\_arm.h 和 src/layer/arm/clip\_arm.cpp，经过处理得到了类名是 Clip\_arm\_arm82，但优化的实现代码和 Clip\_arm 一模一样。

针对这个 arm82 的复制品，添加新的编译参数，以便正确打开对应的宏 \_\_ARM\_FEATURE\_FP16\_VECTOR\_ARITHMETIC，启用 armv8.2 的优化代码

set_source_files_properties(clip_arm_arm82.cpp PROPERTIES COMPILE_FLAGS "-march=armv8.2-a+fp16")

接下来就是如何运行时根据 cpu 特性选择的问题了，先观察 create\_layer 的动作

Layer* create_layer(int index)
{
    layer_creator_func layer_creator = layer_registry[index].creator;
    if (!layer_creator)
        return 0;

    Layer* layer = layer_creator();
    layer->typeindex = index;
    return layer;
}

这里的 layer\_registry 都是通用版本的

class Clip_final : virtual public Clip, virtual public Clip_arm, virtual public Clip_vulkan
{
};

再让 cmake 自动生成个 armv8.2 版本的 layer\_registry\_arm82，用刚才的 Clip\_arm\_arm82 换掉 Clip\_arm。注意这里的 naive C Clip 和 gpu Clip\_vulkan 和通用版本是一样的父类，这样编译后的 clip.o clip\_vulkan.o 也会复用，不会造成二进制体积成倍变大。

class Clip_final_arm82 : virtual public Clip, virtual public Clip_arm_arm82, virtual public Clip_vulkan
{
};

安排好所有的 XXX\_arm 后，layer\_registry\_arm82 都是 armv8.2 优化版本，最后就是根据 CPU 指令集支持情况，创建 layer 的时候，分派不同的 layer\_creator，完成动态切换 layer 的实现

Layer* create_layer(int index)
{
    layer_creator_func layer_creator = 0;
#if __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
    if (ncnn::cpu_support_arm_asimdhp())
    {
        layer_creator = layer_registry_arm82[index].creator;
    }
    else
#endif // __ARM_FEATURE_FP16_VECTOR_ARITHMETIC
    {
        layer_creator = layer_registry[index].creator;
    }
    if (!layer_creator)
        return 0;

    Layer* layer = layer_creator();
    layer->typeindex = index;
    return layer;
}

qaq

推荐阅读

更多嵌入式AI算法部署等请关注极术嵌入式AI专栏。

问题的出现

通常有这样的做法

ncnn 的基础架构

推荐阅读

目录