[Unity] Remove non-deterministic behavior from graph pattern matching #14417

masahi · 2023-03-28T21:54:48Z

Graph pattern matching uses std::map and std::set keyed on pointers, which leads to non-deterministic matching result. This PR removed all non-determinism.

For example, consider matching against QKV projection before attention. Given a sequence of bindings,

with R.dataflow():
    lv0 = R.matmul(x, w0)
    lv1 = R.matmul(x, w1)
    lv2 = R.matmul(x, w2)
    ...

and patterns,

with PatternContext() as ctx:
    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)

intuitively I expect the following matching result:

Q_weight_pat - w0
K_weight_pat - w1
V_weight_pat - w2

But since three matmul patterns are independent of each other, each pattern can match any of three matmuls in the bindings. And since graph pattern matching processes patterns and variables in the bindings in a non-deterministic order, matching result becomes non deterministic as well.

To make graph pattern matching deterministic, I introduced an implicit ordering constraint on the patterns: Patterns are matched in the order that they are declared in the pattern context. So in the example above, we make sure to start pattern matching from the first matmul pattern. Likewise, the variables in the bindings, which are the candidates for matching, are also considered for matching in the order they appear in the bindings. So lv0 is matched first, followed by lv1 etc.

@ganler @sunggg

tvm-bot · 2023-03-28T21:54:53Z

Thanks for contributing to TVM! Please refer to the contributing guidelines https://tvm.apache.org/docs/contribute/ for useful information and tips. Please request code reviews from Reviewers by @-ing them in a comment.

cc @quic-sanirudh _{See #10317 for details}

_{Generated by tvm-bot}

masahi · 2023-03-28T22:03:56Z

tests/python/relax/test_dataflow_pattern.py

+        matmul2 = is_op("relax.matmul")(inp_pat, K_weight_pat)
+        matmul3 = is_op("relax.matmul")(inp_pat, V_weight_pat)
+
+        # TODO(masahi): Automate addition of used_by constraints during is_op


I have a follow-up PR to address this, which removes all used_by stuff below @ganler

masahi · 2023-03-28T22:06:01Z

src/relax/ir/dataflow_matcher.cc


      const auto& uses = def2use.at(rparent->ptr);
-      // skip if `rparent` is not used by `r`.
-      if (uses.cend() == uses.find(r->ptr)) continue;


This has been removed since it seems to be always false.

masahi · 2023-03-28T22:08:10Z

@ganler Is start_hint stuff necessary? I want to remove it, since matching proceeds in a fixed order now. Some tests depend on the use of start_hint.

ganler · 2023-03-28T23:13:06Z

@masahi Thanks. Start hint means the start point of graph matching. It is useful when you have multiple matchable patterns and want to explicitly match one of them which are close to / subsuming certain node.

ganler · 2023-03-29T07:14:38Z

(Traveling) if you can wait I can file a review after 10 hours but definitely within 24 hours. :). Hope it won't get you blocked.

sunggg

Thank you for the improvement!
Removing non-determinism should be a right way to go.
Overall, it looks good to me. We can merge this when @ganler can confirm with the implementation details.

ganler

The changes make sense to me. I don't see any (apparent) logic flaws. Thanks for optimizing the code style and modernizing the code to C++17 standard!

ganler · 2023-03-29T20:10:36Z

src/relax/ir/dataflow_matcher.cc


  const auto commit = [&undo_stack](PNode* p, RNode* r) {
    // match with each other.
+    // TODO(ganler, masahi): Why commit on the same p-r pair happens more than once?


I guess in either way the eventual results are the same. But yeah doing a pre-check could be faster (avoid the overhead of undo_stack.emplace(p, r)).

Here my concern is not about performance. If the same pair is committed more than once, I think there is something odd about the matching algorithm.

Later I'll try to improve the matching algorithm implementation. In particular, I want to remove try_match loop on parents,

tvm/src/relax/ir/dataflow_matcher.cc

Lines 577 to 607 in 278bc9b

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 auto& uses = def2use.at(rparent->ptr);

// 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;

// try all parent R nodes that are not matched yet.

// as long as ppattern can match one node.

if (!pparent->matched && try_match(pparent, rparent, m, def2use, use2def)) {

. This bidirectional matching has been a source of confusion to me (and debugging is very hard) - Since the constraint graph is assumed to be a DAG, I think we can start matching from the root nodes in a topo-sorted order, and matching should be able to proceed purely in one direction.

I see. Because the constraint or relation between nodes may or may not be single-way so in the beginning I made it bidirectional such that the pattern can be matched as long as you can let any node of the matched subgraph be start hint.

For example, for A->B pattern, you can start matching from either A or B (forward or backward).

Specifying certain matching order definitely makes the code logic and debugging easier but I am afraid it also cuts the flexibility in some way. Not sure if it is worth to have some flexibility or if we can keep them both.

Do you have an example in test_dataflow_pattern.py that requires such flexibility? Similarly to how I don't understand the need for start_hint, I don't understand why we might want to start matching from B in A -> B.

I'm also going to look at how similar projects like MLIR, OpenVINO etc implements general graph pattern matching.

Sorry, I need some time to illustrate but my flight is taking off now. I will get you back probably after 8 hours. 🥲

ok :) I'll share my thought in the meantime.

Looking at an example in https://github.com/apache/tvm/blob/unity/tests/python/relax/test_dataflow_pattern.py#L540-L545, I cannot imagine how one would use start_hint in practice. Here, you wrote the input mod by hand, so you know dfb.bindings[0].var is associated to the "left" branch of the CBRx2. But in general, we don't have such information, especially in e2e scenarios for real world models.

Even if there was a good use case for it, I claim that it doesn't justify making API and the implementation more complicated (two additional params + try_match loop on parants). If one has such advanced knowledge of model structure and variables, they may as well have a different way to match & extract such subgraph.

Thanks for bringing it up. When designing the usages it was my bad that I did not know enough use cases and I made it "more capable" by compromising simplicity. masa you have more professional experience with how patterns look in practice so I strongly agree that we should cut off the implementation/complexity for such long-tail use cases.

BTW, also to explain why I was using the code style of CONSTANT logic_op VAR. This is just a practice to avoid writing VAR = CONSTANT which causes silent logic errors. But yeah it is nice to normalize the code style.

masahi · 2023-03-30T10:16:04Z

Just hit an interesting bug in the graph matcher.

The following fake QKV projection mod and pattern should not match. Without this PR, I get segfault. With this PR, it incorrectly matches with the result
{*: x2, *: w1, Op(relax.matmul)(*, *): lv1, *: w2, Op(relax.matmul)(*, *): lv2, Op(relax.matmul)(*, *): lv0, *: w0}

@R.function
def main(...) -> R.Tensor:
    with R.dataflow():
        lv0 = R.matmul(x1, w0) # use different lhs
        lv1 = R.matmul(x2, w1)
        lv2 = R.matmul(x2, w2)
        ...

with PatternContext() as ctx:
    inp_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)

I'll work on a fix in a separate PR.

UPDATE: Found a bug in how undo_stack is cleaned up after match failure,

tvm/src/relax/ir/dataflow_matcher.cc

Lines 565 to 571 in 278bc9b

    
           while (!undo_stack.empty()) { 
        
             auto& top = undo_stack.top(); 
        
             top.first->matched = nullptr; 
        
             top.second->matched = nullptr; 
        
             undo_stack.pop(); 
        
           } 
        
           return false;

. When an intermediate fails to match, undo_stacks for parent and child nodes that have matched tentatively also need to be cleaned up as well.

So the fix is to return undo_stack from try_match after match succeeds, and merge parent / child undo_stack with the undo_stack for the "current" node, see cbf14aa

…#14417) * Remove all non-determinsm from graph matching * add test * typo * try fixing compile error for gcc * more style update * suppress compile warning

Remove all non-determinsm from graph matching

0e09ddc

add test

e9a116a

masahi force-pushed the graph-match-disamb branch from 481f10b to e9a116a Compare March 28, 2023 22:02

masahi commented Mar 28, 2023

View reviewed changes

typo

ccd23dd

masahi commented Mar 28, 2023

View reviewed changes

try fixing compile error for gcc

7174cbc

more style update

6032281

sunggg reviewed Mar 29, 2023

View reviewed changes

suppress compile warning

278bc9b

ganler approved these changes Mar 29, 2023

View reviewed changes

psrivas2 approved these changes Mar 30, 2023

View reviewed changes

tqchen approved these changes Mar 30, 2023

View reviewed changes

vinx13 merged commit 0695cfe into apache:unity Mar 30, 2023

masahi mentioned this pull request Mar 30, 2023

[Unity][Graph matching] Clean up undo stack for parent and child nodes properly #14440

Merged

masahi mentioned this pull request Apr 4, 2023

[Unity][Graph matching] Improved matching algorithm and implementation #14501

Merged

	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 auto& uses = def2use.at(rparent->ptr);

	// 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;

	// try all parent R nodes that are not matched yet.
	// as long as ppattern can match one node.
	if (!pparent->matched && try_match(pparent, rparent, m, def2use, use2def)) {

[Unity] Remove non-deterministic behavior from graph pattern matching #14417

[Unity] Remove non-deterministic behavior from graph pattern matching #14417

Uh oh!

Conversation

masahi commented Mar 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

tvm-bot commented Mar 28, 2023

Uh oh!

masahi Mar 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

masahi Mar 28, 2023

Choose a reason for hiding this comment

Uh oh!

masahi commented Mar 28, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

ganler commented Mar 28, 2023

Uh oh!

ganler commented Mar 29, 2023

Uh oh!

sunggg left a comment • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ganler left a comment

Choose a reason for hiding this comment

Uh oh!

ganler Mar 29, 2023

Choose a reason for hiding this comment

Uh oh!

masahi Mar 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ganler Mar 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

masahi Mar 29, 2023

Choose a reason for hiding this comment

Uh oh!

masahi Mar 29, 2023

Choose a reason for hiding this comment

Uh oh!

ganler Mar 29, 2023

Choose a reason for hiding this comment

Uh oh!

masahi Mar 29, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

ganler Mar 30, 2023

Choose a reason for hiding this comment

Uh oh!

ganler Mar 30, 2023

Choose a reason for hiding this comment

Uh oh!

masahi commented Mar 30, 2023 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

7 participants

masahi commented Mar 28, 2023 •

edited

Loading

masahi Mar 28, 2023 •

edited

Loading

masahi commented Mar 28, 2023 •

edited

Loading

sunggg left a comment •

edited

Loading

masahi Mar 29, 2023 •

edited

Loading

ganler Mar 29, 2023 •

edited

Loading

masahi Mar 29, 2023 •

edited

Loading

masahi commented Mar 30, 2023 •

edited

Loading