Function test_recursive_rewrite_rules

tests/python/relax/test_dataflow_rewriter.py:328–380

Rewrite rules are applied until convergence In this test, both the `RewriteAdd` and `RewriteMultiply` patterns must be applied in order to produce the expected output. However, the `RewriteMultiply` pattern relies on the expression produced by the `RewriteAdd` pass.

()

def test_recursive_rewrite_rules():
    """Rewrite rules are applied until convergence

    In this test, both the `RewriteAdd` and `RewriteMultiply` patterns
    must be applied in order to produce the expected output.  However,
    the `RewriteMultiply` pattern relies on the expression produced by
    the `RewriteAdd` pass.

    """

    @R.rewriter
    class RewriteAdd:
        @R.function
        def pattern(A: R.Tensor([16], "float32")):
            return A + A

        @R.function
        def replacement(A: R.Tensor([16], "float32")):
            return A * R.const(2.0, "float32")

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

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

    @R.function(private=True)
    def before(A: R.Tensor([16], "float32")):
        B = A + A
        return B

    @R.function(private=True)
    def expected(A: R.Tensor([16], "float32")):
        B = R.call_pure_packed(
            "my_optimized_mul_impl",
            A,
            R.const(2.0, "float32"),
            sinfo_args=R.Tensor([16], "float32"),
        )
        return B

    rewriter = RewriteAdd | RewriteMultiply

    after = rewriter(before)
    tvm.ir.assert_structural_equal(expected, after)


def test_rewrite_of_arbitrary_dtype():