(tvm) branch main updated: [Relax][Transform] Add SelectNode handling in SymbolicMatcher (#17368)

[ruihangl]

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

ruihangl 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 4bc61a1445 [Relax][Transform] Add SelectNode handling in 
SymbolicMatcher (#17368)
4bc61a1445 is described below

commit 4bc61a14452cdae09231f1085d40a4b04fbe1f75
Author: Mengshiun Yu <mengs...@gmail.com>
AuthorDate: Sat Sep 14 23:07:06 2024 -0400

    [Relax][Transform] Add SelectNode handling in SymbolicMatcher (#17368)
    
    This PR added support for handling SelectNode in the SymbolicMatcher
    class by modifying the VisitExpr_ function to match the true_value
    and false_value expressions between the current SelectNode and the
    other expression. If the other expression is not a SelectNode, the
    matching condition is updated to ensure the current SelectNode
    expression is equivalent to the other expression.
---
 src/relax/transform/fuse_tir.cc | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/src/relax/transform/fuse_tir.cc b/src/relax/transform/fuse_tir.cc
index 612e1459c8..fe247645dc 100644
--- a/src/relax/transform/fuse_tir.cc
+++ b/src/relax/transform/fuse_tir.cc
@@ -139,6 +139,16 @@ class SymbolicMatcher : ExprFunctor<void(const PrimExpr& 
n, const PrimExpr& othe
     }
   }
 
+  void VisitExpr_(const SelectNode* op, const PrimExpr& other) {
+    const auto* rhs = other.as<SelectNode>();
+    if (rhs) {
+      VisitExpr(op->true_value, rhs->true_value);
+      VisitExpr(op->false_value, rhs->false_value);
+    } else {
+      must_prove_ = must_prove_ && (GetRef<PrimExpr>(op) == other);
+    }
+  }
+
   arith::Analyzer* analyzer_;
   Map<tir::Var, PrimExpr>* var_remap_;
   PrimExpr must_prove_ = Bool(true);