Sahar/contrib ops branch

cmake/onnxruntime_providers.cmake

@@ -15,8 +15,10 @@ function(substitute_op_reduction_srcs all_srcs)
   set(original_srcs
     "${ONNXRUNTIME_ROOT}/contrib_ops/cpu/cpu_contrib_kernels.cc"
     "${ONNXRUNTIME_ROOT}/contrib_ops/cuda/cuda_contrib_kernels.cc"
+    "${ONNXRUNTIME_ROOT}/contrib_ops/openvino/openvino_contrib_kernels.cc"
     "${ONNXRUNTIME_ROOT}/core/providers/cpu/cpu_execution_provider.cc"
     "${ONNXRUNTIME_ROOT}/core/providers/cuda/cuda_execution_provider.cc"
+    "${ONNXRUNTIME_ROOT}/core/providers/openvino/openvino_execution_provider.cc"
     "${ONNXRUNTIME_ROOT}/core/providers/op_kernel_type_control_overrides.inc"
     "${ORTTRAINING_SOURCE_DIR}/training_ops/cpu/cpu_training_kernels.cc"
     "${ORTTRAINING_SOURCE_DIR}/training_ops/cuda/cuda_training_kernels.cc"
@@ -93,6 +95,11 @@ file(GLOB_RECURSE onnxruntime_rocm_contrib_ops_cu_srcs CONFIGURE_DEPENDS
   "${ONNXRUNTIME_ROOT}/contrib_ops/rocm/*.cuh"
 )
 
+file(GLOB_RECURSE onnxruntime_openvino_contrib_ops_cc_srcs CONFIGURE_DEPENDS
+  "${ONNXRUNTIME_ROOT}/contrib_ops/openvino/*.h"
+  "${ONNXRUNTIME_ROOT}/contrib_ops/openvino/*.cc"
+)
+
 file(GLOB onnxruntime_providers_common_srcs CONFIGURE_DEPENDS
   "${ONNXRUNTIME_ROOT}/core/providers/*.h"
   "${ONNXRUNTIME_ROOT}/core/providers/*.cc"
@@ -824,6 +831,7 @@ if (onnxruntime_USE_OPENVINO)
     unset(CMAKE_MAP_IMPORTED_CONFIG_RELWITHDEBINFO)
   endif()
 
+
   if ((DEFINED ENV{OPENCL_LIBS}) AND (DEFINED ENV{OPENCL_INCS}))
     add_definitions(-DIO_BUFFER_ENABLED=1)
     list(APPEND OPENVINO_LIB_LIST $ENV{OPENCL_LIBS} ${OV_20_LIBS} ${InferenceEngine_LIBRARIES} ${NGRAPH_LIBRARIES} ngraph::onnx_importer ${PYTHON_LIBRARIES})
@@ -832,6 +840,9 @@ if (onnxruntime_USE_OPENVINO)
   endif()
 
   source_group(TREE ${ONNXRUNTIME_ROOT}/core FILES ${onnxruntime_providers_openvino_cc_srcs})
+  source_group(TREE ${ONNXRUNTIME_ROOT} FILES ${onnxruntime_openvino_contrib_ops_cc_srcs})
+  list(APPEND onnxruntime_providers_openvino_cc_srcs ${onnxruntime_openvino_contrib_ops_cc_srcs})
+
   onnxruntime_add_shared_library_module(onnxruntime_providers_openvino ${onnxruntime_providers_openvino_cc_srcs} "${ONNXRUNTIME_ROOT}/core/dll/onnxruntime.rc")
   onnxruntime_add_include_to_target(onnxruntime_providers_openvino onnxruntime_common onnx)
   install(DIRECTORY ${PROJECT_SOURCE_DIR}/../include/onnxruntime/core/providers/openvino  DESTINATION ${CMAKE_INSTALL_INCLUDEDIR}/onnxruntime/core/providers)

include/onnxruntime/core/framework/op_kernel.h

@@ -199,6 +199,14 @@ KernelCreateInfo BuildKernelCreateInfo();
 }  // namespace snpe
 }  // namespace contrib
 
+namespace contrib {
+namespace openvino_ep {
+template <typename T>
+KernelCreateInfo BuildKernelCreateInfo();
+}  // namespace openvino_ep
+}  // namespace contrib
+
+
 using BuildKernelCreateInfoFn = KernelCreateInfo (*)();
 
 // Naming convention for operator kernel classes

onnxruntime/contrib_ops/cpu/transformers/subgraph_base.cc

@@ -128,8 +128,9 @@ const IExecutionProvider* Subgraph::GetProvider() const {
   const IExecutionProvider* cpu_provider = providers.Get(onnxruntime::kCpuExecutionProvider);
   const IExecutionProvider* cuda_provider = providers.Get(onnxruntime::kCudaExecutionProvider);
   const IExecutionProvider* rocm_provider = providers.Get(onnxruntime::kRocmExecutionProvider);
+  const IExecutionProvider* openvino_provider = providers.Get(onnxruntime::kOpenVINOExecutionProvider);
   const IExecutionProvider* gpu_provider = cuda_provider ? cuda_provider : rocm_provider;
-  const IExecutionProvider* provider = gpu_provider ? gpu_provider : cpu_provider;
+  const IExecutionProvider* provider = openvino_provider ? openvino_provider : cpu_provider;
   return provider;
 }
 

onnxruntime/contrib_ops/cuda/bert/add_bias_transpose.cu

@@ -2,7 +2,8 @@
 // Licensed under the MIT License.
 #include "core/providers/cuda/cuda_common.h"
 #include "core/providers/cuda/cu_inc/common.cuh"
-#include "contrib_ops/cuda/bert/add_bias_transpose.h"
+#include "
+/cuda/bert/add_bias_transpose.h"
 #include "contrib_ops/cuda/bert/rotary_embedding_util.h"
 
 using namespace onnxruntime::cuda;

onnxruntime/contrib_ops/cuda/transformers/generation_device_helper.h

@@ -44,7 +44,7 @@ Status AddToFeeds(const IExecutionProvider* execution_provider,
                   std::initializer_list<OrtValue> inputs,
                   std::vector<OrtValue>& feeds,
                   IAllocatorUniquePtr<char>& buffer);
-
+.
 template <typename T>
 void InitBeamState(transformers::IBeamSearchState<T>* beam_state,
                    gsl::span<int32_t>& sequence_lengths,

onnxruntime/contrib_ops/openvino/beam_search.cc

@@ -0,0 +1,39 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/openvino/beam_search.h"
+
+namespace onnxruntime {
+namespace contrib {
+namespace openvino_ep {
+
+#define REGISTER_KERNEL_TYPED(T)                                  \
+  ONNX_OPERATOR_TYPED_KERNEL_EX(                                  \
+      BeamSearch,                                                 \
+      kMSDomain,                                                  \
+      1,                                                          \
+      T,                                                          \
+      kOpenVINOExecutionProvider,                                 \
+      (*KernelDefBuilder::Create())                               \
+          .TypeConstraint("T", DataTypeImpl::GetTensorType<T>()), \
+      BeamSearch);
+
+REGISTER_KERNEL_TYPED(float)
+
+BeamSearch::BeamSearch(const OpKernelInfo& info)
+    : onnxruntime::contrib::transformers::BeamSearch(info) {
+
+}
+
+Status BeamSearch::ComputeInternal(OpKernelContext* context) const {
+  return onnxruntime::contrib::transformers::BeamSearch::Compute(context);
+}
+
+Status BeamSearch::Compute(OpKernelContext* context) const {
+  auto s = ComputeInternal(context);
+  return s;
+}
+
+}  // namespace cuda
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/openvino/beam_search.h

@@ -0,0 +1,26 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+
+#include "contrib_ops/cpu/transformers/beam_search.h"
+
+namespace onnxruntime {
+class SessionState;
+
+namespace contrib {
+namespace openvino_ep {
+
+class BeamSearch final : public onnxruntime::contrib::transformers::BeamSearch {
+ public:
+  BeamSearch(const OpKernelInfo& info);
+
+  Status Compute(OpKernelContext* context) const override;
+
+ private:
+  Status ComputeInternal(OpKernelContext* context) const;
+};
+
+}  // namespace cuda.
+}  // namespace contrib
+}  // namespace onnxruntime

onnxruntime/contrib_ops/openvino/openvino_contrib_kernels.cc

@@ -0,0 +1,42 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#include "contrib_ops/openvino/openvino_contrib_kernels.h"
+#include "core/providers/shared_library/provider_api.h"
+#include "core/common/status.h"
+
+using namespace onnxruntime::common;
+
+namespace onnxruntime {
+namespace contrib {
+namespace openvino_ep {
+
+class ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kOpenVINOExecutionProvider, kMSDomain, 1, float, BeamSearch);
+
+template <>
+KernelCreateInfo BuildKernelCreateInfo<void>() {
+  KernelCreateInfo info;
+  return info;
+}
+
+Status RegisterOpenVINOContribKernels(KernelRegistry& kernel_registry) {
+  static const BuildKernelCreateInfoFn function_table[] = {
+    BuildKernelCreateInfo<void>,  // default entry to avoid the list become empty after ops-reducing
+    BuildKernelCreateInfo<ONNX_OPERATOR_TYPED_KERNEL_CLASS_NAME(kOpenVINOExecutionProvider, kMSDomain, 1, float, BeamSearch)>
+  };
+
+  for (auto& function_table_entry : function_table) {
+    KernelCreateInfo info = function_table_entry();
+    if (info.kernel_def != nullptr) {  // filter disabled entries where type is void
+      ORT_RETURN_IF_ERROR(kernel_registry.Register(std::move(info)));
+      //return kernel_registry.Register(std::move(info));
+    }
+  }
+
+  return Status::OK();
+}
+
+}  // namespace openvino 
+}  // namespace contrib
+}  // namespace onnxruntime
+

onnxruntime/contrib_ops/openvino/openvino_contrib_kernels.h

@@ -0,0 +1,14 @@
+// Copyright (c) Microsoft Corporation. All rights reserved.
+// Licensed under the MIT License.
+
+#pragma once
+#include "core/common/status.h"
+
+namespace onnxruntime {
+   class KernelRegistry;
+namespace contrib {
+namespace openvino_ep {
+Status RegisterOpenVINOContribKernels(KernelRegistry& kernel_registry);
+} // namespace OpenVINO
+} // namespace contrib
+}  // namespace onnxruntime

onnxruntime/core/providers/openvino/openvino_execution_provider.cc

@@ -2,13 +2,18 @@
 // Licensed under the MIT License
 
 #include "core/providers/shared_library/provider_api.h"
+#include "backend_manager.h"
+#define ORT_API_MANUAL_INIT
 #include "openvino_execution_provider.h"
 #include "contexts.h"
 #include "backend_manager.h"
 #include "ov_versions/capabilities.h"
+#include "contrib_ops/openvino/openvino_contrib_kernels.h"
 
 #define MEMCPY_S(dest, src, destsz, srcsz) memcpy(dest, src, std::min(destsz, srcsz))
 
+using namespace onnxruntime::common;
+
 namespace onnxruntime {
 
 OpenVINOExecutionProvider::OpenVINOExecutionProvider(const OpenVINOExecutionProviderInfo& info)
@@ -188,4 +193,32 @@ common::Status OpenVINOExecutionProvider::Compile(
   return Status::OK();
 }
 
+namespace openvino_ep {
+static Status RegisterOpenVINOKernels(KernelRegistry& kernel_registry) {
+  ORT_RETURN_IF_ERROR(::onnxruntime::contrib::openvino_ep::RegisterOpenVINOContribKernels(kernel_registry));
+  return Status::OK();
+}
+}  // namespace openvino_ep
+
+static std::shared_ptr<KernelRegistry> o_kernel_registry;
+
+void InitializeRegistry() {
+  o_kernel_registry = KernelRegistry::Create();
+  ORT_THROW_IF_ERROR(openvino_ep::RegisterOpenVINOKernels(*o_kernel_registry));
+}
+
+void DeleteRegistry() {
+  o_kernel_registry.reset();
+}
+
+std::shared_ptr<KernelRegistry> OpenVINOExecutionProvider::GetKernelRegistry() const {
+  return o_kernel_registry;
+}
+
+
 }  // namespace onnxruntime
+
+
+
+
+

onnxruntime/core/providers/openvino/openvino_execution_provider.h

@@ -3,10 +3,15 @@
 
 #pragma once
 
-#include "backend_manager.h"
 #include <map>
 #include <algorithm>
 #include <iostream>
+#include <vector>
+#include <sstream>
+#include "core/providers/shared_library/provider_api.h"
+#include "contexts.h"
+#include "backend_manager.h"
+#include "ov_versions/capabilities.h"
 
 namespace onnxruntime {
 
@@ -176,6 +181,8 @@ class OpenVINOExecutionProvider : public IExecutionProvider {
   const void* GetExecutionHandle() const noexcept override {
     return nullptr;
   }
+
+  std::shared_ptr<KernelRegistry> GetKernelRegistry() const override;
 };
 
 }  // namespace onnxruntime

onnxruntime/core/providers/openvino/openvino_provider_factory.cc

@@ -7,6 +7,10 @@
 #include "openvino_provider_factory_creator.h"
 
 namespace onnxruntime {
+
+void InitializeRegistry();
+void DeleteRegistry();
+
 struct OpenVINOProviderFactory : IExecutionProviderFactory {
   OpenVINOProviderFactory(const char* device_type, bool enable_vpu_fast_compile,
                           const char* device_id, size_t num_of_threads,
@@ -72,10 +76,12 @@ struct OpenVINO_Provider : Provider {
   }
 
   void Initialize() override {
+    InitializeRegistry();
   }
 
   void Shutdown() override {
     openvino_ep::BackendManager::ReleaseGlobalContext();
+    DeleteRegistry();
   }
 
 } g_provider;

onnxruntime/core/providers/openvino/ov_versions/data_ops.cc

@@ -57,13 +57,14 @@ std::set<std::string> ops_supported_only_in_model = {
     "Slice",
     "Split",
     "Tile",
-    "TopK"};
+    "TopK",
+    "BeamSearch"};
 
 // Ops which are supported as functions (as composite ops)
 std::set<std::string> ops_supported_as_function = {
     "LessOrEqual",
     "GreaterOrEqual",
-};
+    "BeamSearch"};
 
 std::vector<SupportedOp> supported_op_mode = {
     {"Abs", V_2020_4, {"CPU", "GPU"}},