New extension classes

openvinotoolkit · Jan 30, 2025 · 64b536e · 64b536e
1 parent 864d45a
commit 64b536e
Show file tree

Hide file tree

Showing 3 changed files with 70 additions and 11 deletions.
diff --git a/src/core/include/openvino/core/extension.hpp b/src/core/include/openvino/core/extension.hpp
@@ -29,16 +29,56 @@ class OPENVINO_API Extension {
     virtual ~Extension();
 };
 
-class OPENVINO_API FrontendExtension : public Extension {
-public:
-    virtual std::string get_frontend_name() = 0;
-};
-}  // namespace ov
-
 #ifndef OV_CREATE_EXTENSION
 #    define OV_CREATE_EXTENSION create_extensions
 #endif
 
+template <class T, class... FEs>
+class OPENVINO_API FeExtension : public T, public FEs... {
+public:
+    using T::T;
+};
+
+// All FEs should be added here as IsFeName structures
+struct OPENVINO_API IsIr {
+    virtual ~IsIr();
+};
+
+struct OPENVINO_API IsOnnx {
+    virtual ~IsOnnx();
+};
+
+struct OPENVINO_API IsTf {
+    virtual ~IsTf();
+};
+
+// These classes can be defined in customer code
+template <class T>
+class OPENVINO_API IrFeExtension : public FeExtension<T, IsIr> {
+public:
+    using FeExtension<T, IsIr>::FeExtension;
+};
+
+template <class T>
+class OPENVINO_API OnnxFeExtension : public FeExtension<T, IsOnnx> {
+public:
+    using FeExtension<T, IsOnnx>::FeExtension;
+};
+
+template <class T>
+class OPENVINO_API TfFeExtension : public FeExtension<T, IsTf> {
+public:
+    using FeExtension<T, IsTf>::FeExtension;
+};
+
+// As example. It is possible to create extensions for several FEs
+template <class T>
+class OPENVINO_API IrAndOnnxFeExtension : public FeExtension<T, IsIr, IsOnnx> {
+public:
+    using FeExtension<T, IsIr, IsOnnx>::FeExtension;
+};
+
+}  // namespace ov
 /**
  * @brief The entry point for library with OpenVINO extensions
  *

diff --git a/src/core/src/extension.cpp b/src/core/src/extension.cpp
@@ -15,3 +15,6 @@ using namespace ov;
 
 ov::Extension::~Extension() = default;
 ov::BaseOpExtension::~BaseOpExtension() = default;
+ov::IsIr::~IsIr() = default;
+ov::IsOnnx::~IsOnnx() = default;
+ov::IsTf::~IsTf() = default;
diff --git a/src/inference/src/model_reader.cpp b/src/inference/src/model_reader.cpp
@@ -5,6 +5,8 @@
 #include "model_reader.hpp"
 
 #include "itt.hpp"
+#include "openvino/core/extension.hpp"
+#include "openvino/core/so_extension.hpp"
 #include "openvino/core/model.hpp"
 #include "openvino/core/preprocess/pre_post_process.hpp"
 #include "openvino/frontend/manager.hpp"
@@ -111,13 +113,27 @@ namespace {
     std::vector<ov::Extension::Ptr> filter_extensions_for_fronend(const ov::frontend::FrontEnd::Ptr& fe, const std::vector<ov::Extension::Ptr>& ov_exts) {
         auto fe_name = fe->get_name();
         std::vector<ov::Extension::Ptr> result;
-        for (auto& ext : ov_exts) {
-            if (auto fe_ext = ov::as_type_ptr<FrontendExtension>(ext)){
-                if (!fe_name.compare(fe_ext->get_frontend_name())){
+        for (const auto& ext : ov_exts) {
+            ov::Extension::Ptr ext_to_check;
+
+            if (auto so_ext = ov::as_type_ptr<ov::detail::SOExtension>(ext)) {
+                ext_to_check = so_ext->extension();
+            } else {
+                ext_to_check = ext;
+            }
+
+            if (auto fe_ext = ov::as_type_ptr<ov::IsIr>(ext_to_check)){
+                if (!fe_name.compare("ir"))
+                    result.push_back(ext);
+            } else if (auto fe_ext = ov::as_type_ptr<ov::IsOnnx>(ext_to_check)){
+                if (!fe_name.compare("onnx"))
+                    result.push_back(ext);
+            } else if (auto fe_ext = ov::as_type_ptr<ov::IsTf>(ext_to_check)){
+                if (!fe_name.compare("tf"))
                     result.push_back(ext);
-                }
+            } else {
+                result.push_back(ext);
             }
-            result.push_back(ext);
         }
         return result;
     }