[nGraph] Utilize CommonOptimizations pass with custom transformations callback

Author: dmitrygo

inference-engine/src/cldnn_engine/cldnn_engine.cpp

@@ -83,7 +83,7 @@ InferenceEngine::ICNNNetwork::Ptr clDNNEngine::CloneNetwork(const InferenceEngin
         ::ngraph::op::GenericIE::DisableReshape noReshape(nGraphFunc);
 
         // Note: instead of running all Conversion Transformations you can make up your own transformation pipeline
-        ngraph::pass::CommonOptimizations().run_on_function(nGraphFunc);
+        ngraph::pass::CommonOptimizations(transformations_callback).run_on_function(nGraphFunc);
         ngraph::pass::ConvertOpSet3ToOpSet2(transformations_callback).run_on_function(nGraphFunc);
         ngraph::pass::ConvertOpSet2ToOpSet1(transformations_callback).run_on_function(nGraphFunc);
         ngraph::pass::ConvertOpSet1ToLegacy(transformations_callback).run_on_function(nGraphFunc);

inference-engine/src/mkldnn_plugin/mkldnn_plugin.cpp

@@ -102,7 +102,7 @@ Engine::LoadExeNetworkImpl(const InferenceEngine::ICNNNetwork &network, const st
         ::ngraph::op::GenericIE::DisableReshape noReshape(nGraphFunc);
 
         // Note: instead of running all Conversion Transformations you can make up your own transformation pipeline
-        ngraph::pass::CommonOptimizations().run_on_function(nGraphFunc);
+        ngraph::pass::CommonOptimizations(transformations_callback).run_on_function(nGraphFunc);
         ngraph::pass::ConvertOpSet3ToOpSet2(transformations_callback).run_on_function(nGraphFunc);
         ngraph::pass::ConvertOpSet2ToOpSet1(transformations_callback).run_on_function(nGraphFunc);
         ngraph::pass::ConvertOpSet1ToLegacy(transformations_callback).run_on_function(nGraphFunc);

inference-engine/src/transformations/include/transformations/common_optimizations/common_optimizations.hpp

@@ -21,9 +21,10 @@ class TRANSFORMATIONS_API CommonOptimizations;
 }  // namespace pass
 }  // namespace ngraph
 
-class ngraph::pass::CommonOptimizations: public ngraph::pass::FunctionPass {
+class ngraph::pass::CommonOptimizations: public ngraph::pass::FunctionPass, public ngraph::pass::PassParam {
 public:
-    explicit CommonOptimizations() : FunctionPass() {}
+    explicit CommonOptimizations(const PassParam::param_callback & callback = PassParam::getDefaultCallback())
+            : FunctionPass(), PassParam(callback) {}
 
     bool run_on_function(std::shared_ptr<ngraph::Function> f) override;
 };

inference-engine/src/transformations/src/transformations/common_optimizations/common_optimizations.cpp

@@ -20,11 +20,17 @@
 
 bool ngraph::pass::CommonOptimizations::run_on_function(std::shared_ptr<ngraph::Function> f) {
     ngraph::pass::Manager CommonOptimizations;
+    std::vector<std::shared_ptr<ngraph::pass::PassBase> > transforms;
 
-#define NGRAPH_PASS(NAME, NAMESPACE) CommonOptimizations.register_pass<NAMESPACE::NAME>();
+#define NGRAPH_PASS(NAME, NAMESPACE) transforms.push_back(CommonOptimizations.register_pass<NAMESPACE::NAME>());
 #include <transformations/common_optimizations/common_optimizations_tbl.hpp>
 #undef NGRAPH_PASS
 
+    for (auto & t : transforms) {
+        if (auto t_param = std::dynamic_pointer_cast<PassParam>(t)) {
+            t_param->setCallback(transformation_callback);
+        }
+    }
     CommonOptimizations.run_passes(f);
     return true;
 }

inference-engine/src/transformations/src/transformations/depth_to_space_fusion.cpp

@@ -91,10 +91,6 @@ void ngraph::pass::DepthToSpaceFusion::depth_to_space_fusion() {
     auto reshape_after = std::make_shared<ngraph::opset3::Reshape> (permute, input3, false);
 
     ngraph::graph_rewrite_callback callback = [this](pattern::Matcher& m) {
-        if (!transformation_callback(std::make_shared<ngraph::opset3::DepthToSpace>())) {
-            return false;
-        }
-
         auto reshape_after = std::dynamic_pointer_cast<ngraph::opset3::Reshape>(m.get_match_root());
         if (!reshape_after) {
             return false;
@@ -157,6 +153,11 @@ void ngraph::pass::DepthToSpaceFusion::depth_to_space_fusion() {
                 std::make_shared<ngraph::opset3::DepthToSpace>(reshape_before->input_value(0), mode, block_size);
         depth_to_space->set_friendly_name(reshape_after->get_friendly_name());
         ngraph::copy_runtime_info({reshape_before, permute, reshape_after}, depth_to_space);
+
+        if (!transformation_callback(depth_to_space)) {
+            return false;
+        }
+
         ngraph::replace_node(reshape_after, depth_to_space);
         return true;
     };