Function test_inline_simple

tests/python/relax/test_inline_functions.py:29–67

Simple case of inlining Inlining can be done either by providing a string name or a GlobalVar.

(key_type)

@pytest.mark.parametrize("key_type", [tvm.ir.GlobalVar, str])
def test_inline_simple(key_type):
    """Simple case of inlining

    Inlining can be done either by providing a string name or a
    GlobalVar.
    """

    @I.ir_module
    class Before:
        @R.function(private=True)
        def main(A: R.Tensor([16, 16], "int32")) -> R.Tensor([16, 32], "int32"):
            B = A * A
            C = Before.subroutine(B)
            D = C + C
            return D

        @R.function(private=True)
        def subroutine(B: R.Tensor([16, 16], "int32")) -> R.Tensor([16, 32], "int32"):
            C = R.concat([B, B], axis=1)
            return C

    @R.function(private=True)
    def expected(A: R.Tensor([16, 16], "int32")) -> R.Tensor([16, 32], "int32"):
        B = A * A
        C = R.concat([B, B], axis=1)
        D = C + C
        return D

    gvar = Before.get_global_var("subroutine")
    if key_type is tvm.ir.GlobalVar:
        key = gvar
    elif key_type is str:
        key = gvar.name_hint
    else:
        raise TypeError(f"Unknown key_type: {key_type}")

    after = Before["main"].inline_functions({key: Before[gvar]})

    tvm.ir.assert_structural_equal(expected, after)


def test_ambiguous_function_name():