From 2ac3acb768c95d905ab2507fc93b126f258d533b Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Tue, 10 Jan 2023 18:39:50 -0500 Subject: [PATCH 1/8] [PyConverter] support returning lone global var --- python/tvm/relay/testing/py_converter.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py index 1ec85faea619..9b8a697e4955 100644 --- a/python/tvm/relay/testing/py_converter.py +++ b/python/tvm/relay/testing/py_converter.py @@ -98,7 +98,11 @@ def optimize(self, prog: Expr): # unwrap tuple wrappers (some op calls produce them) unwrapped = prog.astuple() if isinstance(prog, relay.TupleWrapper) else prog assert relay.analysis.well_formed(unwrapped) - mod = self.mod.from_expr(unwrapped, self.mod.functions, self.mod.type_definitions) + # For a lone global var, there is nothing we need to do + if isinstance(unwrapped, relay.GlobalVar): + return unwrapped + else: + mod = self.mod.from_expr(unwrapped, self.mod.functions, self.mod.type_definitions) # necessary pass: SimplifyInference (otherwise we can't generate code for some operators) # and fusion (to get primitive functions) From 4cf205ec574ec97047cf0dea4ff2d2d0b98b7a77 Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Wed, 11 Jan 2023 21:43:02 -0500 Subject: [PATCH 2/8] [py_converter] Cast PackedFuncs into objects so they can appear inside containers --- python/tvm/relay/testing/py_converter.py | 55 ++++++++++++++++++++---- python/tvm/runtime/container.py | 19 ++++++++ src/runtime/container.cc | 13 ++++++ tests/python/relay/test_py_converter.py | 48 ++++++++++++++++++++- 4 files changed, 125 insertions(+), 10 deletions(-) diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py index 9b8a697e4955..10ffbdd2347a 100644 --- a/python/tvm/relay/testing/py_converter.py +++ b/python/tvm/relay/testing/py_converter.py @@ -38,7 +38,7 @@ # import tvm # from tvm import relay # from tvm import nd -# from tvm.runtime import import container as _container +# from tvm.runtime import container as _container # from tvm.relay.backend.interpreter import RefValue, ConstructorValue PROLOGUE = [ ast.Import([alias("numpy", None)]), @@ -201,7 +201,7 @@ def convert_func_node(self, func: Function, name_var=None): var_names = [self.get_var_name(var) for var in func.params] body, defs = self.visit(func.body) - ret = self.create_def(func_name, var_names, defs + [Return(body)]) + ret = self.create_def(func_name, var_names, defs + [Return(body)], register_packed=True) return (ret, func_name) def convert_module(self): @@ -223,10 +223,25 @@ def create_call(self, func_name: str, arguments): """Creates a simple function call.""" return ast.Call(self.parse_name(func_name), arguments, []) - def create_def(self, func_name: str, arguments: [str], body): - """Wrapper over function definition AST node, whose constructor is inconvenient.""" + def create_def(self, func_name: str, arguments: [str], body, register_packed: bool = False): + """ + Wrapper over function definition AST node, whose constructor is inconvenient. + + register_packed includes a tvm.register_func decorator on the generated function if true. + This option should be used for Relay functions (warning: clobbers registry!) + """ inner_args = [ast.arg(argument, None) for argument in arguments] + # add a decorator to register as a PackedFunc so the function will be an ObjectRef + # and will allow for putting functions into tuples or refs + decorator_list = [ + ast.Call( + self.parse_name("tvm.register_func"), + [ast.Constant(value=func_name)], + [ast.keyword(arg="override", value=ast.Constant(value=True))], + ) + ] + global __MAJOR__, __MINOR__ if __MAJOR__ == 3 and __MINOR__ >= 8: arguments = ast.arguments([], inner_args, None, [], [], None, []) @@ -237,7 +252,7 @@ def create_def(self, func_name: str, arguments: [str], body): func_name, arguments, body, - [], + decorator_list if register_packed else [], None, ) @@ -422,7 +437,15 @@ def visit_var(self, var: Expr): def visit_global_var(self, gvar: Expr): # we don't need to add numbers to global var names because # the *names* are checked for uniqueness in the mod - return (Name(str(gvar.name_hint), Load()), []) + func_name = str(gvar.name_hint) + # load in the packed func and cast to object so it can be put in containers + return ( + self.create_call( + "_container._cast_packed_func", + [self.create_call("tvm.get_global_func", [ast.Constant(value=func_name)])], + ), + [], + ) def visit_let(self, letexp: Expr): # To properly account for scoping and ensure that the entire node produces an expression, @@ -494,7 +517,14 @@ def visit_constant(self, constant: Expr): def visit_function(self, func: Expr): # Python's lambdas are very restrictive, so we do "name" inline functions converted_func, func_name = self.convert_func_node(func) - return (Name(func_name, Load()), [converted_func]) + # load in the PackedFunc and cast to object so it can be stored in other containers + return ( + self.create_call( + "_container._cast_packed_func", + [self.create_call("tvm.get_global_func", [ast.Constant(value=func_name)])], + ), + [converted_func], + ) def visit_call(self, call: Expr): """For calls, we must distinguish between ordinary functions, @@ -524,7 +554,10 @@ def visit_call(self, call: Expr): # ordinary function converted_func, defs = self.visit(func) defs += field_defs - return (ast.Call(converted_func, fields, []), defs) + # have to convert func back from an object to be able to call it + recast_func = self.create_call("_container._uncast_packed_func", [converted_func]) + + return (ast.Call(recast_func, fields, []), defs) def visit_ref_create(self, ref: Expr): val, defs = self.visit(ref.value) @@ -606,7 +639,11 @@ def to_python(expr: Expr, mod=None, target=tvm.target.Target("llvm")): def run_as_python(expr: Expr, mod=None, target=tvm.target.Target("llvm")): """Converts the given Relay expression into a Python script and - executes it.""" + executes it. + + Note that closures will be returned as PackedFuncs cast to Objects; + to execute them, they need to be cast back using runtime.container._uncast_packed_func + """ mod = mod if mod is not None else tvm.IRModule() py_ast = to_python(expr, mod, target) code = compile(py_ast, "", "exec") diff --git a/python/tvm/runtime/container.py b/python/tvm/runtime/container.py index cae408b6121d..a93589b34360 100644 --- a/python/tvm/runtime/container.py +++ b/python/tvm/runtime/container.py @@ -18,6 +18,7 @@ import tvm._ffi from .object import Object, PyNativeObject from .object_generic import ObjectTypes +from .packed_func import PackedFunc from . import _ffi_api @@ -60,6 +61,24 @@ def getitem_helper(obj, elem_getter, length, idx): return elem_getter(obj, idx) +# helper functions for putting PackedFuncs inside other containers +def _cast_packed_func(pf: PackedFunc) -> Object: + """ + Helper function for testing: Casts the given PackedFunc into a TVM object + so that it can be stored inside other runtime containers + """ + return _ffi_api.CastPackedFuncToObject(pf) + + +def _uncast_packed_func(obj: Object) -> PackedFunc: + """ + Helper function for testing: Casts the passed Object to a PackedFunc + so that it can be called. This allows for a PackedFunc to be cast to an + object, stored inside a runtime container, and then cast back to be called + """ + return _ffi_api.CastToPackedFunc(obj) + + @tvm._ffi.register_object("runtime.ADT") class ADT(Object): """Algebatic data type(ADT) object. diff --git a/src/runtime/container.cc b/src/runtime/container.cc index adcaecbc64cf..5d4be3dd5ae5 100644 --- a/src/runtime/container.cc +++ b/src/runtime/container.cc @@ -202,5 +202,18 @@ TVM_REGISTER_GLOBAL("runtime.GetShapeTupleElem").set_body_typed([](ShapeTuple sh ICHECK_LT(idx, shape.size()); return shape[idx]; }); + +// Helper functions for testing: PackedFunc is an ObjectRef in C++ and +// can be put inside runtime containers, but it isn't considered an Object type in Python, +// so the Python counterparts of these containers do not accept it. +// These functions cast the PackedFunc into a generic Object so it can be treated as one +TVM_REGISTER_GLOBAL("runtime.CastPackedFuncToObject").set_body_typed([](const PackedFunc& pf) { + return Downcast(pf); +}); + +TVM_REGISTER_GLOBAL("runtime.CastToPackedFunc").set_body_typed([](const ObjectRef& obj) { + return Downcast(obj); +}); + } // namespace runtime } // namespace tvm diff --git a/tests/python/relay/test_py_converter.py b/tests/python/relay/test_py_converter.py index bd5635e8cf09..5fbf1401019d 100644 --- a/tests/python/relay/test_py_converter.py +++ b/tests/python/relay/test_py_converter.py @@ -20,7 +20,7 @@ from tvm import relay from tvm.relay.testing import to_python, run_as_python from tvm.relay.prelude import Prelude -from tvm.runtime.container import ADT +from tvm.runtime.container import ADT, _uncast_packed_func from tvm.relay.backend.interpreter import RefValue, ConstructorValue # helper: uses a dummy let binding to sequence a list @@ -611,3 +611,49 @@ def reference(x, gamma, beta, moving_mean, moving_var): verify_batch_norm([(20, 10), (10,), (10,), (10,), (10,)]) verify_batch_norm([(10, 50), (50,), (50,), (50,), (50,)]) verify_batch_norm([(30, 40), (40,), (40,), (40,), (40,)]) + + +def test_return_global_var(): + tt = relay.TensorType([1], "float32") + x = relay.Var("x", type_annotation=tt) + identity = relay.Function([x], x, ret_type=tt) + mod = tvm.IRModule() + mod["main"] = identity + main_var = mod.get_global_var("main") + main_func = _uncast_packed_func(run_as_python(main_var, mod=mod)) + + arg = tvm.nd.array(np.array([0.0], dtype="float32")) + res = main_func(arg) + assert arg.numpy() == res.numpy() + + +def test_closure_in_tuple(): + tt = relay.TensorType([1], "float32") + x = relay.Var("x", type_annotation=tt) + identity = relay.Function([x], x, ret_type=tt) + tup = relay.Tuple([identity, identity]) + index = relay.TupleGetItem(tup, 0) + + func = _uncast_packed_func(run_as_python(index)) + arg = tvm.nd.array(np.array([0.0], dtype="float32")) + res = func(arg) + assert arg.numpy() == res.numpy() + + +def test_closure_in_ref(): + tt = relay.TensorType([1], "float32") + x = relay.Var("x", type_annotation=tt) + identity = relay.Function([x], x, ret_type=tt) + gv = relay.GlobalVar("id") + + r = relay.Var("r") + seq = relay.Let( + r, + relay.RefCreate(gv), + relay.Call(relay.RefRead(r), [relay.const(np.array([0.0], dtype="float32"))]), + ) + + mod = tvm.IRModule() + mod[gv] = identity + res = run_as_python(seq, mod=mod) + assert res.numpy() == np.array([0.0], dtype="float32") From dc16c527fd2489b1b0f2e08121de13f3bbc39479 Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Wed, 11 Jan 2023 21:56:50 -0500 Subject: [PATCH 3/8] [py_converter] Convert target strings into TVM Target --- python/tvm/relay/testing/py_converter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py index 10ffbdd2347a..4df25c8af39a 100644 --- a/python/tvm/relay/testing/py_converter.py +++ b/python/tvm/relay/testing/py_converter.py @@ -60,7 +60,7 @@ class PythonConverter(ExprFunctor): def __init__(self, mod, target) -> None: super().__init__() self.mod = mod - self.tgt = target + self.tgt = target if isinstance(target, tvm.target.Target) else tvm.target.Target(target) self.tec = te_compiler.get() self.fun_no = 0 self.var_no = 0 From 6897f496067a70eb772faaf049f736f6def9408c Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Wed, 11 Jan 2023 23:01:02 -0500 Subject: [PATCH 4/8] [py_converter] Do not rely on the name "main" --- python/tvm/relay/testing/py_converter.py | 20 ++++++++++++++++---- tests/python/relay/test_py_converter.py | 19 +++++++++++++++++++ 2 files changed, 35 insertions(+), 4 deletions(-) diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py index 4df25c8af39a..eac91f3562eb 100644 --- a/python/tvm/relay/testing/py_converter.py +++ b/python/tvm/relay/testing/py_converter.py @@ -101,16 +101,28 @@ def optimize(self, prog: Expr): # For a lone global var, there is nothing we need to do if isinstance(unwrapped, relay.GlobalVar): return unwrapped - else: - mod = self.mod.from_expr(unwrapped, self.mod.functions, self.mod.type_definitions) + + # main might be in the mod already and from_expr will not override it if it's there, + # so we need a new name + target_name = self.generate_function_name("target") + + wrapped = unwrapped + if not isinstance(unwrapped, relay.Function): + wrapped = relay.Function(relay.analysis.free_vars(unwrapped), unwrapped) + + # easiest way to make a deep copy -- note that main will not be overridden if it's present + copy_mod = tvm.IRModule.from_expr( + relay.Tuple([]), self.mod.functions, self.mod.type_definitions + ) + copy_mod[target_name] = wrapped # necessary pass: SimplifyInference (otherwise we can't generate code for some operators) # and fusion (to get primitive functions) opts = tvm.transform.Sequential( [relay.transform.SimplifyInference(), relay.transform.FuseOps(fuse_opt_level=0)] ) - mod = opts(mod) - optimized = mod["main"] + copy_mod = opts(copy_mod) + optimized = copy_mod[target_name] return optimized if isinstance(unwrapped, Function) else optimized.body def sanitize(self, name: str) -> str: diff --git a/tests/python/relay/test_py_converter.py b/tests/python/relay/test_py_converter.py index 5fbf1401019d..48aba2905c8c 100644 --- a/tests/python/relay/test_py_converter.py +++ b/tests/python/relay/test_py_converter.py @@ -657,3 +657,22 @@ def test_closure_in_ref(): mod[gv] = identity res = run_as_python(seq, mod=mod) assert res.numpy() == np.array([0.0], dtype="float32") + + +def test_compiling_with_main(): + unit_type = relay.TupleType([]) + unit = relay.Function([], relay.Tuple([]), ret_type=unit_type) + + x = relay.Var("x", type_annotation=unit_type) + identity = relay.Function([x], x, ret_type=unit_type) + + mod = tvm.IRModule() + mod["unit"] = unit + mod["main"] = identity + + import astor + print(astor.to_source(to_python(mod.get_global_var("main")(mod.get_global_var("unit")()), mod=mod))) + + res = run_as_python(mod.get_global_var("main")(mod.get_global_var("unit")()), mod=mod) + assert isinstance(res, ADT) + assert len(res) == 0 From ecb98dfd16912545cfe828aad07ad692b557804e Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Wed, 11 Jan 2023 23:01:41 -0500 Subject: [PATCH 5/8] [py_converter] Use numpy instead of asnumpy (deprecated) --- python/tvm/relay/testing/py_converter.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py index eac91f3562eb..67098696e6cb 100644 --- a/python/tvm/relay/testing/py_converter.py +++ b/python/tvm/relay/testing/py_converter.py @@ -510,7 +510,7 @@ def visit_if(self, if_block: Expr): # need to get the value out of a NDArray to check the condition # equvialent to: val.numpy() - cond_check = ast.Call(ast.Attribute(cond_body, "asnumpy", Load()), [], []) + cond_check = ast.Call(ast.Attribute(cond_body, "numpy", Load()), [], []) ret = ast.IfExp(cond_check, true_body, false_body) return (ret, cond_defs + true_defs + false_defs) From 5a50596db6081b3f04dbf674cfbb07136a356e6a Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Wed, 11 Jan 2023 23:38:55 -0500 Subject: [PATCH 6/8] Remove debug prints from tests --- tests/python/relay/test_py_converter.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/python/relay/test_py_converter.py b/tests/python/relay/test_py_converter.py index 48aba2905c8c..9002a3da6762 100644 --- a/tests/python/relay/test_py_converter.py +++ b/tests/python/relay/test_py_converter.py @@ -70,7 +70,6 @@ def test_create_empty_tuple(): def test_create_scalar(): scalar = relay.const(1) tensor_val = run_as_python(scalar) - print(type(tensor_val)) assert_tensor_value(tensor_val, 1) @@ -670,9 +669,6 @@ def test_compiling_with_main(): mod["unit"] = unit mod["main"] = identity - import astor - print(astor.to_source(to_python(mod.get_global_var("main")(mod.get_global_var("unit")()), mod=mod))) - res = run_as_python(mod.get_global_var("main")(mod.get_global_var("unit")()), mod=mod) assert isinstance(res, ADT) assert len(res) == 0 From b70b9f14b589ccc704a46de898335e276444fd22 Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Thu, 12 Jan 2023 22:21:07 -0500 Subject: [PATCH 7/8] Use FFI_API tuple constructor directly so that it will accept PackedFuncs --- python/tvm/relay/testing/py_converter.py | 42 ++++++++++-------------- python/tvm/runtime/container.py | 18 ---------- src/runtime/container.cc | 12 ------- tests/python/relay/test_py_converter.py | 8 ++--- 4 files changed, 22 insertions(+), 58 deletions(-) diff --git a/python/tvm/relay/testing/py_converter.py b/python/tvm/relay/testing/py_converter.py index 67098696e6cb..44489aa9cf7d 100644 --- a/python/tvm/relay/testing/py_converter.py +++ b/python/tvm/relay/testing/py_converter.py @@ -268,6 +268,15 @@ def create_def(self, func_name: str, arguments: [str], body, register_packed: bo None, ) + def create_tuple(self, fields): + """ + Given the ASTs for tuple fields, produce an AST that creates a + tuple value with those fields + """ + # Use the FFI API directly so that PackedFuncs will be correctly converted to ObjectRef. + # Using tvm.runtime.container.tuple_object fails to convert PackedFuncs in Python + return self.create_call("_container._ffi_api.Tuple", fields) + def create_op_call(self, op: Function, relay_args, py_args): """Lowers the passed primitive function, registers it in TVM's global compiler, and produces a call to the lowered function in @@ -321,8 +330,7 @@ def convert_output(ret_type): assignments += inner_assignments extra_args += inner_args fields.append(inner_output) - fields = [ast.List(fields, Load())] - return (assignments, extra_args, self.create_call("_container.tuple_object", fields)) + return (assignments, extra_args, self.create_tuple(fields)) # create a function to wrap the call of the lowered op and return # a call to that function @@ -450,14 +458,8 @@ def visit_global_var(self, gvar: Expr): # we don't need to add numbers to global var names because # the *names* are checked for uniqueness in the mod func_name = str(gvar.name_hint) - # load in the packed func and cast to object so it can be put in containers - return ( - self.create_call( - "_container._cast_packed_func", - [self.create_call("tvm.get_global_func", [ast.Constant(value=func_name)])], - ), - [], - ) + # load in the packed func + return (self.create_call("tvm.get_global_func", [ast.Constant(value=func_name)]), []) def visit_let(self, letexp: Expr): # To properly account for scoping and ensure that the entire node produces an expression, @@ -495,8 +497,7 @@ def let_thunk(var): def visit_tuple(self, tup: Expr): fields, ret_defs = self.convert_fields(tup.fields) - fields = [ast.List(fields, Load())] - return (self.create_call("_container.tuple_object", fields), ret_defs) + return (self.create_tuple(fields), ret_defs) def visit_tuple_getitem(self, tgi: Expr): tup, tup_defs = self.visit(tgi.tuple_value) @@ -529,12 +530,9 @@ def visit_constant(self, constant: Expr): def visit_function(self, func: Expr): # Python's lambdas are very restrictive, so we do "name" inline functions converted_func, func_name = self.convert_func_node(func) - # load in the PackedFunc and cast to object so it can be stored in other containers + # load in the PackedFunc return ( - self.create_call( - "_container._cast_packed_func", - [self.create_call("tvm.get_global_func", [ast.Constant(value=func_name)])], - ), + self.create_call("tvm.get_global_func", [ast.Constant(value=func_name)]), [converted_func], ) @@ -566,10 +564,7 @@ def visit_call(self, call: Expr): # ordinary function converted_func, defs = self.visit(func) defs += field_defs - # have to convert func back from an object to be able to call it - recast_func = self.create_call("_container._uncast_packed_func", [converted_func]) - - return (ast.Call(recast_func, fields, []), defs) + return (ast.Call(converted_func, fields, []), defs) def visit_ref_create(self, ref: Expr): val, defs = self.visit(ref.value) @@ -595,7 +590,7 @@ def visit_ref_write(self, write: Expr): + val_defs + [ Assign([ast.Attribute(ref, "value", Store())], val), - Return(self.create_call("_container.tuple_object", [])), + Return(self.create_tuple([])), ], ) return (self.create_call(thunk_name, []), [thunk]) @@ -653,8 +648,7 @@ def run_as_python(expr: Expr, mod=None, target=tvm.target.Target("llvm")): """Converts the given Relay expression into a Python script and executes it. - Note that closures will be returned as PackedFuncs cast to Objects; - to execute them, they need to be cast back using runtime.container._uncast_packed_func + Note that closures will be returned as PackedFuncs """ mod = mod if mod is not None else tvm.IRModule() py_ast = to_python(expr, mod, target) diff --git a/python/tvm/runtime/container.py b/python/tvm/runtime/container.py index a93589b34360..e960392a0f0d 100644 --- a/python/tvm/runtime/container.py +++ b/python/tvm/runtime/container.py @@ -61,24 +61,6 @@ def getitem_helper(obj, elem_getter, length, idx): return elem_getter(obj, idx) -# helper functions for putting PackedFuncs inside other containers -def _cast_packed_func(pf: PackedFunc) -> Object: - """ - Helper function for testing: Casts the given PackedFunc into a TVM object - so that it can be stored inside other runtime containers - """ - return _ffi_api.CastPackedFuncToObject(pf) - - -def _uncast_packed_func(obj: Object) -> PackedFunc: - """ - Helper function for testing: Casts the passed Object to a PackedFunc - so that it can be called. This allows for a PackedFunc to be cast to an - object, stored inside a runtime container, and then cast back to be called - """ - return _ffi_api.CastToPackedFunc(obj) - - @tvm._ffi.register_object("runtime.ADT") class ADT(Object): """Algebatic data type(ADT) object. diff --git a/src/runtime/container.cc b/src/runtime/container.cc index 5d4be3dd5ae5..7b5105a3fc94 100644 --- a/src/runtime/container.cc +++ b/src/runtime/container.cc @@ -203,17 +203,5 @@ TVM_REGISTER_GLOBAL("runtime.GetShapeTupleElem").set_body_typed([](ShapeTuple sh return shape[idx]; }); -// Helper functions for testing: PackedFunc is an ObjectRef in C++ and -// can be put inside runtime containers, but it isn't considered an Object type in Python, -// so the Python counterparts of these containers do not accept it. -// These functions cast the PackedFunc into a generic Object so it can be treated as one -TVM_REGISTER_GLOBAL("runtime.CastPackedFuncToObject").set_body_typed([](const PackedFunc& pf) { - return Downcast(pf); -}); - -TVM_REGISTER_GLOBAL("runtime.CastToPackedFunc").set_body_typed([](const ObjectRef& obj) { - return Downcast(obj); -}); - } // namespace runtime } // namespace tvm diff --git a/tests/python/relay/test_py_converter.py b/tests/python/relay/test_py_converter.py index 9002a3da6762..d43ec5861b10 100644 --- a/tests/python/relay/test_py_converter.py +++ b/tests/python/relay/test_py_converter.py @@ -18,9 +18,9 @@ import tvm from tvm import te from tvm import relay -from tvm.relay.testing import to_python, run_as_python +from tvm.relay.testing import run_as_python from tvm.relay.prelude import Prelude -from tvm.runtime.container import ADT, _uncast_packed_func +from tvm.runtime.container import ADT from tvm.relay.backend.interpreter import RefValue, ConstructorValue # helper: uses a dummy let binding to sequence a list @@ -619,7 +619,7 @@ def test_return_global_var(): mod = tvm.IRModule() mod["main"] = identity main_var = mod.get_global_var("main") - main_func = _uncast_packed_func(run_as_python(main_var, mod=mod)) + main_func = run_as_python(main_var, mod=mod) arg = tvm.nd.array(np.array([0.0], dtype="float32")) res = main_func(arg) @@ -633,7 +633,7 @@ def test_closure_in_tuple(): tup = relay.Tuple([identity, identity]) index = relay.TupleGetItem(tup, 0) - func = _uncast_packed_func(run_as_python(index)) + func = run_as_python(index) arg = tvm.nd.array(np.array([0.0], dtype="float32")) res = func(arg) assert arg.numpy() == res.numpy() From edd3ee33d3e613e8d342f9136ad8e508e64bf215 Mon Sep 17 00:00:00 2001 From: Steven Lyubomirsky Date: Thu, 12 Jan 2023 22:32:38 -0500 Subject: [PATCH 8/8] Remove unused import in container.py --- python/tvm/runtime/container.py | 1 - 1 file changed, 1 deletion(-) diff --git a/python/tvm/runtime/container.py b/python/tvm/runtime/container.py index e960392a0f0d..cae408b6121d 100644 --- a/python/tvm/runtime/container.py +++ b/python/tvm/runtime/container.py @@ -18,7 +18,6 @@ import tvm._ffi from .object import Object, PyNativeObject from .object_generic import ObjectTypes -from .packed_func import PackedFunc from . import _ffi_api