This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch unity-staging in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 9c369e7f4a903513e4d223c788aca4c3222ecfff Author: Lesheng Jin <[email protected]> AuthorDate: Fri Feb 17 16:25:37 2023 -0800 [Unity][Pass] Normalize Pass (#14031) This PR implements relax `Normalize` Pass, which allows users to transform Relax IR to normal form, i.e., the expressions are normalized (no nesting and hence the AST is in ANF), and all `checked_type_` and `shape_` of expressions are available. (tests are added). Co-Authored-by: Yuchen Jin <[email protected]> Co-Authored-by: Ruihang Lai <[email protected]> Co-Authored-by: Siyuan Feng <[email protected]> Co-Authored-by: Tianqi Chen <[email protected]> --- include/tvm/relax/transform.h | 9 + python/tvm/relax/transform/transform.py | 11 + src/relax/transform/normalize.cc | 186 +++++++++ tests/python/relax/test_transform_normalize.py | 554 +++++++++++++++++++++++++ 4 files changed, 760 insertions(+) diff --git a/include/tvm/relax/transform.h b/include/tvm/relax/transform.h index e9f63ee9db..7a4054d414 100644 --- a/include/tvm/relax/transform.h +++ b/include/tvm/relax/transform.h @@ -133,6 +133,15 @@ TVM_DLL Pass BindParams(String func_name, Map<String, runtime::NDArray> params); * \return The Pass. */ TVM_DLL Pass FoldConstant(); + +/*! + * \brief Transform Relax IR to normal form: transform AST to A-normal form, and fill the + * checked_type_ and shape_ of expressions. + * + * \return The Pass. + */ +TVM_DLL Pass Normalize(); + } // namespace transform } // namespace relax } // namespace tvm diff --git a/python/tvm/relax/transform/transform.py b/python/tvm/relax/transform/transform.py index c0ac180ff1..7fcf0b1121 100644 --- a/python/tvm/relax/transform/transform.py +++ b/python/tvm/relax/transform/transform.py @@ -68,6 +68,17 @@ def CallTIRRewrite() -> tvm.ir.transform.Pass: return _ffi_api.CallTIRRewrite() # type: ignore +def Normalize() -> tvm.ir.transform.Pass: + """Transforming Relax IR to normal form, i.e., the expressions are normalized(no nesting + and hence the AST is in ANF), and all checked_type_ and shape_ of expressions are available. + + Returns + ------- + ret: tvm.ir.transform.Pass + """ + return _ffi_api.Normalize() # type: ignore + + def RewriteDataflowReshape() -> tvm.ir.transform.Pass: """Convert all reshape-like call_tir to VM reshape operator call. The VM reshape operator calls will be further lowered to a CreateView diff --git a/src/relax/transform/normalize.cc b/src/relax/transform/normalize.cc new file mode 100644 index 0000000000..915498178f --- /dev/null +++ b/src/relax/transform/normalize.cc @@ -0,0 +1,186 @@ +/* + * 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. + */ + +/*! + * \file tvm/relax/transform/normalize.cc + * \brief Pass for transforming Relax IR to normal form, i.e., the expressions are normalized(no + * nesting and hence the AST is in ANF), and all checked_type_ and shape_ of expressions are + * available. + */ + +#include <tvm/relax/expr.h> +#include <tvm/relax/expr_functor.h> +#include <tvm/relax/struct_info.h> +#include <tvm/relax/transform.h> + +namespace tvm { +namespace relax { + +// TODO(@altanh): LCA binding lifting +class NormalizeMutator : public ExprMutatorBase { + public: + NormalizeMutator() { builder_ = BlockBuilder::Create(NullOpt); } + + Expr VisitExpr(const Expr& expr) override { + return builder_->Normalize(ExprMutatorBase::VisitExpr(expr)); + } + + Expr VisitExpr_(const FunctionNode* op) final { + Expr body = this->VisitWithNewScope(op->body, op->params); + + if (body.same_as(op->body)) { + return GetRef<Expr>(op); + } else { + return Function(op->params, body, op->ret_struct_info, op->attrs); + } + } + + Expr VisitExpr_(const IfNode* op) final { + Expr guard = this->VisitExpr(op->cond); + Expr true_b = this->VisitWithNewScope(op->true_branch); + Expr false_b = this->VisitWithNewScope(op->false_branch); + if (op->cond.same_as(guard) && op->true_branch.same_as(true_b) && + op->false_branch.same_as(false_b)) { + return GetRef<Expr>(op); + } else { + return If(guard, true_b, false_b, op->span); + } + } + + Expr VisitWithNewScope(const Expr& expr, Optional<Array<Var>> params = NullOpt) { + builder_->BeginBindingBlock(); + builder_->BeginScope(params); + Expr ret = this->VisitExpr(expr); + BindingBlock prologue = builder_->EndBlock(); + if (!prologue->bindings.empty()) { + ret = SeqExpr({prologue}, ret); + } + builder_->EndScope(); + return ret; + } + + Expr VisitExpr_(const SeqExprNode* op) final { + bool all_blocks_unchanged = true; + Array<BindingBlock> blocks; + for (auto block : op->blocks) { + BindingBlock new_block = this->VisitBindingBlock(block); + if (!new_block->bindings.empty()) { + blocks.push_back(new_block); + } + all_blocks_unchanged &= block.same_as(new_block); + } + + builder_->BeginBindingBlock(); + Expr body = this->VisitExpr(op->body); + BindingBlock prologue = builder_->EndBlock(); + if (!prologue->bindings.empty()) { + blocks.push_back(prologue); + all_blocks_unchanged = false; + } + + if (all_blocks_unchanged && body.same_as(op->body)) { + return GetRef<Expr>(op); + } else { + return SeqExpr(blocks, body); + } + } + + BindingBlock VisitBindingBlock(const BindingBlock& block) final { + BindingBlock ret; + if (const auto* node = block.as<DataflowBlockNode>()) { + ret = VisitBindingBlock_(node); + } else if (const auto* node = block.as<BindingBlockNode>()) { + ret = VisitBindingBlock_(node); + } else { + LOG(FATAL) << "TypeError: Invalid type: " << block->GetTypeKey(); + } + return ret; + } + + BindingBlock VisitBindingBlock_(const BindingBlockNode* block) { + builder_->BeginBindingBlock(); + for (Binding binding : block->bindings) { + this->VisitBinding(binding); + } + return builder_->EndBlock(); + } + + BindingBlock VisitBindingBlock_(const DataflowBlockNode* block) { + builder_->BeginDataflowBlock(); + for (Binding binding : block->bindings) { + this->VisitBinding(binding); + } + return builder_->EndBlock(); + } + + void VisitBinding(const Binding& binding) { + if (const auto* node = binding.as<VarBindingNode>()) { + VisitBinding_(node); + } else if (const auto* node = binding.as<MatchCastNode>()) { + VisitBinding_(node); + } else { + LOG(FATAL) << "TypeError: Invalid type: " << binding->GetTypeKey(); + } + } + + void VisitBinding_(const VarBindingNode* binding) { + Expr new_value = this->VisitExpr(binding->value); + if (!binding->var->struct_info_.defined()) { + UpdateStructInfo(binding->var, GetStructInfo(new_value)); + } + + if (new_value.same_as(binding->value)) { + builder_->EmitNormalized(GetRef<VarBinding>(binding)); + } else { + builder_->EmitNormalized(VarBinding(binding->var, new_value)); + } + } + + void VisitBinding_(const MatchCastNode* binding) { + Expr new_value = this->VisitExpr(binding->value); + + if (new_value.same_as(binding->value)) { + builder_->EmitNormalized(GetRef<MatchCast>(binding)); + } else { + builder_->EmitNormalized( + MatchCast(binding->var, builder_->NormalizeArgument(new_value), binding->struct_info)); + } + } + + private: + /*! \brief Internal block builder to emit bindings during rewriting. */ + BlockBuilder builder_; +}; // namespace relax + +Expr Normalize(const Expr& e) { return NormalizeMutator().VisitExpr(e); } + +namespace transform { + +Pass Normalize() { + runtime::TypedPackedFunc<Function(Function, IRModule, PassContext)> pass_func = + [=](Function f, IRModule m, PassContext pc) { return Downcast<Function>(Normalize(f)); }; + return CreateFunctionPass(pass_func, 1, "Normalize", {}); +} + +TVM_REGISTER_GLOBAL("relax.transform.Normalize").set_body_typed(Normalize); + +} // namespace transform + +} // namespace relax +} // namespace tvm diff --git a/tests/python/relax/test_transform_normalize.py b/tests/python/relax/test_transform_normalize.py new file mode 100644 index 0000000000..9e9533a5ed --- /dev/null +++ b/tests/python/relax/test_transform_normalize.py @@ -0,0 +1,554 @@ +# 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 + + +def test_normalize_function(): + m = tir.Var("m", "int64") + n = tir.Var("n", "int64") + x = relax.Var("x", R.Tensor([m, n], "float16")) + + # Note: the parser automatically normalize the IR written in TVMScript, + # so we manually construct the function here. + mul_add = relax.Function( + [x], + relax.op.multiply(relax.op.add(x, x), relax.op.add(x, x)), + ret_struct_info=R.Tensor("float16", ndim=2), + ) + + # Note: from_expr api names private function (function without global_symbol) as "main" + before_mod = tvm.IRModule.from_expr(mul_add) + + after_mod = relax.transform.Normalize()(before_mod) + + @R.function + def expected(x: R.Tensor(("m", "n"), "float16")) -> R.Tensor(dtype="float16", ndim=2): + gv = R.add(x, x) + gv1 = R.add(x, x) + return R.multiply(gv, gv1) + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_if(): + cond = relax.Var("cond", R.Tensor([], "bool")) + x = relax.Var("x", R.Tensor([1], "float32")) + # TODO(relax-team): add type and shape inference for IfNode + y = relax.Var("y") + + # Note: the parser automatically normalize the IR written in TVMScript, + # so we manually construct the function and If here. + f = relax.Function( + [cond, x], + relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding( + y, + relax.If( + cond, + relax.op.multiply(relax.op.add(x, x), relax.op.add(x, x)), + relax.op.add(relax.op.multiply(x, x), relax.op.multiply(x, x)), + ), + ) + ] + ) + ], + y, + ), + ret_struct_info=R.Tensor("float32", ndim=1), + ) + + before_mod = tvm.IRModule.from_expr(f) + after_mod = relax.transform.Normalize()(before_mod) + + @R.function + def expected( + cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32") + ) -> R.Tensor(dtype="float32", ndim=1): + if cond: + gv = R.add(x, x) + gv1 = R.add(x, x) + y = R.multiply(gv, gv1) + else: + gv = R.multiply(x, x) + gv1 = R.multiply(x, x) + y = R.add(gv, gv1) + return y + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_no_op(): + # the normalize pass should be no-op for IR in ANF + @tvm.script.ir_module + class ANFMod1: + @R.function + def f(x: R.Tensor(dtype="float32")): + gv = R.add(x, x) + gv1 = R.add(gv, gv) + gv2 = R.add(gv, gv1) + return (gv, gv2) + + before_mod = ANFMod1 + after_mod = relax.transform.Normalize()(before_mod) + assert_structural_equal(before_mod, after_mod, map_free_vars=True) + + @tvm.script.ir_module + class ANFMod2: + @R.function + def foo(x: R.Tensor(("m", "n"), "float32")): + m, n = T.var("int64"), T.var("int64") + with R.dataflow(): + lv0 = R.call_tir("test.op.identity", (x,), R.Tensor((m, n), dtype="float32")) + gv0 = R.call_tir("test.op.identity", (lv0,), R.Tensor((m, n), dtype="float32")) + R.output(gv0) + return gv0 + + mod = ANFMod2 + mod_post = relax.transform.Normalize()(mod) + + assert_structural_equal(mod, mod_post) + + +def test_normalize_seq_body(): + # a seq expression with a non-leaf body should bind the body to a var as well + x = relax.Var("x", R.Tensor([], "int32")) + y = relax.Var("y", R.Tensor([], "int32")) + seq = relax.SeqExpr([], relax.op.add(x, y)) + f = relax.Function( + [x, y], + seq, + ret_struct_info=R.Tensor([], "int32"), + ) + + before_mod = tvm.IRModule.from_expr(f) + after_mod = relax.transform.Normalize()(before_mod) + + @R.function + def expected( + x: R.Tensor((), dtype="int32"), y: R.Tensor((), dtype="int32") + ) -> R.Tensor(ndim=0, dtype="int32"): + # normalization inserts a binding like this + z = R.add(x, y) + return z + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_func_body(): + # a function with a body that is not a seq expr should have it wrapped in a seq expr + x = relax.Var("x", R.Tensor([], "int32")) + y = relax.Var("y", R.Tensor([], "int32")) + f = relax.Function( + [x, y], + relax.op.add(x, y), + ret_struct_info=R.Tensor([], "int32"), + ) + + before_mod = tvm.IRModule.from_expr(f) + after_mod = relax.transform.Normalize()(before_mod) + + @R.function + def expected( + x: R.Tensor((), dtype="int32"), y: R.Tensor((), dtype="int32") + ) -> R.Tensor(ndim=0, dtype="int32"): + # result will be a seq expr where the body is a var + z = R.add(x, y) + return z + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_if_branches(): + # an if node's branches must be seq exprs + x = relax.Var("x", R.Tensor([], "int32")) + y = relax.Var("y", R.Tensor([], "int32")) + # TODO(@relax-team): z has a shape of () and type of DynTensorType(ndim=0), + # but normalization fails to infer these even though it should + z = relax.Var("z") + cond = relax.Var("cond", R.Tensor([], "bool")) + plus = relax.op.add(x, y) + mult = relax.op.multiply(x, y) + if_node = relax.If(cond, plus, mult) + seq = relax.SeqExpr([relax.BindingBlock([relax.VarBinding(z, if_node)])], z) + f = relax.Function( + [cond, x, y], + seq, + ret_struct_info=R.Tensor([], "int32"), + ) + + before_mod = tvm.IRModule.from_expr(f) + after_mod = relax.transform.Normalize()(before_mod) + + @R.function + def expected( + cond: R.Tensor((), dtype="bool"), + x: R.Tensor((), dtype="int32"), + y: R.Tensor((), dtype="int32"), + ) -> R.Tensor(ndim=0, dtype="int32"): + # the bodies of the branches will be seq exprs with a binding + if cond: + w = R.add(x, y) + z = w + else: + w = R.multiply(x, y) + z = w + return z + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_if_condition(): + cond = relax.Var("cond", R.Tensor([], "bool")) + x = relax.Var("x", R.Tensor([1], "float32")) + # TODO(relax-team): add type and shape inference for IfNode + y = relax.Var("y") + + # The condition is wrapped in a tuple and then indexed + f = relax.Function( + [cond, x], + relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding( + y, + relax.If( + relax.TupleGetItem(relax.Tuple([cond]), 0), + relax.op.add(x, x), + relax.op.multiply(x, x), + ), + ) + ] + ) + ], + y, + ), + ret_struct_info=R.Tensor("float32", ndim=1), + ) + + before_mod = tvm.IRModule.from_expr(f) + after_mod = relax.transform.Normalize()(before_mod) + + @R.function + def expected( + cond: R.Tensor((), "bool"), x: R.Tensor((1,), "float32") + ) -> R.Tensor(dtype="float32", ndim=1): + c = R.TupleGetItem(R.tuple(cond), 0) + if c: + gv = R.add(x, x) + y = gv + else: + gv = R.multiply(x, x) + y = gv + return y + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_tuple_get_item(): + x = relax.Var("x", R.Tensor([], "int32")) + f = relax.Function( + [x], + relax.TupleGetItem( + relax.TupleGetItem( + relax.Tuple([relax.Tuple([x])]), + 0, + ), + 0, + ), + ret_struct_info=R.Tensor([], "int32"), + ) + + before_mod = tvm.IRModule.from_expr(f) + after_mod = relax.transform.Normalize()(before_mod) + + # TODO: Revisit once we canonicalize SeqExprs (part of normalization?) + # Not using the parser this time because writing it out correctly results in + # *one* binding block, whereas the normalized version has *two* + idx_var = relax.Var("idx_var", R.Tuple([R.Tensor([], "int32")])) + ret_var = relax.Var("ret", R.Tensor([], "int32")) + expected_f = relax.Function( + [x], + relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding( + idx_var, relax.TupleGetItem(relax.Tuple([relax.Tuple([x])]), 0) + ) + ] + ), + relax.BindingBlock([relax.VarBinding(ret_var, relax.TupleGetItem(idx_var, 0))]), + ], + ret_var, + ), + ret_struct_info=R.Tensor([], "int32"), + ) + expected_mod = tvm.IRModule.from_expr(expected_f) + # apply normalization to fill in type and shape annotations (tedious otherwise) + final_mod = relax.transform.Normalize()(expected_mod) + + assert_structural_equal(after_mod, final_mod) + + +def test_normalize_combine_nearby_blocks(): + x = relax.Var("x", R.Tensor([], "int32")) + v0 = relax.Var("v0", R.Tensor([], "int32")) + v1 = relax.Var("v1", R.Tensor([], "int32")) + v2 = relax.Var("v2", R.Tensor([], "int32")) + v3 = relax.Var("v3", R.Tensor([], "int32")) + f = relax.Function( + [x], + relax.SeqExpr( + [ + relax.DataflowBlock([relax.VarBinding(v0, x)]), + relax.DataflowBlock([relax.VarBinding(v1, v0)]), + relax.BindingBlock([relax.VarBinding(v2, v1)]), + relax.BindingBlock([relax.VarBinding(v3, v2)]), + ], + v3, + ), + ret_struct_info=R.Tensor([], "int32"), + ) + + after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f)) + + @R.function + def expected(x: R.Tensor((), "int32")): + with R.dataflow(): + v0 = x + v1 = v0 + R.output(v0, v1) + v2 = v1 + v3 = v2 + return v3 + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_nested_seq(): + x = relax.Var("x", R.Tensor([], "int32")) + y = relax.Var("y", R.Tensor([], "int32")) + z = relax.Var("z", R.Tensor([], "int32")) + seq = relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding(x, relax.const(1)), + relax.VarBinding( + y, + relax.SeqExpr( + [relax.BindingBlock([relax.VarBinding(z, relax.const(2))])], + z, + ), + ), + ] + ) + ], + y, + ) + + f = relax.Function( + [], + seq, + ret_struct_info=R.Tensor([], "int32"), + ) + after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f)) + + @R.function + def expected(): + x = relax.const(1) + z = relax.const(2) + y = z + return y + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_nested_seq_dataflow(): + x = relax.Var("x", R.Tensor([], "int32")) + y = relax.Var("y", R.Tensor([], "int32")) + z = relax.Var("z", R.Tensor([], "int32")) + q = relax.Var("u", R.Tensor([], "int32")) + w = relax.DataflowVar("w", R.Tensor([], "int32")) + u = relax.Var("u", R.Tensor([], "int32")) + seq = relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding(x, relax.const(1)), + relax.VarBinding( + y, + relax.SeqExpr( + [ + relax.BindingBlock([relax.VarBinding(q, relax.const(2))]), + relax.DataflowBlock( + [ + relax.VarBinding(w, q), + relax.VarBinding(u, w), + ] + ), + relax.BindingBlock([relax.VarBinding(z, u)]), + ], + z, + ), + ), + ] + ) + ], + y, + ) + + f = relax.Function( + [], + seq, + ret_struct_info=R.Tensor([], "int32"), + ) + after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f)) + + @R.function + def expected(): + x = relax.const(1) + q = relax.const(2) + with R.dataflow(): + w = q + u = w + R.output(u) + z = u + y = z + return y + + assert_structural_equal(after_mod["main"], expected) + + +def test_normalize_deeply_nested_seq(): + x = relax.Var("x", R.Tensor([], "int32")) + y = relax.Var("y", R.Tensor([], "int32")) + z = relax.Var("z", R.Tensor([], "int32")) + u = relax.Var("u", R.Tensor([], "int32")) + v = relax.Var("v", R.Tensor([], "int32")) + w = relax.Var("w", R.Tensor([], "int32")) + _ = relax.Var("w", R.Tensor([], "int32")) + seq = relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding(x, relax.const(1)), + relax.VarBinding( + y, + relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding( + z, + relax.SeqExpr( + [ + relax.BindingBlock( + [ + relax.VarBinding(u, relax.const(2)), + relax.MatchCast( + _, u, R.Tensor([], "int32") + ), + relax.VarBinding(v, u), + relax.MatchCast( + w, v, R.Tensor([], "int32") + ), + ] + ) + ], + w, + ), + ) + ] + ) + ], + z, + ), + ), + ] + ) + ], + y, + ) + + f = relax.Function( + [], + seq, + ret_struct_info=R.Tensor([], "int32"), + ) + after_mod = relax.transform.Normalize()(tvm.IRModule.from_expr(f)) + + @R.function + def expected(): + x = relax.const(1) + u = relax.const(2) + _ = R.match_cast(u, R.Tensor((), "int32")) + v = u + w = R.match_cast(v, R.Tensor((), "int32")) + z = w + y = z + return y + + assert_structural_equal(after_mod["main"], expected) + + [email protected]() +def test_nesting_non_dataflow_in_dataflow_error(): + x = relax.DataflowVar("x", R.Tensor([], "int32")) + y = relax.Var("y", R.Tensor([], "int32")) + z = relax.Var("z", R.Tensor([], "int32")) + seq = relax.SeqExpr( + [ + relax.DataflowBlock( + [ + relax.VarBinding(x, relax.const(1)), + relax.VarBinding( + y, + relax.SeqExpr( + [relax.BindingBlock([relax.VarBinding(z, relax.const(2))])], + z, + ), + ), + ] + ) + ], + y, + ) + f = relax.Function( + [], + seq, + ret_struct_info=R.Tensor([], "int32"), + ) + relax.transform.Normalize()(tvm.IRModule.from_expr(f)) + # should fail due to a normal binding block being inside a dataflowblock + + +if __name__ == "__main__": + tvm.testing.main()
