[Relax] Implement Rewriter class for pattern-rewrite

[Lunderberg]

Prior to this commit, the pattern to be matched and the rewrite to be performed were provided as separate arguments to rewrite_call. This commit introduces a new class ExprRewriter, which contains both parts. This abstraction makes it easier to combine multiple independent rewrite rules, then apply them all in a IRModule pass.

This implementation also allows pattern rewrite rule to be written using TVMScript, by representing both the pattern to be matched, and the replacement to be applied as functions. For example:

@R.rewriter
class Rewriter:
    @R.function
    def pattern(
        A: R.Tensor([16], "float32"),
        B: R.Tensor([16], "float32"),
    ):
        C = R.add(A, B)
        return C

    @R.function
    def replacement(
        A: R.Tensor([16], "float32"),
        B: R.Tensor([16], "float32"),
    ):
        C = R.call_pure_packed(
            "my_optimized_add_impl", A, B,
            sinfo_args=R.Tensor([16], "float32"),
        )
        return C

[Lunderberg]

@sunggg @csullivan

[sunggg]

Hi, @Lunderberg! Thank you for the cool addition!
I haven't taken a deeper look at the detailed mechanics yet, but I love how this PR allows us to manage the pattern and rewrite conveniently with better readability :)

My main question for you is if this PR is aiming the full support for the existing pattern language.
It is okay if there are some rough edges that we don't currently support, but I would like to clearly document them somewhere to avoid the potential misusage.

Here are some of the cases I wonder how these will be supported with your PR:

• ~ not pattern: https://github.com/apache/tvm/blob/main/tests/python/relax/test_dataflow_pattern.py#L222
• If we want to match with any placeholder (wildcard()), which could be tensor, constant, etc., how should we express this?
• Can we match with any number of function arguments?

[Lunderberg]

My main question for you is if this PR is aiming the full support for the existing pattern language. It is okay if there are some rough edges that we don't currently support, but I would like to clearly document them somewhere to avoid the potential misusage.

I've been aiming to support as much of the existing pattern language as possible, but not the entire syntax. There are some portions of the pattern language that don't have clear Relax equivalents, and those wouldn't be supported.

Here are some of the cases I wonder how these will be supported with your PR:

* `~` not pattern: https://github.com/apache/tvm/blob/main/tests/python/relax/test_dataflow_pattern.py#L222

This is probably the biggest one that doesn't have a relax-level equivalent. The main usage would be to avoid aliasing in a pattern-match. (e.g. An implementation that is valid so long as two wildcard patterns do not refer to the same tensor.)

* If we want to match with any placeholder (`wildcard()`), which could be tensor, constant, etc., how should we express this?

If the pattern has an argument var_name: R.Object, then it would match against any Relax value. This is converted to StructInfoPattern(WildcardPattern(), ObjectStructInfo()) when producing the pattern rule. When the StructInfoPattern validates a match, it does so using StructInfoBaseChecker. Since ObjectStructInfo() is the base of all other relax types, this validation always passes.

There are some edge cases where this doesn't work, mostly arising from struct inference applied within the pattern or replacement. For example, suppose the pattern has wildcard A: R.Object and returns A + A. This would raise an error, because Relax requires binary operations to have either R.Prim or R.Tensor struct info. Resolving this edge case would require re-visiting this discussion on whether the struct inference should throw an error for provably-incorrect usage, or should throw an error for not-provably-correct usage.

* Can we match with any number of function arguments?

We can have arbitrarily many function arguments for the pattern/replacement. However, each pattern must have a fixed number of arguments.

[sunggg]

Thanks, @Lunderberg for the clarification! I think it is absolutely okay to have some rough edges that we cannot support yet, but I hope we clearly document in the code so that people won't get confused.

[Lunderberg]

I think it is absolutely okay to have some rough edges that we cannot support yet, but I hope we clearly document in the code so that people won't get confused.

Agreed. Any unsupported functionality should raise an error during the conversion, which (hopefully) will make it as clear as possible. I'm a big fan of error-checking as early as possible, especially when the input is relatively unstructured.

[sunggg]

Thank you @Lunderberg for addressing my comments and pushing this direction! :)