Move external codegen test helpers into utils

This is so they can be re-used as part of other tests which don't extend test_external_codegen.py

tests/python/relay/test_external_codegen.py

@@ -16,88 +16,22 @@
 # under the License.
 """Unit tests for graph partitioning."""
 
-import os
 import sys
 from collections import OrderedDict
 import numpy as np
 import pytest
 
 import tvm
 from tvm import relay, runtime
-from tvm.contrib import utils
 from tvm.relay.build_module import bind_params_by_name
 from tvm.relay.op.annotation import compiler_begin, compiler_end
-
-
-def update_lib(lib):
-    test_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
-    source_dir = os.path.join(test_dir, "..", "..", "..")
-    contrib_path = os.path.join(source_dir, "src", "runtime", "contrib")
-
-    kwargs = {}
-    kwargs["options"] = ["-O2", "-std=c++14", "-I" + contrib_path]
-    tmp_path = utils.tempdir()
-    lib_name = "lib.so"
-    lib_path = tmp_path.relpath(lib_name)
-    lib.export_library(lib_path, fcompile=False, **kwargs)
-    lib = tvm.runtime.load_module(lib_path)
-
-    return lib
-
-
-def check_vm_result(mod, map_inputs, out_shape, result, tol=1e-5, target="llvm", device=tvm.cpu()):
-    with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
-        exe = relay.vm.compile(mod, target=target)
-    code, lib = exe.save()
-    lib = update_lib(lib)
-    exe = runtime.vm.Executable.load_exec(code, lib)
-    vm = runtime.vm.VirtualMachine(exe, device)
-    out = vm.run(**map_inputs)
-    tvm.testing.assert_allclose(out.numpy(), result, rtol=tol, atol=tol)
-
-
-def check_graph_executor_result(
-    mod, map_inputs, out_shape, result, tol=1e-5, target="llvm", device=tvm.cpu()
-):
-    with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
-        json, lib, _ = relay.build(mod, target=target)
-    lib = update_lib(lib)
-    rt_mod = tvm.contrib.graph_executor.create(json, lib, device)
-
-    for name, data in map_inputs.items():
-        rt_mod.set_input(name, data)
-    rt_mod.run()
-    out = tvm.nd.empty(out_shape, device=device)
-    out = rt_mod.get_output(0, out)
-
-    tvm.testing.assert_allclose(out.numpy(), result, rtol=tol, atol=tol)
-
-
-def check_aot_executor_result(
-    mod, map_inputs, out_shape, result, tol=1e-5, target="llvm", device=tvm.cpu()
-):
-    if tvm.support.libinfo().get("USE_MICRO", "OFF") != "ON":
-        pytest.skip("MicroTVM support not enabled. Set USE_MICRO=ON in config.cmake to enable.")
-
-    # Late import to avoid breaking test with USE_MICRO=OFF.
-    from aot.aot_test_utils import AOTTestModel, AOT_DEFAULT_RUNNER, compile_and_run
-
-    interface_api = "packed"
-    use_unpacked_api = False
-    test_runner = AOT_DEFAULT_RUNNER
-    compile_and_run(
-        AOTTestModel(module=mod, inputs=map_inputs, outputs=[result]),
-        test_runner,
-        interface_api,
-        use_unpacked_api,
-    )
-
-
-def set_external_func_attr(func, compiler, ext_symbol):
-    func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
-    func = func.with_attr("Compiler", compiler)
-    func = func.with_attr("global_symbol", ext_symbol)
-    return func
+from utils.external_codegen import (
+    update_lib,
+    set_external_func_attr,
+    check_aot_executor_result,
+    check_graph_executor_result,
+    check_vm_result,
+)
 
 
 @pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now")

tests/python/relay/utils/external_codegen.py

@@ -0,0 +1,94 @@
+# 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.
+"""Utilities for testing external code generation"""
+
+import os
+
+import pytest
+
+import tvm
+from tvm import relay, runtime
+from tvm.contrib import utils
+from tests.python.relay.aot.aot_test_utils import AOTTestModel, compile_and_run
+
+
+def update_lib(lib):
+    test_dir = os.path.dirname(os.path.realpath(os.path.expanduser(__file__)))
+    source_dir = os.path.join(test_dir, "..", "..", "..")
+    contrib_path = os.path.join(source_dir, "src", "runtime", "contrib")
+
+    kwargs = {}
+    kwargs["options"] = ["-O2", "-std=c++14", "-I" + contrib_path]
+    tmp_path = utils.tempdir()
+    lib_name = "lib.so"
+    lib_path = tmp_path.relpath(lib_name)
+    lib.export_library(lib_path, fcompile=False, **kwargs)
+    lib = tvm.runtime.load_module(lib_path)
+
+    return lib
+
+
+def check_vm_result(mod, map_inputs, out_shape, result, tol=1e-5, target="llvm", device=tvm.cpu()):
+    with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
+        exe = relay.vm.compile(mod, target=target)
+    code, lib = exe.save()
+    lib = update_lib(lib)
+    exe = runtime.vm.Executable.load_exec(code, lib)
+    vm = runtime.vm.VirtualMachine(exe, device)
+    out = vm.run(**map_inputs)
+    tvm.testing.assert_allclose(out.numpy(), result, rtol=tol, atol=tol)
+
+
+def check_graph_executor_result(
+    mod, map_inputs, out_shape, result, tol=1e-5, target="llvm", device=tvm.cpu()
+):
+    with tvm.transform.PassContext(opt_level=3, disabled_pass=["AlterOpLayout"]):
+        json, lib, _ = relay.build(mod, target=target)
+    lib = update_lib(lib)
+    rt_mod = tvm.contrib.graph_executor.create(json, lib, device)
+
+    for name, data in map_inputs.items():
+        rt_mod.set_input(name, data)
+    rt_mod.run()
+    out = tvm.nd.empty(out_shape, device=device)
+    out = rt_mod.get_output(0, out)
+
+    tvm.testing.assert_allclose(out.numpy(), result, rtol=tol, atol=tol)
+
+
+def check_aot_executor_result(
+    mod, map_inputs, out_shape, result, tol=1e-5, target="llvm", device=tvm.cpu()
+):
+    if tvm.support.libinfo().get("USE_MICRO", "OFF") != "ON":
+        pytest.skip("MicroTVM support not enabled. Set USE_MICRO=ON in config.cmake to enable.")
+
+    interface_api = "packed"
+    use_unpacked_api = False
+    use_calculated_workspaces = True
+    compile_and_run(
+        AOTTestModel(module=mod, inputs=map_inputs, outputs=[result]),
+        interface_api,
+        use_unpacked_api,
+        use_calculated_workspaces,
+    )
+
+
+def set_external_func_attr(func, compiler, ext_symbol):
+    func = func.with_attr("Primitive", tvm.tir.IntImm("int32", 1))
+    func = func.with_attr("Compiler", compiler)
+    func = func.with_attr("global_symbol", ext_symbol)
+    return func