From 8bbed387d2153e41103fddb3d0ec8fe308b5d480 Mon Sep 17 00:00:00 2001 From: Christopher Sidebottom Date: Wed, 15 Sep 2021 04:37:51 +0100 Subject: [PATCH] Move external codegen test helpers into utils (#9008) This is so they can be re-used as part of other tests which don't extend test_external_codegen.py I've identified that `test_json_runtime.py` and `test_pass_partition_graph.py` use a very similar but slightly different variant of these functions for future iterations. --- tests/python/relay/test_external_codegen.py | 105 ++-------------- tests/python/relay/utils/external_codegen.py | 125 +++++++++++++++++++ 2 files changed, 137 insertions(+), 93 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 30db5facc208..ad5f2aa9d4fa 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,86 +23,17 @@ 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 +from utils.external_codegen import ( + update_lib, + set_external_func_attr, + parametrize_external_codegen_checks, + parametrize_external_json_codegen_checks, +) -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 - - -@pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now") -@pytest.mark.parametrize( - "check_result", [check_vm_result, check_graph_executor_result, check_aot_executor_result] -) +@parametrize_external_codegen_checks def test_multi_node_subgraph(check_result): x = relay.var("x", shape=(10, 10)) w0 = relay.var("w0", shape=(10, 10)) @@ -170,10 +100,7 @@ def test_multi_node_subgraph(check_result): ) -@pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now") -@pytest.mark.parametrize( - "check_result", [check_vm_result, check_graph_executor_result, check_aot_executor_result] -) +@parametrize_external_codegen_checks def test_extern_gcc_single_op(check_result): x = relay.var("x", shape=(8, 8)) y = relay.var("y", shape=(8, 8)) @@ -191,10 +118,7 @@ def test_extern_gcc_single_op(check_result): check_result(mod, {"x": x_data, "y": y_data}, (8, 8), x_data + y_data) -@pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now") -@pytest.mark.parametrize( - "check_result", [check_vm_result, check_graph_executor_result, check_aot_executor_result] -) +@parametrize_external_codegen_checks def test_extern_gcc_single_op_int(check_result): x = relay.var("x", shape=(8, 8), dtype="int32") y = relay.var("y", shape=(8, 8), dtype="int32") @@ -212,10 +136,7 @@ def test_extern_gcc_single_op_int(check_result): check_result(mod, {"x": x_data, "y": y_data}, (8, 8), x_data + y_data) -@pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now") -@pytest.mark.parametrize( - "check_result", [check_vm_result, check_graph_executor_result, check_aot_executor_result] -) +@parametrize_external_codegen_checks def test_extern_gcc(check_result): x = relay.var("x", shape=(2, 2)) y = relay.var("y", shape=(2, 2)) @@ -292,12 +213,11 @@ def constant_updater(expr, symbol): tvm._ffi.registry.remove_global_func("relay.ext.ccompiler.constant_updater") -@pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now") @pytest.mark.skipif( not tvm.get_global_func("relay.ext.dnnl", True), reason="skip because DNNL codegen is not available", ) -@pytest.mark.parametrize("check_result", [check_vm_result, check_graph_executor_result]) +@parametrize_external_json_codegen_checks def test_extern_dnnl(check_result): dtype = "float32" ishape = (1, 32, 14, 14) @@ -335,12 +255,11 @@ def test_extern_dnnl(check_result): ) -@pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now") @pytest.mark.skipif( not tvm.get_global_func("relay.ext.dnnl", True), reason="skip because DNNL codegen is not available", ) -@pytest.mark.parametrize("check_result", [check_vm_result, check_graph_executor_result]) +@parametrize_external_json_codegen_checks def test_extern_dnnl_const(check_result): dtype = "float32" ishape = (1, 32, 14, 14) diff --git a/tests/python/relay/utils/external_codegen.py b/tests/python/relay/utils/external_codegen.py new file mode 100644 index 000000000000..85583f6ccc5d --- /dev/null +++ b/tests/python/relay/utils/external_codegen.py @@ -0,0 +1,125 @@ +# 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 sys + +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 + + +skip_windows = pytest.mark.skipif(sys.platform == "win32", reason="Skip test on Windows for now") +skip_micro = pytest.mark.skipif( + tvm.support.libinfo().get("USE_MICRO", "OFF") != "ON", + reason="MicroTVM support not enabled. Set USE_MICRO=ON in config.cmake to enable.", +) + + +def parametrize_external_codegen_checks(test): + """Parametrize over the various check_result functions which are available""" + return pytest.mark.parametrize( + "check_result", + [ + pytest.param(check_aot_executor_result, marks=[skip_windows, skip_micro]), + pytest.param(check_graph_executor_result, marks=[skip_windows]), + pytest.param(check_vm_result, marks=[skip_windows]), + ], + )(test) + + +def parametrize_external_json_codegen_checks(test): + """Parametrize over the various check_result functions which are available for JSON""" + return pytest.mark.parametrize( + "check_result", + [ + pytest.param(check_graph_executor_result, marks=[skip_windows]), + pytest.param(check_vm_result, marks=[skip_windows]), + ], + )(test) + + +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"]): + executor_factory = relay.build(mod, target=target) + lib = update_lib(executor_factory.lib) + rt_mod = tvm.contrib.graph_executor.create(executor_factory.graph_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() +): + # 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