[EP ABI] Initial support for kernel-based EPs #26206

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged

adrianlizarraga merged 79 commits into main from adrianl/ep-abi-kernel-based-eps

Dec 5, 2025

cmake/onnxruntime_unittests.cmake

-Original file line number
+Diff line change
@@ Expand Up / @@ -2094,6 +2094,50 @@ if (onnxruntime_BUILD_SHARED_LIB AND @@
       set_target_properties(example_plugin_ep_virt_gpu PROPERTIES FOLDER "ONNXRuntimeTest")
       source_group(TREE ${TEST_SRC_DIR} FILES ${onnxruntime_autoep_test_example_plugin_ep_virt_gpu_src})
+      #
+      # example_plugin_ep_kernel_registry
+      #
+      set(onnxruntime_autoep_test_example_plugin_ep_kernel_registry_src
+              "${TEST_SRC_DIR}/autoep/library/plugin_ep_utils.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_lib_entry.cc"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_factory.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_factory.cc"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep.cc"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_kernel_registration.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_kernel_registration.cc"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/utils.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/base.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/base.cc"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/squeeze.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/squeeze.cc"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/relu.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/relu.cc"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/mul.h"
+              "${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/kernels/mul.cc")
+      onnxruntime_add_shared_library_module(example_plugin_ep_kernel_registry ${onnxruntime_autoep_test_example_plugin_ep_kernel_registry_src})
+      target_include_directories(example_plugin_ep_kernel_registry PRIVATE ${REPO_ROOT}/include/onnxruntime/core/session)
+      target_link_libraries(example_plugin_ep_kernel_registry PRIVATE onnxruntime ${GSL_TARGET})
+      if(UNIX)
+        if (APPLE)
+    	    set(ONNXRUNTIME_EXAMPLE_PLUGIN_EP_KERNEL_REGISTRY_LINK_FLAG "-Xlinker -dead_strip")
+        elseif (NOT CMAKE_SYSTEM_NAME MATCHES "AIX")
+          string(CONCAT ONNXRUNTIME_EXAMPLE_PLUGIN_EP_KERNEL_REGISTRY_LINK_FLAG
+                 "-Xlinker --version-script=${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_lib.lds "
+                 "-Xlinker --no-undefined -Xlinker --gc-sections -z noexecstack")
+        endif()
+      else()
+        set(ONNXRUNTIME_EXAMPLE_PLUGIN_EP_KERNEL_REGISTRY_LINK_FLAG
+            "-DEF:${TEST_SRC_DIR}/autoep/library/example_plugin_ep_kernel_registry/ep_lib.def")
+      endif()
+      set_property(TARGET example_plugin_ep_kernel_registry APPEND_STRING PROPERTY LINK_FLAGS
+                   ${ONNXRUNTIME_EXAMPLE_PLUGIN_EP_KERNEL_REGISTRY_LINK_FLAG})
+      set_target_properties(example_plugin_ep_kernel_registry PROPERTIES FOLDER "ONNXRuntimeTest")
+      source_group(TREE ${TEST_SRC_DIR} FILES ${onnxruntime_autoep_test_example_plugin_ep_kernel_registry_src})
       #
       # test library
       #
@@ Expand Down Expand Up / @@ -2129,7 +2173,7 @@ if (onnxruntime_BUILD_SHARED_LIB AND @@
               TARGET onnxruntime_autoep_test
               SOURCES ${onnxruntime_autoep_test_SRC} ${onnxruntime_unittest_main_src}
               LIBS ${onnxruntime_autoep_test_LIBS}
-              DEPENDS ${all_dependencies} example_plugin_ep example_plugin_ep_virt_gpu
+              DEPENDS ${all_dependencies} example_plugin_ep example_plugin_ep_virt_gpu example_plugin_ep_kernel_registry
       )
     endif()
@@ Expand Down @@

include/onnxruntime/core/session/onnxruntime_cxx_api.h

-Original file line number
+Diff line change
@@ Expand Up / @@ -644,6 +644,9 @@ ORT_DEFINE_RELEASE(ValueInfo); @@
     ORT_DEFINE_RELEASE_FROM_API_STRUCT(ModelCompilationOptions, GetCompileApi);
     ORT_DEFINE_RELEASE_FROM_API_STRUCT(EpDevice, GetEpApi);
+    ORT_DEFINE_RELEASE_FROM_API_STRUCT(KernelDef, GetEpApi);
+    ORT_DEFINE_RELEASE_FROM_API_STRUCT(KernelDefBuilder, GetEpApi);
+    ORT_DEFINE_RELEASE_FROM_API_STRUCT(KernelRegistry, GetEpApi);
     // This is defined explicitly since OrtTensorRTProviderOptionsV2 is not a C API type,
     // but the struct has V2 in its name to indicate that it is the second version of the options.
@@ Expand Down Expand Up / @@ -3286,5 +3289,89 @@ struct Model : detail::ModelImpl<OrtModel> { @@
       explicit Model(const std::vector<DomainOpsetPair>& opsets);
     #endif
     };
+    namespace detail {
+    template <typename T>
+    struct ConstKernelDefImpl : Base<T> {
+      using B = Base<T>;
+      using B::B;
+      ///< Wraps OrtEpApi::KernelDef_GetOperatorType
+      const char* GetOperatorType() const;
+      ///< Wraps OrtEpApi::KernelDef_GetDomain
+      const char* GetDomain() const;
+      ///< Wraps OrtEpApi::KernelDef_GetSinceVersion
+      std::pair<int, int> GetSinceVersion() const;
+      ///< Wraps OrtEpApi::KernelDef_GetExecutionProvider
+      const char* GetExecutionProvider() const;
+      ///< Wraps OrtEpApi::KernelDef_GetInputMemType
+      OrtMemType GetInputMemType(size_t input_index) const;
+      ///< Wraps OrtEpApi::KernelDef_GetOutputMemType
+      OrtMemType GetOutputMemType(size_t output_index) const;
+    };
+    }  // namespace detail
+    using ConstKernelDef = detail::ConstKernelDefImpl<detail::Unowned<const OrtKernelDef>>;
+    struct KernelDef : detail::ConstKernelDefImpl<OrtKernelDef> {
+      using Base = detail::ConstKernelDefImpl<OrtKernelDef>;
+      using Base::Base;
+      explicit KernelDef(std::nullptr_t) {}
+      explicit KernelDef(OrtKernelDef* p) : detail::ConstKernelDefImpl<OrtKernelDef>{p} {}
+      ConstKernelDef GetConst() const { return ConstKernelDef{this->p_}; }
+    };
+    /** \brief Builder for OrtKernelDef.
+     *
+     * Used by plugin EPs to build a kernel definition.
+     */
+    struct KernelDefBuilder : detail::Base<OrtKernelDefBuilder> {
+      KernelDefBuilder();                           ///< Wraps OrtEpApi::CreateKernelDefBuilder
+      explicit KernelDefBuilder(std::nullptr_t) {}  ///< Create an empty object, must be assigned a valid one to be used
+      explicit KernelDefBuilder(OrtKernelDefBuilder* ort_kernel_def_builder);
+      KernelDefBuilder& SetOperatorType(const char* op_type);
+      KernelDefBuilder& SetDomain(const char* domain);
+      KernelDefBuilder& SetSinceVersion(int since_version_start, int since_version_end);
+      KernelDefBuilder& SetExecutionProvider(const char* ep_name);
+      KernelDefBuilder& SetInputMemType(size_t input_index, OrtMemType mem_type);
+      KernelDefBuilder& SetOutputMemType(size_t output_index, OrtMemType mem_type);
+      KernelDefBuilder& AddTypeConstraint(const char* arg_name, const OrtDataType* data_type);
+      KernelDefBuilder& AddTypeConstraint(const char* arg_name, const std::vector<const OrtDataType*>& data_types);
+      KernelDefBuilder& AddInputOutputAlias(int input_index, int output_index);
+      KernelDefBuilder& AddInputOutputAliases(const std::vector<int>& input_indices,
+                                              const std::vector<int>& output_indices);
+      KernelDefBuilder& AddInputOutputMutableAlias(int input_index, int output_index);
+      KernelDefBuilder& AddInputOutputMutableAliases(const std::vector<int>& input_indices,
+                                                     const std::vector<int>& output_indices);
+      KernelDef Build();
+    };
+    /** \brief Registry for kernels supported by an EP.
+     *
+     * Used by plugin EPs to register definitions for supported kernels.
+     */
+    struct KernelRegistry : detail::Base<OrtKernelRegistry> {
+      ///< Wrapper around OrtEpApi::CreateKernelRegistry
+      KernelRegistry();
+      ///< Create an empty object, must be assigned a valid one to be used
+      explicit KernelRegistry(std::nullptr_t) {}
+      ///< Take ownership of a pointer created with the C API.
+      explicit KernelRegistry(OrtKernelRegistry* ort_kernel_registry);
+      ///< Wraps KernelRegistry_AddKernel
+      Status AddKernel(const OrtKernelDef* kernel_def, OrtKernelCreateFunc kernel_create_func,
+                       void* kernel_create_func_state);
+    };
     }  // namespace Ort
     #include "onnxruntime_cxx_inline.h"

include/onnxruntime/core/session/onnxruntime_cxx_inline.h

-Original file line number
+Diff line change
@@ Expand Up / @@ -12,6 +12,7 @@ @@
     #include <iterator>
     #include <string>
     #include <type_traits>
+    #include <utility>
     #include <vector>
     // Convert OrtStatus to Ort::Status and return
@@ Expand Down Expand Up @@
     }
     #endif
+    namespace detail {
+    template <typename T>
+    inline const char* ConstKernelDefImpl<T>::GetOperatorType() const {
+      return GetEpApi().KernelDef_GetOperatorType(this->p_);
+    }
+    template <typename T>
+    inline const char* ConstKernelDefImpl<T>::GetDomain() const {
+      return GetEpApi().KernelDef_GetDomain(this->p_);
+    }
+    template <typename T>
+    inline std::pair<int, int> ConstKernelDefImpl<T>::GetSinceVersion() const {
+      int start = 0;
+      int end = 0;
+      ThrowOnError(GetEpApi().KernelDef_GetSinceVersion(this->p_, &start, &end));
+      return std::pair<int, int>(start, end);
+    }
+    template <typename T>
+    inline const char* ConstKernelDefImpl<T>::GetExecutionProvider() const {
+      return GetEpApi().KernelDef_GetExecutionProvider(this->p_);
+    }
+    template <typename T>
+    inline OrtMemType ConstKernelDefImpl<T>::GetInputMemType(size_t input_index) const {
+      OrtMemType mem_type{};
+      ThrowOnError(GetEpApi().KernelDef_GetInputMemType(this->p_, input_index, &mem_type));
+      return mem_type;
+    }
+    template <typename T>
+    inline OrtMemType ConstKernelDefImpl<T>::GetOutputMemType(size_t output_index) const {
+      OrtMemType mem_type{};
+      ThrowOnError(GetEpApi().KernelDef_GetOutputMemType(this->p_, output_index, &mem_type));
+      return mem_type;
+    }
+    }  // namespace detail
+    inline KernelDefBuilder::KernelDefBuilder() {
+      ThrowOnError(GetEpApi().CreateKernelDefBuilder(&p_));
+    }
+    inline KernelDefBuilder::KernelDefBuilder(OrtKernelDefBuilder* p) : detail::Base<OrtKernelDefBuilder>{p} {
+    }
+    inline KernelDefBuilder& KernelDefBuilder::SetOperatorType(const char* op_type) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_SetOperatorType(p_, op_type));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::SetDomain(const char* domain) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_SetDomain(p_, domain));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::SetSinceVersion(int since_version_start, int since_version_end) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_SetSinceVersion(p_, since_version_start, since_version_end));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::SetExecutionProvider(const char* ep_name) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_SetExecutionProvider(p_, ep_name));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::SetInputMemType(size_t input_index, OrtMemType mem_type) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_SetInputMemType(p_, input_index, mem_type));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::SetOutputMemType(size_t output_index, OrtMemType mem_type) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_SetOutputMemType(p_, output_index, mem_type));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::AddTypeConstraint(const char* arg_name,
+                                                                 const OrtDataType* data_type) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_AddTypeConstraint(p_, arg_name, &data_type, 1));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::AddTypeConstraint(const char* arg_name,
+                                                                 const std::vector<const OrtDataType*>& data_types) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_AddTypeConstraint(p_, arg_name, data_types.data(), data_types.size()));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::AddInputOutputAlias(int input_index, int output_index) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputAliases(p_, &input_index, &output_index, 1));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::AddInputOutputAliases(const std::vector<int>& input_indices,
+                                                                     const std::vector<int>& output_indices) {
+      if (input_indices.size() != output_indices.size()) {
+        ORT_CXX_API_THROW("Expecting input and output indices to have the same element count", ORT_INVALID_ARGUMENT);
+      }
+      ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputAliases(p_, input_indices.data(), output_indices.data(),
+                                                                     input_indices.size()));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::AddInputOutputMutableAlias(int input_index, int output_index) {
+      ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputMutableAliases(p_, &input_index, &output_index, 1));
+      return *this;
+    }
+    inline KernelDefBuilder& KernelDefBuilder::AddInputOutputMutableAliases(const std::vector<int>& input_indices,
+                                                                            const std::vector<int>& output_indices) {
+      if (input_indices.size() != output_indices.size()) {
+        ORT_CXX_API_THROW("Expecting input and output indices to have the same element count", ORT_INVALID_ARGUMENT);
+      }
+      ThrowOnError(GetEpApi().KernelDefBuilder_AddInputOutputMutableAliases(p_, input_indices.data(), output_indices.data(),
+                                                                            input_indices.size()));
+      return *this;
+    }
+    inline KernelDef KernelDefBuilder::Build() {
+      OrtKernelDef* kernel_def = nullptr;
+      ThrowOnError(GetEpApi().KernelDefBuilder_Build(p_, &kernel_def));
+      return KernelDef(kernel_def);
+    }
+    inline KernelRegistry::KernelRegistry() {
+      ThrowOnError(GetEpApi().CreateKernelRegistry(&p_));
+    }
+    inline KernelRegistry::KernelRegistry(OrtKernelRegistry* p) : detail::Base<OrtKernelRegistry>{p} {
+    }
+    inline Status KernelRegistry::AddKernel(const OrtKernelDef* kernel_def, OrtKernelCreateFunc kernel_create_func,
+                                            void* kernel_create_func_state) {
+      return Status{GetEpApi().KernelRegistry_AddKernel(p_, kernel_def, kernel_create_func, kernel_create_func_state)};
+    }
     }  // namespace Ort

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[EP ABI] Initial support for kernel-based EPs #26206

Uh oh!

Diff view

Diff view

There are no files selected for viewing

Uh oh!

Uh oh!

Uh oh!

Uh oh!