- Fix bug with pattern var visit

Author: mbs-octoml

src/relay/ir/indexed_graph.cc

@@ -19,7 +19,9 @@
 
 /*!
  * \file src/relay/ir/indexed_graph.cc
- * \brief Utilities for creating Indexed Graphs.
+ * \brief A graph representation of the dataflow in a Relay expression or Relay dataflow
+ * pattern. Nodes in this graph capture generic dataflow inputs, outputs, dominators and control
+ * flow scopes.
  */
 #include "indexed_graph.h"
 
@@ -34,33 +36,30 @@ namespace tvm {
 namespace relay {
 
 std::string RefToSummary(const Expr& expr) {
-  if (const auto* var_node = expr.as<VarNode>()) {
-    return "%" + var_node->name_hint();
-  } else if (const auto* global_var_node = expr.as<GlobalVarNode>()) {
-    return "@" + global_var_node->name_hint;
-  } else if (const auto* op_node = expr.as<OpNode>()) {
-    return op_node->name;
-  } else if (expr.as<ConstantNode>()) {
-    return "const";
-  } else if (expr.as<FunctionNode>()) {
-    return "fn";
-  } else if (const auto* tuple_node = expr.as<TupleNode>()) {
-    return "tuple(" + std::to_string(tuple_node->fields.size()) + ")";
-  } else if (const auto* tuple_get_item_node = expr.as<TupleGetItemNode>()) {
-    return "." + std::to_string(tuple_get_item_node->index);
-  } else if (const auto* call_node = expr.as<CallNode>()) {
-    return RefToSummary(call_node->op) + "(" + std::to_string(call_node->args.size()) + ")";
-  } else if (expr.as<IfNode>()) {
-    return "if";
-  } else if (expr.as<ConstructorNode>()) {
-    return "ctor";
-  } else if (expr.as<MatchNode>()) {
-    return "match";
-  } else if (expr.as<LetNode>()) {
-    return "let";
-  } else {
-    return "";
-  }
+  class Visitor : public ExprFunctor<std::string(const Expr&)> {
+    std::string VisitExpr_(const VarNode* op) final { return "%" + op->name_hint(); }
+    std::string VisitExpr_(const GlobalVarNode* op) final { return "@" + op->name_hint; }
+    std::string VisitExpr_(const ConstantNode* op) final { return "const"; }
+    std::string VisitExpr_(const TupleNode* op) final {
+      return "tuple(" + std::to_string(op->fields.size()) + ")";
+    }
+    std::string VisitExpr_(const FunctionNode* op) final { return "fn"; }
+    std::string VisitExpr_(const CallNode* op) final {
+      return VisitExpr(op->op) + "(" + std::to_string(op->args.size()) + ")";
+    }
+    std::string VisitExpr_(const LetNode* op) final { return "let"; }
+    std::string VisitExpr_(const IfNode* op) final { return "if"; }
+    std::string VisitExpr_(const OpNode* op) final { return op->name; }
+    std::string VisitExpr_(const TupleGetItemNode* op) final {
+      return "." + std::to_string(op->index);
+    }
+    std::string VisitExpr_(const RefCreateNode* op) final { return "ref_create"; }
+    std::string VisitExpr_(const RefReadNode* op) final { return "ref_read"; }
+    std::string VisitExpr_(const RefWriteNode* op) final { return "ref_write"; }
+    std::string VisitExpr_(const ConstructorNode* op) final { return "ctor"; }
+    std::string VisitExpr_(const MatchNode* op) final { return "match"; }
+  };
+  return Visitor().VisitExpr(expr);
 }
 
 std::string RefToSummary(const DFPattern& pattern) { return ""; }
@@ -140,7 +139,7 @@ std::unique_ptr<IndexedGraph<Expr>> CreateIndexedGraph(const Expr& expr) {
 
      private:
       void VisitPattern_(const PatternVarNode* pattern_var_node) final {
-        creator_->graph_->AddNode(pattern_var_node->var);
+        creator_->VisitLeaf(pattern_var_node->var);
       }
 
       Creator* creator_;
@@ -173,9 +172,9 @@ std::unique_ptr<IndexedGraph<Expr>> CreateIndexedGraph(const Expr& expr) {
    */
   class Annotator : public ExprFunctor<void(const Expr&)> {
    public:
-    Annotator(std::pair<std::unique_ptr<IndexedGraph<Expr>>,
-                        std::unique_ptr<std::unordered_set<const CallNode*>>>
-                  args)
+    explicit Annotator(std::pair<std::unique_ptr<IndexedGraph<Expr>>,
+                                 std::unique_ptr<std::unordered_set<const CallNode*>>>
+                           args)
         : graph_(std::move(args.first)), rec_calls_(std::move(args.second)) {}
 
     std::unique_ptr<IndexedGraph<Expr>> Annotate() {

src/relay/ir/indexed_graph.h

@@ -22,6 +22,8 @@
  * \brief A graph representation of the dataflow in a Relay expression or Relay dataflow
  * pattern. Nodes in this graph capture generic dataflow inputs, outputs, dominators and control
  * flow scopes.
+ *
+ * TODO(mbs): Rename to 'DataflowGraph'.
  */
 #ifndef TVM_RELAY_IR_INDEXED_GRAPH_H_
 #define TVM_RELAY_IR_INDEXED_GRAPH_H_
@@ -264,7 +266,7 @@ class IndexedGraph {
       ICHECK(node->node_ref_);
       auto itr = node_map_.find(node->node_ref_);
       ICHECK(itr != node_map_.end());
-      ICHECK_EQ(itr->second, node);
+      ICHECK_EQ(itr->second, node) << "at index " << index << " in:" << std::endl << ToString();
       // Inputs come before.
       for (size_t i = 0; i < node->inputs_.size(); ++i) {
         const Node* input = node->inputs_[i];

tests/python/relay/test_dataflow_pattern.py

@@ -1458,18 +1458,6 @@ def test_partition_overused():
     assert pattern.partition(out) == out
 
 
-def test_partition_not_overused():
-    pattern = is_op("nn.relu")(is_op("nn.conv2d")(wildcard(), wildcard()))
-
-    x = relay.var("input")
-    w = relay.var("weight")
-    conv2d = relay.op.nn.conv2d(x, w)
-    relu = relay.op.nn.relu(conv2d)
-    out = relu + relu
-
-    assert pattern.partition(out) == out
-
-
 def test_partition_fuzzy_tuple():
     x = relay.var("x")
     y = relay.var("y")
@@ -1496,6 +1484,7 @@ def concat(*args):
 
 
 def test_partition_fuzzy_function_args():
+
     func_pattern = FunctionPattern(None, wildcard() + wildcard())(None) + wildcard()
     x = relay.var("x")
     y = relay.var("y")