跳到主要内容

K230 RVV优化性能说明

1. 概述

近些年AI领域的快速发展,衍生出各种各样的神经网络模型,不断出现新的算子。但是AI芯片的迭代周期相比AI模型会长很多,这些新出现的算子多数不能直接使用AI芯片进行推理加速,同时旧的算子中也有一部分不适合使用AI芯片进行推理加速。因此CPU成为这部分算子的执行载体,这也意味着CPU的性能会成为模型部署时影响最终性能的一个因素,而RVV扩展就是RISC-V CPU提升性能的一个重要手段。K230 中采用的玄铁C908双核处理器中大核具备RVV1.0扩展的特性,能够大幅度提升CPU算子推理时的表现。

K230进行模型推理时需要使用 .kmodel格式的模型,.kmodel是由nncaseONNXTFLite模型进行编译后的模型格式,适用于本公司及相关合作企业生产的开发板。nncase支持目前常见的神经网络算子,但是部分算子无法通过K230进行推理加速,这部分算子只能使用CPU进行推理。

2. RVV应用场景

目前研究应用最广泛的神经网络 Transformer中,模型结构与 CNN存在较大差异,很多基于 CNN设计的AI芯片无法完全对 Transformer进行加速。以下为 TransformerDecoder模型在开启RVV优化和不开启RVV优化的算子执行情况。

2.1 未开启RVV优化

stackvm tensor opcounttime consumption(ms)percentage(%)
softmax51749.6188.6574
where4199.43210.1058
EXTCALL6516.0990.815779
layer_norm75.810.294408
gather20.3930.0199144
STLOCAL2120.3910.019813
LDC_I42410.3880.019661
reduce_arg10.3360.017026
reshape260.2810.014239
LDLOCAL1490.260.0131749
LDNULL1060.1660.00841166
LDTENSOR290.1030.00521929
LEA_GP580.0970.00491525
LDDATATYPE290.070.00354709
LDARG50.0080.000405381
RET10.0040.000202691
LDTUPLE10.0030.000152018
total9411973.45100

2.2 开启RVV优化

stackvm tensor opcounttime consumption(ms)percentage(%)
softmax525.72255.6175
EXTCALL6516.17934.9831
layer_norm70.9672.0909
where40.9121.97198
gather20.390.84328
LDC_I42410.3860.834631
STLOCAL2120.3790.819495
reduce_arg10.340.735167
LDLOCAL1490.2590.560024
reshape260.2430.525428
LDNULL1060.170.367583
LEA_GP580.1030.222712
LDTENSOR290.1030.222712
LDDATATYPE290.0760.164331
LDARG50.0110.0237848
RET10.0050.0108113
LDTUPLE10.0030.00648677
total94146.248100

2.3 性能分析及说明

上述模型推理中,K230的KPU单元不支持对 softmaxlayer_normwheregatherreduce_argreshape 进行硬件推理加速,因此需要使用C908实现推理,目前已经完成了对 softmaxlayer_normwhere的RVV优化,性能提升明显。

以下为使用RVV优化前后各算子在模型推理时间中的占比图。

image-20240719162322440

image-20240719162343905

以下为使用RVV优化前后相关算子的性能对比。

RVV

从以上的对比结果可以看出,在开启RVV优化后能够极大的提升CPU算子的推理性能,缩短整个模型的推理时间(1973–> 46)ms,RVV优化后占据大部分时间的 softmax算子时间减少到25ms,layer_norm算子时间减少到0.97ms,where算子时间减少到0.91ms,整个模型的推理时间缩短了97.6%,在实际模型部署时具有很高的应用价值。

2.4 RVV优化示例

2.4.1 RVV代码

具体实现请参考 nncase中的layer_norm,需要具备一定RV指令和V扩展指令知识。

layer_norm的计算公式如下:

y= (x−E[x])/sqrt(Var[x]+ϵ)∗γ+β

整体计算流程详见 layernorm_impl函数,为了代码具备更高的可读性,该流程中将RVV优化代码拆分成三个部分:

  1. 计算 E[x],具体请参考 get_mean函数。
  2. 计算 Var[x],具体请参考 get_var函数。
  3. 按照上述公式进行layer_norm的计算,具体请参考 layer_norm_update1函数。

由于乘法相比除法耗时更短,第3步的计算中进行了公式变换,使用 rsqrt代替 sqrt,再用乘法替代除法。

2.4.2 核心代码说明

以下为 get_mean中核心代码的说明,这段代码实现了对a1处数组的循环加载求和,求和结果存储在v0,最后求平均值保存在ret中。它利用RVV的向量加载,向量累加指令来实现求和,从而提高计算性能。

"vle32.v v8, (a1);"   // 加载a1地址处的32bit向量到v8寄存器
"sub a0,a0, t0;" // a0 -= t0,用于循环控制计数
"slli t1, t0, 2;" // t1 = t0 << 2,因为每个float32是4字节,所以地址增加4*t0
"vfredsum.vs v0,v8,v0;" // v0 += v8,向量累加求和到v0

"add a1, a1, t1;" // a1 += t1,更新加载地址
"bnez a0, XXXXXX%=;" // 如果a0!=0,跳转至循环开始地址
"vfmv.f.s f0, v0;" // 把v0向量累加结果移动到f0
"fcvt.s.w f1, %[avl];" // 将avl转换为float保存到f1
"fdiv.s %[ret], f0, f1;" // ret = f0/f1,即求平均值

2.4.3 添加RVV算子流程

以下流程中路径均以nncase作为根目录

  1. 函数声明:src/Native/src/kernels/stackvm/optimized/opt_ops.h
  2. 算子实现:
    • 通用优化 src/Native/src/kernels/stackvm/optimized
    • x86优化 src/Native/src/kernels/stackvm/optimized/x86_64
    • RVV优化 src/Native/src/kernels/stackvm/optimized/riscv64
  3. 逻辑调用:src/Native/src/kernels/stackvm/tensor_ops.cpp
  4. 修改CMakeLists:src/Native/src/kernels/stackvm/optimized/CMakeLists.txt
    • 通用优化:15行增加源文件名
    • 特定平台优化:44行增加源文件名

2.5 tips

如果遇到尚未支持RVV优化的算子,且需要进行支持的,欢迎在nncase提issue和PR