ONNX Runtime 插件 EP 参考实现全解:从三个示例插件看懂 Plugin Execution Provider 的加载、注册与内核开发
【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime
ONNX Runtime 支持以动态库形式加载的插件式执行提供程序(Plugin Execution Provider),开发者无需修改 ORT 主库即可向运行时注入自定义的算子加速能力。onnxruntime/test/autoep/library/目录下的三个示例插件(readme.md)是理解这套基础设施的最佳入口:它们既是插件 EP 开发者的参考实现,也被 ORT 单测套件用于持续验证插件加载、设备注册与内核执行行为。读完本文,你能掌握插件 EP 的导出符号约定、工厂/EP 虚函数表填充方式、设备与分配器注册流程,以及“编译型 EP”“虚拟设备 EP”“内核注册型 EP”三种实现范式的差异。
目录定位与用途
该目录的定位在文档中被明确为两类受众(readme.md):
- 插件 EP 开发者:以这里的文件为参考创建自己的自定义 EP 插件;
- ORT 开发者:这些文件被 ONNX Runtime 的单测套件调用,用于确保插件支持按预期工作,功能更新时需同步更新示例代码。
需要牢记文档中的两点声明:这些代码仅用于演示和测试目的,未针对生产环境优化;其中展示的 API 用法对应“当前版本”的插件 EP 接口,接口演进时需对照源码更新。
目录结构与文档描述一致:
| 路径 | 说明 |
|---|---|
onnxruntime/test/autoep/library/example_plugin_ep/ | 基础“编译型”插件 EP,可编译单个 Mul 算子 |
onnxruntime/test/autoep/library/example_plugin_ep_virt_gpu/ | 编译型插件 EP,且注册自己的虚拟硬件设备,可用于为不同目标交叉编译模型 |
onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/ | 基础插件 EP,直接注册算子内核(kernel registration)而非编译节点 |
onnxruntime/test/autoep/library/plugin_ep_utils.h | 三个示例共用的公共工具头文件 |
插件 EP 的动态库导出约定
每个示例插件最终都是一个只导出两个 C 符号的共享库。从 example_plugin_ep.cc 可以看到入口实现:
extern "C" { // 公开符号 EXPORT_SYMBOL OrtStatus* CreateEpFactories(const char* registration_name, const OrtApiBase* ort_api_base, const OrtLogger* default_logger, OrtEpFactory** factories, size_t max_factories, size_t* num_factories) { const OrtApi* ort_api = ort_api_base->GetApi(ORT_API_VERSION); const OrtEpApi* ep_api = ort_api->GetEpApi(); const OrtModelEditorApi* model_editor_api = ort_api->GetModelEditorApi(); // C++ API 手动初始化 Ort::InitApi(ort_api); // 工厂可以使用传入的 registration_name,也可以自定义 EP 名称 std::unique_ptr<OrtEpFactory> factory = std::make_unique<ExampleEpFactory>(registration_name, ApiPtrs{*ort_api, *ep_api, *model_editor_api}, *default_logger); if (max_factories < 1) { return ort_api->CreateStatus(ORT_INVALID_ARGUMENT, "Not enough space to return EP factory. Need at least one."); } factories[0] = factory.release(); *num_factories = 1; return nullptr; } EXPORT_SYMBOL OrtStatus* ReleaseEpFactory(OrtEpFactory* factory) { delete static_cast<ExampleEpFactory*>(factory); return nullptr; } } // extern "C"关键要点:
CreateEpFactories是 ORT 加载动态库后查找的入口点,插件通过ort_api_base->GetApi(ORT_API_VERSION)获取与自身编译版本匹配的OrtApi,再用GetEpApi()/GetModelEditorApi()拿到插件 EP 专用与模型编辑 API;registration_name是 ORT 侧注册插件时使用的名称,工厂可以照搬它作为 EP 名称,也可以自行定义;- 一个插件库可以返回多个工厂(
max_factories控制容量上限),示例只返回一个; EXPORT_SYMBOL宏在 Apple 平台展开为__attribute__((visibility("default"))),保证符号在 macOS/iOS 上可见。
符号可见性由 Windows 的.def与 Linux 的.lds链接脚本双重限定,例如 example_plugin_ep_library.def:
LIBRARY "example_plugin_ep.dll" EXPORTS CreateEpFactories @1 ReleaseEpFactory @2即插件库对外只暴露这两个符号,其余实现细节全部隐藏——这也是插件与主库解耦的关键设计。
公共工具头 plugin_ep_utils.h
三个示例共享的工具层在 plugin_ep_utils.h,它示范了插件 EP 中非常常见的几类模式:
- 错误处理宏:
RETURN_IF_ERROR(把返回OrtStatus*的调用转换为早期返回)、EP_ENFORCE(条件断言,失败抛std::runtime_error)、IGNORE_ORTSTATUS(接管并释放不需要的OrtStatus*防止内存泄漏)、EXCEPTION_TO_RETURNED_STATUS_BEGIN/END(将 C++ 异常统一转换为OrtStatus*返回,适配 C API 的无异常边界); - 日志宏
LOG(level, ...):通过api_.Logger_LogMessage写入 ORT 日志系统,自动附带文件、行号与函数名; - 配置读取:
GetSessionConfigEntryOrDefault(session_options, key, default),从OrtSessionOptions的 session config 中取键值,供 EP 在创建时读取应用层配置(后文 EPContext 配置即此用法); - 类型检查辅助:
IsFloatTensor、GetTensorShape(返回std::optional<std::vector<int64_t>>,非张量时为空)、AreShapesStaticAndEqual(判断两个 shape 是否均为静态且相等,用于快速排除动态维度与广播场景); - 数据读取模板:
GetKernelInputDataAndShape<T>校验元素类型后取出张量数据 span 与 shape,是插件 EP 内核读取输入的惯用写法。
头文件开头使用#define ORT_API_MANUAL_INIT包含 C++ API 再取消,这是插件场景的固定写法——插件不能依赖 ORT 静态初始化的全局 API 指针,必须由入口函数手动初始化(对应上文Ort::InitApi(ort_api))。
示例一:example_plugin_ep —— 编译型 EP 的完整样板
这是三者中最完整的参考实现,覆盖了插件 EP 生命周期中的大部分回调。
工厂层:虚函数表填充与元数据
ep_factory.cc 的构造函数展示了工厂必须填充的函数指针组:
ort_version_supported = ORT_API_VERSION; // 记录编译时的 ORT 版本 GetName = GetNameImpl; GetVendor = GetVendorImpl; GetVendorId = GetVendorIdImpl; GetVersion = GetVersionImpl; GetSupportedDevices = GetSupportedDevicesImpl; CreateEp = CreateEpImpl; ReleaseEp = ReleaseEpImpl; CreateAllocator = CreateAllocatorImpl; ReleaseAllocator = ReleaseAllocatorImpl; CreateDataTransfer = CreateDataTransferImpl; IsStreamAware = IsStreamAwareImpl; CreateSyncStreamForDevice = CreateSyncStreamForDeviceImpl; GetHardwareDeviceIncompatibilityDetails = GetHardwareDeviceIncompatibilityDetailsImpl; CreateExternalResourceImporterForDevice = CreateExternalResourceImporterForDeviceImpl; GetNumCustomOpDomains = GetNumCustomOpDomainsImpl; GetCustomOpDomains = GetCustomOpDomainsImpl; ValidateCompiledModelCompatibilityInfo = ValidateCompiledModelCompatibilityInfoImpl; SelectBestModelCandidate = SelectBestModelCandidateImpl;同时,工厂为 EP 预置了三类Ort::MemoryInfo,分别对应三种内存语义(ep_factory.cc#L69-L99):
default_memory_info_:设备默认内存(OrtDeviceAllocator),注释里明确OrtArenaAllocator是 ORT 内部 Arena 实现保留的类型,插件不可用;readonly_memory_info_:OrtReadOnlyAllocator,用于初始化器(权重);host_accessible_memory_info_:OrtDeviceMemoryType_HOST_ACCESSIBLE,暴露 EP 设备上 CPU 可访问的内存(如 pinned memory)。
工厂成员中还有共享的 arena 分配器(arena_allocator_+ 互斥锁 + 使用计数)与数据搬运实现ExampleDataTransfer。
设备注册:GetSupportedDevices 与设备元数据键
GetSupportedDevicesImpl(ep_factory.cc#L143-L215)遍历 ORT 传入的硬件设备,对 CPU 设备构造OrtEpDevice。这个函数是插件 EP 向 ORT“自我介绍”的核心位置,展示了几个标准元数据键:
factory->ort_api.AddKeyValuePair(ep_metadata, "supported_devices", "CrackGriffin 7+"); // 示例 os_driver_version,格式为 4 段点分版本 factory->ort_api.AddKeyValuePair(ep_metadata, kOrtEpDevice_EpMetadataKey_OSDriverVersion, "31.0.101.1000"); // GroupQueryAttention Value 缓存布局偏好 factory->ort_api.AddKeyValuePair(ep_metadata, kOrtEpDevice_EpMetadataKey_GqaPreferredValueLayout, "BNSH"); // 报告对全部初始化器的 weightless 支持 factory->ort_api.AddKeyValuePair(ep_metadata, kOrtEpDevice_EpMetadataKey_WeightlessSupport, "all"); factory->ort_api.AddKeyValuePair(ep_options, "run_really_fast", "true");注意源码注释中一段很有价值的自述:如果 EP 上报了BNHS布局但GetCapability并未实现对应的 Transpose 融合序列,会误导应用层选择无法更快执行的布局——元数据声明的能力必须与能力查询实现一致。
随后通过ep_api.EpDevice_AddAllocatorInfo把三类分配器信息挂到OrtEpDevice上(read-only 与 host-accessible 为可选)。函数尾部还保留了一段被注释掉的 C++ API 等价写法,展示用 RAII 包装的Ort::EpDevice完成同样的注册。
EP 实例:从 session options 到 EPContext 配置
CreateEpImpl(ep_factory.cc#L218-L276)中示范了如何把 session 配置转成 EP 内部Config:
RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpContextEnable, "0", ep_context_enable)); RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpContextEmbedMode, "0", ep_context_embed_mode)); RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpContextFilePath, "", ep_context_output_model_path)); RETURN_IF_ERROR(GetSessionConfigEntryOrDefault(*session_options, kOrtSessionOptionEpEnableWeightlessEpContextNodes, "0", weightless_ep_context_nodes_enable)); ExampleEp::Config config = {}; config.enable_ep_context = ep_context_enable == "1"; config.embed_ep_context_in_model = ep_context_embed_mode == "1"; config.ep_context_output_model_path = std::move(ep_context_output_model_path); config.enable_weightless_ep_context_nodes = weightless_ep_context_nodes_enable == "1";源码中特别警告:不要保存对OrtSessionOptions的引用,因为其生命周期不受保证——只能提取值。整个函数体包裹在EXCEPTION_TO_RETURNED_STATUS_BEGIN/END中,Ort::Experimental::EpContextConfig构造函数失败(实验性 API 不可用)也会被统一转成OrtStatus。
能力查询与编译:MulKernel 与 EPContextKernel 双内核
ep.h 中ExampleEp继承OrtEp,虚函数表覆盖GetName、GetWeightlessSupport、CreateAllocator、CreateSyncStreamForDevice、GetCapability、Compile、ReleaseNodeComputeInfos、GetCompiledModelCompatibilityInfo、Sync、GetDefaultMemoryDevice等。EP 内部持有两种内核:
MulKernel:示例性的 Mul 实现(不处理广播),在编译阶段生成,计算时通过float_initializers_使用 EP 自行保存的常量初始化器副本;EpContextKernel:专门处理从已编译模型加载的 EPContext 节点。示例中其Compute()返回 NOT_IMPLEMENTED——注释说明生产级 EP 会反序列化ep_cache_context属性并恢复编译状态,示例把它拆出来只是为了“清晰分离 EPContext 处理与 MulKernel”。
能力查询侧(ep.cc#L350-L397)展示了两种上报方式的分歧点:
- 具体实现的自定义节点(如
Custom_Mul)调用EpGraphSupportInfo_AddSingleNode,告知 ORT 该节点不参与融合/编译,直接由内核执行; - 需要编译的节点(
Mul、EPContext)调用EpGraphSupportInfo_AddNodesToFuse,并设置node_fusion_options.drop_constant_initializers。源码注释解释了其含义:为 true 表示 EP 不需要 ORT 在推理时提供常量初始化器输入(因为 EP 已自行拷贝并管理权重),这给 ORT 释放未使用初始化器的机会;反之若应用要求生成 weightless EPContext 模型,则设为 false,让 ORT 把权重作为 EPContext 节点输入提供(对应ep.enable_weightless_ep_context_nodes配置)。
CompileImpl(ep.cc#L400-L469)则示范了编译侧的防御性校验:拒绝编译单节点之外的输入、拒绝重新编译已含 EPContext 节点的模型(配置冲突时给出可操作的报错文本)、校验 fused node 的 EP 名称归属、按embed_mode属性选择从属性内嵌上下文或外部文件/回调读取 EPContext 二进制数据的流程。
编译产物兼容性:ValidateCompiledModelCompatibilityInfo 与候选选择
ep_factory.cc#L491-L634 实现了一套完整的兼容性协商逻辑,值得逐段学习:
- 兼容性信息串格式为
"<EP名>;version=<EP版本>;ort_api_version=<N>[;hardware_architecture=<arch>]",由 EP 的GetCompiledModelCompatibilityInfo生成; - 解析时逐项比对:EP 名不匹配或格式非法 →
OrtCompiledModelCompatibility_EP_UNSUPPORTED;EP 版本、ORT API 版本或硬件架构任一项不同 →EP_SUPPORTED_PREFER_RECOMPILATION;全部一致 →EP_SUPPORTED_OPTIMAL; SelectBestModelCandidateImpl对模型包中的多个候选各取ep_compatibility_info键,复用上面的校验并排名:
int CompatibilityRank(OrtCompiledModelCompatibility c) { switch (c) { case OrtCompiledModelCompatibility_EP_SUPPORTED_OPTIMAL: return 3; case OrtCompiledModelCompatibility_EP_SUPPORTED_PREFER_RECOMPILATION: return 2; case OrtCompiledModelCompatibility_EP_NOT_APPLICABLE: return 1; case OrtCompiledModelCompatibility_EP_UNSUPPORTED: return 0; default: return 0; } }即“最优匹配 > 偏好重编译 > 不适用 > 不支持”,全部不支持时返回SIZE_MAX表示无可用候选。
其他可选回调的示例覆盖
- 流感知:
IsStreamAwareImpl返回true,CreateSyncStreamForDeviceImpl仅对OrtDeviceMemoryType_DEFAULT的设备内存创建StreamImpl(其他内存类型无需设备流同步); - 数据搬运:
CreateDataTransferImpl直接返回工厂构造时创建的共享ExampleDataTransfer; - 自定义算子域:注册
test/test2两个域,各挂一个ExampleEpCustomOp(Custom_Mul、Custom_Mul2),通过CreateKernelV2/KernelComputeV2提供 V2 内核路径; - 设备不兼容原因上报:
GetHardwareDeviceIncompatibilityDetailsImpl对非 CPU 设备返回OrtDeviceEpIncompatibility_DEVICE_INCOMPATIBLE及描述文本; - 外部资源导入器:
CreateExternalResourceImporterForDeviceImpl创建ExampleExternalResourceImporter,注释提示生产级多 GPU EP 应在导入器中捕获ep_device以支持多物理设备的校验。
示例二:example_plugin_ep_virt_gpu —— 虚拟设备与交叉编译
该示例(ep_factory.cc)的核心差异在于GetSupportedDevicesImpl无视 ORT 传入的物理设备列表,主动创建一个虚拟 GPU:
// 若应用允许(如交叉编译场景),创建虚拟 OrtHardwareDevice。 // 此示例 EP 创建一个虚拟 GPU OrtHardwareDevice,并添加基于它的 OrtEpDevice。 if (factory->allow_virtual_devices_ && num_ep_devices < max_ep_devices) { OrtKeyValuePairs* hw_metadata = nullptr; factory->ort_api_.CreateKeyValuePairs(&hw_metadata); factory->ort_api_.AddKeyValuePair(hw_metadata, kOrtHardwareDevice_MetadataKey_IsVirtual, "1"); auto* status = factory->ep_api_.CreateHardwareDevice( OrtHardwareDeviceType::OrtHardwareDeviceType_GPU, factory->vendor_id_, /*device_id*/ 0, factory->vendor_.c_str(), hw_metadata, &factory->virtual_hw_device_); ... status = factory->ort_api_.GetEpApi()->CreateEpDevice( factory, factory->virtual_hw_device_, ep_metadata, ep_options, &virtual_ep_device); }要点:
- 通过硬件设备元数据键
kOrtHardwareDevice_MetadataKey_IsVirtual显式标记虚拟设备,allow_virtual_devices_构造参数决定开关——这正是文档中“虚拟设备可用于为不同目标交叉编译模型”的落地方式:在没有真实目标硬件的开发机上也能完成针对该虚拟目标的编译流程; - 该示例刻意保持精简:
CreateAllocatorImpl、CreateDataTransferImpl、CreateSyncStreamForDeviceImpl均返回空(注释说明“GPU EP 通常会支持,但示例从简”),IsStreamAware返回false; CreateEpImpl从 session 选项读取ep.context_enable配置决定是否启用 EPContext 输出。
示例三:example_plugin_ep_kernel_registry —— 内核注册范式
与“编译型 EP”相对,该示例展示直接注册算子内核的路线(ep.h):
ep.cc 的构造函数中:
ort_version_supported = ORT_API_VERSION; // 初始化 EP 函数表 GetName = GetNameImpl; GetCapability = GetCapabilityImpl; GetKernelRegistry = GetKernelRegistryImpl; CreateProfiler = CreateProfilerImpl; // 这不是编译型 EP,因此不需要以下项 Compile = nullptr; ReleaseNodeComputeInfos = nullptr;即通过GetKernelRegistry暴露OrtKernelRegistry,而不是Compile。GetCapabilityImpl(ep.cc#L55-L112)的能力查询逻辑:
- 遍历图中所有节点,按算子类型收集候选:
Relu、Squeeze、If、Loop、Scan直接候选;Mul/Sub额外要求两个输入张量 shape 均为静态且相等(不支持广播与动态维度),这正好用到了plugin_ep_utils.h中的GetTensorShape与AreShapesStaticAndEqual; - 对每个候选节点调用
EpGraphSupportInfo_LookUpKernel查询注册表里是否确有对应内核,只有查到才用EpGraphSupportInfo_AddSingleNode上报为支持。
这一“先看注册表再上报”的次序避免了“声明支持但无内核”的能力误报。GetKernelRegistryImpl不就地创建注册表,而是向工厂取缓存实例(factory_.GetKernelRegistryForEp),避免每个 EP 实例重复构建。此外该示例还实现了CreateProfilerImpl,返回ExampleKernelEpProfiler以支持 EP 级性能剖析(见 ep_profiling.h)。内核文件位于 kernels/ 子目录,按算子一文件一对组织(relu、squeeze、if、loop、scan、binary_op)。
构建与测试集成:示例如何被单测驱动
从 onnxruntime_unittests.cmake 可以看到三个示例各自被构建成共享库模块:
# example_plugin_ep file(GLOB onnxruntime_autoep_test_library_src "${TEST_SRC_DIR}/autoep/library/example_plugin_ep/*.h" "${TEST_SRC_DIR}/autoep/library/example_plugin_ep/*.cc" "${TEST_SRC_DIR}/autoep/library/ep_context_data_utils.h" "${TEST_SRC_DIR}/autoep/library/plugin_ep_utils.h") onnxruntime_add_shared_library_module(example_plugin_ep ${onnxruntime_autoep_test_library_src}) target_include_directories(example_plugin_ep PRIVATE ${REPO_ROOT}/include/onnxruntime/core/session) target_link_libraries(example_plugin_ep PRIVATE onnxruntime ${GSL_TARGET}) # Linux 侧用 --version-script 链接脚本限定导出符号,Windows 侧用 .def即每个示例库依赖主库onnxruntime与 GSL,导出面由.def/.lds严格限定为CreateEpFactories/ReleaseEpFactory。
测试侧由 test_autoep_utils.h 提供统一入口:静态的example_ep_info、example_ep_virt_gpu_info、example_plugin_ep_kernel_registry_info三个ExamplePluginInfo(库路径 + 注册名 + EP 名),以及RegisterAndGetExampleEp(注册插件库、取回OrtEpDevice并在析构时自动反注册)。具体测试分布在 test/ 下各文件:test_registration.cc 验证设备元数据(如示例插件声明的os_driver_version、GQA 布局偏好、weightless 支持,以及虚拟 GPU 的加载与反注册行为),test_execution.cc 覆盖端到端执行(并链接示例库内的测试钩头 ep_test_hooks.h 以检查同步计数等内部行为),另有 test_data_transfer.cc、test_ep_compatibility.cc(编译产物兼容性)、test_selection.cc(候选选择)、test_external_resource_importer.cc、test_handle_leak.cc 等,恰好与三个示例覆盖的回调一一对应。
小结:三种范式的适用选择
| 维度 | example_plugin_ep | example_plugin_ep_virt_gpu | example_plugin_ep_kernel_registry |
|---|---|---|---|
| 执行模型 | 编译型(Compile+ 融合节点) | 编译型 + 虚拟 GPU 设备 | 内核注册(GetKernelRegistry) |
| 目标算子 | Mul、EPContext、Custom_Mul(自定义域) | 可编译节点(配置ep.context_enable) | Relu、Squeeze、If、Loop、Scan、Mul/Sub(静态同形) |
| 设备 | 只支持 CPU 设备 | 主动注册虚拟 GPU | 工厂默认路径 |
| 额外能力 | arena 分配器、流同步、数据搬运、外部资源导入、自定义算子域、兼容性校验与候选选择、Custom Op | 虚拟硬件设备(交叉编译场景) | 工厂级内核注册表缓存、EP Profiler |
| 参考文件 | example_plugin_ep/ | example_plugin_ep_virt_gpu/ | example_plugin_ep_kernel_registry/ |
如果你的插件需要对子图做图级编译/代码生成(如针对自家硬件的图编译器),以第一个示例为骨架,重点研究GetCapability的融合上报与Compile的上下文产物管理;如果目标是“算子级替换”且不想引入编译环节,参考第三个示例的内核注册路线;如果需要支持在无真实硬件环境下完成编译验证或多目标构建,则第二个示例的虚拟设备机制是现成范式。所有示例均声明仅用于演示与测试,生产实现需在此基础上补齐分配器、数据搬运、流同步等路径,并对照当前版本的插件 EP 接口逐一核对回调签名。
【免费下载链接】onnxruntimeONNX Runtime: cross-platform, high performance ML inferencing and training accelerator项目地址: https://gitcode.com/GitHub_Trending/on/onnxruntime
创作声明:本文部分内容由AI辅助生成(AIGC),仅供参考