Function test_inline_multiple_instances

tests/python/relax/test_inline_functions.py:236–267

A subroutine may be inlined multiple times When inlining, SSA should still be respected.

()

def test_inline_multiple_instances():
    """A subroutine may be inlined multiple times

    When inlining, SSA should still be respected.
    """

    @I.ir_module
    class Before:
        @R.function(private=True)
        def main(A: R.Tensor):
            B = Before.subroutine(A)
            C = Before.subroutine(B)
            return C

        @R.function(private=True)
        def subroutine(A0: R.Tensor) -> R.Tensor:
            A1 = A0 * A0
            A2 = A1 + A1
            return A2

    @R.function(private=True)
    def expected(A: R.Tensor):
        # First call
        B = A * A
        C = B + B
        # Second call
        D = C * C
        E = D + D
        return E

    after = Before["main"].inline_functions({"subroutine": Before["subroutine"]})
    tvm.ir.assert_structural_equal(expected, after)


def test_inline_multiple_instances_with_distinct_static_shapes():