This is an automated email from the ASF dual-hosted git repository.
lunderberg pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 58e622b74d [Unity][Transform] Implement Relax function inlining
(#16194)
58e622b74d is described below
commit 58e622b74dd2da6138005b865e0e0b2d6bd4a548
Author: Eric Lunderberg <[email protected]>
AuthorDate: Sat Dec 9 12:25:13 2023 -0600
[Unity][Transform] Implement Relax function inlining (#16194)
* [Unity][Transform] Implement Relax function inlining
This commit adds the ability to inline one Relax function into
another. This is provided with two APIs, one which allows explicit
specification of the functions to be inlined, and one which inlines
every private Relax function in an `IRModule`. Function inlining with
explicit replacements is intended for use by an end-user when building
a model function (e.g. inlining of a rotary embedding), and the second
is intended for use as part of a generic optimization pipeline.
* lint fix
* Use DetectRecursion from relax/analysis.h
* Added test case for error on mutually-recursive functions
---
python/tvm/relax/expr.py | 5 +
python/tvm/relax/transform/__init__.py | 1 +
python/tvm/relax/transform/transform.py | 10 +
src/relax/transform/inline_functions.cc | 228 ++++++++++++
tests/python/relax/test_inline_functions.py | 404 +++++++++++++++++++++
.../test_transform_inline_private_functions.py | 105 ++++++
6 files changed, 753 insertions(+)
diff --git a/python/tvm/relax/expr.py b/python/tvm/relax/expr.py
index 71f23577e7..c9780bea7e 100644
--- a/python/tvm/relax/expr.py
+++ b/python/tvm/relax/expr.py
@@ -818,6 +818,11 @@ class Function(BaseFunc, Scriptable):
return _ffi_api.FunctionBindParams(self, binding_map) # type: ignore
+ def inline_functions(
+ self, function_map: Mapping[Union[str, tvm.ir.GlobalVar], "Function"]
+ ) -> "Function":
+ return _ffi_api.FunctionInlineFunctions(self, function_map) # type:
ignore
+
@tvm._ffi.register_object("relax.expr.ExternFunc")
class ExternFunc(BaseFunc):
diff --git a/python/tvm/relax/transform/__init__.py
b/python/tvm/relax/transform/__init__.py
index c3f037da5f..4e0b2a5fa7 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/__init__.py
@@ -42,6 +42,7 @@ from .transform import (
FuseTIR,
FusionPattern,
Gradient,
+ InlinePrivateFunctions,
KillAfterLastUse,
LambdaLift,
LegalizeOps,
diff --git a/python/tvm/relax/transform/transform.py
b/python/tvm/relax/transform/transform.py
index 0af89b7d9e..b543cac24e 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -588,6 +588,16 @@ def RemoveUnusedOutputs() -> tvm.ir.transform.Pass:
return _ffi_api.RemoveUnusedOutputs() # type: ignore
+def InlinePrivateFunctions() -> tvm.ir.transform.Pass:
+ """Inline all private relax functions
+
+ Returns
+ -------
+ ret: tvm.ir.transform.Pass
+ """
+ return _ffi_api.InlinePrivateFunctions() # type: ignore
+
+
def AnnotateTIROpPattern() -> tvm.ir.transform.Pass:
"""Annotate Op Pattern Kind for TIR functions
diff --git a/src/relax/transform/inline_functions.cc
b/src/relax/transform/inline_functions.cc
new file mode 100644
index 0000000000..457d6d9ead
--- /dev/null
+++ b/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.
+
+ Var param_var(func->params[i]->name_hint(),
args[i]->struct_info_.as<StructInfo>());
+ param_bindings.push_back(VarBinding(param_var, args[i]));
+ param_map.Set(func->params[i], param_var);
+ }
+
+ DataflowBlock binding_block(param_bindings);
+ Expr body = Bind(func, param_map).as<FunctionNode>()->body;
+
+ return SeqExpr({binding_block}, body);
+ }
+
+ const Map<Variant<String, GlobalVar>, Function>& replacements_;
+ support::OrderedSet<GlobalVar> inline_stack_;
+};
+} // namespace
+
+/*!
+ * \brief Bind params to function by using name
+ * \param func Relax function
+ * \param params params dict
+ * \return Function
+ */
+Function FunctionInlineFunctions(Function func,
+ const Map<Variant<String, GlobalVar>,
Function>& replacements) {
+ for (const auto& [key, func] : replacements) {
+ if (auto ptr = key.as<GlobalVarNode>()) {
+ CHECK(!replacements.count(ptr->name_hint))
+ << "ValueError: "
+ << "Map of functions to inline must be unambiguous. "
+ << "However, the map provided contains both the GlobalVar " << key
<< " and the string \'"
+ << ptr->name_hint << "'";
+ }
+ }
+
+ FunctionInliner mutator(replacements);
+ return Downcast<Function>(mutator(std::move(func)));
+}
+
+TVM_REGISTER_GLOBAL("relax.FunctionInlineFunctions").set_body_typed(FunctionInlineFunctions);
+
+namespace transform {
+
+Pass InlinePrivateFunctions() {
+ auto pass_func = [=](IRModule mod, PassContext pc) {
+ Map<Variant<String, GlobalVar>, Function> replacements;
+ for (const auto& [gvar, base_func] : mod->functions) {
+ if (auto opt = base_func.as<relax::Function>()) {
+ auto func = opt.value();
+ bool is_private =
!func->GetAttr<String>(tvm::attr::kGlobalSymbol).defined();
+ if (is_private) {
+ replacements.Set(gvar, func);
+ }
+ }
+ }
+
+ if (replacements.empty()) {
+ // Early bail-out if there are no private functions.
+ return mod;
+ }
+
+ for (const auto& recursive_set : DetectRecursion(mod)) {
+ for (const auto& recursive_func : recursive_set) {
+ replacements.erase(recursive_func);
+ }
+ }
+
+ if (replacements.empty()) {
+ // Early bail-out if all private functions are recursive.
+ return mod;
+ }
+
+ IRModule updates;
+ for (const auto& [gvar, base_func] : mod->functions) {
+ if (!replacements.count(gvar)) {
+ if (auto opt = base_func.as<relax::Function>()) {
+ auto func = FunctionInlineFunctions(opt.value(), replacements);
+ if (!base_func.same_as(func)) {
+ updates->Add(gvar, func);
+ }
+ }
+ }
+ }
+
+ auto write_ptr = mod.CopyOnWrite();
+ for (const auto& [key, func] : replacements) {
+ write_ptr->Remove(Downcast<GlobalVar>(key));
+ }
+ write_ptr->Update(updates);
+ return mod;
+ };
+ return tvm::transform::CreateModulePass(pass_func, 0,
"InlinePrivateFunctions", {});
+}
+
+TVM_REGISTER_GLOBAL("relax.transform.InlinePrivateFunctions")
+ .set_body_typed(InlinePrivateFunctions);
+
+} // namespace transform
+
+} // namespace relax
+} // namespace tvm
diff --git a/tests/python/relax/test_inline_functions.py
b/tests/python/relax/test_inline_functions.py
new file mode 100644
index 0000000000..c6cf4c1303
--- /dev/null
+++ b/tests/python/relax/test_inline_functions.py
@@ -0,0 +1,404 @@
+# 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"]})
+
+
+def test_error_when_inlining_mutually_recursive_functions():
+ """Inlining a recursive function call should raise an error"""
+
+ @I.ir_module
+ class Before:
+ @R.function(private=True)
+ def main():
+ B = Before.subroutine_a()
+ return B
+
+ @R.function(private=True)
+ def subroutine_a() -> 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_b()
+ else:
+ Out = R.const(0, "int64")
+
+ return Out
+
+ @R.function(private=True)
+ def subroutine_b() -> 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_a()
+ else:
+ Out = R.const(0, "int64")
+
+ return Out
+
+ with pytest.raises(Exception):
+ Before["main"].inline_functions(
+ {
+ "subroutine_a": Before["subroutine_a"],
+ "subroutine_b": Before["subroutine_b"],
+ }
+ )
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
diff --git a/tests/python/relax/test_transform_inline_private_functions.py
b/tests/python/relax/test_transform_inline_private_functions.py
new file mode 100644
index 0000000000..21321ebc0a
--- /dev/null
+++ b/tests/python/relax/test_transform_inline_private_functions.py
@@ -0,0 +1,105 @@
+# 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.script
+import tvm.testing
+from tvm import relax
+from tvm.script import ir as I
+from tvm.script import relax as R
+from tvm.script import tir as T
+
+
+def test_inline_simple():
+ """Simple case of inlining
+
+ Inlining applies to all private functions
+ """
+
+ @I.ir_module
+ class Before:
+ @R.function
+ 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
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def main(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 = tvm.relax.transform.InlinePrivateFunctions()(Before)
+ tvm.ir.assert_structural_equal(Expected, After)
+
+
+def test_skip_inline_of_recursive_functions():
+ """Recursively-defined functions
+
+ This behavior is deliberately different between the
+ `relax.transform.InlinePrivateFunctions` pass, and the
+ `relax.Function.inline_functions` utility.
+
+ For a user-facing utility, such as `func.inline_functions(...)`,
+ the functions to be inlined are specifically listed, and must not
+ be ignored. If it is unable to inline the user-requested
+ function, it should return an appropriate error.
+
+ For a generic utility to be used in optimization pipelines, the
+ framework is tasked with selecting the functions to be inlined,
+ and should avoid selecting any function that cannot be inlined.
+ This includes recursively-defined functions.
+ """
+
+ @I.ir_module
+ class Before:
+ @R.function
+ 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
+
+ Expected = Before
+
+ After = tvm.relax.transform.InlinePrivateFunctions()(Before)
+ tvm.ir.assert_structural_equal(Expected, After)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
