From 60a4db6616e3e6863ab85d59f62824abd16e72b7 Mon Sep 17 00:00:00 2001 From: Chris Sidebottom Date: Thu, 29 Jul 2021 18:09:34 +0000 Subject: [PATCH] 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 | 80 ++--------------- tests/python/relay/utils/external_codegen.py | 94 ++++++++++++++++++++ 2 files changed, 101 insertions(+), 73 deletions(-) create mode 100644 tests/python/relay/utils/external_codegen.py diff --git a/tests/python/relay/test_external_codegen.py b/tests/python/relay/test_external_codegen.py index 30db5facc2085..14a2f234ff5df 100644 --- a/tests/python/relay/test_external_codegen.py +++ b/tests/python/relay/test_external_codegen.py @@ -16,7 +16,6 @@ # under the License. """Unit tests for graph partitioning.""" -import os import sys from collections import OrderedDict import numpy as np @@ -24,80 +23,15 @@ 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") diff --git a/tests/python/relay/utils/external_codegen.py b/tests/python/relay/utils/external_codegen.py new file mode 100644 index 0000000000000..15211f970a1bb --- /dev/null +++ b/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