This is an automated email from the ASF dual-hosted git repository.
syfeng 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 2dcb8716e8 [Unity][BlockBuilder] Depracate `BlockBuilder.get()` and
change it to `BlockBuilder.finalize()` (#16090)
2dcb8716e8 is described below
commit 2dcb8716e8f391e3bb395f6942bdc5a422adbab3
Author: Yixin Dong <[email protected]>
AuthorDate: Mon Nov 27 18:41:46 2023 -0800
[Unity][BlockBuilder] Depracate `BlockBuilder.get()` and change it to
`BlockBuilder.finalize()` (#16090)
* 1108
* fix ci
* 1109
* finished
* fix ci
---
include/tvm/relax/block_builder.h | 9 ++
include/tvm/relax/transform.h | 9 ++
python/tvm/relax/block_builder.py | 46 +++++----
python/tvm/relax/ir/instrument.py | 2 +-
python/tvm/relax/transform/__init__.py | 1 +
python/tvm/relax/transform/transform.py | 15 +++
src/relax/analysis/well_formed.cc | 5 +-
src/relax/ir/block_builder.cc | 6 ++
src/relax/transform/normalize.cc | 103 +++++++++++++++++++++
tests/python/relax/test_blockbuilder_core.py | 103 +++++++++++++++------
.../relax/test_transform_normalize_global_var.py | 98 ++++++++++++++++++++
11 files changed, 350 insertions(+), 47 deletions(-)
diff --git a/include/tvm/relax/block_builder.h
b/include/tvm/relax/block_builder.h
index 9b339babb4..4272b3f75e 100644
--- a/include/tvm/relax/block_builder.h
+++ b/include/tvm/relax/block_builder.h
@@ -82,6 +82,15 @@ class BlockBuilderNode : public Object {
*/
virtual IRModule GetContextIRModule() const = 0;
+ /*!
+ * \brief Finalize the building process and return the result IRModule.
Possibly rename
+ * GlobalVars in the IRModule to ensure name uniqueness and the invariant:
+ * every public function has the same name as its "global_symbol" attribute.
+ *
+ * \return The IRModule in this BlockBuilder.
+ */
+ virtual IRModule Finalize() = 0;
+
/*!
* \brief Add a Relax function or a TIR PrimFunc to internal context module.
* \param func The function to be added.
diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h
index f99d74b7aa..ae6d4059a6 100644
--- a/include/tvm/relax/transform.h
+++ b/include/tvm/relax/transform.h
@@ -154,6 +154,15 @@ TVM_DLL Pass AttachGlobalSymbol();
*/
TVM_DLL Pass Normalize();
+/*!
+ * \brief Possibly rename the GlobalVar in an IRModule to ensure these
properties:
+ * 1. (Invariant) First ensure every public function has the same name as its
"global_symbol"
+ * attribute;
+ * 2. To ensure 1., we may need to rename private functions with conflicting
names;
+ * 3. Finally, the name of every GlobalVar is unique in the IRModule.
+ */
+TVM_DLL Pass NormalizeGlobalVar();
+
/*!
* \brief Simplify a Relax module by folding var bindings and match shape
nodes,
* as well as tuple indices.
diff --git a/python/tvm/relax/block_builder.py
b/python/tvm/relax/block_builder.py
index 5bf36ce330..ab807203a3 100644
--- a/python/tvm/relax/block_builder.py
+++ b/python/tvm/relax/block_builder.py
@@ -14,26 +14,22 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
-# pylint: disable=no-else-return, invalid-name, unused-argument
+# pylint: disable=no-else-return, invalid-name, unused-argument,
import-outside-toplevel
"""Developer API of constructing Relax AST."""
-from typing import Dict, List, Optional, Union, Any, Callable, Sequence
+from typing import Any, Callable, Dict, List, Optional, Sequence, Union
+
+import tvm
+from tvm import relax as rx
+from tvm import tir
+from tvm.ir.base import deprecated
from tvm.ir.module import IRModule
from tvm.runtime import Object
-from tvm import relax as rx, tir
-import tvm
-from .expr import (
- Expr,
- Var,
- GlobalVar,
- BindingBlock,
- Tuple,
- BaseFunc,
- Binding,
-)
-from .struct_info import StructInfo
-from .op.base import call_tir, call_tir_with_grad
+
from . import _ffi_api
+from .expr import BaseFunc, Binding, BindingBlock, Expr, GlobalVar, Tuple, Var
+from .op.base import call_tir, call_tir_with_grad
+from .struct_info import StructInfo
from .utils import gen_call_tir_inputs
@@ -658,6 +654,7 @@ class BlockBuilder(Object):
"""
return _ffi_api.BlockBuilderNormalize(self, expr) # type: ignore
+ @deprecated("tvm.relax.BlockBuilder.get",
"tvm.relax.BlockBuilder.finalize")
def get(self) -> tvm.IRModule:
"""Return the IRModule being built.
@@ -666,7 +663,24 @@ class BlockBuilder(Object):
ret : tvm.IRModule
An IRModule with Relax and TIR functions being built.
"""
- return _ffi_api.BlockBuilderGetContextIRModule(self) # type: ignore
+ return self.finalize()
+
+ def finalize(self) -> tvm.IRModule:
+ """Finalize the building process and return the result IRModule.
+
+ Possibly rename GlobalVars in the IRModule to ensure name uniqueness
and the invariant:
+ every public function has the same name as its "global_symbol"
attribute.
+
+ Note this call may invalidate global vars previously returned by this
builder
+ (see tvm.relax.transform.NormalizeGlobalVar), so it can only be called
once at the end of
+ the building process.
+
+ Returns
+ -------
+ ret : tvm.IRModule
+ An IRModule with Relax and TIR functions being built.
+ """
+ return _ffi_api.BlockBuilderFinalize(self) # type: ignore
def get_unique_name(self, name_prefix: str) -> str:
"""Generate a unique name with a specified prefix.
diff --git a/python/tvm/relax/ir/instrument.py
b/python/tvm/relax/ir/instrument.py
index 1ecd87fe1b..281a8af511 100644
--- a/python/tvm/relax/ir/instrument.py
+++ b/python/tvm/relax/ir/instrument.py
@@ -39,7 +39,7 @@ class WellFormedInstrument:
"""
def __init__(self, check_struct_info: bool = True,
validate_before_transform: bool = True):
- self.skip_pass_name = ["Normalize", "ResolveGlobals"]
+ self.skip_pass_name = ["Normalize", "NormalizeGlobalVar",
"ResolveGlobals"]
self.check_struct_info = check_struct_info
self.validate_before_transform = validate_before_transform
diff --git a/python/tvm/relax/transform/__init__.py
b/python/tvm/relax/transform/__init__.py
index 35be359e76..4d841a8e7b 100644
--- a/python/tvm/relax/transform/__init__.py
+++ b/python/tvm/relax/transform/__init__.py
@@ -51,6 +51,7 @@ from .transform import (
MetaScheduleTuneIRMod,
MetaScheduleTuneTIR,
Normalize,
+ NormalizeGlobalVar,
PatternCheckContext,
RealizeVDevice,
RemovePurityChecking,
diff --git a/python/tvm/relax/transform/transform.py
b/python/tvm/relax/transform/transform.py
index 5bdea37b50..049bce7428 100644
--- a/python/tvm/relax/transform/transform.py
+++ b/python/tvm/relax/transform/transform.py
@@ -284,6 +284,21 @@ def Normalize() -> tvm.ir.transform.Pass:
return _ffi_api.Normalize() # type: ignore
+def NormalizeGlobalVar() -> tvm.ir.transform.Pass:
+ """Possibly rename the GlobalVar in an IRModule to ensure these properties:
+
+ 1. (Invariant) First ensure every public function has the same name as its
"global_symbol"
+ attribute
+ 2. To ensure 1., we may need to rename private functions with conflicting
names;
+ 3. Finally, the name of every GlobalVar is unique in the IRModule.
+
+ Returns
+ -------
+ ret: tvm.ir.transform.Pass
+ """
+ return _ffi_api.NormalizeGlobalVar() # type: ignore
+
+
def CanonicalizeBindings() -> tvm.ir.transform.Pass:
"""
Canonicalizes variable definitions
diff --git a/src/relax/analysis/well_formed.cc
b/src/relax/analysis/well_formed.cc
index 9eed78d270..12fa30b21b 100644
--- a/src/relax/analysis/well_formed.cc
+++ b/src/relax/analysis/well_formed.cc
@@ -26,7 +26,7 @@
* This pass will check:
* 1. Each Expr should have `struct_info_` field already populated, when
* `check_struct_info` is true.
- * 2. GlobalVars are defined before use.
+ * 2. GlobalVars are defined before use. And all GlobalVars have different
names.
* 3. When a Function has a corresponding GlobalVar and a `global_symbol`
* attribute, the name of the GlobalVar must equal the value of the
* `global_symbol` attribute value.
@@ -126,6 +126,9 @@ class WellFormedChecker : public relax::ExprVisitor,
}
void CheckGlobalVarAndGsymbolConsistency(GlobalVar var, Function func) {
+ // the uniqueness of all global vars are ensured by
IRModule->global_var_map_, so do not need
+ // to check again
+
// check name in global var and gsymbol
Optional<String> gsymbol = func->GetAttr<String>(tvm::attr::kGlobalSymbol);
if (gsymbol.defined() && gsymbol != var->name_hint) {
diff --git a/src/relax/ir/block_builder.cc b/src/relax/ir/block_builder.cc
index 210af591b5..fda31e44a9 100644
--- a/src/relax/ir/block_builder.cc
+++ b/src/relax/ir/block_builder.cc
@@ -27,6 +27,7 @@
#include <tvm/relax/op_attr_types.h>
#include <tvm/relax/struct_info.h>
#include <tvm/relax/struct_info_functor.h>
+#include <tvm/relax/transform.h>
#include <tvm/relax/type.h>
#include <tvm/relay/op.h>
#include <tvm/runtime/registry.h>
@@ -72,6 +73,8 @@ class BlockBuilderImpl : public BlockBuilderNode {
IRModule GetContextIRModule() const final { return context_mod_; }
+ IRModule Finalize() final { return
transform::NormalizeGlobalVar()(context_mod_); }
+
GlobalVar AddFunction(const BaseFunc& func, String func_name_hint) final {
LazyInitCtxFuncDedupMap();
auto it = ctx_func_dedup_map_->find(func);
@@ -1002,6 +1005,9 @@ TVM_REGISTER_GLOBAL("relax.BlockBuilderUpdateFunction")
TVM_REGISTER_GLOBAL("relax.BlockBuilderGetContextIRModule")
.set_body_method<BlockBuilder>(&BlockBuilderNode::GetContextIRModule);
+TVM_REGISTER_GLOBAL("relax.BlockBuilderFinalize")
+ .set_body_method<BlockBuilder>(&BlockBuilderNode::Finalize);
+
TVM_REGISTER_GLOBAL("relax.BlockBuilderCurrentBlockIsDataFlow")
.set_body_method<BlockBuilder>(&BlockBuilderNode::CurrentBlockIsDataFlow);
diff --git a/src/relax/transform/normalize.cc b/src/relax/transform/normalize.cc
index fdd2ccc17e..6b45a4c8e9 100644
--- a/src/relax/transform/normalize.cc
+++ b/src/relax/transform/normalize.cc
@@ -170,6 +170,99 @@ class NormalizeMutator : public ExprMutatorBase {
Expr Normalize(const Expr& e) { return NormalizeMutator().VisitExpr(e); }
+class GlobalVarNormalizer : private ExprMutator {
+ public:
+ static IRModule Normalize(const IRModule& m) {
+ GlobalVarNormalizer renamer(m);
+ return renamer.RenameModule();
+ }
+
+ private:
+ explicit GlobalVarNormalizer(const IRModule& m) : ExprMutator(), module_(m),
name_supply_("") {}
+
+ using ExprMutator::VisitExpr_;
+
+ IRModule RenameModule() {
+ // Step 1. Add public functions (functions with global_symbol attributes)
+ auto name_changes = AddPublicFunctions();
+ if (!name_changes) {
+ return module_;
+ }
+
+ // Step 2. Rename private functions
+ AddPrivateFunctions();
+
+ // Step 3. Substitute global vars in functions
+ for (auto [gvar, func] : module_->functions) {
+ if (!func->IsInstance<FunctionNode>()) {
+ continue;
+ }
+ auto new_func = Downcast<BaseFunc>(this->VisitExpr(func));
+ builder_->UpdateFunction(gvar_map_[gvar], new_func);
+ }
+
+ // Step 4. Update the original module (because we do not want to copy all
metadata to the new
+ // module)
+ auto after_module = builder_->GetContextIRModule();
+ auto module_node = module_.CopyOnWrite();
+ module_node->functions = after_module->functions;
+ module_node->global_var_map_ = after_module->global_var_map_;
+ return module_;
+ }
+
+ /**
+ * \brief Add public functions to the builder, and update the name supplier.
+ * \return true if any name changes are made.
+ */
+ bool AddPublicFunctions() {
+ bool name_changes = false;
+ for (const auto& [gvar, func] : module_->functions) {
+ auto global_symbol = func->GetAttr<String>("global_symbol");
+ if (!global_symbol) {
+ continue;
+ }
+
+ auto global_symbol_value = global_symbol.value();
+ CHECK(!name_supply_->ContainsName(global_symbol_value))
+ << "IRModule contains duplicate global symbol: " <<
global_symbol_value;
+ name_supply_->ReserveName(global_symbol_value);
+ auto new_gvar = builder_->AddFunction(func, global_symbol_value);
+ gvar_map_.Set(gvar, new_gvar);
+
+ if (global_symbol.value() != gvar->name_hint) {
+ name_changes = true;
+ }
+ }
+ return name_changes;
+ }
+
+ /**
+ * \brief Add private functions to the builder with names provided by name
supplier. Renaming may
+ * happen if the name of any function conflicts with the name of a public
function.
+ */
+ void AddPrivateFunctions() {
+ for (auto [gvar, func] : module_->functions) {
+ auto global_symbol = func->GetAttr<String>("global_symbol");
+ if (global_symbol) {
+ continue;
+ }
+
+ auto new_name = name_supply_->FreshName(gvar->name_hint, false, false);
+ auto new_gvar = builder_->AddFunction(func, new_name);
+ gvar_map_.Set(gvar, new_gvar);
+ }
+ }
+
+ Expr VisitExpr_(const GlobalVarNode* op) final {
+ ICHECK(gvar_map_.count(GetRef<GlobalVar>(op)));
+ return gvar_map_[GetRef<GlobalVar>(op)];
+ }
+
+ IRModule module_;
+ NameSupply name_supply_;
+ Map<GlobalVar, GlobalVar> gvar_map_;
+};
+
namespace transform {
Pass Normalize() {
@@ -180,6 +273,16 @@ Pass Normalize() {
TVM_REGISTER_GLOBAL("relax.transform.Normalize").set_body_typed(Normalize);
+Pass NormalizeGlobalVar() {
+ runtime::TypedPackedFunc<IRModule(IRModule, PassContext)> pass_func =
+ [=](IRModule mod, PassContext pc) { return
GlobalVarNormalizer::Normalize(mod); };
+ return CreateModulePass(/*pass_function=*/pass_func,
+ /*opt_level=*/0,
+ /*pass_name=*/"NormalizeGlobalVar",
+ /*required=*/{});
+}
+TVM_REGISTER_GLOBAL("relax.transform.NormalizeGlobalVar").set_body_typed(NormalizeGlobalVar);
+
} // namespace transform
} // namespace relax
diff --git a/tests/python/relax/test_blockbuilder_core.py
b/tests/python/relax/test_blockbuilder_core.py
index f09b7bc6f5..e1cbe37b18 100644
--- a/tests/python/relax/test_blockbuilder_core.py
+++ b/tests/python/relax/test_blockbuilder_core.py
@@ -28,9 +28,11 @@ from tvm.script import relax as R, tir as T
from tvm.tir.function import PrimFunc
[email protected]_func("test.blockbuilder.nop")
-def nop():
- pass
[email protected](scope="module")
+def register_nop():
+ @tvm.register_func("test.blockbuilder.nop")
+ def nop():
+ pass
def test_block_builder():
@@ -91,7 +93,7 @@ def test_function_single_block():
assert gv0.name_hint == "gv"
bb.emit_func_output(gv0)
- func = bb.get()["func"]
+ func = bb.finalize()["func"]
assert func.params[0] == x
assert func.params[1] == y
assert func.body.body == gv0
@@ -121,7 +123,7 @@ def test_function_multi_blocks():
gv2 = bb.emit_output(gv1)
bb.emit_func_output(gv2)
- func = bb.get()["func"]
+ func = bb.finalize()["func"]
assert_structural_equal(gv2.struct_info, rx.TensorStructInfo([m, n],
"float16"))
assert func.params[0] == x
@@ -134,35 +136,41 @@ def test_function_multi_blocks():
def test_multi_functions():
- m = tir.Var("m", "int64")
- n = tir.Var("n", "int64")
- x = rx.Var("x", rx.TensorStructInfo([m, n], "float16"))
- y = rx.Var("y", rx.TensorStructInfo([n], "float16"))
bb = rx.BlockBuilder()
- with bb.function("func1", [x, y]):
+ m_1 = tir.Var("m", "int64")
+ n_1 = tir.Var("n", "int64")
+ x_1 = rx.Var("x", rx.TensorStructInfo([m_1, n_1], "float16"))
+ y_1 = rx.Var("y", rx.TensorStructInfo([n_1], "float16"))
+
+ with bb.function("func1", [x_1, y_1]):
with bb.dataflow():
- lv0 = bb.emit(rx.op.add(x, y))
+ lv0 = bb.emit(rx.op.add(x_1, y_1))
assert lv0.name_hint == "lv"
gv0 = bb.emit_output(lv0)
bb.emit_func_output(gv0)
- with bb.function("func2", [x, y]):
+ m_2 = tir.Var("m", "int64")
+ n_2 = tir.Var("n", "int64")
+ x_2 = rx.Var("x", rx.TensorStructInfo([m_2, n_2], "float16"))
+ y_2 = rx.Var("y", rx.TensorStructInfo([n_2], "float16"))
+
+ with bb.function("func2", [x_2, y_2]):
with bb.dataflow():
- lv0 = bb.emit(rx.op.add(y, x))
+ lv0 = bb.emit(rx.op.add(y_2, x_2))
# TODO(@yuchen): enable block builder to reset local var unique
name map
assert lv0.name_hint == "lv1"
gv0 = bb.emit_output(lv0)
bb.emit_func_output(gv0)
- mod = bb.get()
+ mod = bb.finalize()
func1 = mod["func1"]
- assert func1.params[0] == x
- assert func1.params[1] == y
+ assert func1.params[0] == x_1
+ assert func1.params[1] == y_1
assert len(func1.body.blocks) == 1
func2 = mod["func2"]
- assert func2.params[0] == x
- assert func2.params[1] == y
+ assert func2.params[0] == x_2
+ assert func2.params[1] == y_2
assert len(func2.body.blocks) == 1
@@ -223,7 +231,7 @@ def test_emit_match_cast():
gv0 = bb.emit_output(lv1)
bb.emit_func_output(gv0)
- func = bb.get()["func"]
+ func = bb.finalize()["func"]
block = func.body.blocks[0]
b0, b1 = block.bindings[:2]
assert isinstance(b0, rx.MatchCast)
@@ -252,7 +260,7 @@ def test_emit_match_cast_binding_in_dataflow_block():
bb.emit_output(gv)
bb.emit_func_output(x)
- func = bb.get()["main"]
+ func = bb.finalize()["main"]
block = func.body.blocks[0]
b0 = block.bindings[0]
assert isinstance(b0, rx.MatchCast)
@@ -353,7 +361,7 @@ def test_call_te():
out = bb.emit_output(bb.call_te(te_func, [x, y], {"C": z},
msg="hello"))
bb.emit_func_output(out)
- mod = bb.get()
+ mod = bb.finalize()
rx_func = mod["rx_func"]
assert rx_func.params[0] == x
@@ -372,7 +380,7 @@ def test_call_te_unique_tensor_name():
gv = bb.emit_te(topi.nn.matmul, x, y)
bb.emit_func_output(gv)
- f_matmul = bb.get()["matmul"]
+ f_matmul = bb.finalize()["matmul"]
param_A = f_matmul.params[0]
param_B = f_matmul.params[1]
buffer_A = f_matmul.buffer_map[param_A]
@@ -411,7 +419,7 @@ def test_emit_te():
out = bb.emit_te(te_func, [x, y], {"C": z}, msg="hello")
bb.emit_func_output(out)
- mod = bb.get()
+ mod = bb.finalize()
rx_func = mod["rx_func"]
def get_tir_func():
@@ -453,13 +461,13 @@ def test_emit_te_multiple():
B = te.compute((128, 128), lambda i, j: A[i, j] + 1)
return B
- with bb.function("rx_func", [x, y]):
+ with bb.function("rx_func", [x, y, z]):
x1 = bb.emit_te(te_func, x)
y1 = bb.emit_te(te_func, y)
z1 = bb.emit_te(te_func, z)
bb.emit_func_output(z1)
- mod = bb.get()
+ mod = bb.finalize()
rx_func = mod["rx_func"]
prim_func = []
@@ -488,7 +496,7 @@ def test_emit_te_multiple_output():
z = rx.TupleGetItem(y, 0)
bb.emit_func_output([y, z])
- rx_func = bb.get()["rx_func"]
+ rx_func = bb.finalize()["rx_func"]
# check call tir output shape is a Tuple of ShapeExpr
assert rx_func.params[0] == x
@@ -511,7 +519,7 @@ def test_emit_te_extern():
out = bb.emit_te(tvm.contrib.cblas.matmul, x, y, transa=False,
transb=False)
bb.emit_func_output(out)
- mod = bb.get()
+ mod = bb.finalize()
rx_func = mod["rx_cblas_matmul"]
# check Relax function calls TIR function with call_tir call
@@ -540,7 +548,7 @@ def test_emit_te_prim_value():
out = bb.emit_te(topi.clip, x, a_min, a_max)
bb.emit_func_output(out)
- rx_func = bb.get()["rx_clip"]
+ rx_func = bb.finalize()["rx_clip"]
# check Relax function calls TIR function with call_tir call
assert rx_func.params[0] == x
@@ -649,7 +657,7 @@ def test_emit_nested_tuple(emit_nested_tuple):
output = (scalars, x, y)
bb.emit_func_output(output)
- return bb.get()["func"]
+ return bb.finalize()["func"]
def make_expected(emit_nested_tuple: bool):
if emit_nested_tuple:
@@ -683,5 +691,42 @@ def test_emit_nested_tuple(emit_nested_tuple):
tvm.ir.assert_structural_equal(expected, actual)
[email protected]_well_formed_check_before_transform
+def test_finalize_public_private_name_conflict():
+ # tir call
+ bb = rx.BlockBuilder()
+
+ def te_zero():
+ return topi.full((), "int64", tir.IntImm("int64", 0))
+
+ def te_one():
+ return topi.full((), "int64", tir.IntImm("int64", 1))
+
+ with bb.function("func", []):
+ gv0 = bb.emit_te(te_zero, primfunc_name_hint="func")
+ gv1 = bb.emit_te(te_one, primfunc_name_hint="func")
+ bb.emit_func_output((gv0, gv1))
+
+ mod = bb.finalize()
+ assert rx.analysis.well_formed(mod)
+
+ # relax function call
+ bb = rx.BlockBuilder()
+
+ with bb.function("func", [], private=True):
+ gvar = bb.emit_func_output(rx.const(0, "int64"))
+
+ with bb.function("func", [], private=True):
+ gv0 = bb.emit(rx.Call(gvar, []))
+ gvar1 = bb.emit_func_output(gv0)
+
+ with bb.function("func", []):
+ gv0 = bb.emit(rx.Call(gvar1, []))
+ bb.emit_func_output(gv0)
+
+ mod = bb.finalize()
+ assert rx.analysis.well_formed(mod)
+
+
if __name__ == "__main__":
tvm.testing.main()
diff --git a/tests/python/relax/test_transform_normalize_global_var.py
b/tests/python/relax/test_transform_normalize_global_var.py
new file mode 100644
index 0000000000..0a26ffc8e6
--- /dev/null
+++ b/tests/python/relax/test_transform_normalize_global_var.py
@@ -0,0 +1,98 @@
+# 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 import relax
+from tvm import tir
+from tvm.ir.base import assert_structural_equal
+
+import tvm.script
+from tvm.script import tir as T, relax as R, ir as I
+
+
[email protected]_well_formed_check_before_transform
+def test_normalize_relax_function():
+ @I.ir_module
+ class Before:
+ @R.function(private=True)
+ def f():
+ return R.const(1, "int32")
+
+ @R.function
+ def f1():
+ R.func_attr({"global_symbol": "f"})
+ cls = Before
+ gv: R.Tensor((), dtype="int32") = cls.f()
+ return gv
+
+ @I.ir_module
+ class Expected:
+ @R.function
+ def f():
+ cls = Expected
+ gv: R.Tensor((), dtype="int32") = cls.f1()
+ return gv
+
+ @R.function(private=True)
+ def f1():
+ return R.const(1, "int32")
+
+ After = relax.transform.NormalizeGlobalVar()(Before)
+
+ assert not relax.analysis.well_formed(Before)
+ assert relax.analysis.well_formed(After)
+ assert_structural_equal(After, Expected)
+
+
[email protected]_well_formed_check_before_transform
+def test_normalize_tir_function():
+ @I.ir_module
+ class Before:
+ @T.prim_func(private=True)
+ def f(x: T.Buffer((1,), "int32")):
+ x[0] = T.int32(0)
+
+ @R.function
+ def f1():
+ R.func_attr({"global_symbol": "f"})
+ cls = Before
+ gv: R.Tensor((), dtype="int32") = R.call_tir(cls.f, (),
R.Tensor((1,), dtype="int32"))
+ return gv
+
+ @I.ir_module
+ class Expected:
+ @T.prim_func(private=True)
+ def f1(x: T.Buffer((1,), "int32")):
+ x[0] = 0
+
+ @R.function
+ def f() -> R.Tensor((1,), dtype="int32"):
+ cls = Expected
+ gv = R.call_tir(cls.f1, R.tuple(), out_sinfo=R.Tensor((1,),
dtype="int32"))
+ return gv
+
+ After = relax.transform.NormalizeGlobalVar()(Before)
+
+ assert not relax.analysis.well_formed(Before)
+ assert relax.analysis.well_formed(After)
+ assert_structural_equal(After, Expected)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
