[Unity][Hexagon] Enable Relax VM for Hexagon (#14415)

This is attempt to port PR (tlc-pack/relax#167)
(submitted by @psrivas2 and @YuchenJin) from tlc-pack/relax to enable
Hexagon tests with Relax VM.

Author: ibsidorenko

python/tvm/contrib/hexagon/session.py

@@ -23,7 +23,9 @@
 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 @@ def get_graph_debug_executor(
             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 @@ def get_executor_from_factory(self, module: ExecutorFactoryModule):
             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 @@ def _graph_executor_from_factory(
         """
         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],

python/tvm/runtime/module.py

@@ -498,7 +498,7 @@ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=No
                             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}")

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_;

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)
+
+
+@pytest.mark.skip("takes too long (~20min)")
+@tvm.testing.requires_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)
+
+
+@pytest.mark.skip("takes too long (~20min)")
+@tvm.testing.requires_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)
+
+
+@pytest.mark.skip("takes too long (~20min)")
+@tvm.testing.requires_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()