masahi commented on code in PR #14501:
URL: https://github.com/apache/tvm/pull/14501#discussion_r1159163415
##########
src/relax/ir/dataflow_matcher.cc:
##########
@@ -523,109 +523,117 @@ bool MatchExpr(DFPattern pattern, Expr expr,
Optional<Map<Var, Expr>> bindings_o
TVM_REGISTER_GLOBAL("relax.dpl.match_expr").set_body_typed(MatchExpr);
+class MatcherUseDefAnalysis : public relax::ExprVisitor {
+ public:
+ std::vector<const VarNode*> vars;
+ std::map<const VarNode*, std::vector<const VarNode*>> def2use;
+ // caller -> callee table.
+ std::map<const VarNode*, std::vector<const VarNode*>> caller2callees;
+
+ const VarNode* cur_user_;
+
+ void VisitBinding_(const VarBindingNode* binding) override {
+ // init
+ cur_user_ = binding->var.get();
+ this->VisitVarDef(binding->var);
+ this->VisitExpr(binding->value);
+ cur_user_ = nullptr;
+ }
+
+ void VisitExpr_(const VarNode* op) override {
+ if (nullptr == cur_user_) return;
+
+ auto check_and_push = [](std::vector<const VarNode*>& vec, const VarNode*
var) {
+ if (std::find(vec.begin(), vec.end(), var) == vec.end()) {
+ vec.push_back(var);
+ }
+ };
+
+ check_and_push(def2use[op], cur_user_);
+ check_and_push(vars, op);
+
+ caller2callees[cur_user_].push_back(op);
+ }
+
+ void VisitExpr_(const DataflowVarNode* op) override {
+ VisitExpr_(static_cast<const VarNode*>(op));
+ }
+};
+
struct PNode {
const DFPatternNode* ptr;
- const VarNode* matched = nullptr;
std::vector<std::pair<PNode*, const std::vector<PairCons>&>> children;
std::vector<std::pair<PNode*, const std::vector<PairCons>&>> parents;
};
struct RNode {
const VarNode* ptr;
- const DFPatternNode* matched = nullptr;
std::vector<RNode*> children;
std::vector<RNode*> parents;
};
-/**
- * \brief This method try to match a real node and a pattern node along with
its neighbors.
- */
-using UndoItems = std::vector<std::pair<PNode*, RNode*>>;
-static std::optional<UndoItems> try_match(
- PNode* p, RNode* r, DFPatternMatcher* m,
- const std::map<const VarNode*, std::vector<const VarNode*>>& def2use,
- const std::map<const VarNode*, std::vector<const VarNode*>>& use2def) {
- if (p->matched != nullptr && p->matched == r->ptr) return {}; // matched
before.
- if (!m->Match(GetRef<DFPattern>(p->ptr), GetRef<Var>(r->ptr))) return
std::nullopt;
-
- UndoItems undo;
-
- const auto commit = [&undo](PNode* p, RNode* r) {
- // match with each other.
- // TODO(ganler, masahi): Why commit on the same p-r pair happens more than
once?
- if (p->ptr == r->matched) {
- ICHECK_EQ(p->matched, r->ptr);
- return;
- }
- p->matched = r->ptr;
- r->matched = p->ptr;
- undo.emplace_back(p, r);
- };
-
- const auto quit = [&undo] {
- for (auto& [p_node, r_node] : undo) {
- p_node->matched = nullptr;
- r_node->matched = nullptr;
- }
- return std::nullopt;
- };
+struct MatchState {
+ void add(const PNode* p, const RNode* r) {
+ match_p_r[p] = r;
+ match_r_p[r] = p;
+ }
- const auto try_match_update_undo = [&](PNode* p, RNode* r) {
- if (auto undo_more = try_match(p, r, m, def2use, use2def)) {
- undo.insert(undo.end(), undo_more->begin(), undo_more->end());
- return true;
+ void add(const MatchState& other) {
+ for (const auto& [p, r] : other.match_p_r) {
+ add(p, r);
}
- return false;
- };
+ }
- commit(p, r);
+ const VarNode* matched(const PNode* p) const {
+ if (!match_p_r.count(p)) return nullptr;
+ return match_p_r.at(p)->ptr;
+ }
- // match parent patterns.
- for (auto& [pparent, constraints] : p->parents) {
- bool any_cons_sat = false;
- for (auto& rparent : r->parents) {
- // skip if mismatch.
- if (rparent->matched && rparent->matched != pparent->ptr) continue;
+ const DFPatternNode* matched(const RNode* r) const {
+ if (!match_r_p.count(r)) return nullptr;
+ return match_r_p.at(r)->ptr;
+ }
- const auto& uses = def2use.at(rparent->ptr);
+ const VarNode* matched(const PNode& p) const { return matched(&p); }
+ const DFPatternNode* matched(const RNode& r) const { return matched(&r); }
- // check edge constraints.
- bool cons_sat = true;
- for (const auto& cons : constraints) {
- if (cons.type == PairCons::kOnlyUsedBy && uses.size() != 1) {
- cons_sat = false;
- break;
- }
-
- if (cons.index != -1) {
- const auto& callees = use2def.at(r->ptr);
- if (callees.size() <= static_cast<size_t>(cons.index) ||
- callees[cons.index] != rparent->ptr) {
- cons_sat = false;
- break;
- }
- }
- }
- if (!cons_sat) continue;
- any_cons_sat = true;
+ private:
+ std::unordered_map<const PNode*, const RNode*> match_p_r;
+ std::unordered_map<const RNode*, const PNode*> match_r_p;
+};
- // try all parent R nodes that are not matched yet.
- // as long as ppattern can match one node.
- if (!pparent->matched && try_match_update_undo(pparent, rparent)) {
- commit(pparent, rparent);
- break;
+/**
+ * \brief This method try to match a real node and a pattern node along with
its neighbors.
+ */
+static std::optional<MatchState> TryMatch(const PNode& p, const RNode& r,
DFPatternMatcher* m,
+ const MatcherUseDefAnalysis&
ud_analysis) {
+ if (!m->Match(GetRef<DFPattern>(p.ptr), GetRef<Var>(r.ptr))) return
std::nullopt;
+
+ MatchState result;
+
+ for (size_t i = 0; i < p.parents.size(); ++i) {
+ const auto p_node_parent = p.parents[i].first;
+ if (p_node_parent->ptr->IsInstance<WildcardPatternNode>()) {
+ ICHECK_EQ(p.parents.size(), r.parents.size());
+ if (auto v = result.matched(p_node_parent); v && v != r.parents[i]->ptr)
{
+ // A parent wildcard pattern is already matched to other variable.
+ return std::nullopt;
Review Comment:
This is for preventing false matching below:
```
lv0 = R.matmul(x1, w0)
lv1 = R.matmul(x2, w1)
lv2 = R.matmul(x2, w2)
```
```
inp_pat = wildcard()
Q_weight_pat = wildcard()
K_weight_pat = wildcard()
V_weight_pat = wildcard()
matmul1 = is_op("relax.matmul")(inp_pat, Q_weight_pat)
matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat)
matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat)
```
The pattern adds a constraint that the LHS of three matmuls be the same.
After matching the first matmul, we learn that `inp_pat` matches with `x1`. On
the second matmul pattern + the binding (`lv1`), `m->Match(...)` would return
true , since it simply matches a matmul pattern with two wildcards. This
additional check tells us that `inp_pat` has already been matched to `x1`, so
the second matching where `inp_pat` would match to `x2` must be disallowed.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
