slyubomirsky commented on code in PR #16194: URL: https://github.com/apache/tvm/pull/16194#discussion_r1418121101
########## src/relax/transform/inline_functions.cc: ########## @@ -0,0 +1,228 @@ +/* + * 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. + */ + +#include <tvm/relax/analysis.h> +#include <tvm/relax/expr.h> +#include <tvm/relax/expr_functor.h> +#include <tvm/relax/transform.h> + +#include <utility> + +#include "../../support/ordered_set.h" +#include "utils.h" + +namespace tvm { +namespace relax { + +namespace { + +class FunctionInliner : public ExprMutator { + public: + explicit FunctionInliner(const Map<Variant<String, GlobalVar>, Function>& replacements) + : replacements_(replacements) {} + + using ExprMutator::VisitExpr_; + + Expr VisitExpr_(const FunctionNode* op) override { + auto node = ExprMutator::VisitExpr_(op); + if (node.get() != op) { + node = CanonicalizeBindings(node); + node = RemoveAllUnused(node); + } + return node; + } + + Expr VisitExpr_(const CallNode* op) override { + auto node = Downcast<Call>(ExprMutator::VisitExpr_(op)); + + if (auto opt = node->op.as<GlobalVar>()) { + auto gvar = opt.value(); + if (auto opt = GetFunction(gvar)) { + auto callee = opt.value(); + CHECK_EQ(callee->params.size(), node->args.size()) + << "Attempted to inline call to " << gvar << ", which accepts " << callee->params.size() + << " parameters. " + << "However, it was called with " << node->args.size() << " arguments in expression " + << node; + + Expr inlined = InlinedCall(callee, node->args); + + CHECK(!inline_stack_.count(gvar)) + << "Relax function inlining does not support recursive functions. " + << "However, recursive function " << gvar << " was requested to be inlined."; + + inline_stack_.insert(gvar); + inlined = VisitExpr(std::move(inlined)); + inline_stack_.erase(gvar); + + return inlined; + } + } + + return std::move(node); + } + + private: + Optional<Function> GetFunction(const GlobalVar& gvar) const { + if (auto opt = replacements_.Get(gvar)) { + return opt; + } else if (auto opt = replacements_.Get(gvar->name_hint)) { + return opt; + } else { + return NullOpt; + } + } + + Expr InlinedCall(Function func, const Array<Expr>& args) const { + // Ensures that the inlined instance does not have duplicate usage + // with other inlined copies, or with the original callee. + func = CopyWithNewVars(std::move(func)); + + Array<Binding> param_bindings; + + Map<Var, Expr> param_map; + for (size_t i = 0; i < args.size(); i++) { + // Option 1: Use tvm::relax::Bind to substitute arguments into + // the body. If the arguments contain DataflowVar instances, + // but the subroutine does not use DataflowBlock, this would + // result in invalid AST. + // + // Option 2: Define a VarBinding `param[i] = args[i]` for each + // parameter, then rely on CanonicalizeBindings to replace with + // DataflowVar where possible. This would solve the invalid use + // of DataflowVar, but wouldn't handle symbolic variables. If + // the subroutine has symbolic variables defined by its + // arguments, the VarBinding would leave them undefined. + // + // Option 3: Define a MatchCast `param[i] = args[i]` for each + // parameter, followed by CanonicalizeBindings. This is the + // first option that would result in well-formed AST, but it + // wouldn't be optimal. Symbolic variables would have two + // copies, one from the initial definition, and one + // from the MatchCast inlined portion. + // + // Option 4: Define a VarBinding `param[i] = args[i]`, with + // CanonicalizeBindings to handle conversion of Var to + // DataflowVar, and tvm::relax::Bind to handle substitution of + // symbolic variables. This would result in a well-formed Relax + // function, with no duplicate definitions of symbolic + // variables. + // + // This implementation uses Option 4. Review Comment: Very clever solution to avoid dealing with dataflow blocks. It's good that `CanonicalizeBindings` has become a single source of truth for dealing with these issues. ########## 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: What logic in the pass handles this replacement? The use of CanonicalizeBindings per "option 4"? ########## 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. + """ + + @I.ir_module + class Before: + @R.function(private=True) + def main(A: R.Tensor(["n", 16], "int32")) -> R.Tensor(["n", 32], "int32"): + B = A * A + C = Before.subroutine(B) + D = C + C + return D + + @R.function(private=True) + def subroutine(B: R.Tensor(["n", 16], "int32")) -> R.Tensor(["n", 32], "int32"): + C = R.concat([B, B], axis=1) + return C + + @R.function(private=True) + def expected(A: R.Tensor(["n", 16], "int32")) -> R.Tensor(["n", 32], "int32"): + B = A * A + C = R.concat([B, B], axis=1) + D = C + C + return D + + after = Before["main"].inline_functions({"subroutine": Before["subroutine"]}) + tvm.ir.assert_structural_equal(expected, after) + + +def test_subroutine_with_symbolic_vars_and_static_argument(): + """Inlined subroutines should use the caller's static shape + + Before inlining, the subroutine has symbolic variables, and the + caller have static shape. After inlining, no symbolic variables + should remain. + """ + + @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(["n", 16], "int32")) -> R.Tensor(["n", 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 + + after = Before["main"].inline_functions({"subroutine": Before["subroutine"]}) + tvm.ir.assert_structural_equal(expected, after) + + +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(): + """A subroutine may be inlined multiple times + + When inlining, each instance of the inlined function may have a + different value for the symbolic variables it uses. + """ + + @I.ir_module + class Before: + @R.function(private=True) + def main(A: R.Tensor([16, 16]), B: R.Tensor([32, 32])): + A_out: R.Tensor([16, 16]) = Before.subroutine(A) + B_out: R.Tensor([32, 32]) = Before.subroutine(B) + return (A_out, B_out) + + @R.function(private=True) + def subroutine(Input: R.Tensor(["n", "m"])) -> R.Tensor(["n", "m"]): + Output = Input + Input + return Output + + @R.function(private=True) + def expected(A: R.Tensor([16, 16]), B: R.Tensor([32, 32])): + A_out: R.Tensor([16, 16]) = A + A + B_out: R.Tensor([32, 32]) = B + B + return (A_out, B_out) + + after = Before["main"].inline_functions({"subroutine": Before["subroutine"]}) + tvm.ir.assert_structural_equal(expected, after) + + +def test_inline_nested_subroutine_calls(): + """A private function may itself require inlining""" + + @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 + D = Before.subroutine(B) + E = D + D + return E + + @R.function(private=True) + def subroutine(B: R.Tensor([16, 16], "int32")) -> R.Tensor([16, 32], "int32"): + C = R.concat([B, B], axis=1) + D = Before.subsubroutine(C) + return D + + @R.function(private=True) + def subsubroutine(C: R.Tensor([16, 32], "int32")) -> R.Tensor([16, 32], "int32"): + D = C * C * C + return D + + @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 * C + E = D + D + return E + + after = Before["main"].inline_functions( + { + "subroutine": Before["subroutine"], + "subsubroutine": Before["subsubroutine"], + } + ) + tvm.ir.assert_structural_equal(expected, after) + + +def test_error_when_inlining_recursive_function(): + """Inlining a recursive function call should raise an error""" + + @I.ir_module + class Before: + @R.function(private=True) + def main(): + B = Before.subroutine() + return B + + @R.function(private=True) + def subroutine() -> R.Tensor([], "int64"): + R.func_attr({"relax.force_pure": True}) + cond = R.call_packed("dummy_function", sinfo_args=R.Tensor([], "bool")) + if cond: + Out = Before.subroutine() + else: + Out = R.const(0, "int64") + + return Out + + with pytest.raises(Exception): + Before["main"].inline_functions({"subroutine": Before["subroutine"]}) + Review Comment: A couple of cases that are also worth testing: 1. Inlining within an inlined function 2. Making sure that mutually recursive functions can't be inlined either (since if you can inline within an inlined function, mutual recursion would also cause infinite inlining) -- 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]
