ONNX Runtime 源码阅读：模型分区的实现

相关阅读：
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