From aa649760f942b540ebf1e856906fe64043b979a4 Mon Sep 17 00:00:00 2001
From: Mengshiun Yu <mengshyu@gmail.com>
Date: Thu, 12 Sep 2024 08:50:57 -0400
Subject: [PATCH] [Relax][Transform] Add SelectNode handling in SymbolicMatcher

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 612e1459c826..fe247645dc24 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);