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(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並關注，及時獲取更多最新文章！

ONNX Runtime 源碼閱讀：模型分區的實現

前言

說明

涉及文件

正文

總結

認知提升的方法

C#開源的兩款功能強大的錄屏神器

螞蟻面試：Springcloud核心組件的底層原理，你知道多少？

前端 Vue yarn.lock文件：詳解和使用指南

使用Github+Markdown搭鍵自己的筆記本

PyTorch框架Tensor底層原理

ONNX Runtime 源碼閱讀：模型分區的實現

冰山之下：從extern C到程序鏈接

交叉編譯和ABI簡介

https://yachay.unat.edu.pe/blog/index.php?comment_area=format_blog&comment_component=blog&comment_co

linux以太網驅動總結