[EP ABI] Add APIs to create control flow kernels for plugin EPs - #26927
Conversation
…Impl returned by ORT with an OrtKernelImpl created by the EP
… along with the helpers
There was a problem hiding this comment.
Pull request overview
This pull request adds APIs to enable plugin execution providers (EPs) to create and register kernels for control flow operators (If, Loop, and Scan). The implementation provides ORT-managed kernel implementations that handle the complex subgraph execution while allowing EPs to provide device-specific helper functions for operations like tensor concatenation and transposition.
Key changes:
- Adds three new OrtEpApi functions:
CreateIfKernel,CreateLoopKernel, andCreateScanKernelfor creating control flow kernel implementations - Introduces helper structures (
OrtLoopKernelHelper,OrtScanKernelHelper) that EPs implement to provide device-specific operations - Adds comprehensive test coverage with new ONNX model test cases and execution tests for all three control flow operators
Reviewed changes
Copilot reviewed 28 out of 31 changed files in this pull request and generated 10 comments.
Show a summary per file
| File | Description |
|---|---|
| onnxruntime/test/testdata/scan_mul.py | Python script to generate test model for Scan operator |
| onnxruntime/test/testdata/scan_mul.onnx | Binary ONNX test model for Scan operator |
| onnxruntime/test/testdata/loop_sub_one.py | Python script to generate test model for Loop operator |
| onnxruntime/test/testdata/loop_sub_one.onnx | Binary ONNX test model for Loop operator |
| onnxruntime/test/testdata/if_mul.py | Python script to generate test model for If operator |
| onnxruntime/test/testdata/if_mul.onnx | Binary ONNX test model for If operator |
| onnxruntime/test/autoep/test_execution.cc | Adds execution tests for If, Loop, and Scan operators with the plugin EP |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/kernels/utils.h | Updates kernel creation macro to match new Create signature returning OrtKernelImpl* |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/kernels/squeeze.h/cc | Updates Create method signature and adds creator field initialization |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/kernels/relu.h/cc | Updates Create method signature and adds creator field initialization |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/kernels/binary_op.h/cc | Updates Create method signature and adds creator field initialization |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/kernels/scan.h/cc | Implements ScanHelper for Scan operator with transpose helper function |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/kernels/loop.h/cc | Implements LoopHelper for Loop operator with concatenation helper function |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/kernels/if.h/cc | Implements IfHelper for If operator |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/ep_kernel_registration.cc | Registers If, Loop, and Scan kernels with the plugin EP |
| onnxruntime/test/autoep/library/example_plugin_ep_kernel_registry/ep.cc | Updates GetCapability to claim If, Loop, and Scan operators |
| onnxruntime/core/session/plugin_ep/ep_kernel_registration.cc | Updates PluginEpOpKernel to derive from IControlFlowKernel and implement SetupSubgraphExecutionInfo |
| onnxruntime/core/session/plugin_ep/ep_control_flow_kernel_impls.h | Defines control flow kernel implementation classes that wrap CPU EP kernels |
| onnxruntime/core/session/plugin_ep/ep_control_flow_kernel_impls.cc | Implements control flow kernel wrappers and helper function adapters |
| onnxruntime/core/session/plugin_ep/ep_api.h | Declares new control flow kernel creation functions |
| onnxruntime/core/session/plugin_ep/ep_api.cc | Implements control flow kernel creation APIs with validation |
| onnxruntime/core/providers/cpu/controlflow/loop.h | Changes SetConcatOutputFunc visibility to public |
| onnxruntime/core/framework/session_state.cc | Updates CreateSubgraphSessionState to handle plugin EPs with control flow operators |
| include/onnxruntime/core/session/onnxruntime_ep_c_api.h | Adds OrtKernelImplCreator enum, helper structures, and new API function declarations |
| cmake/onnxruntime_unittests.cmake | Adds new control flow kernel source files to the test build |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
There was a problem hiding this comment.
Pull request overview
Copilot reviewed 28 out of 31 changed files in this pull request and generated 4 comments.
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Description
Adds APIs to enable plugin EPs to create and register kernels for control flow operators (If, Loop, and Scan). The implementation provides ORT-managed kernel implementations that handle subgraph execution while allowing EPs to provide device-specific helper functions for operations like tensor concatenation and transposition.
Key changes:
CreateIfKernel,CreateLoopKernel,CreateScanKernel, andReleaseKernelImplfor creating control flow kernel implementationsOrtLoopKernelHelper,OrtScanKernelHelper) that EPs implement to provide device-specific operationsMotivation and Context