ONNXruntime添加自己EP的详细步骤

Question

微软的ONNXruntime文档中有涉及EP的添加，但内容很简略，不知道有没有朋友自己亲自实现过添加自己或者别人的加速器硬件的抽象到ONNXruntime中

Answer 1

引自免费微信小程序：皆我百晓生

为了在ONNXRuntime中添加一个新的执行提供者（Execution Provider，EP），你需要实现一系列接口并将其与运行时集成。以下是一个超详细步骤的概述：

1. 定义新EP类： 首先，你需要创建一个继承自onnxruntime::Provider基类的新类，并实现必要的接口。例如：

   namespace MyAccelerator {
     class MyExecutionProvider : public onnxruntime::IExecutionProvider {
     public:
       explicit MyExecutionProvider(const OrtProviderOptions& options) : IExecutionProvider(kMyExecutionProvider, options) {}
   
       // 实现必须的方法
       virtual Status Initialize() override;
       virtual void Shutdown() override;
       virtual std::unique_ptr<onnxruntime::IDataTransfer> GetDataTransfer() const override;
       virtual std::shared_ptr<IDataAllocator> CreateDataAllocator(int device_id, const OrtMemoryInfo& info) override;
       virtual common::Status GetCapability(const onnxruntime::GraphViewer& graph_viewer, ExecutionProvider::Capabilities& out) override;
       virtual std::vector<std::unique_ptr<ComputeCapability>> GetCapability(
         const onnxruntime::GraphViewer& graph,
         const std::vector<const KernelRegistry*>& kernel_registries) override;
   
       // 添加对模型节点进行优化和执行的方法
       virtual Status Compile(const std::vector<Node*>& fused_nodes, std::vector<NodeComputeInfo>& node_compute_info) override;
       virtual Status Execute(const std::vector<OrtValue>& inputs, std::vector<OrtValue>& outputs,
                             const std::vector<NodeComputeInfo>& node_compute_info,
                             const std::string& provider_name) override;
     };
   }
   

2. 实现核心功能：

   • Initialize()：初始化硬件或软件环境。
   • GetDataTransfer()：如果需要数据传输层，则实现此方法以处理跨设备的数据移动。
   • CreateDataAllocator()：为该EP创建内存分配器。
   • GetCapability()：分析模型图，确定EP可以支持哪些运算符。
   • Compile() 和 Execute()：分别负责编译和执行由ONNX Runtime调度的计算图部分。

3. 注册EP到ONNXRuntime： 在ONNXRuntime启动前，通过调用RegisterExecutionProvider()函数将你的EP注册到全局上下文中：

   using namespace onnxruntime;
   using namespace MyAccelerator;
   
   int main() {
     OrtEnv env(ORT_LOGGING_LEVEL_WARNING, "my_ort_session");
     OrtSessionOptions session_options;
     auto status = OrtSessionOptionsAppendExecutionProvider_MyExecutionProvider(session_options, /* your options */);
     if (!status.IsOK()) {
       // handle error
     }
   
     // 创建会话时使用session_options
     Ort::Session session(env, model_path.c_str(), session_options);
   
     // ...其余代码...
   }
   
   // 注册EP到ONNXRuntime的核心库
   extern "C" ONNX_EXPORT void RegisterCustomExecutionProvider() {
     auto creator = [](const OrtProviderOptions& options) {
       return std::make_unique<MyExecutionProvider>(options);
     };
     onnxruntime::RegisterExecutionProviderFactory(MyExecutionProvider::Type(), creator);
   }
   

4. 构建和链接： 将包含新EP实现的源文件添加到ONNXRuntime项目的构建系统中，并确保正确链接所需的库或框架。

5. 测试： 使用一些简单的ONNX模型来测试新EP是否能成功编译并执行模型中的操作。

注意：上述代码示例是简化的，实际开发中需要根据具体硬件加速器特性进行适配。此外，ONNXRuntime的API可能会随着版本更新而变化，请参考最新版的官方文档和源码进行开发。