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

[masahi]

This is a fix for the issue described in #14417 (comment)

In graph pattern matching, backtracking behavior is implemented by undoing matchings in undo_stack. 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. This corner-case handling is missing in the current implementation.

The fix is to return undo_stack from recursive try_match calls when matching succeeds, and merge parent / child undo_stack with the undo_stack for the "current" node

@ganler @sunggg @cyx-6

[tvm-bot]

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}

[ganler]

LGTM.

[masahi]

Hey I'm happy to report that I have pattern-based rewriter for dataflow block working. Compared to the existing call node based matching & rewriting that requires a common post-dominator in a pattern, it lets us match a tree structure and replace leaf nodes or branches with new expression. The branch is here which I'll send after this one.

This can be immediately used for combining any number of multiple matmuls sharing the same LHS into one matmul. In Relay, we have a dedicated pass for that purpose ( CombineParallelDense ), but we can achieve the same thing via graph (tree) matching and rewrite. The same story for CombineParallelConv2D and CombineParallelBatchMatmul.

For example, in SD UNet we have many three parallel matmul for QKV projections. In addition, there are also highly non-obvious parallel matmuls consisting of 32 or 22 of them. Those patterns can all be matched and rewritten via the following generic pattern and rewriter. I'm very excited about this!
https://github.com/masahi/web-stable-diffusion/blob/unet-opt/test.py#L37-L68

[psrivas2]

Fantastic, thanks!