diff --git a/tests/python/unittest/test_link_params.py b/tests/python/unittest/test_link_params.py index 0e56f6c00ae2..b2e7398c0ee0 100644 --- a/tests/python/unittest/test_link_params.py +++ b/tests/python/unittest/test_link_params.py @@ -20,7 +20,6 @@ import os import re import sys -import tempfile import numpy as np import pytest @@ -39,15 +38,17 @@ # The data types that are linkable. -LINKABLE_DTYPES = ( - [f"uint{b}" for b in (8, 16, 32, 64)] - + [f"int{b}" for b in (8, 16, 32, 64)] - + ["float32", "float64"] +linkable_dtype = tvm.testing.parameter( + *( + [f"uint{b}" for b in (8, 16, 32, 64)] + + [f"int{b}" for b in (8, 16, 32, 64)] + + ["float32", "float64"] + ) ) def dtype_info(dtype): - """Lookup numpy type info for the given string dtype (of LINKABLE_DTYPES above).""" + """Lookup numpy type info for the given string dtype (of linkable_dtype params above).""" if "int" in dtype: return np.iinfo(getattr(np, dtype)) else: @@ -181,58 +182,56 @@ def _add_decl(name, dtype): @tvm.testing.requires_llvm -def test_llvm_link_params(): - for dtype in LINKABLE_DTYPES: - ir_mod, param_init = _make_mod_and_params(dtype) - rand_input = _make_random_tensor(dtype, INPUT_SHAPE) - target = "llvm" - runtime = Runtime("crt", {"system-lib": True}) - executor = Executor("graph", {"link-params": True}) - with tvm.transform.PassContext(opt_level=3): - lib = tvm.relay.build( - ir_mod, target, runtime=runtime, executor=executor, params=param_init - ) - - # NOTE: Need to export_library() and load_library() to link all the Module(llvm, ...) - # against one another. - temp_dir = tempfile.mkdtemp() - export_file = os.path.join(temp_dir, "lib.so") - lib.lib.export_library(export_file) - mod = tvm.runtime.load_module(export_file) - assert set(lib.params.keys()) == {"p0", "p1"} # NOTE: op folded - assert mod.get_function("TVMSystemLibEntryPoint") != None - - graph = json.loads(lib.graph_json) - for p in lib.params: - _verify_linked_param(dtype, lib, mod, graph, p) or found_one - - # Wrap in function to explicitly deallocate the runtime. - def _run_linked(lib, mod): - graph_json, _, _ = lib - graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0)) - graph_rt.set_input("rand_input", rand_input) # NOTE: params not required. - graph_rt.run() - return graph_rt.get_output(0) - - linked_output = _run_linked(lib, mod) - - runtime = Runtime("cpp", {"system-lib": True}) - with tvm.transform.PassContext(opt_level=3): - lib = tvm.relay.build(ir_mod, "llvm", runtime=runtime, params=param_init) - - def _run_unlinked(lib): - graph_json, mod, lowered_params = lib - graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0)) - graph_rt.set_input("rand_input", rand_input, **lowered_params) - graph_rt.run() - return graph_rt.get_output(0) - - unlinked_output = _run_unlinked(lib) - - if "int" in dtype: - np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy()) - else: - np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy()) +def test_llvm_link_params(linkable_dtype): + ir_mod, param_init = _make_mod_and_params(linkable_dtype) + rand_input = _make_random_tensor(linkable_dtype, INPUT_SHAPE) + main_func = ir_mod["main"] + target = "llvm" + runtime = Runtime("crt", {"system-lib": True}) + executor = Executor("graph", {"link-params": True}) + with tvm.transform.PassContext(opt_level=3): + lib = tvm.relay.build(ir_mod, target, runtime=runtime, executor=executor, params=param_init) + + # NOTE: Need to export_library() and load_library() to link all the Module(llvm, ...) + # against one another. + temp_dir = utils.TempDirectory() + export_file = temp_dir / "lib.so" + lib.lib.export_library(export_file) + mod = tvm.runtime.load_module(export_file) + assert set(lib.params.keys()) == {"p0", "p1"} # NOTE: op folded + assert mod.get_function("TVMSystemLibEntryPoint") != None + + graph = json.loads(lib.graph_json) + for p in lib.params: + _verify_linked_param(linkable_dtype, lib, mod, graph, p) or found_one + + # Wrap in function to explicitly deallocate the runtime. + def _run_linked(lib, mod): + graph_json, _, _ = lib + graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0)) + graph_rt.set_input("rand_input", rand_input) # NOTE: params not required. + graph_rt.run() + return graph_rt.get_output(0) + + linked_output = _run_linked(lib, mod) + + runtime = Runtime("cpp", {"system-lib": True}) + with tvm.transform.PassContext(opt_level=3): + lib = tvm.relay.build(ir_mod, "llvm", runtime=runtime, params=param_init) + + def _run_unlinked(lib): + graph_json, mod, lowered_params = lib + graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0)) + graph_rt.set_input("rand_input", rand_input, **lowered_params) + graph_rt.run() + return graph_rt.get_output(0) + + unlinked_output = _run_unlinked(lib) + + if "int" in linkable_dtype: + np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy()) + else: + np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy()) def _get_c_datatype(dtype): @@ -267,141 +266,140 @@ def _format_c_value(dtype, width, x): HEX_NUM_RE = re.compile(r"[+\-]?(?:(?:0x[0-9A-Fa-f.p+-]+)|(?:INFINITY)|(?:NAN))") -def test_c_link_params(): +def test_c_link_params(linkable_dtype): temp_dir = utils.tempdir() - for dtype in LINKABLE_DTYPES: - mod, param_init = _make_mod_and_params(dtype) - rand_input = _make_random_tensor(dtype, INPUT_SHAPE) - target = "c" - executor = Executor("graph", {"link-params": True}) - with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): - lib = tvm.relay.build(mod, target, executor=executor, params=param_init) - assert set(lib.params.keys()) == {"p0", "p1"} # NOTE: op folded - - src = lib.lib.get_source() - lib.lib.save(temp_dir.relpath("test.c"), "c") - c_dtype = _get_c_datatype(dtype) - src_lines = src.split("\n") - param = lib.params["p0"].numpy().reshape(np.prod(KERNEL_SHAPE)) - param_def = f'static const {c_dtype} __attribute__((section(".rodata.tvm"), aligned(16))) __tvm_param__p0[{np.prod(param.shape)}] = {{' - for i, line in enumerate(src_lines): - if line == param_def: - i += 1 - break - else: - assert False, f'did not find parameter definition "{param_def}":\n{src}' - - cursor = 0 - width = dtype_info(dtype).bits // 4 + 2 - if dtype.startswith("int"): - width += 1 # Account for sign - - while "};" not in src_lines[i]: - for match in HEX_NUM_RE.finditer(src_lines[i]): - assert match.group() == _format_c_value(dtype, width, param[cursor]), ( - f'p0 byte {cursor}: want "{_format_c_value(dtype, width, param[cursor])}" got ' - f'"{match.group(0)}"; full p0 follows:\n{src}' - ) - cursor += 1 + mod, param_init = _make_mod_and_params(linkable_dtype) + rand_input = _make_random_tensor(linkable_dtype, INPUT_SHAPE) + main_func = mod["main"] + target = "c" + executor = Executor("graph", {"link-params": True}) + with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): + lib = tvm.relay.build(mod, target, executor=executor, params=param_init) + assert set(lib.params.keys()) == {"p0", "p1"} # NOTE: op folded + + src = lib.lib.get_source() + lib.lib.save(temp_dir.relpath("test.c"), "c") + c_dtype = _get_c_datatype(linkable_dtype) + src_lines = src.split("\n") + param = lib.params["p0"].numpy().reshape(np.prod(KERNEL_SHAPE)) + param_def = f'static const {c_dtype} __attribute__((section(".rodata.tvm"), aligned(16))) __tvm_param__p0[{np.prod(param.shape)}] = {{' + + for i, line in enumerate(src_lines): + if line == param_def: i += 1 - - assert cursor == np.prod(param.shape) - - # Need a unique name per library to avoid dlopen caching the lib load. - lib_path = temp_dir.relpath(f"test-{dtype}-linked.so") - lib["remove_params"]().export_library(lib_path) - lib_mod = tvm.runtime.load_module(lib_path) - - # lib_mod = lib_factory['default']() - graph = json.loads(lib.graph_json) - for p in lib.params: - _verify_linked_param(dtype, lib, lib_mod, graph, p) - - # Wrap in function to explicitly deallocate the runtime. - def _run_linked(lib_mod): - graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0))) - graph_rt.set_input("rand_input", rand_input) # NOTE: params not required. - graph_rt.run() - - return graph_rt.get_output(0) - - linked_output = _run_linked(lib_mod) - - linked_params = lib.params - with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): - lib = tvm.relay.build(mod, "c", params=param_init) - _, _, params = lib - # Need a unique name per library to avoid dlopen caching the lib load. - lib_path = temp_dir.relpath(f"test-{dtype}-unlinked.so") - lib.export_library(lib_path) - lib_mod = tvm.runtime.load_module(lib_path) - - def _run_unlinked(lib_mod): - graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0))) - graph_rt.set_input("rand_input", rand_input, **params) - graph_rt.run() - return graph_rt.get_output(0) - - unlinked_output = _run_unlinked(lib_mod) - - if "int" in dtype: - np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy()) + break else: - np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy()) + assert False, f'did not find parameter definition "{param_def}":\n{src}' + + cursor = 0 + width = dtype_info(linkable_dtype).bits // 4 + 2 + if linkable_dtype.startswith("int"): + width += 1 # Account for sign + + while "};" not in src_lines[i]: + for match in HEX_NUM_RE.finditer(src_lines[i]): + assert match.group() == _format_c_value(linkable_dtype, width, param[cursor]), ( + f'p0 byte {cursor}: want "{_format_c_value(linkable_dtype, width, param[cursor])}" got ' + f'"{match.group(0)}"; full p0 follows:\n{src}' + ) + cursor += 1 + i += 1 + + assert cursor == np.prod(param.shape) + + # Need a unique name per library to avoid dlopen caching the lib load. + lib_path = temp_dir.relpath(f"test-{linkable_dtype}-linked.so") + lib["remove_params"]().export_library(lib_path) + lib_mod = tvm.runtime.load_module(lib_path) + + # lib_mod = lib_factory['default']() + graph = json.loads(lib.graph_json) + for p in lib.params: + _verify_linked_param(linkable_dtype, lib, lib_mod, graph, p) + + # Wrap in function to explicitly deallocate the runtime. + def _run_linked(lib_mod): + graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0))) + graph_rt.set_input("rand_input", rand_input) # NOTE: params not required. + graph_rt.run() + + return graph_rt.get_output(0) + + linked_output = _run_linked(lib_mod) + + linked_params = lib.params + with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): + lib = tvm.relay.build(mod, "c", params=param_init) + _, _, params = lib + # Need a unique name per library to avoid dlopen caching the lib load. + lib_path = temp_dir.relpath(f"test-{linkable_dtype}-unlinked.so") + lib.export_library(lib_path) + lib_mod = tvm.runtime.load_module(lib_path) + + def _run_unlinked(lib_mod): + graph_rt = tvm.contrib.graph_executor.GraphModule(lib_mod["default"](tvm.cpu(0))) + graph_rt.set_input("rand_input", rand_input, **params) + graph_rt.run() + return graph_rt.get_output(0) + + unlinked_output = _run_unlinked(lib_mod) + + if "int" in linkable_dtype: + np.testing.assert_equal(unlinked_output.numpy(), linked_output.numpy()) + else: + np.testing.assert_allclose(unlinked_output.numpy(), linked_output.numpy()) @tvm.testing.requires_micro -def test_crt_link_params(): +def test_crt_link_params(linkable_dtype): from tvm import micro - for dtype in LINKABLE_DTYPES: - mod, param_init = _make_mod_and_params(dtype) - rand_input = _make_random_tensor(dtype, INPUT_SHAPE) - target = "c" - runtime = Runtime("crt", {"system-lib": True}) - executor = Executor("graph", {"link-params": True}) - with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): - factory = tvm.relay.build( - mod, target, runtime=runtime, executor=executor, params=param_init + mod, param_init = _make_mod_and_params(linkable_dtype) + rand_input = _make_random_tensor(linkable_dtype, INPUT_SHAPE) + main_func = mod["main"] + target = "c" + runtime = Runtime("crt", {"system-lib": True}) + executor = Executor("graph", {"link-params": True}) + with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): + factory = tvm.relay.build( + mod, target, runtime=runtime, executor=executor, params=param_init + ) + assert set(factory.get_params().keys()) == {"p0", "p1"} # NOTE: op folded + + temp_dir = tvm.contrib.utils.tempdir() + template_project_dir = os.path.join(tvm.micro.get_standalone_crt_dir(), "template", "host") + project = tvm.micro.generate_project( + template_project_dir, factory, temp_dir / "project", {"verbose": 1} + ) + project.build() + project.flash() + with tvm.micro.Session(project.transport()) as sess: + graph_rt = tvm.micro.session.create_local_graph_executor( + factory.get_graph_json(), sess.get_system_lib(), sess.device ) - assert set(factory.get_params().keys()) == {"p0", "p1"} # NOTE: op folded - temp_dir = tvm.contrib.utils.tempdir() - template_project_dir = os.path.join( - tvm.micro.get_standalone_crt_dir(), "template", "host" - ) - project = tvm.micro.generate_project( - template_project_dir, factory, temp_dir / "project", {"verbose": 1} - ) - project.build() - project.flash() - with tvm.micro.Session(project.transport()) as sess: - graph_rt = tvm.micro.session.create_local_graph_executor( - factory.get_graph_json(), sess.get_system_lib(), sess.device - ) - - # NOTE: not setting params here. - graph_rt.set_input("rand_input", rand_input) - graph_rt.run() - linked_output = graph_rt.get_output(0).numpy() + # NOTE: not setting params here. + graph_rt.set_input("rand_input", rand_input) + graph_rt.run() + linked_output = graph_rt.get_output(0).numpy() - runtime = Runtime("cpp", {"system-lib": True}) - with tvm.transform.PassContext(opt_level=3): - lib = tvm.relay.build(mod, "llvm", runtime=runtime, params=param_init) + runtime = Runtime("cpp", {"system-lib": True}) + with tvm.transform.PassContext(opt_level=3): + lib = tvm.relay.build(mod, "llvm", runtime=runtime, params=param_init) - def _run_unlinked(lib): - graph_json, mod, lowered_params = lib - graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0)) - graph_rt.set_input("rand_input", rand_input, **lowered_params) - graph_rt.run() - return graph_rt.get_output(0).numpy() + def _run_unlinked(lib): + graph_json, mod, lowered_params = lib + graph_rt = tvm.contrib.graph_executor.create(graph_json, mod, tvm.cpu(0)) + graph_rt.set_input("rand_input", rand_input, **lowered_params) + graph_rt.run() + return graph_rt.get_output(0).numpy() - unlinked_output = _run_unlinked(lib) + unlinked_output = _run_unlinked(lib) - if "int" in dtype: - np.testing.assert_equal(unlinked_output, linked_output) - else: - np.testing.assert_allclose(unlinked_output, linked_output) + if "int" in linkable_dtype: + np.testing.assert_equal(unlinked_output, linked_output) + else: + np.testing.assert_allclose(unlinked_output, linked_output) if __name__ == "__main__":