Skip to content

Commit

Permalink
[Relax][Transform] Add SelectNode handling in SymbolicMatcher (#17368)
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
mengshyu authored Sep 15, 2024
1 parent e0105e4 commit 4bc61a1
Showing 1 changed file with 10 additions and 0 deletions.
10 changes: 10 additions & 0 deletions src/relax/transform/fuse_tir.cc
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand Down

0 comments on commit 4bc61a1

Please sign in to comment.