Add provenance tags to opset transformation passes (#4198)

* Add provenance tag in opset0 downgrade pass

* Add provenance tag in opset1 upgrade pass

* Type name as string

* Change op_cast to return replacement_node instead of bool

* Add provenance tags to all nodes created while downgrade

* Add provenance tags to all nodes created while upgrade

* Comments

* Style apply

* Update ONNX import provenance test function

* Add const statement

* Add upgrade/downgrade provenance tag tests

* Update tests

* Style apply

* Provenance enabled check

* Removed redundant add_tag

* Test for add_provenance_tags above

* Add graph test for provenance tags in transformation pass

* Use EXPECT_TRUE and EXPECT_FALSE instead of EXPECT_EQ

* Return replacement node directly

* Style apply

* Test downgrade provenance tag with ONNX importer

* Update test/onnx/onnx_import_provenance.in.cpp
Co-Authored-By: Michał Karzyński <postrational@users.noreply.github.com>

* Test update

* Update provenance test to check node type occurence

* Style apply
Co-authored-by: Michał Karzyński <postrational@users.noreply.github.com>
Co-authored-by: Sang Ik Lee <sang.ik.lee@intel.com>
Co-authored-by: Scott Cyphers <diyessi@users.noreply.github.com>

test/models/onnx/provenance_downgrade_topk.prototxt

+ir_version: 4
+producer_name: "nGraph ONNX Importer"
+graph {
+  node {
+    input: "x"
+    input: "k"
+    output: "values"
+    output: "indices"
+    op_type: "TopK"
+    name: "TOPK"
+  }
+  name: "test_graph"
+  input {
+    name: "x"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+            dim_value: 3
+          }
+          dim {
+            dim_value: 4
+          }
+        }
+      }
+    }
+  }
+  input {
+    name: "k"
+    type {
+      tensor_type {
+        elem_type: 7
+        shape {
+          dim {
+            dim_value: 1
+          }
+        }
+      }
+    }
+  }
+  output {
+    name: "values"
+    type {
+      tensor_type {
+        elem_type: 1
+        shape {
+          dim {
+            dim_value: 3
+          }
+          dim {
+            dim_value: 3
+          }
+        }
+      }
+    }
+  }
+  output {
+    name: "indices"
+    type {
+      tensor_type {
+        elem_type: 7
+        shape {
+          dim {
+            dim_value: 3
+          }
+          dim {
+            dim_value: 3
+          }
+        }
+      }
+    }
+  }
+}
+opset_import {
+  version: 10
+}

test/onnx/onnx_import_provenance.in.cpp

@@ -18,6 +18,10 @@
 #include "ngraph/file_util.hpp"
 #include "ngraph/frontend/onnx_import/default_opset.hpp"
 #include "ngraph/frontend/onnx_import/onnx.hpp"
+#include "ngraph/opsets/opset0.hpp"
+#include "ngraph/pass/manager.hpp"
+#include "ngraph/pass/opset0_downgrade.hpp"
+#include "ngraph/provenance.hpp"
 #include "util/test_control.hpp"
 #include "util/type_prop.hpp"
 
@@ -44,52 +48,73 @@ NGRAPH_TEST(onnx_${BACKEND_NAME}, provenance_tag_text)
 // the NodeToCheck parameter of this template is used to find a node in the whole subgraph
 // that a particular unit test is supposed to check against the expected provenance tag
 template <typename NodeToCheck>
-void test_provenance_tags(const std::string& model_path, const std::string& expected_provenance_tag)
+void test_provenance_tags(const std::shared_ptr<Function> function,
+                          const std::string& expected_provenance_tag)
 {
-    const auto function =
-        onnx_import::import_onnx_model(file_util::path_join(SERIALIZED_ZOO, model_path));
-
+    int node_count = 0;
     for (const auto ng_node : function->get_ordered_ops())
     {
         if (as_type_ptr<NodeToCheck>(ng_node))
         {
+            ++node_count;
             const auto tags = ng_node->get_provenance_tags();
-            ASSERT_EQ(tags.size(), 1) << "There should be exactly one provenance tag set for "
-                                      << ng_node;
-
-            EXPECT_EQ(*(tags.cbegin()), expected_provenance_tag);
+            ASSERT_TRUE(tags.size() > 0) << "Node " << ng_node->get_friendly_name()
+                                         << " should have at least one provenance tag.";
+            EXPECT_TRUE(tags.find(expected_provenance_tag) != tags.end());
         }
     }
+    EXPECT_TRUE(node_count > 0) << "Expected type of node doesn't exist in graph.";
 }
 
 NGRAPH_TEST(onnx_${BACKEND_NAME}, provenance_only_output)
 {
     // the Add node in the model does not have a name,
     // only its output name should be found in the provenance tags
-    test_provenance_tags<default_opset::Add>("onnx/provenance_only_outputs.prototxt",
-                                             "<ONNX Add (-> output_of_add)>");
+    const auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/provenance_only_outputs.prototxt"));
+    test_provenance_tags<default_opset::Add>(function, "<ONNX Add (-> output_of_add)>");
 }
 
 NGRAPH_TEST(onnx_${BACKEND_NAME}, provenance_node_name_and_outputs)
 {
-    test_provenance_tags<default_opset::Add>("onnx/provenance_node_name_and_outputs.prototxt",
-                                             "<ONNX Add (Add_node -> output_of_add)>");
+    const auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/provenance_node_name_and_outputs.prototxt"));
+    test_provenance_tags<default_opset::Add>(function, "<ONNX Add (Add_node -> output_of_add)>");
 }
 
 NGRAPH_TEST(onnx_${BACKEND_NAME}, provenance_multiple_outputs_op)
 {
-    test_provenance_tags<default_opset::TopK>("onnx/provenance_multiple_outputs_op.prototxt",
-                                              "<ONNX TopK (TOPK -> values, indices)>");
+    const auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/provenance_multiple_outputs_op.prototxt"));
+    test_provenance_tags<default_opset::TopK>(function, "<ONNX TopK (TOPK -> values, indices)>");
 }
 
 NGRAPH_TEST(onnx_${BACKEND_NAME}, provenance_tagging_constants)
 {
-    test_provenance_tags<default_opset::Constant>("onnx/provenance_input_tags.prototxt",
+    const auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/provenance_input_tags.prototxt"));
+    test_provenance_tags<default_opset::Constant>(function,
                                                   "<ONNX Input (initializer_of_A) Shape:{1}>");
 }
 
 NGRAPH_TEST(onnx_${BACKEND_NAME}, provenance_tagging_parameters)
 {
-    test_provenance_tags<default_opset::Parameter>("onnx/provenance_input_tags.prototxt",
-                                                   "<ONNX Input (input_B) Shape:{}>");
+    const auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/provenance_input_tags.prototxt"));
+    test_provenance_tags<default_opset::Parameter>(function, "<ONNX Input (input_B) Shape:{}>");
+}
+
+NGRAPH_TEST(onnx_${BACKEND_NAME}, provenance_tag_downgrade_pass)
+{
+    set_provenance_enabled(true);
+
+    const auto function = onnx_import::import_onnx_model(
+        file_util::path_join(SERIALIZED_ZOO, "onnx/provenance_downgrade_topk.prototxt"));
+
+    ngraph::pass::Manager pass_manager;
+    pass_manager.register_pass<pass::Opset0Downgrade>();
+    pass_manager.run_passes(function);
+
+    test_provenance_tags<op::v0::TopK>(function, "<ONNX TopK (TOPK -> values, indices)>");
+    test_provenance_tags<op::v0::TopK>(function, "<Opset0_Downgrade (v1 TopK)>");
 }

test/provenance.cpp

@@ -27,6 +27,8 @@
 #include "ngraph/ngraph.hpp"
 #include "ngraph/pass/fused_op_decomposition.hpp"
 #include "ngraph/pass/manager.hpp"
+#include "ngraph/pass/opset0_downgrade.hpp"
+#include "ngraph/pass/opset1_upgrade.hpp"
 #include "ngraph/provenance.hpp"
 
 using namespace std;
@@ -333,6 +335,56 @@ TEST(provenance, add_group_above)
     EXPECT_EQ(m1->get_provenance_tags(), (ProvSet{"m1"}));
 }
 
+TEST(provenance, add_tags_above)
+{
+    auto x = make_shared<op::Parameter>(element::i32, PartialShape{2, 3, 4});
+    auto y = make_shared<op::Parameter>(element::i32, PartialShape{2, 3, 4});
+
+    auto a = make_shared<op::Add>(x, y);
+    auto b = make_shared<op::Multiply>(x, y);
+    auto c = make_shared<op::Subtract>(a, b);
+    auto d = make_shared<op::Abs>(c);
+
+    // Add tags to Subtract and all nodes until Parameters (all above c, until params x, y)
+    c->add_provenance_tags_above(OutputVector{x, y}, {"tag_above_c - until_params"});
+    // Add tags to Abs and Subtract (above d, until c inputs)
+    d->add_provenance_tags_above(c->input_values(), {"tag_above_d - until_c_inputs"});
+    // Add tags to Abs and all nodes above
+    d->add_provenance_tags_above(OutputVector{}, {"tag_all_above_d"});
+
+    auto x_tags = x->get_provenance_tags();
+    EXPECT_EQ(x_tags.size(), 1);
+    EXPECT_TRUE(x_tags.find("tag_all_above_d") != x_tags.end());
+
+    auto y_tags = y->get_provenance_tags();
+    EXPECT_EQ(y_tags.size(), 1);
+    EXPECT_TRUE(y_tags.find("tag_all_above_d") != y_tags.end());
+
+    auto a_tags = a->get_provenance_tags();
+    EXPECT_EQ(a_tags.size(), 2);
+    EXPECT_TRUE(a_tags.find("tag_above_c - until_params") != a_tags.end());
+    EXPECT_FALSE(a_tags.find("tag_above_d - until_c_inputs") != a_tags.end());
+    EXPECT_TRUE(a_tags.find("tag_all_above_d") != a_tags.end());
+
+    auto b_tags = b->get_provenance_tags();
+    EXPECT_EQ(b_tags.size(), 2);
+    EXPECT_TRUE(b_tags.find("tag_above_c - until_params") != b_tags.end());
+    EXPECT_FALSE(b_tags.find("tag_above_d - until_c_inputs") != b_tags.end());
+    EXPECT_TRUE(b_tags.find("tag_all_above_d") != b_tags.end());
+
+    auto c_tags = c->get_provenance_tags();
+    EXPECT_EQ(c_tags.size(), 3);
+    EXPECT_TRUE(c_tags.find("tag_above_c - until_params") != c_tags.end());
+    EXPECT_TRUE(c_tags.find("tag_above_d - until_c_inputs") != c_tags.end());
+    EXPECT_TRUE(c_tags.find("tag_all_above_d") != c_tags.end());
+
+    auto d_tags = d->get_provenance_tags();
+    EXPECT_EQ(d_tags.size(), 2);
+    EXPECT_FALSE(d_tags.find("tag_above_c - until_params") != d_tags.end());
+    EXPECT_TRUE(d_tags.find("tag_above_d - until_c_inputs") != d_tags.end());
+    EXPECT_TRUE(d_tags.find("tag_all_above_d") != d_tags.end());
+}
+
 TEST(provenance, builder)
 {
     auto p1 = make_shared<op::Parameter>(element::i32, PartialShape{2, 3, 4});
@@ -501,3 +553,155 @@ TEST(provenance, scaled_quantize_concat_unsigned)
         }
     }
 }
+
+TEST(provenance, opset1_upgrade_pass_topk)
+{
+    set_provenance_enabled(true);
+
+    const size_t axis = 2;
+    const size_t k = 10;
+    const auto data = make_shared<op::Parameter>(element::i32, Shape{5, 10, 15});
+
+    const auto topk_v0 = make_shared<op::v0::TopK>(data, axis, element::i32, k);
+    const auto result = make_shared<op::Result>(topk_v0);
+    auto f = make_shared<Function>(ResultVector{result}, ParameterVector{data});
+
+    ngraph::pass::Manager pass_manager;
+    pass_manager.register_pass<pass::Opset1Upgrade>();
+    pass_manager.run_passes(f);
+
+    const auto pass_replacement_node =
+        f->get_result()->input(0).get_source_output().get_node_shared_ptr();
+    const auto topk_v1 = as_type_ptr<op::v1::TopK>(pass_replacement_node);
+
+    const std::string tag = "<Opset1_Upgrade (v0 TopK)>";
+    auto tag_check = [&tag](std::shared_ptr<ngraph::Node> node) {
+        auto tags = node->get_provenance_tags();
+        EXPECT_TRUE(tags.find(tag) != tags.end());
+    };
+    traverse_nodes(as_node_vector(topk_v1->outputs()),
+                   tag_check,
+                   false,
+                   as_node_vector(topk_v0->input_values()));
+}
+
+TEST(provenance, opset0_downgrade_pass_topk)
+{
+    set_provenance_enabled(true);
+
+    const auto data = make_shared<op::Parameter>(element::i32, Shape{5, 10, 15});
+    const int32_t k = 10;
+    const auto k_node = op::Constant::create(element::i64, Shape{}, {k});
+    const size_t axis = 2;
+    const auto mode = op::v1::TopK::Mode::MAX;
+    const auto sort = op::v1::TopK::SortType::SORT_INDICES;
+    const auto elem_type = element::i64;
+
+    const auto topk_v1 = make_shared<op::v1::TopK>(data, k_node, axis, mode, sort, elem_type);
+    const auto result = make_shared<op::Result>(topk_v1);
+    auto f = make_shared<Function>(ResultVector{result}, ParameterVector{data});
+
+    ngraph::pass::Manager pass_manager;
+    pass_manager.register_pass<pass::Opset0Downgrade>();
+    pass_manager.run_passes(f);
+
+    const auto pass_replacement_node =
+        f->get_result()->input(0).get_source_output().get_node_shared_ptr();
+    const auto topk_v0 = as_type_ptr<op::v0::TopK>(pass_replacement_node);
+
+    const std::string tag = "<Opset0_Downgrade (v1 TopK)>";
+    auto tag_check = [&tag](std::shared_ptr<ngraph::Node> node) {
+        auto tags = node->get_provenance_tags();
+        EXPECT_TRUE(tags.find(tag) != tags.end());
+    };
+    traverse_nodes(as_node_vector(topk_v0->outputs()),
+                   tag_check,
+                   false,
+                   as_node_vector(topk_v1->input_values()));
+}
+
+TEST(provenance, opset1_upgrade_pass_graph)
+{
+    set_provenance_enabled(true);
+
+    auto x = make_shared<op::Parameter>(element::i32, PartialShape{2, 3, 4});
+    auto y = make_shared<op::Parameter>(element::i32, PartialShape{2, 3, 4});
+
+    auto a = make_shared<op::v0::Add>(x, y);
+    auto b = make_shared<op::v0::Subtract>(x, y);
+    auto c = make_shared<op::v0::Abs>(b);
+    auto d = make_shared<op::v0::Multiply>(a, b);
+
+    auto f = make_shared<Function>(d, ParameterVector{x, y});
+
+    ngraph::pass::Manager pass_manager;
+    pass_manager.register_pass<pass::Opset1Upgrade>();
+    pass_manager.run_passes(f);
+
+    for (auto node : f->get_ordered_ops())
+    {
+        auto tags = node->get_provenance_tags();
+        if (as_type_ptr<op::v1::Add>(node))
+        {
+            EXPECT_EQ(tags.size(), 1);
+            EXPECT_TRUE(tags.find("<Opset1_Upgrade (v0 Add)>") != tags.end());
+        }
+        else if (as_type_ptr<op::v1::Multiply>(node))
+        {
+            EXPECT_EQ(tags.size(), 1);
+            EXPECT_TRUE(tags.find("<Opset1_Upgrade (v0 Multiply)>") != tags.end());
+        }
+        else if (as_type_ptr<op::v1::Subtract>(node))
+        {
+            EXPECT_EQ(tags.size(), 1);
+            EXPECT_TRUE(tags.find("<Opset1_Upgrade (v0 Subtract)>") != tags.end());
+        }
+        else if (as_type_ptr<op::v0::Abs>(node))
+        {
+            EXPECT_TRUE(tags.empty());
+        }
+    }
+}
+
+TEST(provenance, opset0_downgrade_pass_graph)
+{
+    set_provenance_enabled(true);
+
+    auto x = make_shared<op::Parameter>(element::i32, PartialShape{2, 3, 4});
+    auto y = make_shared<op::Parameter>(element::i32, PartialShape{2, 3, 4});
+
+    auto a = make_shared<op::v1::Add>(x, y);
+    auto b = make_shared<op::v1::Subtract>(x, y);
+    auto c = make_shared<op::v0::Abs>(b);
+    auto d = make_shared<op::v1::Multiply>(a, b);
+
+    auto f = make_shared<Function>(d, ParameterVector{x, y});
+
+    ngraph::pass::Manager pass_manager;
+    pass_manager.register_pass<pass::Opset0Downgrade>();
+    pass_manager.run_passes(f);
+
+    for (auto node : f->get_ordered_ops())
+    {
+        auto tags = node->get_provenance_tags();
+        if (as_type_ptr<op::v0::Add>(node))
+        {
+            EXPECT_EQ(tags.size(), 1);
+            EXPECT_TRUE(tags.find("<Opset0_Downgrade (v1 Add)>") != tags.end());
+        }
+        else if (as_type_ptr<op::v0::Multiply>(node))
+        {
+            EXPECT_EQ(tags.size(), 1);
+            EXPECT_TRUE(tags.find("<Opset0_Downgrade (v1 Multiply)>") != tags.end());
+        }
+        else if (as_type_ptr<op::v0::Subtract>(node))
+        {
+            EXPECT_EQ(tags.size(), 1);
+            EXPECT_TRUE(tags.find("<Opset0_Downgrade (v1 Subtract)>") != tags.end());
+        }
+        else if (as_type_ptr<op::v0::Abs>(node))
+        {
+            EXPECT_TRUE(tags.empty());
+        }
+    }
+}