[WIP]EP interface

cmake/CMakeLists.txt

@@ -142,6 +142,7 @@ option(onnxruntime_TVM_CUDA_RUNTIME "Build TVM with CUDA support" OFF)
 option(onnxruntime_TVM_USE_LLVM "Build TVM with LLVM. Set customized path to llvm-config.exe here if need" OFF)
 option(onnxruntime_TVM_USE_HASH "Build ipp-crypto library for support hash algorithm. It is defined for TVM only")
 option(onnxruntime_USE_XNNPACK "Build with XNNPACK support. Provides an alternative math library on ARM, WebAssembly and x86." OFF)
+option(onnxruntime_USE_INTREE "InTree EP support" OFF)
 option(onnxruntime_USE_WEBNN "Build with WebNN support. Enable hardware acceleration in web browsers." OFF)
 
 # Options related to reducing the binary size produced by the build
@@ -800,6 +801,11 @@ if (onnxruntime_USE_XNNPACK)
   list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_XNNPACK=1)
   list(APPEND ONNXRUNTIME_PROVIDER_NAMES xnnpack)
 endif()
+if (onnxruntime_USE_INTREE)
+  list(APPEND ORT_PROVIDER_FLAGS -DUSE_INTREE=1)
+  list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_INTREE=1)
+  list(APPEND ONNXRUNTIME_PROVIDER_NAMES intree)
+endif()
 if (onnxruntime_USE_WEBNN)
   list(APPEND ORT_PROVIDER_FLAGS -DUSE_WEBNN=1)
   list(APPEND ORT_PROVIDER_CMAKE_FLAGS -Donnxruntime_USE_WEBNN=1)

cmake/onnxruntime_providers.cmake

@@ -108,6 +108,9 @@ endif()
 if (onnxruntime_USE_XNNPACK)
   set(PROVIDERS_XNNPACK onnxruntime_providers_xnnpack)
 endif()
+if (onnxruntime_USE_INTREE)
+  set(PROVIDERS_INTREE onnxruntime_providers_intree)
+endif()
 if(onnxruntime_USE_WEBNN)
   set(PROVIDERS_WEBNN onnxruntime_providers_webnn)
 endif()
@@ -196,6 +199,10 @@ if (onnxruntime_USE_XNNPACK)
   include(onnxruntime_providers_xnnpack.cmake)
 endif()
 
+if (onnxruntime_USE_INTREE)
+  include(onnxruntime_providers_intree.cmake)
+endif()
+
 if (onnxruntime_USE_CANN)
   include(onnxruntime_providers_cann.cmake)
 endif()

cmake/onnxruntime_providers_intree.cmake

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+  add_compile_definitions(USE_INTREE=1)
+
+  file(GLOB_RECURSE onnxruntime_providers_intree_cc_srcs CONFIGURE_DEPENDS
+    "${ONNXRUNTIME_INCLUDE_DIR}/core/providers/intree/*.h"
+    "${ONNXRUNTIME_ROOT}/core/providers/intree/*.h"
+    "${ONNXRUNTIME_ROOT}/core/providers/intree/*.cc"
+  )
+
+  source_group(TREE ${REPO_ROOT} FILES ${onnxruntime_providers_intree_cc_srcs})
+  onnxruntime_add_static_library(onnxruntime_providers_intree ${onnxruntime_providers_intree_cc_srcs})
+  onnxruntime_add_include_to_target(onnxruntime_providers_intree
+    onnxruntime_common onnxruntime_framework onnx pthreadpool Boost::mp11 safeint_interface
+  )
+
+  add_dependencies(onnxruntime_providers_intree onnx ${onnxruntime_EXTERNAL_DEPENDENCIES})
+  set_target_properties(onnxruntime_providers_intree PROPERTIES FOLDER "ONNXRuntime")
+
+  set_target_properties(onnxruntime_providers_intree PROPERTIES LINKER_LANGUAGE CXX)
+  #target_include_directories(onnxruntime_providers_intree PUBLIC "/bert_ort/leca/code/onnxruntime2/include/onnxruntime")
+
+  if (NOT onnxruntime_BUILD_SHARED_LIB)
+    install(TARGETS onnxruntime_providers_intree
+            ARCHIVE   DESTINATION ${CMAKE_INSTALL_LIBDIR}
+            LIBRARY   DESTINATION ${CMAKE_INSTALL_LIBDIR}
+            RUNTIME   DESTINATION ${CMAKE_INSTALL_BINDIR}
+            FRAMEWORK DESTINATION ${CMAKE_INSTALL_BINDIR})
+  endif()

cmake/onnxruntime_python.cmake

@@ -71,9 +71,9 @@ onnxruntime_add_shared_library_module(onnxruntime_pybind11_state ${onnxruntime_p
 
 if(MSVC)
   target_compile_options(onnxruntime_pybind11_state PRIVATE "$<$<COMPILE_LANGUAGE:CUDA>:SHELL:--compiler-options /utf-8>" "$<$<NOT:$<COMPILE_LANGUAGE:CUDA>>:/utf-8>")
-  if(onnxruntime_ENABLE_TRAINING)
-    target_compile_options(onnxruntime_pybind11_state PRIVATE "/bigobj")
-  endif()
+  #if(onnxruntime_ENABLE_TRAINING)
+  target_compile_options(onnxruntime_pybind11_state PRIVATE "/bigobj")
+  #endif()
 endif()
 if(HAS_CAST_FUNCTION_TYPE)
   target_compile_options(onnxruntime_pybind11_state PRIVATE "-Wno-cast-function-type")
@@ -179,6 +179,7 @@ target_link_libraries(onnxruntime_pybind11_state PRIVATE
     ${PROVIDERS_ACL}
     ${PROVIDERS_ARMNN}
     ${PROVIDERS_XNNPACK}
+    ${PROVIDERS_INTREE}
     ${PROVIDERS_AZURE}
     ${PROVIDERS_QNN}
     onnxruntime_optimizer

include/onnxruntime/core/framework/allocator.h

@@ -4,11 +4,12 @@
 #pragma once
 
 #include "core/common/common.h"
-#include "core/framework/allocator_stats.h"
 #include "core/session/onnxruntime_c_api.h"
 #include "ortdevice.h"
 #include "ortmemoryinfo.h"
 #include <map>
+#include <string>
+#include <sstream>
 
 // This configures the arena based allocator used by ORT
 // See docs/C_API.md for details on what these mean and how to choose these values
@@ -38,6 +39,49 @@
 };
 
 namespace onnxruntime {
+// Runtime statistics collected by an allocator.
+struct AllocatorStats {
+  int64_t num_allocs;             // Number of allocations.
+  int64_t num_reserves;           // Number of reserves. (Number of calls to Reserve() in arena-based allocators)
+  int64_t num_arena_extensions;   // Number of arena extensions (Relevant only for arena based allocators)
+  int64_t num_arena_shrinkages;   // Number of arena shrinkages (Relevant only for arena based allocators)
+  int64_t bytes_in_use;           // Number of bytes in use.
+  int64_t total_allocated_bytes;  // The total number of allocated bytes by the allocator.
+  int64_t max_bytes_in_use;       // The maximum bytes in use.
+  int64_t max_alloc_size;         // The max single allocation seen.
+                                  // The upper limit what the allocator can allocate, if such a limit
+                                  // is known. Certain allocator may return 0 to indicate the limit is
+                                  // unknown.
+  int64_t bytes_limit;
+
+  AllocatorStats() { Clear(); }
+
+  void Clear() {
+    this->num_allocs = 0;
+    this->num_reserves = 0;
+    this->num_arena_extensions = 0;
+    this->num_arena_shrinkages = 0;
+    this->bytes_in_use = 0;
+    this->max_bytes_in_use = 0;
+    this->max_alloc_size = 0;
+    this->bytes_limit = 0;
+    this->total_allocated_bytes = 0;
+  }
+
+  std::string DebugString() const {
+    std::ostringstream ss;
+    ss << "Limit:                    " << this->bytes_limit << "\n"
+       << "InUse:                    " << this->bytes_in_use << "\n"
+       << "TotalAllocated:           " << this->total_allocated_bytes << "\n"
+       << "MaxInUse:                 " << this->max_bytes_in_use << "\n"
+       << "NumAllocs:                " << this->num_allocs << "\n"
+       << "NumReserves:              " << this->num_reserves << "\n"
+       << "NumArenaExtensions:       " << this->num_arena_extensions << "\n"
+       << "NumArenaShrinkages:       " << this->num_arena_shrinkages << "\n"
+       << "MaxAllocSize:             " << this->max_alloc_size << "\n";
+    return ss.str();
+  }
+};
 constexpr const char* CPU = "Cpu";
 constexpr const char* CUDA = "Cuda";
 constexpr const char* CUDA_PINNED = "CudaPinned";

include/onnxruntime/core/framework/execution_provider.h

@@ -27,30 +27,14 @@ class Node;
 #include "core/common/basic_types.h"
 #include "core/common/profiler_common.h"
 #include "core/framework/allocator_utils.h"
-#include "core/framework/func_api.h"
+#include "core/framework/node_compute_info.h"
 #include "core/framework/provider_options.h"
 #include "core/framework/framework_provider_common.h"
 #include "core/framework/stream_handles.h"
 #include "core/framework/tuning_context.h"
 
 namespace onnxruntime {
 
-/**
-   Logical device representation.
-*/
-
-// if we are export the fused function to dll, the function will still in the same binary as onnxruntime
-// use std function to give execution provider some chance to capture some state.
-using CreateFunctionStateFunc = std::function<int(ComputeContext*, FunctionState*)>;
-using ComputeFunc = std::function<Status(FunctionState, const OrtApi*, OrtKernelContext*)>;
-using DestroyFunctionStateFunc = std::function<void(FunctionState)>;
-
-struct NodeComputeInfo {
-  CreateFunctionStateFunc create_state_func;
-  ComputeFunc compute_func;
-  DestroyFunctionStateFunc release_state_func;
-};
-
 enum class DataLayout {
   NCHW,
   NHWC,

include/onnxruntime/core/framework/kernel_def_builder.h

@@ -56,6 +56,10 @@ class KernelDef {
     return type_constraints_;
   }
 
+  std::unordered_map<std::string, std::vector<MLDataType>>& TypeConstraints() {
+    return type_constraints_;
+  }
+
   const std::vector<std::pair<int, int>>& MayInplace() const {
     return inplace_map_;
   }

include/onnxruntime/core/framework/kernel_registry.h

@@ -78,6 +78,10 @@ class KernelRegistry {
     return kernel_creator_fn_map_;
   }
 
+  KernelCreateMap& GetKernelCreateMap() {
+    return kernel_creator_fn_map_;
+  }
+
  private:
   // TryFindKernel implementation. Either kernel_type_str_resolver or type_constraints is provided.
   Status TryFindKernelImpl(const Node& node, ProviderType exec_provider,

include/onnxruntime/core/framework/node_compute_info.h

@@ -0,0 +1,17 @@
+#pragma once
+
+#include "core/framework/func_api.h"
+
+namespace onnxruntime {
+// if we are export the fused function to dll, the function will still in the same binary as onnxruntime
+// use std function to give execution provider some chance to capture some state.
+using CreateFunctionStateFunc = std::function<int(ComputeContext*, FunctionState*)>;
+using ComputeFunc = std::function<Status(FunctionState, const OrtApi*, OrtKernelContext*)>;
+using DestroyFunctionStateFunc = std::function<void(FunctionState)>;
+
+struct NodeComputeInfo {
+  CreateFunctionStateFunc create_state_func;
+  ComputeFunc compute_func;
+  DestroyFunctionStateFunc release_state_func;
+};
+}

include/onnxruntime/core/framework/op_kernel.h

@@ -26,11 +26,15 @@
 #include "core/graph/graph_viewer.h"
 #include "core/graph/onnx_protobuf.h"
 #include "core/common/gsl.h"
+
 namespace onnxruntime {
 class OpKernelContext;
 }
 #endif
 
+#include "interface/framework/kernel.h"
+
+
 namespace onnxruntime {
 
 std::unique_ptr<OpKernelInfo> CopyOpKernelInfo(const OpKernelInfo& info);
@@ -39,7 +43,7 @@ class OpKernel {
  public:
   using DoneCallback = std::function<void()>;
 
-  explicit OpKernel(const OpKernelInfo& info) : op_kernel_info_(CopyOpKernelInfo(info)) {}
+  explicit OpKernel(const OpKernelInfo& info): op_kernel_info_(CopyOpKernelInfo(info)) {}
   virtual ~OpKernel() = default;
 
   const onnxruntime::Node& Node() const;

include/onnxruntime/core/framework/op_kernel_context.h

@@ -1,14 +1,16 @@
 // Copyright (c) Microsoft Corporation. All rights reserved.
 // Licensed under the MIT License.
 
+#include "interface/framework/kernel.h"
+
 namespace onnxruntime {
 class IExecutionFrame;
 class Stream;
 namespace concurrency {
 class ThreadPool;
 }
 
-class OpKernelContext {
+class OpKernelContext : public interface::IKernelContext {
  public:
   using ArgMap = std::unordered_map<std::string, size_t>;
 
@@ -43,7 +45,15 @@
     }
   }
 
+  const void* InputData(int index) const override {
+    //todo - check tensor type
+    const auto* tensor = Input<onnxruntime::Tensor>(index);
+    return tensor->DataRaw();
+  }
+
   // Fetch a required input, enforcing that it is present.
+  // Fetch a required input, enforcing that it is present. Fetch a required input, enforcing that it is present.
+  // Fetch a required input, enforc Fetch a required input, enforcing that it is present.
   template <typename T>
   const T& RequiredInput(int index) const {
     const T* input_ptr = Input<T>(index);
@@ -68,6 +78,20 @@
   Tensor* Output(int index, const std::vector<int64_t>& shape);
   Tensor* Output(int index, const std::initializer_list<int64_t>& shape);
 
+  void* AllocateOutput(int index, const interface::TensorShape& shape) override {
+    auto* tensor = Output(index, shape);
+    ORT_ENFORCE(tensor);
+    return tensor->MutableDataRaw();
+  }
+
+  const int64_t* InputShape(int index, size_t* num_dims) const override {
+    const auto* tensor = Input<onnxruntime::Tensor>(index);
+    const auto& shape = tensor->Shape();
+    auto dims = shape.GetDims();
+    *num_dims = dims.size();
+    return dims.data();
+  };
+
   // Fetch a required tensor output, enforcing that it is present.
   Tensor& RequiredOutput(int index, const TensorShape& shape) {
     Tensor* output_ptr = Output(index, shape);

include/onnxruntime/core/framework/op_kernel_info.h

@@ -9,6 +9,7 @@
 #include "core/framework/op_node_proto_helper.h"
 #include "core/graph/graph_viewer.h"
 #include "core/common/gsl.h"
+#include "interface/framework/kernel.h"
 
 namespace onnxruntime {
 
@@ -20,7 +21,7 @@ struct AllocPlanPerValue;
 // A very light-weight class, which works as an aggregated
 // view of all data needed for constructing a Kernel instance.
 // NOTE: it does not own/hold any objects.
-class OpKernelInfo : public OpNodeProtoHelper<ProtoHelperNodeContext> {
+class OpKernelInfo : public OpNodeProtoHelper<ProtoHelperNodeContext>, public interface::IKernelInfo {
  public:
   explicit OpKernelInfo(const onnxruntime::Node& node,
                         const KernelDef& kernel_def,

include/onnxruntime/core/graph/constants.h

@@ -51,6 +51,7 @@ constexpr const char* kXnnpackExecutionProvider = "XnnpackExecutionProvider";
 constexpr const char* kWebNNExecutionProvider = "WebNNExecutionProvider";
 constexpr const char* kCannExecutionProvider = "CANNExecutionProvider";
 constexpr const char* kAzureExecutionProvider = "AzureExecutionProvider";
+constexpr const char* kInTreeExecutionProvider = "InTreeExecutionProvider";
 
 constexpr const char* kExecutionProviderSharedLibraryPath = "shared_lib_path";
 constexpr const char* kExecutionProviderSharedLibraryEntry = "provider_factory_entry_point";

include/onnxruntime/core/session/environment.h

@@ -10,9 +10,29 @@
 #include "core/platform/threadpool.h"
 #include "core/common/logging/logging.h"
 #include "core/framework/allocator.h"
+#include "interface/provider/provider.h"
 
 struct OrtThreadingOptions;
 namespace onnxruntime {
+struct ProviderLibrary2 {
+  ProviderLibrary2(const ORTCHAR_T* library_path) : library_path_{library_path} {}
+  ~ProviderLibrary2() {
+    // assert(!handle_); // We should already be unloaded at this point (disabled until Python shuts down deterministically)
+  }
+
+  void Load();
+
+  void Unload();
+
+  interface::ExecutionProvider* CreateExternalEPInstance(const std::unordered_map<std::string, std::string>& provider_options);
+
+ private:
+  const ORTCHAR_T* library_path_;
+  void* handle_{};
+
+  ORT_DISALLOW_COPY_AND_ASSIGNMENT(ProviderLibrary2);
+};
+
 /** TODO: remove this class
    Provides the runtime environment for onnxruntime.
    Create one instance for the duration of execution.
@@ -88,6 +108,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);
 
+  Status LoadExternalExecutionProvider(const std::string& provider_type, const std::string& library_path);
+
+  interface::ExecutionProvider* CreateExternalEPInstance(const std::string& provider_type, const std::unordered_map<std::string, std::string>& provider_options);
+
  private:
   ORT_DISALLOW_COPY_ASSIGNMENT_AND_MOVE(Environment);
   Status Initialize(std::unique_ptr<logging::LoggingManager> logging_manager,
@@ -99,5 +123,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<ProviderLibrary2>> external_ep_libs_;
 };
 }  // namespace onnxruntime