[WIP] Out-Tree EP feature

include/onnxruntime/core/session/environment.h

@@ -88,6 +88,10 @@
    */
   Status CreateAndRegisterAllocatorV2(const std::string& provider_type, const OrtMemoryInfo& mem_info, const std::unordered_map<std::string, std::string>& options, const OrtArenaCfg* arena_cfg = nullptr);
 
+  void InsertCustomEp(const char* ep_name, OrtExecutionProviderFactory* ep_factory);
+
+  const std::unordered_map<std::string, std::unique_ptr<OrtExecutionProviderFactory>>& GetCustomEpFactories() const { return custom_ep_factories_; }
+
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Environment);
   Status Initialize(std::unique_ptr<logging::LoggingManager> logging_manager,
@@ -99,5 +103,6 @@
   std::unique_ptr<onnxruntime::concurrency::ThreadPool> inter_op_thread_pool_;
   bool create_global_thread_pools_{false};
   std::vector<AllocatorPtr> shared_allocators_;
+  std::unordered_map<std::string, std::unique_ptr<OrtExecutionProviderFactory>> custom_ep_factories_;
 };
 }  // namespace onnxruntime

include/onnxruntime/core/session/onnxruntime_c_api.h

@@ -304,6 +304,10 @@
 ORT_RUNTIME_CLASS(OpAttr);
 ORT_RUNTIME_CLASS(Logger);
 ORT_RUNTIME_CLASS(ShapeInferContext);
+ORT_RUNTIME_CLASS(ExecutionProvider);
+ORT_RUNTIME_CLASS(ExecutionProviderFactory);
+ORT_RUNTIME_CLASS(Node);
+ORT_RUNTIME_CLASS(GraphViewer);
 
 #ifdef _WIN32
 typedef _Return_type_success_(return == 0) OrtStatus* OrtStatusPtr;
@@ -689,6 +693,50 @@
  */
 ORT_EXPORT const OrtApiBase* ORT_API_CALL OrtGetApiBase(void) NO_EXCEPTION;
 
+typedef struct OrtMetaDef {
+  const char* name;
+  const char* domain;
+  int since_version;
+
+  const char** inputs;
+  size_t input_len;
+  const char** outputs;
+  size_t output_len;
+  const char** constant_initializers;
+  size_t initializer_len;
+
+  const char* doc_string;
+} OrtMetaDef;
+
+typedef struct OrtIndexedSubGraph {
+  OrtMetaDef* meta_def; // TODO(leca): how to define a nested structure pointer?
+  size_t* node_index;
+  size_t node_index_len;
+} OrtIndexedSubGraph;
+
+typedef struct OrtComputeContext {
+  void*(ORT_API_CALL* AllocateFunc)(void*, size_t, size_t);
+  void(ORT_API_CALL* DestroyFunc)(void*, void*);
+  void* allocator_handle;
+  const char* node_name;
+} OrtComputeContext;
+
+typedef struct OrtNodeComputeInfo {
+  int(ORT_API_CALL* CreateFunctionStateFunc)(OrtComputeContext*, void**);
+  OrtStatusPtr(ORT_API_CALL* ComputeFunc)(void*, const OrtApi*, OrtKernelContext*);
+  void(ORT_API_CALL* DestroyFunctionStateFunc)(void*);
+} OrtNodeComputeInfo;
+
+typedef struct OrtExecutionProvider {
+  void(ORT_API_CALL* GetCapability)(const OrtExecutionProvider* this_, const OrtGraphViewer* graph, size_t* cnt, OrtIndexedSubGraph***);
+  void(ORT_API_CALL* Compile)(OrtExecutionProvider* this_, const OrtGraphViewer** graph, const OrtNode** node, size_t cnt, OrtNodeComputeInfo*** node_compute_info);
+  const char* type;
+} OrtExecutionProvider;
+
+typedef struct OrtExecutionProviderFactory {
+  void*(ORT_API_CALL* CreateExecutionProvider)(OrtExecutionProviderFactory* this_, const char* const* ep_option_keys, const char* const* ep_option_values, size_t option_size);
+} OrtExecutionProviderFactory;
+
 /** \brief Thread work loop function
  *
  * Onnxruntime will provide the working loop on custom thread creation
@@ -4665,7 +4713,28 @@
                   _In_reads_(num_external_initializer_files) char* const* external_initializer_file_buffer_array,
                   _In_reads_(num_external_initializer_files) const size_t* external_initializer_file_lengths,
                   size_t num_external_initializer_files);
-};
+
+  ORT_API2_STATUS(RegisterOrtExecutionProviderLibrary, _In_ const ORTCHAR_T* lib_path, _In_ OrtEnv* env, _In_ const char* ep_name);
+
+  ORT_API2_STATUS(SessionOptionsAppendOrtExecutionProvider, _In_ OrtSessionOptions* options, _In_ const char* ep_name,
+                   _In_reads_(num_keys) const char* const* provider_options_keys, _In_reads_(num_keys) const char* const* provider_options_values, _In_ size_t num_keys);
+
+  ORT_API2_STATUS(OrtGraph_IsConstantInitializer, const OrtGraphViewer* graph, const char* name, bool check_outer_scope, _Out_ bool* ret);
+
+  ORT_API2_STATUS(OrtGraph_GetNodesIndexInTopologicalOrder, const OrtGraphViewer* graph, _Out_ size_t* len, _Out_ const size_t** nodes_index_in_topological_order);
+
+  ORT_API2_STATUS(OrtGraph_GetOrtNode, const OrtGraphViewer* graph, size_t node_index, _Outptr_ const OrtNode** node);
+
+  ORT_API2_STATUS(OrtNode_GetOpType, const OrtNode* node, _Out_ const char** op_type);
+
+  ORT_API2_STATUS(OrtNode_GetInputSize, const OrtNode* node, _Out_ size_t* input_size);
+
+  ORT_API2_STATUS(OrtNode_GetIthInputName, const OrtNode* node, size_t i, _Out_ const char** ith_input_name);
+
+  ORT_API2_STATUS(OrtNode_GetOutputSize, const OrtNode* node, _Out_ size_t* output_size);
+
+  ORT_API2_STATUS(OrtNode_GetIthOutputName, const OrtNode* node, size_t i, _Out_ const char** ith_output_name);
+};  // struct OrtApi
 
 /*
  * Steps to use a custom op:

onnxruntime/core/framework/provider_adapter.h

@@ -0,0 +1,96 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/session/onnxruntime_c_api.h"
+#include "core/framework/compute_capability.h"
+
+namespace onnxruntime {
+class ExecutionProviderAdapter : public IExecutionProvider {
+public:
+  ExecutionProviderAdapter(OrtExecutionProvider* ep) : IExecutionProvider(ep->type), ep_impl_(ep) {}
+  virtual std::vector<std::unique_ptr<ComputeCapability>> GetCapability(const GraphViewer& graph_viewer, const IKernelLookup& kernel_lookup) const override {
+    size_t cnt = 0;
+    OrtIndexedSubGraph** indexed_subgraph = nullptr;
+    ep_impl_->GetCapability(ep_impl_, reinterpret_cast<const OrtGraphViewer*>(&graph_viewer), &cnt, &indexed_subgraph);
+
+    if (cnt == 0) return IExecutionProvider::GetCapability(graph_viewer, kernel_lookup);
+
+    std::vector<std::unique_ptr<ComputeCapability>> ret;
+    for (size_t i = 0; i < cnt; i++) {
+        std::unique_ptr<IndexedSubGraph> sb = std::make_unique<IndexedSubGraph>();
+        sb->nodes.reserve(indexed_subgraph[i]->node_index_len);
+        for (size_t j = 0; j < indexed_subgraph[i]->node_index_len; j++) sb->nodes.push_back((indexed_subgraph[i]->node_index)[j]);
+        if (indexed_subgraph[i]->meta_def != nullptr) {
+            std::unique_ptr<IndexedSubGraph::MetaDef> meta_def = std::make_unique<IndexedSubGraph::MetaDef>();
+            meta_def->name = indexed_subgraph[i]->meta_def->name ? indexed_subgraph[i]->meta_def->name : "";
+            meta_def->doc_string = indexed_subgraph[i]->meta_def->doc_string ? indexed_subgraph[i]->meta_def->doc_string : "";
+            meta_def->domain = indexed_subgraph[i]->meta_def->domain ? indexed_subgraph[i]->meta_def->domain : "";
+            meta_def->since_version = indexed_subgraph[i]->meta_def->since_version;
+
+            meta_def->inputs.reserve(indexed_subgraph[i]->meta_def->input_len);
+            for (size_t j = 0; j < indexed_subgraph[i]->meta_def->input_len; j++) meta_def->inputs.push_back(indexed_subgraph[i]->meta_def->inputs[j]);
+
+            meta_def->outputs.reserve(indexed_subgraph[i]->meta_def->output_len);
+            for (size_t j = 0; j < indexed_subgraph[i]->meta_def->output_len; j++) meta_def->outputs.push_back(indexed_subgraph[i]->meta_def->outputs[j]);
+
+            meta_def->constant_initializers.reserve(indexed_subgraph[i]->meta_def->initializer_len);
+            for (size_t j = 0; j < indexed_subgraph[i]->meta_def->initializer_len; j++) meta_def->constant_initializers.push_back(indexed_subgraph[i]->meta_def->constant_initializers[j]);
+
+            sb->SetMetaDef(std::move(meta_def));
+        }
+
+        ret.push_back(std::make_unique<ComputeCapability>(std::move(sb)));
+    }
+    return ret;
+  }
+
+  virtual common::Status Compile(const std::vector<FusedNodeAndGraph>& fused_nodes_and_graphs, std::vector<NodeComputeInfo>& node_compute_funcs) override {
+    std::vector<const OrtGraphViewer*> ortGraphs;
+    std::vector<const OrtNode*> ortNodes;
+    for (auto& fused_node_graph : fused_nodes_and_graphs) {
+      const GraphViewer& graph_viewer = fused_node_graph.filtered_graph;
+      const Node& fused_node = fused_node_graph.fused_node;
+      ortGraphs.push_back(reinterpret_cast<const OrtGraphViewer*>(&graph_viewer));
+      ortNodes.push_back(reinterpret_cast<const OrtNode*>(&fused_node));
+    }
+    size_t count = fused_nodes_and_graphs.size();
+    node_compute_info_ = new OrtNodeComputeInfo* [count];
+    ep_impl_->Compile(ep_impl_, ortGraphs.data(), ortNodes.data(), count, &node_compute_info_);
+
+    node_compute_funcs.reserve(count);
+    for (size_t i = 0; i < count; i++) {
+        NodeComputeInfo compute_info;
+        compute_info.create_state_func = [&, i](ComputeContext* context, void** state) {
+            if (node_compute_info_[i]->CreateFunctionStateFunc) {
+                OrtComputeContext occ;
+                occ.AllocateFunc = context->allocate_func;
+                occ.DestroyFunc = context->release_func;
+                occ.allocator_handle = context->allocator_handle;
+                occ.node_name = context->node_name;
+                return node_compute_info_[i]->CreateFunctionStateFunc(&occ, state);  // TODO(leca): reinterpret_cast<OrtComputeContext*>(context)?
+            }
+            return 0;
+        };
+        compute_info.compute_func = [&, i](void* state, const OrtApi* api, OrtKernelContext* context) {
+            return ToStatus(node_compute_info_[i]->ComputeFunc(state, api, context));
+        };
+        compute_info.release_state_func = [&, i](void* state) {
+            if (node_compute_info_[i]->DestroyFunctionStateFunc) {
+                node_compute_info_[i]->DestroyFunctionStateFunc(state);
+            }
+        };
+        node_compute_funcs.push_back(compute_info);
+    }
+
+/*    node_compute_funcs.resize(count);
+    NodeComputeInfo*
+    ep_impl_->Compile(ep_impl_, ortGraphs.data(), ortNodes.data(), count, reinterpret_cast<>(&node_compute_funcs.data()));
+*/
+    return Status::OK();
+  }
+private:
+  OrtExecutionProvider* ep_impl_;
+  OrtNodeComputeInfo** node_compute_info_;
+};
+}

onnxruntime/core/framework/session_options.h

@@ -15,6 +15,7 @@
 #include "core/session/onnxruntime_c_api.h"
 #include "core/optimizer/graph_transformer_level.h"
 #include "core/util/thread_utils.h"
+#include "core/framework/provider_options.h"
 
 #if !defined(ORT_MINIMAL_BUILD) || defined(ORT_MINIMAL_BUILD_CUSTOM_OPS)
 #include "core/framework/library_handles.h"
@@ -184,6 +185,8 @@ struct SessionOptions {
   // User specified logging func and param
   OrtLoggingFunction user_logging_function = nullptr;
   void* user_logging_param = nullptr;
+
+  ProviderOptionsMap custom_ep_options;
 };
 
 inline std::ostream& operator<<(std::ostream& os, const SessionOptions& session_options) {

onnxruntime/core/session/environment.cc

@@ -348,4 +348,9 @@ Status Environment::CreateAndRegisterAllocatorV2(const std::string& provider_typ
   return Status{ONNXRUNTIME, common::INVALID_ARGUMENT, provider_type + " is not implemented in CreateAndRegisterAllocatorV2()"};
 }
 
+void Environment::InsertCustomEp(const char* ep_name, OrtExecutionProviderFactory* ep_factory) {
+  std::unique_ptr<OrtExecutionProviderFactory> p(ep_factory);
+  custom_ep_factories_.insert({ep_name, std::move(p)});  // TODO(leca): review
+}
+
 }  // namespace onnxruntime

onnxruntime/core/session/inference_session.cc

@@ -89,6 +89,7 @@
 #include "core/framework/stream_execution_context.h"
 #include "orttraining/core/optimizer/memory_optimizer/memory_optimizer.h"
 #endif
+#include "core/framework/provider_adapter.h"
 
 using namespace ONNX_NAMESPACE;
 using namespace onnxruntime::common;
@@ -1655,6 +1656,22 @@ common::Status InferenceSession::Initialize() {
     const Env& env = Env::Default();
     env.GetTelemetryProvider().LogSessionCreationStart();
 
+    const std::unordered_map<std::string, std::unique_ptr<OrtExecutionProviderFactory>>& custom_ep_factories = environment_.GetCustomEpFactories();
+    if (custom_ep_factories.size() > 0) {
+      for (auto const& [ep_name, ep_factory] : custom_ep_factories) {
+        if (session_options_.custom_ep_options.find(ep_name) != session_options_.custom_ep_options.end()) {
+          std::vector<const char*> keys, values;
+          for (auto const& [op_k, op_v] : session_options_.custom_ep_options[ep_name]) {
+            keys.push_back(op_k.c_str());
+            values.push_back(op_v.c_str());
+          }
+          OrtExecutionProvider* ep = reinterpret_cast<OrtExecutionProvider*>(ep_factory->CreateExecutionProvider(ep_factory.get(), keys.data(), values.data(), keys.size()));
+          std::unique_ptr<ExecutionProviderAdapter> ep_adapter = std::make_unique<ExecutionProviderAdapter>(ep);
+          ORT_RETURN_IF_ERROR(RegisterExecutionProvider(std::move(ep_adapter)));
+        }
+      }
+    }
+
     bool have_cpu_ep = false;
 
     {

onnxruntime/core/session/onnxruntime_c_api.cc

@@ -22,6 +22,7 @@
 #include "core/common/safeint.h"
 #include "core/graph/constants.h"
 #include "core/graph/graph.h"
+#include "core/graph/graph_viewer.h"
 #include "core/framework/allocator.h"
 #include "core/framework/tensor.h"
 #include "core/framework/ort_value.h"
@@ -2353,6 +2354,82 @@ ORT_API(const OrtTrainingApi*, OrtApis::GetTrainingApi, uint32_t version) {
 #endif
 }
 
+ORT_API_STATUS_IMPL(OrtApis::RegisterOrtExecutionProviderLibrary, _In_ const char* lib_path, _In_ OrtEnv* env, _In_ const char* ep_name) {
+  API_IMPL_BEGIN
+  void* handle = nullptr;
+  ORT_THROW_IF_ERROR(Env::Default().LoadDynamicLibrary(ToPathString(lib_path), false, &handle));
+  if (handle) {
+    OrtExecutionProviderFactory* (*symbol)();
+    ORT_THROW_IF_ERROR(Env::Default().GetSymbolFromLibrary(handle, "RegisterCustomEp", (void**)&symbol));
+    env->InsertCustomEp(ep_name, symbol());
+    return nullptr;
+  }
+  return CreateStatus(ORT_RUNTIME_EXCEPTION, "cannot load the shared library for out-tree EP");
+  API_IMPL_END
+}
+
+ORT_API_STATUS_IMPL(OrtApis::SessionOptionsAppendOrtExecutionProvider, _In_ OrtSessionOptions* options, _In_ const char* ep_name,
+                    _In_reads_(num_keys) const char* const* provider_options_keys, _In_reads_(num_keys) const char* const* provider_options_values, _In_ size_t num_keys) {
+  std::unordered_map<std::string, std::string> kv;
+  for (size_t i = 0; i < num_keys; i++) {
+    kv.insert({provider_options_keys[i], provider_options_values[i]});
+  }
+  options->value.custom_ep_options.insert({ep_name, kv});
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtGraph_IsConstantInitializer, const OrtGraphViewer* graph, const char* name, bool check_outer_scope, _Out_ bool* ret) {
+  const ::onnxruntime::GraphViewer* graph_viewer = reinterpret_cast<const ::onnxruntime::GraphViewer*>(graph);
+  *ret = graph_viewer->IsConstantInitializer(name, check_outer_scope);
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtGraph_GetNodesIndexInTopologicalOrder, const OrtGraphViewer* graph, _Out_ size_t* len, _Out_ const size_t** nodes_index_in_topological_order) {
+  const ::onnxruntime::GraphViewer* graph_viewer = reinterpret_cast<const ::onnxruntime::GraphViewer*>(graph);
+  const std::vector<size_t>& nodes = graph_viewer->GetNodesInTopologicalOrder();
+  *len = nodes.size();
+  *nodes_index_in_topological_order = nodes.data();
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtGraph_GetOrtNode, const OrtGraphViewer* graph, size_t node_index, _Outptr_ const OrtNode** node) {
+  const ::onnxruntime::GraphViewer* graph_viewer = reinterpret_cast<const ::onnxruntime::GraphViewer*>(graph);
+  *node = reinterpret_cast<const OrtNode*>(graph_viewer->GetNode(node_index));
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtNode_GetOpType, const OrtNode* node, _Out_ const char** op_type) {
+  const ::onnxruntime::Node* n = reinterpret_cast<const ::onnxruntime::Node*>(node);
+  *op_type = n->OpType().c_str();
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtNode_GetInputSize, const OrtNode* node, _Out_ size_t* input_size) {
+  const ::onnxruntime::Node* n = reinterpret_cast<const ::onnxruntime::Node*>(node);
+  *input_size = n->InputDefs().size();
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtNode_GetIthInputName, const OrtNode* node, size_t i, _Out_ const char** ith_input_name) {
+  const ::onnxruntime::Node* n = reinterpret_cast<const ::onnxruntime::Node*>(node);
+  assert(i < n->InputDefs().size());
+  *ith_input_name = n->InputDefs()[i]->Name().c_str();
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtNode_GetOutputSize, const OrtNode* node, _Out_ size_t* output_size) {
+  const ::onnxruntime::Node* n = reinterpret_cast<const ::onnxruntime::Node*>(node);
+  *output_size = n->OutputDefs().size();
+  return nullptr;
+}
+
+ORT_API_STATUS_IMPL(OrtApis::OrtNode_GetIthOutputName, const OrtNode* node, size_t i, _Out_ const char** ith_output_name) {
+  const ::onnxruntime::Node* n = reinterpret_cast<const ::onnxruntime::Node*>(node);
+  assert(i < n->OutputDefs().size());
+  *ith_output_name = n->OutputDefs()[i]->Name().c_str();
+  return nullptr;
+}
+
 static constexpr OrtApiBase ort_api_base = {
     &OrtApis::GetApi,
     &OrtApis::GetVersionString};
@@ -2730,6 +2807,18 @@ static constexpr OrtApi ort_api_1_to_19 = {
     &OrtApis::KernelInfoGetAllocator,
     &OrtApis::AddExternalInitializersFromFilesInMemory,
     // End of Version 18 - DO NOT MODIFY ABOVE (see above text for more information)
+
+    &OrtApis::RegisterOrtExecutionProviderLibrary,
+    &OrtApis::SessionOptionsAppendOrtExecutionProvider,
+
+    &OrtApis::OrtGraph_IsConstantInitializer,
+    &OrtApis::OrtGraph_GetNodesIndexInTopologicalOrder,
+    &OrtApis::OrtGraph_GetOrtNode,
+    &OrtApis::OrtNode_GetOpType,
+    &OrtApis::OrtNode_GetInputSize,
+    &OrtApis::OrtNode_GetIthInputName,
+    &OrtApis::OrtNode_GetOutputSize,
+    &OrtApis::OrtNode_GetIthOutputName,
 };
 
 // OrtApiBase can never change as there is no way to know what version of OrtApiBase is returned by OrtGetApiBase.