Adds provenance tag propagation (#2747)

* Add some unit tests for provenance with node replacement

* One more test

* add provenace propagation

* Adds ability to control whether provenance is on or off

* Add some unit tests for provenance with node replacement

* One more test

* add provenace propagation

* Adds ability to control whether provenance is on or off

* Removes extraneous semicolon

* Fixes include header errors

* Reimplements provenance enable flag using global variable

* Fixes style issues.

* Adds newline at the end of file.

* Fixes function return type.

* Adds RAII mechanism for localizing provenance enable state in testing.

* Improves getter/setter method name to match variable name.

* Adds s_ prefix to global variable provenance_enabled.

src/ngraph/CMakeLists.txt

@@ -354,6 +354,8 @@ set (SRC
     pattern/op/skip.hpp
     placement.cpp
     placement.hpp
+    provenance.cpp
+    provenance.hpp
     rank.hpp
     result_vector.hpp
     runtime/aligned_buffer.cpp

src/ngraph/graph_util.cpp

@@ -32,6 +32,7 @@
 #include "ngraph/op/result.hpp"
 #include "ngraph/pass/manager.hpp"
 #include "ngraph/pass/visualize_tree.hpp"
+#include "ngraph/provenance.hpp"
 #include "ngraph/result_vector.hpp"
 #include "ngraph/util.hpp"
 
@@ -150,6 +151,24 @@ void ngraph::replace_node(std::shared_ptr<Node> target, std::shared_ptr<Node> re
     // Fix input/output descriptors
     NGRAPH_CHECK(target->get_output_size() == replacement->get_output_size());
 
+    if (ngraph::get_provenance_enabled())
+    {
+        auto set_replacement_prov = [replacement](std::shared_ptr<Node> node) {
+            replacement->merge_provenance_tags_from(node);
+        };
+
+        traverse_nodes({target}, set_replacement_prov, false, replacement->get_arguments());
+
+        auto propagate_replacement_prov = [replacement](std::shared_ptr<Node> node) {
+            if (is_post_dominated(node.get(), replacement.get()))
+            {
+                node->merge_provenance_tags_from(replacement);
+            }
+        };
+
+        traverse_nodes({replacement}, propagate_replacement_prov, false);
+    }
+
     // For each of target's output O with replacement output O_rep:
     //     For each O's connected downstream input I:
     //         Change I's connected upstream output to O_rep
@@ -160,7 +179,6 @@ void ngraph::replace_node(std::shared_ptr<Node> target, std::shared_ptr<Node> re
             input.replace_source_output(replacement->output(i));
         }
     }
-    replacement->merge_provenance_tags_from(target);
 }
 
 // Check if all paths from X to a result go through Y

src/ngraph/provenance.cpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#include <cstdlib>
+
+#include "provenance.hpp"
+
+namespace ngraph
+{
+    void set_provenance_enabled(bool enabled) { s_provenance_enabled = enabled; }
+    bool get_provenance_enabled() { return s_provenance_enabled; }
+}

src/ngraph/provenance.hpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#pragma once
+
+#include <cstdlib>
+
+namespace ngraph
+{
+    static bool s_provenance_enabled = std::getenv("NGRAPH_PROVENANCE_ENABLE") != nullptr;
+    void set_provenance_enabled(bool enabled);
+    bool get_provenance_enabled();
+}

test/CMakeLists.txt

@@ -58,6 +58,7 @@ set(SRC
     pass_memory_layout.cpp
     pass_shape_specialization.cpp
     pattern.cpp
+    provenance.cpp
     reshape_elimination.cpp
     reshape_sinking.cpp
     serialize.cpp

test/provenance.cpp

+//*****************************************************************************
+// Copyright 2017-2019 Intel Corporation
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//*****************************************************************************
+
+#include <fstream>
+#include <sstream>
+#include <string>
+#include <vector>
+
+#include "gmock/gmock.h"
+#include "gtest/gtest.h"
+
+#include "ngraph/ngraph.hpp"
+#include "ngraph/provenance.hpp"
+
+using namespace std;
+using namespace ngraph;
+using ::testing::Return;
+
+using ProvSet = std::unordered_set<std::string>;
+
+TEST(provenance, provenance)
+{
+    class ProvenanceEnabler
+    {
+    public:
+        ProvenanceEnabler()
+        {
+            saved_enable_state = get_provenance_enabled();
+            set_provenance_enabled(true);
+        }
+        ~ProvenanceEnabler() { set_provenance_enabled(saved_enable_state); }
+    private:
+        bool saved_enable_state;
+    } provenance_enabler;
+
+    //
+    // Before:
+    //
+    //   A{tag_a}  B{tag_b}
+    //         |   |
+    //        C{tag_c}
+    //
+    // Replacement:
+    //
+    //       A{tag_a} B{tag_b}
+    //              | |
+    //         C := D{}
+    //
+    // After:
+    //
+    //   A{tag_a}  B{tag_b}
+    //         |   |
+    //        D{tag_c}
+    //
+    // Comment:
+    //   * D is the replacement root, and its insertion kills C. We should not, however, consider
+    //     A and B to be killed, because they are not post-dominated by D until after C is cut out
+    //     of the graph.
+    //
+    {
+        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);
+        a->add_provenance_tag("tag_a");
+        auto b = make_shared<op::Multiply>(y, x);
+        b->add_provenance_tag("tag_b");
+        auto c = make_shared<op::Subtract>(a, b);
+        c->add_provenance_tag("tag_c");
+
+        auto f = make_shared<Function>(c, ParameterVector{x, y});
+
+        auto new_c = make_shared<op::Subtract>(a, b);
+        replace_node(c, new_c);
+
+        EXPECT_EQ(new_c->get_provenance_tags(), ProvSet{"tag_c"});
+    }
+
+    //
+    // Before:
+    //
+    //   A{tag_a}  B{tag_b}
+    //         |   |
+    //        C{tag_c}
+    //
+    // Replacement:
+    //
+    //
+    //
+    //     A{tag_a}  B{tag_b}
+    //        |      |
+    //   C -> D{tag_d}
+    //
+    // After:
+    //
+    //   A{tag_a}  B{tag_b}
+    //         |   |
+    //        D{tag_c,tag_d}
+    //
+    // Comment:
+    //   * D is the replacement root, and its insertion kills C. We should not, however, consider
+    //     A and B to be killed, because they are not post-dominated by D until after C is cut out
+    //     of the graph.
+    //
+    {
+        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);
+        a->add_provenance_tag("tag_a");
+        auto b = make_shared<op::Multiply>(y, x);
+        b->add_provenance_tag("tag_b");
+        auto c = make_shared<op::Subtract>(a, b);
+        c->add_provenance_tag("tag_c");
+
+        auto f = make_shared<Function>(c, ParameterVector{x, y});
+
+        auto d = make_shared<op::Subtract>(a, b);
+        d->add_provenance_tag("tag_d");
+        replace_node(c, d);
+
+        EXPECT_EQ(d->get_provenance_tags(), (ProvSet{"tag_c", "tag_d"}));
+    }
+
+    //
+    // Before:
+    //
+    //   A{tag_a}  B{tag_b}
+    //         |   |
+    //        C{tag_c}
+    //
+    // Replacement:
+    //
+    //   C -> D{tag_d}
+    //
+    // After:
+    //
+    //   D{tag_a,tag_b,tag_c,tag_d}
+    //
+    // Comment:
+    //   * D is the replacement root, and its insertion kills A, B, and C.
+    //
+    {
+        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);
+        a->add_provenance_tag("tag_a");
+        auto b = make_shared<op::Multiply>(y, x);
+        b->add_provenance_tag("tag_b");
+        auto c = make_shared<op::Subtract>(a, b);
+        c->add_provenance_tag("tag_c");
+
+        auto f = make_shared<Function>(c, ParameterVector{x, y});
+
+        auto d = make_zero(element::i32, Shape{2, 3, 4});
+        d->add_provenance_tag("tag_d");
+        replace_node(c, d);
+
+        EXPECT_EQ(d->get_provenance_tags(), (ProvSet{"tag_a", "tag_b", "tag_c", "tag_d"}));
+    }
+
+    //
+    // Before:
+    //
+    //   A{tag_a}  B{tag_b}
+    //         |   |
+    //        C{tag_c}
+    //
+    // Replacement:
+    //
+    //   C -> D{}
+    //
+    // After:
+    //
+    //   D{tag_a,tag_b,tag_c}
+    //
+    // Comment:
+    //   * D is the replacement root, and its insertion kills A, B, and C.
+    //
+    {
+        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);
+        a->add_provenance_tag("tag_a");
+        auto b = make_shared<op::Multiply>(y, x);
+        b->add_provenance_tag("tag_b");
+        auto c = make_shared<op::Subtract>(a, b);
+        c->add_provenance_tag("tag_c");
+
+        auto f = make_shared<Function>(c, ParameterVector{x, y});
+
+        auto d = make_zero(element::i32, Shape{2, 3, 4});
+        replace_node(c, d);
+
+        EXPECT_EQ(d->get_provenance_tags(), (ProvSet{"tag_a", "tag_b", "tag_c"}));
+    }
+
+    //
+    // Before:
+    //
+    //   A{tag_a}  B{tag_b}
+    //         |   |
+    //        C{tag_c}
+    //
+    // Replacement:
+    //
+    //        D{}
+    //         |
+    //   C -> E{}
+    //
+    // After:
+    //
+    //   D{tag_a,tag_b,tag_c}
+    //   |
+    //   E{tag_a,tag_b,tag_c}
+    //
+    // Comment:
+    //   * E is the replacement root, and its insertion kills A, B, and C.
+    //   * D is post-dominated by E, so the tags inherited by E should also be taken on by D.
+    //
+    {
+        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);
+        a->add_provenance_tag("tag_a");
+        auto b = make_shared<op::Multiply>(y, x);
+        b->add_provenance_tag("tag_b");
+        auto c = make_shared<op::Subtract>(a, b);
+        c->add_provenance_tag("tag_c");
+
+        auto f = make_shared<Function>(c, ParameterVector{x, y});
+
+        auto d = make_zero(element::i32, Shape{2, 3, 4});
+        auto e = make_shared<op::Negative>(d);
+        replace_node(c, e);
+
+        EXPECT_EQ(d->get_provenance_tags(), (ProvSet{"tag_a", "tag_b", "tag_c"}));
+        EXPECT_EQ(e->get_provenance_tags(), (ProvSet{"tag_a", "tag_b", "tag_c"}));
+    }
+}