(tvm) branch main updated: [BugFix][Relay] skip leaf args when matching 'path' part for dominator pattern (#16983)

[wr...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

wrongtest pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git


The following commit(s) were added to refs/heads/main by this push:
     new e6bfaf8d80 [BugFix][Relay] skip leaf args when matching 'path' part for dominator pattern (#16983)
e6bfaf8d80 is described below

commit e6bfaf8d80c553cb626c542d57b36698b066c128
Author: Huibin Wang <wa...@intellif.com>
AuthorDate: Fri Jun 21 18:29:59 2024 +0800

    [BugFix][Relay] skip leaf args when matching 'path' part for dominator pattern (#16983)
    
    * [BugFix][Relay] skip leaf args when matching 'path' part for dominator pattern
    
    * add testcase
---
 src/relay/ir/dataflow_matcher.cc            |  8 +++++++-
 tests/python/relay/test_dataflow_pattern.py | 24 +++++++++++++++++++++++-
 2 files changed, 30 insertions(+), 2 deletions(-)

diff --git a/src/relay/ir/dataflow_matcher.cc b/src/relay/ir/dataflow_matcher.cc
index 8e756a8aa2..0c0ff72901 100644
--- a/src/relay/ir/dataflow_matcher.cc
+++ b/src/relay/ir/dataflow_matcher.cc
@@ -300,11 +300,17 @@ bool DFPatternMatcher::VisitDFPattern_(const CallPatternNode* op, const Expr& ex
 
 // Recursively find the Dominator parent along all inputs paths.
 bool DFPatternMatcher::MatchesPath(const DominatorPatternNode* op, const Expr& expr) {
+  // utilities
+  auto is_leaf_node = [](const Expr& expr) {
+    return expr.as<ConstantNode>() || expr.as<VarNode>();
+  };
+
+  // logic
   auto call_node = expr.as<CallNode>();
   auto index_node = expr_to_node(expr);
   size_t arg_counter{0};
   for (auto node : index_node->inputs_) {
-    if (!(call_node && node->ref() == call_node->op)) {
+    if (!(call_node && (node->ref() == call_node->op || is_leaf_node(node->ref())))) {
       arg_counter += 1;
       memoize_ = true;
       if (!VisitDFPattern(op->parent, node->ref())) {
diff --git a/tests/python/relay/test_dataflow_pattern.py b/tests/python/relay/test_dataflow_pattern.py
index 6942c47491..4031790fc3 100644
--- a/tests/python/relay/test_dataflow_pattern.py
+++ b/tests/python/relay/test_dataflow_pattern.py
@@ -28,7 +28,7 @@ from tvm.relay.testing import run_opt_pass
 # convention.
 K_ELEMWISE = 0
 K_BROADCAST = 1
-
+K_INJECTIVE = 2
 
 ## NODE TESTS
 def test_expr_pattern():
@@ -696,6 +696,28 @@ def test_match_dominator():
     assert diamond.match(out)
 
 
+def test_match_dominator2():
+    # Pattern
+    conv2d_pat = is_op("nn.conv2d")(wildcard(), wildcard())
+    eltwise_pat = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(None)
+    broadcast_pat = (wildcard().has_attr({"TOpPattern": K_BROADCAST}))(None)
+    path_pat = eltwise_pat | broadcast_pat
+    injective_pat = (wildcard().has_attr({"TOpPattern": K_INJECTIVE}))(wildcard())
+    pattern = injective_pat.dominates(conv2d_pat, path_pat)
+
+    # Graph
+    inp = relay.var("input")
+    weight = relay.var("weight")
+    bias = relay.var("bias")
+    conv2d = relay.op.nn.conv2d(inp, weight)
+    bias_add = relay.op.nn.bias_add(conv2d, bias)
+    relu = relay.op.nn.relu(bias_add)
+    reshape = relay.op.reshape(relu, newshape=[-1, 2, 8])
+
+    # Check
+    assert pattern.match(reshape)
+
+
 def test_not_match_dominator():
     is_conv2d = is_op("nn.conv2d")(wildcard(), wildcard())
     is_unary_elemwise = (wildcard().has_attr({"TOpPattern": K_ELEMWISE}))(wildcard())