Function test_inline_dataflow_block

tests/python/relax/test_inline_functions.py:88–119

Functions may be inlined within a dataflow block

()

def test_inline_dataflow_block():
    """Functions may be inlined within a dataflow block"""

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

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

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

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


def test_inline_non_dataflow_block_into_dataflow_block():