[Relax] Allow composition of DFPattern replacements

[Lunderberg]

The rewrite_call function accepts a DFPattern, and a function to
rewrite expressions matching that pattern. Often, the rewriting
function will perform additional validation that cannot be expressed
within the DFPattern itself. If this additional validation fails,
the rewriter function will return the matched expression unmodified.

Prior to this commit, an OrPattern that matches on the first branch,
but whose rewriter function does not apply a modification, would
prevent the second branch from being checked. This commit updates the
ExprPatternRewriter to check both branches of a OrPattern, if the
rewriter function of the first branch does not modify the result.

[Lunderberg]

This PR is currently marked as a draft, because it depends on the refactor done in #16730.

[slyubomirsky]

The changes seem easy to understand given the PR this is built on and I think the new behavior is intuitive as well: If the first branch of the OR does not result in a rewrite, check the second. In principle, I think this should be noted in whatever top-level documentation there is for the pattern matching grammar.

[slyubomirsky]

Looks like a typo. I assume it's supposed to be "handle matching," correct?

[Lunderberg]

Whoops, this was a typo. It should be "handle matching", and I've updated the PR with the correction.

[Lunderberg]

The pre-requisite PR #16730 has landed, so this PR is now rebased on top of main and marked as ready. Thank you @slyubomirsky for the review, and so I think it's just waiting on CI now.

[MasterJH5574]

Hi @Lunderberg, it seems that this PR introduces new regression. The original issue came from CI in MLC-LLM https://ci.mlc.ai/blue/organizations/jenkins/mlc-llm/detail/PR-2068/1/pipeline/. And I spent some time adapting your test test_backtrack_if_rewriter_returns_no_op and reducing it to the minimal I can get.

Running the test below

def test_backtrack_if_rewriter_returns_no_op():
    pat_match_no_rewrite = is_op("relax.add")(wildcard(), wildcard())

    pat_arg = wildcard()
    pat_zeros = is_op("relax.zeros")(wildcard())
    pat_add = is_op("relax.add")(pat_arg, pat_zeros)

    pat = pat_match_no_rewrite | pat_add

    def rewriter(expr, matches):
        print(f"matching {pat} to {matches[pat]}")
        assert isinstance(matches[pat], rx.Call)
        if pat_match_no_rewrite in matches:
            return expr   ## <<== return the expr
        elif pat_add in matches:
            return expr   ## <<== also return the expr
        else:
            raise RuntimeError("Pattern matched, but neither branch matched")

    @R.function(private=True)
    def before():
        with R.dataflow():
            A = R.ones([64, 128], "int32")
            B = R.zeros([64, 128], "int32")
            C = R.add(A, B)

            R.output(C)
        return C

    after = rewrite_call(pat, rewriter, before)

will yield the following output:

matching OrPattern(Op(relax.add)(*, *) | Op(relax.add)(*, Op(relax.zeros)(*))) to R.add(A, B)
matching OrPattern(Op(relax.add)(*, *) | Op(relax.add)(*, Op(relax.zeros)(*))) to R.add(A, B)
matching OrPattern(Op(relax.add)(*, *) | Op(relax.add)(*, Op(relax.zeros)(*))) to C
[22:36:29] /home/ruihang/Workspace/tvm/src/relax/ir/block_builder.cc:65: Warning: BlockBuilder destroyed with remaining blocks!
Traceback (most recent call last):
  File "/home/ruihang/Workspace/tvm/tests/python/relax/test_dataflow_pattern.py", line 1949, in <module>
    test_backtrack_if_rewriter_returns_no_op()
  File "/home/ruihang/Workspace/tvm/tests/python/relax/test_dataflow_pattern.py", line 1945, in test_backtrack_if_rewriter_returns_no_op
    after = rewrite_call(pat, rewriter, before)
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ruihang/Workspace/tvm/python/tvm/relax/dpl/rewrite.py", line 59, in rewrite_call
    return ffi.rewrite_call(pattern, rewriter, func)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "tvm/_ffi/_cython/./packed_func.pxi", line 332, in tvm._ffi._cy3.core.PackedFuncBase.__call__
  File "tvm/_ffi/_cython/./packed_func.pxi", line 263, in tvm._ffi._cy3.core.FuncCall
  File "tvm/_ffi/_cython/./packed_func.pxi", line 252, in tvm._ffi._cy3.core.FuncCall3
  File "tvm/_ffi/_cython/./base.pxi", line 182, in tvm._ffi._cy3.core.CHECK_CALL
  File "/home/ruihang/Workspace/tvm/python/tvm/_ffi/base.py", line 481, in raise_last_ffi_error
    raise py_err
  File "tvm/_ffi/_cython/./packed_func.pxi", line 56, in tvm._ffi._cy3.core.tvm_callback
  File "/home/ruihang/Workspace/tvm/tests/python/relax/test_dataflow_pattern.py", line 1921, in rewriter
    assert isinstance(matches[pat], rx.Call)
AssertionError

As you can see in the error message, with this PR, the pattern matcher mapped an add to a single Var. This breaks our assumption that what matched is always an add. Would you mind helping dig and fix this issue? Thanks in ahead.

[MasterJH5574]

The particular rewrite_call that fails in MLC-LLM CI is this one https://github.com/mlc-ai/mlc-llm/blob/main/python/mlc_llm/compiler_pass/fuse_ft_dequantize_matmul_epilogue.py#L213-L225

[Lunderberg]

Ooh, that's a weird one, and thank you for the test case. I can reproduce it on my side, and will look in more detail tomorrow. For now, the first thing I'm noticing is that the failure occurs on the third time that rewriter is called, where I'd expect it to only be called twice (once for each branch).

matches[pat_match_no_rewrite] is filled during the first call, and matches[pat_add] is filled during the second call. During the third call, which fails the assert, matches[pat_match_no_rewrite] is filled again, which seems rather odd.

[Lunderberg]

This is an interesting one.

1. The DFPatternMatcher will unwrap variables when attempting to find a match (link), as part of ExtractMatchedExpr.
2. The ExprPatternMatcher will save the value of each variable binding is saved in order to (link).
3. The ExprPatternMatcher checks for a match in the override for Expr VisitExpr(const Expr&).

So the call to rewriter with the assert fails is made when ExprPatternMatcher is checking the variable usage in the body of the SeqExpr. This gets unwrapped by DFPatternMatcher, and identified as an expression whose value is equal to the match itself.

This potential bug existed all the way back to #15578. However, until #16732, the relax::Var body of a SeqExpr would never be the first match encountered. Any time a relax::Var would have matched, the expression to which it was bound would have matched first, hiding the later match.

I think the resolution is to update condition (1). Where currently, we unwrap all variable bindings as part of a match, we should instead unwrap all variable bindings except for the top-level of the match. This could be most easily implemented by having the ExprPatternMatcher::VisitExpr skip the call to TryRewrite for relax::Var. I'll try this out and see if it resolves the issue.

Edit: Or even better, this should be handled when filling in the OrPattern branches for the match. I had this step to match sure that the matches[pat] was still populated, even when the OrPattern was unwrapped. However, where previously this would be populated by ExtractMatchedExpr, and would contain the expression after TryGetValOfVar, I incorrectly filled it with the expression before applying TryGetValOfVar.

[Lunderberg]

@MasterJH5574 I have a fix implemented in #16828. Can you try it on your side and verify that it resolves the issue in the end-to-end flow?

[MasterJH5574]

@Lunderberg Thanks for the swift fix! It works for e2e compilation now.