Skip to content

Move external codegen test helpers into utils #9008

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
masahi merged 1 commit into from
Sep 15, 2021
Merged

Move external codegen test helpers into utils #9008

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
105 changes: 12 additions & 93 deletions tests/python/relay/test_external_codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,94 +16,24 @@
# 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
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))
Expand Down Expand Up @@ -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))
Expand All @@ -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")
Expand All @@ -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))
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down
125 changes: 125 additions & 0 deletions tests/python/relay/utils/external_codegen.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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