slyubomirsky commented on code in PR #16194: URL: https://github.com/apache/tvm/pull/16194#discussion_r1419798304
########## tests/python/relax/test_inline_functions.py: ########## @@ -0,0 +1,363 @@ +# Licensed to the Apache Software Foundation (ASF) under one +# or more contributor license agreements. See the NOTICE file +# distributed with this work for additional information +# regarding copyright ownership. The ASF licenses this file +# to you under the Apache License, Version 2.0 (the +# "License"); you may not use this file except in compliance +# with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, +# software distributed under the License is distributed on an +# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +# KIND, either express or implied. See the License for the +# specific language governing permissions and limitations +# under the License. + +import pytest + +import tvm +import tvm.testing +from tvm.script import relax as R, ir as I, tir as T + + [email protected]("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 == tvm.ir.GlobalVar: + key = gvar + elif key_type == 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(): + """Raise an error on ambiguous inputs + + For convenience, the function being replaced can be specified + either as a string, or as a GlobalVar. However, all replacements + must be unambiguous. + """ + + @R.function + def func(): + return R.tuple() + + gvar = tvm.ir.GlobalVar("name") + + with pytest.raises(ValueError): + func.inline_functions({gvar: func, "name": func}) + + +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(): + """Function inlining may not produce invalid Relax IR + + A subroutine call may appear within a DataflowBlock, even if the + subroutine does not itself use a DataflowBlock. In this case, to + avoid inserting a non-dataflow block in the middle of a set of + dataflow bindings, the DataflowBlock in the caller must be split + up. + """ + + @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"): + 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"): + # DataflowBlock before subroutine + with R.dataflow(): + B = A * A + R.output(B) + + # BindingBlock from the inlined subroutine. Because B is used + # here, outside of a DataflowBlock, this requires it to be + # updated from a DataflowVar to a normal Var. + C = R.concat([B, B], axis=1) + + # Resuming the DataflowBlock after the inlined subroutine + with R.dataflow(): + 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_subroutine_with_symbolic_vars(): + """Inlined subroutines should use the caller's symbolic variables + + Before inlining, the subroutine and the caller have distinct + `tir::Var` for each symbolic variables. After inlining, only the + caller's `tir::Var` symbolic variables should remain. + """ Review Comment: Ah, very nice. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected]
