Tengine 算子调度分析

上一篇简述了Tengine如何适配ncnn模型,其实本质上Tengine算是一个新的IR,可对各种框架进行适配,虽然不同框架对不同算子有不同的实现方式,但是经过Tengine团队的不懈努力,可以支持nchw与nhwc的计算格式。

现在Tengine开源部分的算子分别对于不同的平台分为arm32,arm64,ref,x86。其中arm32与arm64是针对端侧平台的算子优化,以arm汇编和arm neon为主来进行计算部分的优化,可以大大提高计算效率。X86则是针对Convolution和FullyConnected层在x86平台上进行加速。优化部分只针对需要大量计算的算子,对于其余的算子则是通过加载reference算子来实现。那么现在问题来了,那么多平台,如何对相同算子进行调度则是一个问题。

Tengine 算子调度流程图

上图则表述了Tengine是如何判断相同算子在不同平台的调度。

在Tengine框架内部,每一个类别都有自己独有的注册列表,例如arm32有自己的init.cpp,arm64也有自己的init.cpp来进行注册。在运行Tengine的时候,Tengine首先会获取平台信息,当获取平台信息后则会进行算子部分的调度。

NodeOps* NodeOpsRegistryManager::RealFindNodeOps(const CPUInfo* cpu_info, Node* node) { NodeOps* ops; if(cpu_info != nullptr) { // search cpu_type int master_cpu = cpu_info->GetMasterCPU(); const std::string& cpu_model = cpu_info->GetCPUModelString(master_cpu); ops = FindNodeOps(cpu_model, cpu_info, node); if(ops) return ops; // search arch std::string cpu_arch; int int_arch = cpu_info->GetCPUArch(master_cpu); if(int_arch == ARCH_ARM_V8) { cpu_arch = "arm64"; } else if(int_arch == ARCH_ARM_V7) { cpu_arch = "arm32"; } ops = FindNodeOps(cpu_arch, cpu_info, node); if(ops) return ops; } // search x86 #if CONFIG_ARCH_X86 ops = FindNodeOps("x86", cpu_info, node); if(ops) return ops; #endif // the final search: reference ops = FindNodeOps(REF_REGISTRY_NAME, cpu_info, node); return ops; }

上述代码则展示了Tengine对于不同平台的的算子调度部分。

在进行Tengine框架运行的时候则可以通过不同的环境变量设置来进行算子的调度。Tengine实现算子的步骤是首先实现Reference算子,在保证能正常工作的情况下在进行针对性的算子优化,例如对arm32与arm64,所以性能算子可以直接在Reference中找到对应算子。变相的可以理解Reference算子是各种平台优化算子的基础。如果出现模型运行错误的情况,则可以首先运行Reference算子来进行debug。对于如何调度如下:

export OPS_REGISTRY=reference export OP_NAME=Convolution

以上两条指令则是可以把Convolution算子设置在Reference,以此类推,如果对性能算子中的任何一个报有怀疑,则可以先用reference来进行验证。同样,对于初学者而言,c++当然比汇编要容易理解啦。

刚刚所说了Reference算子是包含了所有算子,那么如果区分相同算子是调用性能算子还是功能算子呢?这就涉及到两方面,第一则是上述所说的平台信息选择,第二则是算子中设置的优先级部分,以Convolution算子为例,如下为在Reference中的算子选择函数:

const int default_prio = 1500; ... ... ... NodeOps* SelectFunc(const CPUInfo* cpu_info, Node* node) { RefConv* ops = new RefConv(); return ops; } } // namespace RefConvolutionOps void RegisterRefConv2d(void) { NodeOpsRegistryManager::RegisterOPImplementor(REF_REGISTRY_NAME, "Convolution", RefConvolutionOps::SelectFunc, RefConvolutionOps::default_prio); }

如下为在arm32中Convolution的选择函数:

const int default_prio = 500; ... ... void RegisterConv2dFast(void) { if(!NodeOpsRegistryManager::RegisterOPImplementor("arm32", "Convolution", conv_fast::SelectFunc, conv_fast::default_prio)) LOG_ERROR() << __FUNCTION__ << " :Regist OP failed for prio[" << conv_fast::default_prio << "]\n"; }

上述在arm32中的default_prio = 500 优先级设置高与Reference中的default_prio = 1500,则表明在算子调度时优先选择arm32中的Convolution算子。

好啦,Tengine简单的算子调度介绍到此结束啦。

https://github.com/OAID/Tenginegithub.com/OAID/Tengine