相关阅读：
ONNX Runtime 源码阅读：模型推理过程概览
ONNX Runtime 源码阅读：模型结点串行执行顺序的确定
前言
为了实现高效的推理，神经网络推理引擎应该尽可能将主机（Host）上能提供更高效计算的硬件设备（Device）利用上，ONNX Runtime当然不能例外。ONNX Runtime目前已经支持了多种不同设备，移动端的支持也在开发中。一台主机上很可能同时存在多种设备，ONNX Runtime是如何选择在那种设备上运行的呢（也就是怎么分区）？对于某些运行时不支持的操作怎么处理（回落）？不同硬件的运行时优先级怎么确定呢？这些就是本文探究的主要内容。
说明
为了不混淆概念，这里先做一点说明。这里所指的主机是由多个软、硬件组成的一个系统，例如一台电脑、一部手机等；设备是指CPU、GPU等计算单元，有些也叫加速器（Accelerator）；而他们所依赖的驱动等软件，有多种叫法，有些推理引擎上称为运行时（Runtime）：例如高通骁龙神经网络推理引擎（SNPE， Snapdragon Neural network Processing Engine），而在ONNX Runtime中，把它称为执行提供者（Execution Provider，其实翻译成赞助商可能更合适，赞助算力嘛，哈哈）。我想大概是为了和ONNX Runtime这个名字区分开。但是出于习惯，接下来都把它们称为Provider。
涉及文件
onnxruntime\onnxruntime\core\framework\sequential_executor.cc
onnxruntime\onnxruntime\core\session\inference_session.cc
onnxruntime\onnxruntime\core\framework\graph_partitioner.cc
onnxruntime\onnxruntime\python\onnxruntime_pybind_state.cc
正文
ONNX Runtime中，对特定硬件和他们依赖的驱动等进行了抽象，这个抽象统一了调用各种硬件的资源的方法，因此将这种抽象称为执行提供者（Execution Provider）；将某一个特定的操作例如卷积、池化等称为算子（kernel），OpKernel是所有算子的基类。不同的Provider对于同一种操作的实现是不一样的，例如同是卷积操作，使用CPU Provider和Cuda Provider就不一样。为节点分配Provider的过程其实也就是模型分区的过程。
通过前面的文章，我们知道ONNX Runtime运行主要分为三个阶段：实例化、初始化、推理。当我们调用InferenceSession.run()的时候，最终通过层层委托，真正执行推理的是IExecutor.Execute()，也就是IExecutor的子类重写的Execute方法。IExecutor有两个子类，SequentialExecutor和ParallelExecutor。不管那里一个子类，最终都是通过给出的Node信息去SessionState里面找到Node对应的OpKernel。下面通过代码具体说明。由于方法体太长，接下来的代码示例中使用// .......表示此处省略的很多代码，想看完整的代码请根据示例开头标注查看。
// onnxruntime\onnxruntime\core\framework\sequential_executor.cc#SequentialExecutor::Execute()
Status SequentialExecutor::Execute(const SessionState& session_state, const std::vector<int>& feed_mlvalue_idxs,
const std::vector<OrtValue>& feeds, const std::vector<int>& fetch_mlvalue_idxs,
std::vector<OrtValue>& fetches,
const std::unordered_map<size_t, CustomAllocator>& fetch_allocators,
const logging::Logger& logger) {
// ......
auto p_op_kernel = session_state.GetKernel(node_index);
// .....
// if a kernel has been added in the session state, it better be NON-null.
if (p_op_kernel == nullptr)
return ORT_MAKE_STATUS(ONNXRUNTIME, FAIL, "Got nullptr from GetKernel for node: ",
node.Name());
// .....
try {
compute_status = p_op_kernel->Compute(&op_kernel_context);
} catch (const std::exception& ex) {
compute_status = ORT_MAKE_STATUS(ONNXRUNTIME, RUNTIME_EXCEPTION, ex.what());
}
// .......
}
从上面代码可以看到，OpKernel是直接通过SessionState.GetKernel()获取的，并在下一步中直接调用了，中间没有做过转换。也就是说在执行的时候，执行该节点的OpKernel是已经确定了，说明Provider的选择不是在这一阶段做的，而是在实例化或者初始化阶段，凭感觉我们觉得是初始化阶段。
因此我们看一下初始化阶段的代码，初始化的主要代码在InferenceSession.Initialize()中。InferenceSession.Initialize()我们并不陌生，ONNX Runtime 源码阅读：模型结点串行执行顺序的确定一文看到，也是它通过调用SessionStateInitializer.CreatePlan()方法确定了模型中各个节点执行的先后顺序。我们再看一下这个方法：
// onnxruntime\onnxruntime\core\session\inference_session.cc#InferenceSession::Initialize()
common::Status InferenceSession::Initialize() {
// ......
SessionStateInitializer session_initializer(session_options_.enable_mem_pattern, model_location_, graph,
*session_state_, execution_providers_, kernel_registry_manager_);
// create SessionState for subgraphs as it's needed by the transformers
ORT_RETURN_IF_ERROR_SESSIONID_(CreateSubgraphSessionState(graph, *session_state_));
// apply any transformations to the main graph and any subgraphs
ORT_RETURN_IF_ERROR_SESSIONID_(TransformGraph(graph, *graph_transformation_mgr_,
execution_providers_, kernel_registry_manager_,
insert_cast_transformer_,
*session_state_));
// ......
ORT_RETURN_IF_ERROR_SESSIONID_(session_initializer.CreatePlan(nullptr, nullptr, session_options_.execution_mode));
// handle any subgraphs
ORT_RETURN_IF_ERROR_SESSIONID_(InitializeSubgraphSessions(graph, *session_state_));
// ......
}
因为我从执行的代码中已经知道了，OpKernel是从SessionState中取出来的，那我们知道找到OpKernel存入SessionState的过程我们就知道运行时是怎么选择的。上面代码中我们对有对SessionState有操作的地方一个一个去看就应该能找到。就比如警察已经锁定了有个人想要买毒品，那么盯着他，看他和谁接触过然后一个一个排查，就应该能找到给他卖毒品的毒贩。毕竟我们不是作者，只能一个一个看。最终发现，Provider的确定了具体OpKernel的获取是分开的：Provider确定在InferenceSession::TransformGraph()里面，确定了Provider之后，SessionStateInitializer.CreatePlan()在为他们初始化相应的OpKernel。
我们先看InferenceSession::TransformGraph()：
// onnxruntime\onnxruntime\core\session\inference_session.cc#InferenceSession::TransformGraph()
common::Status InferenceSession::TransformGraph(onnxruntime::Graph& graph,
const onnxruntime::GraphTransformerManager& graph_transformer_mgr,
const ExecutionProviders& providers,
KernelRegistryManager& kernel_registry_manager,
const InsertCastTransformer& insert_cast_transformer,
SessionState& session_state) {
// .......
// Do partitioning based on execution providers' capability.
GraphPartitioner partitioner(kernel_registry_manager, providers);
ORT_RETURN_IF_ERROR_SESSIONID_(partitioner.Partition(graph, session_state.ExportDll(), session_state.GetMutableFuncMgr()));
// apply transformers except default transformers
// Default transformers are required for correctness and they are owned and run by inference session
for (int i = static_cast<int>(TransformerLevel::Level1); i <= static_cast<int>(TransformerLevel::MaxLevel); i++) {
ORT_RETURN_IF_ERROR_SESSIONID_(graph_transformer_mgr.ApplyTransformers(graph, static_cast<TransformerLevel>(i), *session_logger_));
}
// .......
}
在InferenceSession::TransformGraph()中，委托GraphPartitioner.Partition()方法去为每个节点分配它们对应的Provider。
// onnxruntime\onnxruntime\core\framework\graph_partitioner.cc#GraphPartitioner::Partition()
Status GraphPartitioner::Partition(Graph& graph, bool export_dll, FuncManager& func_mgr) const {
// .......
for (auto& provider : providers_) {
int count = 0;
std::vector<Node*> nodes_need_compile;
std::vector<std::unique_ptr<ComputeCapability>> capabilities =
provider->GetCapability(graph_viewer, kernel_registry_mgr_.GetKernelRegistriesByProviderType(provider->Type()));
for (auto& capability : capabilities) {
Node* n = PlaceNode(graph, std::move(capability->sub_graph), kernel_registry_mgr_, provider->Type(), count);
if (n != nullptr) {
nodes_need_compile.push_back(n);
}
}
if (!nodes_need_compile.empty()) {
if (export_dll) {
std::string dll_path;
ORT_RETURN_IF_ERROR(provider->Compile(nodes_need_compile, dll_path));
for (auto* node : nodes_need_compile)
ORT_RETURN_IF_ERROR(func_mgr.AddFuncInfo(node->Name(), dll_path));
} else {
std::vector<NodeComputeInfo> node_compute_funcs;
ORT_RETURN_IF_ERROR(provider->Compile(nodes_need_compile, node_compute_funcs));
ORT_ENFORCE(node_compute_funcs.size() == nodes_need_compile.size(),
"Provider doesn't return correct number of compiled functions");
for (size_t j = 0; j < nodes_need_compile.size(); j++)
ORT_RETURN_IF_ERROR(func_mgr.AddFuncInfo(nodes_need_compile[j]->Name(), node_compute_funcs[j].compute_func,
node_compute_funcs[j].create_state_func,
node_compute_funcs[j].release_state_func));
}
for (auto* node : nodes_need_compile) {
//prepare the func kernel
KernelDefBuilder builder;
BuildFusedKernelDef(builder, *node);
ORT_RETURN_IF_ERROR(fused_kernel_registry->Register(
builder, static_cast<KernelCreatePtrFn>([](const OpKernelInfo& info) -> OpKernel* { return new FunctionKernel(info); })));
}
}
}
// .......
}
上面的代码中，providers_是ExecutionProviders的一个实例，InferenceSession中持有它的引用。它实际上是一个std::vector类型的列表，里面存放着所有本次运行中所有可以使用的Provider。GraphPartitioner::Partition()会按照从头到尾按的顺序通过provider->GetCapability()方法询问每一个Provider的处理能力，也就是在这个模型中，哪些节点是这个模型所支持的。最后会得到一个该Provider支持的节点的列表，再使用PlaceNode()函数设置每一个节点中execution_provider_type_属性，也就是表示对应Provider的字符串。当一个一个Provider询问下来，很可能有些节点没有一个特殊的加速器能处理，这些不能跑在其他加速器上的节点，都由CPU Provider的支持，所以要求CPU Provider要支持所有的ONNX Operator。
打个不太恰当的比方，模型里面的节点就是一群高考完的孩子，各大高校就是不同Provider。高考完后志愿系统会询问各个高校，我这里有这么一批学生，你要不要？高校根据自身的资源条件将符合要求的孩子接收了，给了他们一张录取通知书。当然所有高校询问下来，很可能有一些孩子没能被高校录取，怎么办？社会大学给他托底。ONNX Runtime中CPU就是社会大学。GetCapability()的签名是
std::vector<std::unique_ptr<ComputeCapability>>
IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
const std::vector<const KernelRegistry*>& kernel_registries) const
它接收一个GraphViewer&参数和一个const std::vector<const KernelRegistry*>&参数，返回一个std::vector<std::unique_ptr<ComputeCapability>>类型列表，里表里可以获取所支持的结点的索引。其中GraphViewer&的参数里面还有模型的结构和结点信息，包括结点的类型、索引值、参数等。IExecutionProvider的子类根据自己的情况重写这个方法，通过GraphViewer&获取到模型的所有结点并决定是否支持这些结点的一部分或者全部。如果支持某个结点，就将该结点的索引值用ComputeCapability封装起来放入一个列表中返回给调用者。IExecutionProvider提供了默认的实现，默认实现会尝试将Provider已注册的所有Kernel都支持，而不管具体参数。CPU Provider使用的就是默认的实现，而类似Cuda Provider，除了看结点类型，还要考量结点的参数才决定支持不支持，例如Cuda Provider只支持对称填充的卷积操作。下面是GetCapability()默认的实现代码：
// onnxruntime\onnxruntime\core\framework\execution_provider.cc#IExecutionProvider::GetCapability()
std::vector<std::unique_ptr<ComputeCapability>>
IExecutionProvider::GetCapability(const onnxruntime::GraphViewer& graph,
const std::vector<const KernelRegistry*>& kernel_registries) const {
std::vector<std::unique_ptr<ComputeCapability>> result;
for (auto& node : graph.Nodes()) {
for (auto registry : kernel_registries) {
if (registry->TryFindKernel(node, Type()) != nullptr) {
std::unique_ptr<IndexedSubGraph> sub_graph = onnxruntime::make_unique<IndexedSubGraph>();
sub_graph->nodes.push_back(node.Index());
result.push_back(onnxruntime::make_unique<ComputeCapability>(std::move(sub_graph)));
break;
}
}
}
return result;
}
但这里有个小小的疑问，既然是依次询问每个Provider对同一个模型中所有的节点的处理能力，就可能出现对于某个特定的节点，有多个Provider都声明可以处理的这种情况。那怎么解决这种冲突呢？ONNX Runtime的解决方法非常简单粗暴——先到先得。没错，如果某个节点已经被某个Provider声明可以处理并被分配给该Provider了，那么后续即便再有Provider进行声明，也会被忽略。情看下面的代码，通过一个if条件判断，如果某个节点已经被分配，就不再重新分配了：
// onnxruntime\onnxruntime\core\framework\graph_partitioner.cc#PlaceNode()
static Node* PlaceNode(Graph& graph, std::unique_ptr<IndexedSubGraph> capability,
const KernelRegistryManager& kernel_registry_mgr, const std::string& provider_type, int& count) {
if (nullptr == capability) {
return nullptr;
}
if (nullptr == capability->GetMetaDef()) {
// The <provider> can run a single node in the <graph> if not using meta-defs.
// A fused kernel is not supported in this case.
ORT_ENFORCE(1 == capability->nodes.size());
auto* node = graph.GetNode(capability->nodes[0]);
if (nullptr != node && node->GetExecutionProviderType().empty()) {
// The node was not fused or assigned. Assign it to this <provider>.
node->SetExecutionProviderType(provider_type);
}
// .......
}
// .......
}
那么，怎么确定各个Provider谁先谁后？分两种情况：
- 用户指定了Provider：按照客户指定的顺序；
- 用户未指定：使用ONNX Runtime的开发者写死了的顺序。作者们通过某种考量方式，在代码里面已经直接为各种支持的Provider排了顺序。所有支持的Provider和他们的默认顺序如下，从中可以看到，CPU的优先级最低，因此它可以为所有其他Provider托底：
// onnxruntime\onnxruntime\python\onnxruntime_pybind_state.cc#GetAllProviders()
// ordered by default priority. highest to lowest.
const std::vector<std::string>& GetAllProviders() {
static std::vector<std::string> all_providers = {kTensorrtExecutionProvider, kCudaExecutionProvider, kDnnlExecutionProvider,
kNGraphExecutionProvider, kOpenVINOExecutionProvider, kNupharExecutionProvider,
kBrainSliceExecutionProvider, kCpuExecutionProvider};
return all_providers;
}
我们可以看到Provider的注册过程如下：
// onnxruntime\onnxruntime\python\onnxruntime_pybind_state.cc#InitializeSession()
void InitializeSession(InferenceSession* sess, const std::vector<std::string>& provider_types) {
if (provider_types.empty()) {
// use default registration priority.
RegisterExecutionProviders(sess, GetAllProviders());
} else {
RegisterExecutionProviders(sess, provider_types);
}
OrtPybindThrowIfError(sess->Initialize());
}
上面第一种情况是用户没有指定Provider，默认注册所有本机所支持的Provider；第二种情况用户指定了Provider。由于ONNXRuntime要求CPU必须能够托底，因此如果用户指定了Provider单却不包含CPU Provider的时候，系统会自动加上CPU Provider确保所有操作都可能回落到CPU上执行。
// onnxruntime\onnxruntime\core\session\inference_session.cc#InferenceSession::Initialize()
common::Status InferenceSession::Initialize() {
// ......
// Register default CPUExecutionProvider if user didn't provide it through the Register() calls
if (!execution_providers_.Get(onnxruntime::kCpuExecutionProvider)) {
LOGS(*session_logger_, INFO) << "Adding default CPU execution provider.";
CPUExecutionProviderInfo epi{session_options_.enable_cpu_mem_arena};
auto p_cpu_exec_provider = onnxruntime::make_unique<CPUExecutionProvider>(epi);
ORT_RETURN_IF_ERROR_SESSIONID_(RegisterExecutionProvider(std::move(p_cpu_exec_provider)));
}
// ......
}
总结
通过上面的一些代码示例我们知道了ONNX Runtime的分区方法：
- 按照出现在列表中的先后顺序注册支持的Provider，这个列表可以由用户指定亦可以使用系统默认的一个列表，越靠前优先级越高；
- ONNX Runtime按照第一步确定的优先级依次调用Provider的
GetCapability()方法询问不同的Provider的处理能力，GetCapability()返回一个所给模型中所有可以运行在被自己处理的节点列表，也就是模型的子图； - 根据第二步取得的节点列表，依次检查其中的节点，如果节点未分配给其他更高优先级的Provider，则将该节点分配给当前Provider。
- 最终不能被其他专有的Provider所支持的节点都将回落到CPU上执行。
虽然实现方式不一样，但是很多推理引擎的基本的思想是一样的，比如分区。高通神经网络处理引擎分区是离线做的，也就是使用一个专门的编译器将能需要运行在某个特定硬件上的节点进行编译并将相关信息写入到专用模型文件中。而ONNX Runtime的分区实现方法却简单的多，可以在运行时根据模型直接进行分区。
个人微信公众号TensorBoy。微信扫描上方二维码或者微信搜索TensorBoy并关注，及时获取更多最新文章！