This is an automated email from the ASF dual-hosted git repository.
csullivan pushed a commit to branch unity
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/unity by this push:
new 854f2e9bc1 [Unity][Hexagon] Enable Relax VM for Hexagon (#14415)
854f2e9bc1 is described below
commit 854f2e9bc12fb8bc938db0d65f70835e129a48f9
Author: Ivan Sidorenko <[email protected]>
AuthorDate: Tue Mar 28 19:28:32 2023 +0300
[Unity][Hexagon] Enable Relax VM for Hexagon (#14415)
This is attempt to port PR (https://github.com/tlc-pack/relax/pull/167)
(submitted by @psrivas2 and @YuchenJin) from tlc-pack/relax to enable
Hexagon tests with Relax VM.
---
python/tvm/contrib/hexagon/session.py | 37 +++-
python/tvm/runtime/module.py | 2 +-
src/runtime/hexagon/hexagon_module.h | 1 +
.../contrib/test_hexagon/test_relax_integration.py | 236 +++++++++++++++++++++
4 files changed, 273 insertions(+), 3 deletions(-)
diff --git a/python/tvm/contrib/hexagon/session.py
b/python/tvm/contrib/hexagon/session.py
index 506f1d968d..5690fa44c1 100644
--- a/python/tvm/contrib/hexagon/session.py
+++ b/python/tvm/contrib/hexagon/session.py
@@ -23,7 +23,9 @@ import tempfile
from typing import Union
import tvm
+from tvm import relax
from tvm import rpc as _rpc
+from tvm.contrib import utils
import tvm.contrib.hexagon as hexagon
from tvm.relay.backend.executor_factory import (
ExecutorFactoryModule,
@@ -283,13 +285,13 @@ class Session:
graph_json, graph_debug_mod, self.device, dump_root=str(dump_root)
)
- def get_executor_from_factory(self, module: ExecutorFactoryModule):
+ def get_executor_from_factory(self, module: Union[ExecutorFactoryModule,
relax.Executable]):
"""Create a local GraphModule which consumes a remote libmod.
Parameters
----------
- module : ExecutorFactoryModule
+ module : Union[ExecutorFactoryModule, relax.Executable]
The module to upload to the remote
session and load.
@@ -298,6 +300,8 @@ class Session:
return self._aot_executor_from_factory(module)
if isinstance(module, GraphExecutorFactoryModule):
return self._graph_executor_from_factory(module)
+ if isinstance(module, relax.Executable):
+ return self._relax_vm_executable_executor(module)
raise TypeError(f"Unsupported executor type: {type(module)}")
@@ -349,6 +353,35 @@ class Session:
"""
return self.get_graph_executor(module.get_graph_json(),
module.get_lib())
+ def _relax_vm_executable_executor(self, vm_exec: relax.Executable):
+ """Create a local TVM module which consumes a remote vm executable.
+
+ Paramters
+ ---------
+
+ vm_exec : relax.Executable
+ The Relax VM Executable to upload to the remote and load. This
will typically be the
+ output of `relax.build`.
+
+ Returns
+ -------
+ TVMModule :
+ TVM module object
+ """
+ assert self._rpc is not None, "Hexagon session must be started using
__enter__ prior to use"
+
+ temp_dir = utils.tempdir()
+ path_exec = temp_dir.relpath("exec.so")
+
+ vm_exec.mod.export_library(
+ path_exec,
+ fcompile=hexagon.create_aot_shared,
+ hexagon_arch="v68",
+ )
+
+ path = self.upload(path_exec, "exec.so")
+ return self._rpc.get_function("tvm.hexagon.load_module")(str(path))
+
def _aot_executor_from_factory(
self,
module: Union[str, pathlib.Path, AOTExecutorFactoryModule],
diff --git a/python/tvm/runtime/module.py b/python/tvm/runtime/module.py
index db2b704a60..e15bf7507c 100644
--- a/python/tvm/runtime/module.py
+++ b/python/tvm/runtime/module.py
@@ -498,7 +498,7 @@ class Module(object):
object_format = "cu"
has_c_module = True
else:
- assert module.type_key == "llvm" or module.type_key ==
"static_library"
+ assert module.is_dso_exportable
global_object_format = object_format = "o"
path_obj = os.path.join(workspace_dir,
f"lib{index}.{object_format}")
diff --git a/src/runtime/hexagon/hexagon_module.h
b/src/runtime/hexagon/hexagon_module.h
index aac75002c2..160bd1de74 100644
--- a/src/runtime/hexagon/hexagon_module.h
+++ b/src/runtime/hexagon/hexagon_module.h
@@ -64,6 +64,7 @@ class HexagonModuleNode : public runtime::ModuleNode {
const char* type_key() const final { return "hexagon"; }
void SaveToFile(const std::string& file_name, const std::string& format)
override;
void SaveToBinary(dmlc::Stream* stream) override;
+ bool IsDSOExportable() const final { return true; }
protected:
std::string data_;
diff --git a/tests/python/contrib/test_hexagon/test_relax_integration.py
b/tests/python/contrib/test_hexagon/test_relax_integration.py
new file mode 100644
index 0000000000..823f4bdb92
--- /dev/null
+++ b/tests/python/contrib/test_hexagon/test_relax_integration.py
@@ -0,0 +1,236 @@
+# 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.
+"""Relax hexagon test."""
+
+import numpy as np
+import pytest
+import tvm.testing
+from tvm import relay, relax, runtime
+from tvm.relax.testing import relay_translator
+from tvm.contrib.hexagon.session import Session
+from tvm.relay import testing
+
+
+class TestConv2d:
+ """Test conv2d op"""
+
+ n_batch = tvm.testing.parameter(1, relay.Any())
+
+ @tvm.testing.requires_hexagon
+ def test_conv2d(self, hexagon_session: Session, n_batch):
+ """Test Relax conv2d op and compare with Relay"""
+ dtype = "float32"
+ data = relay.var("data", relay.TensorType((n_batch, 64, 64, 3), dtype))
+ weight = relay.var("weight", relay.TensorType((5, 5, 3, 8), dtype))
+ y = relay.nn.conv2d(
+ data,
+ weight,
+ padding=(2, 2),
+ kernel_size=(5, 5),
+ data_layout="NHWC",
+ kernel_layout="HWIO",
+ out_dtype="float32",
+ )
+ f = relay.Function([data, weight], y)
+ relay_mod = tvm.IRModule.from_expr(f)
+
+ target_hexagon = tvm.target.hexagon("v68")
+ target = tvm.target.Target(target_hexagon, host=target_hexagon)
+ relax_mod = relay_translator.from_relay(relay_mod["main"], target)
+
+ exe = relax.build(relax_mod, target)
+ dev = hexagon_session.device
+ vm_mod = hexagon_session.get_executor_from_factory(exe)
+ vm_rt = relax.VirtualMachine(vm_mod, dev)
+
+ data_np = np.random.rand(1, 64, 64, 3).astype(np.float32)
+ weight_np = np.random.rand(5, 5, 3, 8).astype(np.float32)
+
+ # Run on hexagon and get result
+ data = tvm.nd.array(data_np, dev)
+ weight = tvm.nd.array(weight_np, dev)
+ vm_rt.set_input("main", data, weight)
+ vm_rt.invoke_stateful("main")
+ hexagon_res = vm_rt.get_outputs("main")
+
+ # Compile and run on Relay for comparison.
+ dev = tvm.cpu()
+ data = tvm.nd.array(data_np, dev)
+ weight = tvm.nd.array(weight_np, dev)
+
+ target = tvm.target.Target("llvm", host="llvm")
+ vm_exec = relay.vm.compile(relay_mod, target=target)
+ vm_factory = runtime.vm.VirtualMachine(vm_exec, tvm.cpu())
+ relay_res = vm_factory.invoke("main", data, weight)
+ tvm.testing.assert_allclose(hexagon_res.numpy(), relay_res.numpy(),
rtol=1e-3)
+
+
+class TestMLP:
+ """Test MLP"""
+
+ n_batch = tvm.testing.parameter(1, relay.Any())
+
+ @tvm.testing.requires_hexagon
+ def test_mlp(self, hexagon_session: Session, n_batch):
+ """Test Relax MLP and compare with Relay"""
+ relay_mod, params = testing.mlp.get_workload(batch_size=n_batch,
dtype="float32")
+
+ target_hexagon = tvm.target.hexagon("v68")
+ target = tvm.target.Target(target_hexagon, host=target_hexagon)
+ relax_mod = relay_translator.from_relay(relay_mod["main"], target,
params)
+
+ exe = relax.build(relax_mod, target)
+ hexagon_device = hexagon_session.device
+
+ vm_mod = hexagon_session.get_executor_from_factory(exe)
+ vm_rt = relax.VirtualMachine(vm_mod, hexagon_device)
+
+ shape = (1, 1, 28, 28)
+ data_np = np.random.rand(*shape).astype("float32")
+ data = tvm.nd.array(data_np, hexagon_device)
+ vm_rt.set_input("main", data)
+ vm_rt.invoke_stateful("main")
+ hexagon_res = vm_rt.get_outputs("main")
+
+ # Compile and run on Relay for comparison.
+ cpu_dev = tvm.cpu()
+ data = tvm.nd.array(data_np, cpu_dev)
+
+ target = tvm.target.Target("llvm", host="llvm")
+ vm_exec = relay.vm.compile(relay_mod, target=target)
+ vm_factory = runtime.vm.VirtualMachine(vm_exec, cpu_dev)
+ relay_res = vm_factory.invoke("main", data, **params)
+ tvm.testing.assert_allclose(hexagon_res.numpy(), relay_res.numpy(),
rtol=1e-3)
+
+
+def get_onnx_mobilenet():
+ """Download and import mobilenet model with ONNX"""
+ import onnx # pylint: disable=import-outside-toplevel
+
+ # pylint: disable=line-too-long
+ model_url =
"https://github.com/onnx/models/raw/main/vision/classification/mobilenet/model/mobilenetv2-7.onnx";
+ model_path = tvm.contrib.download.download_testdata(
+ model_url, "mobilenetv2-7.onnx", module="onnx"
+ )
+ return onnx.load(model_path)
+
+
[email protected]("takes too long (~20min)")
[email protected]_hexagon
+def test_mobilenet_onnx(hexagon_session: Session):
+ """Test MobileNetV2 ONNX model"""
+ onnx_model = get_onnx_mobilenet()
+ data_np = np.random.rand(1, 3, 224, 224).astype("float32")
+ shape_dict = {"input": data_np.shape}
+ relay_mod, _ = relay.frontend.from_onnx(onnx_model, shape_dict,
freeze_params=True)
+
+ target_hexagon = tvm.target.hexagon("v68")
+ target = tvm.target.Target(target_hexagon, host=target_hexagon)
+ relax_mod = relay_translator.from_relay(relay_mod["main"], target_hexagon)
+
+ # Compile and run on Hexagon.
+ exe = relax.build(relax_mod, target)
+ dev = hexagon_session.device
+
+ vm_mod = hexagon_session.get_executor_from_factory(exe)
+ vm_rt = relax.VirtualMachine(vm_mod, dev)
+ data = tvm.nd.array(data_np, dev)
+ vm_rt.set_input("main", data)
+ vm_rt.invoke_stateful("main")
+ hexagon_res = vm_rt.get_outputs("main")
+
+ # Compile and run on LLVM for comparison.
+ relax_mod = relay_translator.from_relay(relay_mod["main"], "llvm")
+ exe = relax.build(relax_mod, "llvm")
+ dev = tvm.cpu()
+ vm_rt = relax.VirtualMachine(exe, dev)
+ data = tvm.nd.array(data_np, dev)
+ llvm_res = vm_rt["main"](data)
+ tvm.testing.assert_allclose(hexagon_res.numpy(), llvm_res.numpy(),
rtol=1e-3)
+
+
[email protected]("takes too long (~20min)")
[email protected]_hexagon
+def test_mobilenet(hexagon_session: Session):
+ """Test MobileNet workload"""
+ relay_mod, params = testing.mobilenet.get_workload(batch_size=1,
dtype="float32")
+ data_np = np.random.rand(1, 3, 224, 224).astype("float32")
+
+ target_hexagon = tvm.target.hexagon("v68")
+ target = tvm.target.Target(target_hexagon, host=target_hexagon)
+
+ # translate the relay mobilenet and bind params
+ relax_mod = relay_translator.from_relay(relay_mod["main"], target, params)
+
+ # Compile and run on Hexagon.
+ exe = relax.build(relax_mod, target)
+ dev = hexagon_session.device
+
+ vm_mod = hexagon_session.get_executor_from_factory(exe)
+ vm_rt = relax.VirtualMachine(vm_mod, dev)
+ data = tvm.nd.array(data_np, dev)
+ vm_rt.set_input("main", data)
+ vm_rt.invoke_stateful("main")
+ hexagon_res = vm_rt.get_outputs("main")
+
+ # Compile and run on LLVM for comparison.
+ relax_mod = relay_translator.from_relay(relay_mod["main"], "llvm", params)
+ exe = relax.build(relax_mod, "llvm")
+ dev = tvm.cpu()
+ vm_rt = relax.VirtualMachine(exe, dev)
+ data = tvm.nd.array(data_np, dev)
+ llvm_res = vm_rt["main"](data)
+ tvm.testing.assert_allclose(hexagon_res.numpy(), llvm_res.numpy(),
rtol=1e-3)
+
+
[email protected]("takes too long (~20min)")
[email protected]_hexagon
+def test_mobilenet_dyn(hexagon_session: Session):
+ """Test MobileNet workload with dynamic batch size"""
+ relay_mod, params = testing.mobilenet.get_workload(batch_size=relay.Any(),
dtype="float32")
+ data_np = np.random.rand(1, 3, 224, 224).astype("float32")
+
+ target_hexagon = tvm.target.hexagon("v68")
+ target = tvm.target.Target(target_hexagon, host=target_hexagon)
+
+ # translate the relay mobilenet and bind params
+ relax_mod = relay_translator.from_relay(relay_mod["main"], target, params)
+
+ # Compile and run on Hexagon.
+ exe = relax.build(relax_mod, target)
+ dev = hexagon_session.device
+
+ vm_mod = hexagon_session.get_executor_from_factory(exe)
+ vm_rt = relax.VirtualMachine(vm_mod, dev)
+ data = tvm.nd.array(data_np, dev)
+ vm_rt.set_input("main", data)
+ vm_rt.invoke_stateful("main")
+ hexagon_res = vm_rt.get_outputs("main")
+
+ # Compile and run on Relay for comparison.
+ dev = tvm.cpu()
+ data = tvm.nd.array(data_np, dev)
+
+ target = tvm.target.Target("llvm", host="llvm")
+ vm_exec = relay.vm.compile(relay_mod, target=target)
+ vm_factory = runtime.vm.VirtualMachine(vm_exec, tvm.cpu())
+ relay_res = vm_factory.invoke("main", data, **params)
+ tvm.testing.assert_allclose(hexagon_res.numpy(), relay_res.numpy(),
rtol=1e-3)
+
+
+if __name__ == "__main__":
+ tvm.testing.main()
