From 3cd73c48199144b1eab787a891a47f49f207908b Mon Sep 17 00:00:00 2001 From: Mehrdad Hessar Date: Mon, 21 Mar 2022 10:44:53 -0700 Subject: [PATCH 1/6] AOT with LLVM Codegen on Hexagon --- python/tvm/contrib/hexagon/build.py | 4 +- python/tvm/contrib/hexagon/session.py | 38 ++++++++++++++----- src/relay/backend/aot_executor_codegen.cc | 6 +-- .../hexagon/hexagon/hexagon_device_api_v2.cc | 8 +++- src/runtime/hexagon/rpc/hexagon/rpc_server.cc | 26 ++++++++----- .../hexagon/rpc/simulator/rpc_server.cc | 1 + src/runtime/hexagon/rpc/simulator/session.cc | 34 ++++++++++++++++- src/runtime/library_module.cc | 4 +- src/target/llvm/codegen_hexagon.cc | 6 +++ .../contrib/test_hexagon/test_launcher.py | 35 ++++++++++++----- 10 files changed, 121 insertions(+), 41 deletions(-) diff --git a/python/tvm/contrib/hexagon/build.py b/python/tvm/contrib/hexagon/build.py index fd74eb7738cf..60b1ec7a49d2 100644 --- a/python/tvm/contrib/hexagon/build.py +++ b/python/tvm/contrib/hexagon/build.py @@ -182,7 +182,7 @@ def upload(self, local_path: Union[str, pathlib.Path], remote_filename: str): assert self._workspace self._copy_to_remote(local_path, os.path.join(str(self._workspace), remote_filename)) - def start_session(self) -> Session: + def start_session(self, name="hexagon-rpc") -> Session: """Connect to the RPC server. Returns @@ -197,7 +197,7 @@ def start_session(self) -> Session: "timeout": 0, "key": self._device_key, } - return Session(self, hexagon_remote_kw) + return Session(self, hexagon_remote_kw, session_name=name) def load_module(self, module: Union[str, pathlib.Path, tvm.runtime.Module], session: Session): """Load TVM module. diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index 56bd1b79faad..68ea09315fbb 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/python/tvm/contrib/hexagon/session.py @@ -60,10 +60,10 @@ def __init__( rpc_receive_buffer_size_bytes: int = 2 * 1024 * 1024, ): self._launcher = launcher - self._session_name = session_name - self._remote_stack_size_bytes = remote_stack_size_bytes - self._rpc_receive_buffer_size_bytes = rpc_receive_buffer_size_bytes - self._remote_kw = remote_kw + self._session_name: str = session_name + self._remote_stack_size_bytes: int = remote_stack_size_bytes + self._rpc_receive_buffer_size_bytes: int = rpc_receive_buffer_size_bytes + self._remote_kw: dict = remote_kw self._rpc = None self.device = None @@ -86,7 +86,12 @@ def __enter__(self): self._rpc_receive_buffer_size_bytes, ], ) - self.device = self._rpc.hexagon(0) + if self._session_name == "cpu-rpc": + self.device = self._rpc.cpu(0) + elif self._session_name == "hexagon-rpc": + self.device = self._rpc.hexagon(0) + else: + raise RuntimeError(f"Incorrect session name: {self._session_name}") return self except RuntimeError as exception: @@ -286,6 +291,11 @@ def _aot_executor_from_factory( for target in module.target.values() if "hexagon" in target.keys ) + assert len(module.target.values()) == 1 + + for target in module.target.values(): + target_kind = str(target).split()[0] + assert hexagon_arch, "No hexagon target architecture found" assert len(hexagon_arch) == 1, f"Inconsistent hexagon architecture found, {hexagon_arch}" hexagon_arch = hexagon_arch.pop() @@ -295,11 +305,19 @@ def _aot_executor_from_factory( binary_name = "test_binary.so" binary_path = temp_dir / binary_name - module.export_library( - str(binary_path), - fcompile=hexagon.create_aot_shared, - hexagon_arch=hexagon_arch, - ) + if target_kind == "hexagon": + module.export_library( + str(binary_path), + fcompile=hexagon.create_aot_shared, + hexagon_arch=hexagon_arch, + ) + elif target_kind == "llvm": + module.export_library( + str(binary_path), + cc=hexagon.hexagon_clang_plus(), + ) + else: + raise ValueError("Incorrect Target kind.") self.upload(binary_path, binary_name) diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc index c2b2ac0fc5e2..945bbc5b3897 100644 --- a/src/relay/backend/aot_executor_codegen.cc +++ b/src/relay/backend/aot_executor_codegen.cc @@ -1234,12 +1234,10 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode { Target target_host; for (const auto& it : tmp) { auto dev_type = it.first.as(); - if (!target_host.defined() && it.second->kind->device_type == kDLCPU) { + if (!target_host.defined() && ((it.second->kind->device_type == kDLCPU) || + (it.second->kind->device_type == kDLHexagon))) { target_host = it.second; } - if (!target_host.defined() && it.second->kind->device_type == kDLHexagon) { - target_host = *(new Target("c")); - } ICHECK(dev_type); targets[static_cast(dev_type->value)] = it.second; } diff --git a/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc b/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc index 5a7642abeb55..72215abd5b0b 100644 --- a/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc +++ b/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc @@ -107,12 +107,16 @@ struct HexagonWorkspacePool : public WorkspacePool { }; void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType type_hint) { - CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type; + bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) || + (DLDeviceType(dev.device_type) == kDLCPU); + CHECK(is_valid_device) << "dev.device_type: " << dev.device_type; return dmlc::ThreadLocalStore::Get()->AllocWorkspace(dev, size); } void HexagonDeviceAPIv2::FreeWorkspace(Device dev, void* data) { - CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type; + bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) || + (DLDeviceType(dev.device_type) == kDLCPU); + CHECK(is_valid_device) << "dev.device_type: " << dev.device_type; CHECK(hexagon_buffer_map_.count(data) != 0) << "Attempt made to free unknown or already freed workspace allocation"; dmlc::ThreadLocalStore::Get()->FreeWorkspace(dev, data); diff --git a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc index af91dd3b4e6d..df5b5f19c6ab 100644 --- a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc +++ b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc @@ -58,11 +58,10 @@ class HexagonIOHandler { read_buffer_size_bytes_{read_buffer_size_bytes}, write_buffer_available_length_{0} {} - void MessageStart(size_t message_size_bytes) {} + void MessageStart(size_t message_size_bytes) { LOG(INFO) << "MessageStart called."; } ssize_t PosixWrite(const uint8_t* buf, size_t write_len_bytes) { - LOG(INFO) << "INFO: HexagonIOHandler PosixWrite called, write_len_bytes(" << write_len_bytes - << ")"; + LOG(INFO) << "HexagonIOHandler PosixWrite called, write_len_bytes(" << write_len_bytes << ")"; int32_t written_size = write_buffer_.sputn(reinterpret_cast(buf), write_len_bytes); if (written_size != write_len_bytes) { LOG(ERROR) << "written_size(" << written_size << ") != write_len_bytes(" << write_len_bytes @@ -72,10 +71,10 @@ class HexagonIOHandler { return (ssize_t)written_size; } - void MessageDone() { LOG(INFO) << "INFO: Message Done."; } + void MessageDone() { LOG(INFO) << "Message Done."; } ssize_t PosixRead(uint8_t* buf, size_t read_len_bytes) { - LOG(INFO) << "INFO: HexagonIOHandler PosixRead called, read_len_bytes(" << read_len_bytes + LOG(INFO) << "HexagonIOHandler PosixRead called, read_len_bytes(" << read_len_bytes << "), read_buffer_index_(" << read_buffer_index_ << ")"; uint32_t bytes_to_read = 0; @@ -99,7 +98,7 @@ class HexagonIOHandler { * \return The status */ AEEResult SetReadBuffer(const uint8_t* data, size_t data_size_bytes) { - LOG(INFO) << "INFO: HexagonIOHandler SetReadBuffer: data_size_bytes(" << data_size_bytes + LOG(INFO) << "HexagonIOHandler SetReadBuffer: data_size_bytes(" << data_size_bytes << "), read_buffer_index_(" << read_buffer_index_ << "), read_buffer_size_bytes_(" << read_buffer_size_bytes_ << ")"; if (data_size_bytes > read_buffer_size_bytes_) { @@ -121,7 +120,7 @@ class HexagonIOHandler { * \return The size of data that is read in bytes. */ int64_t ReadFromWriteBuffer(uint8_t* buf, size_t read_size_bytes) { - LOG(INFO) << "INFO: HexagonIOHandler ReadFromWriteBuffer called, read_size_bytes: " + LOG(INFO) << "HexagonIOHandler ReadFromWriteBuffer called, read_size_bytes: " << read_size_bytes; int64_t size = (int64_t)write_buffer_.sgetn(reinterpret_cast(buf), read_size_bytes); write_buffer_available_length_ -= size; @@ -133,7 +132,7 @@ class HexagonIOHandler { return size; } - void Close() { LOG(INFO) << "INFO: HexagonIOHandler Close called"; } + void Close() { LOG(INFO) << "HexagonIOHandler Close called"; } void Exit(int code) { exit(code); } @@ -156,13 +155,19 @@ class HexagonRPCServer { * \param data The data pointer * \param data_size_bytes The data size in bytes. * - * \return The size of data written to IOHandler. + * \return The size of data written to IOHandler if no error. + * Otherwise, returns -1; */ int64_t Write(const uint8_t* data, size_t data_size_bytes) { if (io_.SetReadBuffer(data, data_size_bytes) != AEE_SUCCESS) { + LOG(ERROR) << "ERROR: SetReadBuffer failed"; + return -1; + } + + if (!rpc_server_.ProcessOnePacket()) { + LOG(ERROR) << "ERROR: ProcessOnePacket failed"; return -1; } - rpc_server_.ProcessOnePacket(); return (int64_t)data_size_bytes; } @@ -211,6 +216,7 @@ const tvm::runtime::PackedFunc get_runtime_func(const std::string& name) { void reset_device_api() { const tvm::runtime::PackedFunc api = get_runtime_func("device_api.hexagon.v2"); tvm::runtime::Registry::Register("device_api.hexagon", true).set_body(api); + tvm::runtime::Registry::Register("device_api.cpu", true).set_body(api); } int __QAIC_HEADER(hexagon_rpc_open)(const char* uri, remote_handle64* handle) { diff --git a/src/runtime/hexagon/rpc/simulator/rpc_server.cc b/src/runtime/hexagon/rpc/simulator/rpc_server.cc index 76f168cd20ad..dee324ec1cb4 100644 --- a/src/runtime/hexagon/rpc/simulator/rpc_server.cc +++ b/src/runtime/hexagon/rpc/simulator/rpc_server.cc @@ -292,6 +292,7 @@ int main() { const auto* api_v2 = tvm::runtime::Registry::Get("device_api.hexagon.v2"); ICHECK(api_v2 != nullptr); tvm::runtime::Registry::Register("device_api.hexagon", true).set_body(*api_v2); + tvm::runtime::Registry::Register("device_api.cpu", true).set_body(*api_v2); tvm::runtime::hexagon::SimulatorRPCServer server; diff --git a/src/runtime/hexagon/rpc/simulator/session.cc b/src/runtime/hexagon/rpc/simulator/session.cc index b0f71c7bcf8a..2c1f4003f1c1 100644 --- a/src/runtime/hexagon/rpc/simulator/session.cc +++ b/src/runtime/hexagon/rpc/simulator/session.cc @@ -214,6 +214,11 @@ class SimulatorRPCChannel final : public RPCChannel { std::string runmain; // Path to run_main_on_hexagon. }; + struct Message_ { + Message msg; + std::string str() const; + }; + Message SendMsg(Message msg); Message SendMsg(uint32_t code, uint32_t len, uint32_t va); void ReadFromProcess(void* host_dst, HEX_VA_t src, size_t len); @@ -461,6 +466,27 @@ std::string SimulatorRPCChannel::Cpu_::str() const { return default_cpu_; } +std::string SimulatorRPCChannel::Message_::str() const { + switch (msg.code) { + case Message::kNone: + return "kNone"; + case Message::kAck: + return "kAck"; + case Message::kTerminate: + return "kTerminate"; + case Message::kReceiveStart: + return "kReceiveStart"; + case Message::kReceiveEnd: + return "kReceiveEnd"; + case Message::kSendStart: + return "kSendStart"; + case Message::kSendEnd: + return "kSendEnd"; + default: + break; + } +} + SimulatorRPCChannel::SDKInfo_::SDKInfo_(const std::string& sdk_root, const std::string& cpu) : root(sdk_root) { // For v69 chips, still look for v68 in the directory names. @@ -524,6 +550,7 @@ SimulatorRPCChannel::SimulatorRPCChannel(int stack_size, std::string args) { const auto* api_v2 = tvm::runtime::Registry::Get("device_api.hexagon.v2"); ICHECK(api_v2 != nullptr); tvm::runtime::Registry::Register("device_api.hexagon", true).set_body(*api_v2); + tvm::runtime::Registry::Register("device_api.cpu", true).set_body(*api_v2); const char* sdk_root_env = std::getenv("HEXAGON_SDK_ROOT"); ICHECK(sdk_root_env != nullptr) << "Please set HEXAGON_SDK_ROOT"; @@ -651,9 +678,14 @@ Message SimulatorRPCChannel::SendMsg(Message msg) { HEX_4u_t result; core = sim_->Run(&result); - ICHECK_EQ(core, HEX_CORE_BREAKPOINT); + Core_ core_ = {core}; + ICHECK_EQ(core, HEX_CORE_BREAKPOINT) + << "Expecting HEX_CORE_BREAKPOINT, received: " << core_.str(); }; + Message_ msg_ = {msg}; + LOG(INFO) << "Sending message: " << msg_.str(); + WriteToProcess(message_buffer_v_, &msg, sizeof msg); run(); diff --git a/src/runtime/library_module.cc b/src/runtime/library_module.cc index 7efa91d912eb..54fd362387c5 100644 --- a/src/runtime/library_module.cc +++ b/src/runtime/library_module.cc @@ -115,8 +115,8 @@ Module LoadModuleFromBinary(const std::string& type_key, dmlc::Stream* stream) { loaders += name.substr(loadkey.size()); } } - LOG(FATAL) << "Binary was created using " << type_key - << " but a loader of that name is not registered. Available loaders are " << loaders + LOG(FATAL) << "Binary was created using {" << type_key + << "} but a loader of that name is not registered. Available loaders are " << loaders << ". Perhaps you need to recompile with this runtime enabled."; } diff --git a/src/target/llvm/codegen_hexagon.cc b/src/target/llvm/codegen_hexagon.cc index 9f7ee6194117..035f772f8d6c 100644 --- a/src/target/llvm/codegen_hexagon.cc +++ b/src/target/llvm/codegen_hexagon.cc @@ -475,6 +475,12 @@ runtime::Module BuildHexagon(IRModule mod, Target target) { TVM_REGISTER_GLOBAL("target.build.hexagon").set_body_typed(BuildHexagon); +TVM_REGISTER_GLOBAL("tvm.codegen.llvm.target_hexagon") + .set_body([](const TVMArgs& targs, TVMRetValue* rv) { + CodeGenLLVM* cg = new CodeGenHexagon(); + *rv = static_cast(cg); + }); + } // namespace codegen } // namespace tvm diff --git a/tests/python/contrib/test_hexagon/test_launcher.py b/tests/python/contrib/test_hexagon/test_launcher.py index c2152cf62355..4a817a93d672 100644 --- a/tests/python/contrib/test_hexagon/test_launcher.py +++ b/tests/python/contrib/test_hexagon/test_launcher.py @@ -16,22 +16,23 @@ # under the License. import os -import pathlib import sys import pytest import numpy as np -import logging import tvm.testing from tvm import te from tvm import relay from tvm.relay.backend import Executor, Runtime -from tvm.contrib import utils, ndk -from tvm.contrib.hexagon.build import HexagonLauncher import tvm.contrib.hexagon as hexagon from .conftest import requires_hexagon_toolchain +aot_target_kind = tvm.testing.parameter( + "c", + "llvm -keys=hexagon -link-params=0 -mattr=+hvxv68,+hvx-length128b,+hvx-qfloat,-hvx-ieee-fp -mcpu=hexagonv68 -mtriple=hexagon", +) + @requires_hexagon_toolchain def test_add(hexagon_session): @@ -270,8 +271,18 @@ def _workaround_create_aot_shared(): ) +def get_target_and_session(target_kind: str): + if target_kind == "c": + target_hexagon = tvm.target.hexagon("v68") + session_key = "hexagon-rpc" + elif target_kind.startswith("llvm"): + target_hexagon = target_kind + session_key = "cpu-rpc" + return target_hexagon, session_key + + @requires_hexagon_toolchain -def test_aot_executor(hexagon_session): +def test_aot_executor(hexagon_launcher, aot_target_kind): dtype = "float32" input_shape = (1, 128, 128, 3) w_shape = (5, 5, 3, 8) @@ -290,7 +301,7 @@ def test_aot_executor(hexagon_session): relay_mod = tvm.IRModule.from_expr(f) relay_mod = relay.transform.InferType()(relay_mod) - target_hexagon = tvm.target.hexagon("v68") + target_hexagon, session_key = get_target_and_session(aot_target_kind) weight_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2], w_shape[3]).astype(dtype=dtype) input_data = np.random.rand( @@ -304,11 +315,13 @@ def test_aot_executor(hexagon_session): lowered = tvm.relay.build( relay_mod, params=params, - target=tvm.target.Target(target_hexagon, host="c"), + target=tvm.target.Target(target_hexagon, host=aot_target_kind), runtime=Runtime("cpp"), executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) + hexagon_session = hexagon_launcher.start_session(name=session_key) + hexagon_session.__enter__() aot_mod = hexagon_session.get_executor_from_factory(lowered) aot_mod.set_input(**inputs) aot_mod.run() @@ -332,7 +345,7 @@ def test_aot_executor(hexagon_session): @requires_hexagon_toolchain -def test_aot_executor_multiple_conv2d(hexagon_session): +def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind): dtype = "float32" input_shape = (1, 8, 8, 3) w1_shape = (5, 5, 3, 1) @@ -362,7 +375,7 @@ def test_aot_executor_multiple_conv2d(hexagon_session): relay_mod = tvm.IRModule.from_expr(f) relay_mod = relay.transform.InferType()(relay_mod) - target_hexagon = tvm.target.hexagon("v68") + target_hexagon, session_key = get_target_and_session(aot_target_kind) weight1_data = np.random.rand(w1_shape[0], w1_shape[1], w1_shape[2], w1_shape[3]).astype( dtype=dtype @@ -381,11 +394,13 @@ def test_aot_executor_multiple_conv2d(hexagon_session): lowered = tvm.relay.build( relay_mod, params=params, - target=tvm.target.Target(target_hexagon, host="c"), + target=tvm.target.Target(target_hexagon, host=aot_target_kind), runtime=Runtime("cpp"), executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) + hexagon_session = hexagon_launcher.start_session(name=session_key) + hexagon_session.__enter__() aot_mod = hexagon_session.get_executor_from_factory(lowered) aot_mod.set_input(**inputs) aot_mod.run() From 278a78cc02fc137d62a1b180e1e2509ec2f5ef6a Mon Sep 17 00:00:00 2001 From: Mehrdad Hessar Date: Wed, 20 Apr 2022 08:44:19 -0700 Subject: [PATCH 2/6] Address comments --- python/tvm/contrib/hexagon/build.py | 9 +++-- python/tvm/contrib/hexagon/session.py | 12 +++++-- python/tvm/script/tir/__init__.pyi | 2 ++ src/relay/backend/aot_executor_codegen.cc | 2 ++ .../hexagon/hexagon/hexagon_device_api_v2.cc | 4 +++ src/runtime/hexagon/rpc/hexagon/rpc_server.cc | 8 +++-- .../contrib/test_hexagon/test_launcher.py | 34 +++++++++---------- 7 files changed, 46 insertions(+), 25 deletions(-) diff --git a/python/tvm/contrib/hexagon/build.py b/python/tvm/contrib/hexagon/build.py index 60b1ec7a49d2..1664ee4b1184 100644 --- a/python/tvm/contrib/hexagon/build.py +++ b/python/tvm/contrib/hexagon/build.py @@ -182,9 +182,14 @@ def upload(self, local_path: Union[str, pathlib.Path], remote_filename: str): assert self._workspace self._copy_to_remote(local_path, os.path.join(str(self._workspace), remote_filename)) - def start_session(self, name="hexagon-rpc") -> Session: + def start_session(self, session_name: str = "hexagon-rpc") -> Session: """Connect to the RPC server. + Parameters + ---------- + session_name : str + RPC session name. + Returns ------- Session : @@ -197,7 +202,7 @@ def start_session(self, name="hexagon-rpc") -> Session: "timeout": 0, "key": self._device_key, } - return Session(self, hexagon_remote_kw, session_name=name) + return Session(self, hexagon_remote_kw, session_name=session_name) def load_module(self, module: Union[str, pathlib.Path, tvm.runtime.Module], session: Session): """Load TVM module. diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index 68ea09315fbb..627b3ad7774a 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/python/tvm/contrib/hexagon/session.py @@ -44,7 +44,7 @@ class Session: Remote configs for RPC tracker. session_name : str - Hexagon RPC session name. + Hexagon RPC session name. Options are [hexagon-rpc, cpu-rpc] remote_stack_size_bytes : int The stack size of the remote device, to be passed to @@ -91,7 +91,10 @@ def __enter__(self): elif self._session_name == "hexagon-rpc": self.device = self._rpc.hexagon(0) else: - raise RuntimeError(f"Incorrect session name: {self._session_name}") + raise RuntimeError( + f"Incorrect session name: {self._session_name}.\n" + f"Options for session name are [hexagon-rpc, cpu-rpc]." + ) return self except RuntimeError as exception: @@ -317,7 +320,10 @@ def _aot_executor_from_factory( cc=hexagon.hexagon_clang_plus(), ) else: - raise ValueError("Incorrect Target kind.") + raise ValueError( + f"Incorrect Target kind.\n" + f"Target kind should be from these options: [hexagon, llvm]." + ) self.upload(binary_path, binary_name) diff --git a/python/tvm/script/tir/__init__.pyi b/python/tvm/script/tir/__init__.pyi index 3eb383ed9974..9727a8db6316 100644 --- a/python/tvm/script/tir/__init__.pyi +++ b/python/tvm/script/tir/__init__.pyi @@ -226,6 +226,7 @@ def alloc_buffer( """ special_stmt - Reads/Writes """ + @overload def reads(read_regions: List[BufferSlice]) -> None: ... @overload @@ -337,6 +338,7 @@ def Assert(condition: Union[PrimExpr, builtins.bool], message: str) -> PrimExpr: """ Scope handler - Loops """ + @overload def serial( begin: Union[PrimExpr, int], diff --git a/src/relay/backend/aot_executor_codegen.cc b/src/relay/backend/aot_executor_codegen.cc index 945bbc5b3897..9a194965ded4 100644 --- a/src/relay/backend/aot_executor_codegen.cc +++ b/src/relay/backend/aot_executor_codegen.cc @@ -1234,6 +1234,8 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode { Target target_host; for (const auto& it : tmp) { auto dev_type = it.first.as(); + // TODO(tvm-team): AoT only works with kDLCPU device type. We can remove kDLHexagon + // here once we refactored kDLHexagon to kDLCPU. if (!target_host.defined() && ((it.second->kind->device_type == kDLCPU) || (it.second->kind->device_type == kDLHexagon))) { target_host = it.second; diff --git a/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc b/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc index 72215abd5b0b..ebd826b2c7b3 100644 --- a/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc +++ b/src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc @@ -84,6 +84,7 @@ void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, int ndim, const int64_t* sh void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, size_t nbytes, size_t alignment, DLDataType type_hint) { + // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU; bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) || (DLDeviceType(dev.device_type) == kDLCPU); CHECK(is_valid_device) << "dev.device_type: " << dev.device_type; @@ -94,6 +95,7 @@ void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, size_t nbytes, size_t align } void HexagonDeviceAPIv2::FreeDataSpace(Device dev, void* ptr) { + // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU; bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) || (DLDeviceType(dev.device_type) == kDLCPU); CHECK(is_valid_device) << "dev.device_type: " << dev.device_type; @@ -107,6 +109,7 @@ struct HexagonWorkspacePool : public WorkspacePool { }; void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType type_hint) { + // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU; bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) || (DLDeviceType(dev.device_type) == kDLCPU); CHECK(is_valid_device) << "dev.device_type: " << dev.device_type; @@ -114,6 +117,7 @@ void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType typ } void HexagonDeviceAPIv2::FreeWorkspace(Device dev, void* data) { + // Added kDLCPU since we use hexagon as a sub-target of LLVM which by default maps to kDLCPU; bool is_valid_device = (TVMDeviceExtType(dev.device_type) == kDLHexagon) || (DLDeviceType(dev.device_type) == kDLCPU); CHECK(is_valid_device) << "dev.device_type: " << dev.device_type; diff --git a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc index df5b5f19c6ab..f352eb7e0828 100644 --- a/src/runtime/hexagon/rpc/hexagon/rpc_server.cc +++ b/src/runtime/hexagon/rpc/hexagon/rpc_server.cc @@ -58,7 +58,7 @@ class HexagonIOHandler { read_buffer_size_bytes_{read_buffer_size_bytes}, write_buffer_available_length_{0} {} - void MessageStart(size_t message_size_bytes) { LOG(INFO) << "MessageStart called."; } + void MessageStart(size_t message_size_bytes) {} ssize_t PosixWrite(const uint8_t* buf, size_t write_len_bytes) { LOG(INFO) << "HexagonIOHandler PosixWrite called, write_len_bytes(" << write_len_bytes << ")"; @@ -159,8 +159,9 @@ class HexagonRPCServer { * Otherwise, returns -1; */ int64_t Write(const uint8_t* data, size_t data_size_bytes) { - if (io_.SetReadBuffer(data, data_size_bytes) != AEE_SUCCESS) { - LOG(ERROR) << "ERROR: SetReadBuffer failed"; + AEEResult rc = io_.SetReadBuffer(data, data_size_bytes); + if (rc != AEE_SUCCESS) { + LOG(ERROR) << "ERROR: SetReadBuffer failed: " << rc; return -1; } @@ -216,6 +217,7 @@ const tvm::runtime::PackedFunc get_runtime_func(const std::string& name) { void reset_device_api() { const tvm::runtime::PackedFunc api = get_runtime_func("device_api.hexagon.v2"); tvm::runtime::Registry::Register("device_api.hexagon", true).set_body(api); + // Registering device_api.cpu as device_api.hexagon.v2 since we use hexagon as sub-target of LLVM. tvm::runtime::Registry::Register("device_api.cpu", true).set_body(api); } diff --git a/tests/python/contrib/test_hexagon/test_launcher.py b/tests/python/contrib/test_hexagon/test_launcher.py index 4a817a93d672..2d1cc7418e5a 100644 --- a/tests/python/contrib/test_hexagon/test_launcher.py +++ b/tests/python/contrib/test_hexagon/test_launcher.py @@ -274,11 +274,13 @@ def _workaround_create_aot_shared(): def get_target_and_session(target_kind: str): if target_kind == "c": target_hexagon = tvm.target.hexagon("v68") - session_key = "hexagon-rpc" + session_name = "hexagon-rpc" elif target_kind.startswith("llvm"): target_hexagon = target_kind - session_key = "cpu-rpc" - return target_hexagon, session_key + session_name = "cpu-rpc" + else: + assert False, "Incorrect target_kind: {target_kind}. Options are [c, llvm]." + return target_hexagon, session_name @requires_hexagon_toolchain @@ -301,7 +303,7 @@ def test_aot_executor(hexagon_launcher, aot_target_kind): relay_mod = tvm.IRModule.from_expr(f) relay_mod = relay.transform.InferType()(relay_mod) - target_hexagon, session_key = get_target_and_session(aot_target_kind) + target_hexagon, session_name = get_target_and_session(aot_target_kind) weight_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2], w_shape[3]).astype(dtype=dtype) input_data = np.random.rand( @@ -320,12 +322,11 @@ def test_aot_executor(hexagon_launcher, aot_target_kind): executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) - hexagon_session = hexagon_launcher.start_session(name=session_key) - hexagon_session.__enter__() - aot_mod = hexagon_session.get_executor_from_factory(lowered) - aot_mod.set_input(**inputs) - aot_mod.run() - hexagon_output = aot_mod.get_output(0).numpy() + with hexagon_launcher.start_session(session_name=session_name) as hexagon_session: + aot_mod = hexagon_session.get_executor_from_factory(lowered) + aot_mod.set_input(**inputs) + aot_mod.run() + hexagon_output = aot_mod.get_output(0).numpy() target_llvm = tvm.target.Target("llvm") with tvm.transform.PassContext(opt_level=3): @@ -375,7 +376,7 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind): relay_mod = tvm.IRModule.from_expr(f) relay_mod = relay.transform.InferType()(relay_mod) - target_hexagon, session_key = get_target_and_session(aot_target_kind) + target_hexagon, session_name = get_target_and_session(aot_target_kind) weight1_data = np.random.rand(w1_shape[0], w1_shape[1], w1_shape[2], w1_shape[3]).astype( dtype=dtype @@ -399,12 +400,11 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind): executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) - hexagon_session = hexagon_launcher.start_session(name=session_key) - hexagon_session.__enter__() - aot_mod = hexagon_session.get_executor_from_factory(lowered) - aot_mod.set_input(**inputs) - aot_mod.run() - hexagon_output = aot_mod.get_output(0).numpy() + with hexagon_launcher.start_session(session_name=session_name) as hexagon_session: + aot_mod = hexagon_session.get_executor_from_factory(lowered) + aot_mod.set_input(**inputs) + aot_mod.run() + hexagon_output = aot_mod.get_output(0).numpy() target_llvm = tvm.target.Target("llvm") with tvm.transform.PassContext(opt_level=3): From 2ce844511b6d68fc2db7b9f3beb854e8bf5ab812 Mon Sep 17 00:00:00 2001 From: Mehrdad Hessar Date: Wed, 20 Apr 2022 10:51:10 -0700 Subject: [PATCH 3/6] address Chris comments --- python/tvm/contrib/hexagon/session.py | 4 ++- tests/python/contrib/test_hexagon/conftest.py | 26 ++++++++++++++ .../contrib/test_hexagon/test_launcher.py | 35 +++++-------------- 3 files changed, 37 insertions(+), 28 deletions(-) diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index 627b3ad7774a..4697b5dc7fe0 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/python/tvm/contrib/hexagon/session.py @@ -44,7 +44,9 @@ class Session: Remote configs for RPC tracker. session_name : str - Hexagon RPC session name. Options are [hexagon-rpc, cpu-rpc] + Hexagon RPC session name. Options are [hexagon-rpc, cpu-rpc]. + `hexagon-rpc` is used with hexagon target and `cpu-rpc` used with + hexagon as a sub-target of LLVM. remote_stack_size_bytes : int The stack size of the remote device, to be passed to diff --git a/tests/python/contrib/test_hexagon/conftest.py b/tests/python/contrib/test_hexagon/conftest.py index 009150b1081c..1a917724b521 100644 --- a/tests/python/contrib/test_hexagon/conftest.py +++ b/tests/python/contrib/test_hexagon/conftest.py @@ -202,3 +202,29 @@ def terminate_rpc_servers(): yield [] if serial == "simulator": os.system("ps ax | grep tvm_rpc_x86 | awk '{print $1}' | xargs kill") + + +aot_host_target = tvm.testing.parameter( + "c", + "llvm -keys=hexagon -link-params=0 -mattr=+hvxv68,+hvx-length128b,+hvx-qfloat,-hvx-ieee-fp -mcpu=hexagonv68 -mtriple=hexagon", +) + + +@pytest.fixture() +def aot_target(aot_host_target): + if aot_host_target == "c": + yield tvm.target.hexagon("v68") + elif aot_host_target.startswith("llvm"): + yield aot_host_target + else: + assert False, "Incorrect AoT host target: {aot_host_target}. Options are [c, llvm]." + + +@pytest.fixture() +def rpc_session_name(aot_host_target): + if aot_host_target == "c": + yield "hexagon-rpc" + elif aot_host_target.startswith("llvm"): + yield "cpu-rpc" + else: + assert False, "Incorrect AoT host target: {aot_host_target}. Options are [c, llvm]." diff --git a/tests/python/contrib/test_hexagon/test_launcher.py b/tests/python/contrib/test_hexagon/test_launcher.py index 2d1cc7418e5a..a6624969e7e9 100644 --- a/tests/python/contrib/test_hexagon/test_launcher.py +++ b/tests/python/contrib/test_hexagon/test_launcher.py @@ -28,11 +28,6 @@ from .conftest import requires_hexagon_toolchain -aot_target_kind = tvm.testing.parameter( - "c", - "llvm -keys=hexagon -link-params=0 -mattr=+hvxv68,+hvx-length128b,+hvx-qfloat,-hvx-ieee-fp -mcpu=hexagonv68 -mtriple=hexagon", -) - @requires_hexagon_toolchain def test_add(hexagon_session): @@ -271,20 +266,8 @@ def _workaround_create_aot_shared(): ) -def get_target_and_session(target_kind: str): - if target_kind == "c": - target_hexagon = tvm.target.hexagon("v68") - session_name = "hexagon-rpc" - elif target_kind.startswith("llvm"): - target_hexagon = target_kind - session_name = "cpu-rpc" - else: - assert False, "Incorrect target_kind: {target_kind}. Options are [c, llvm]." - return target_hexagon, session_name - - @requires_hexagon_toolchain -def test_aot_executor(hexagon_launcher, aot_target_kind): +def test_aot_executor(hexagon_launcher, aot_host_target, aot_target, rpc_session_name): dtype = "float32" input_shape = (1, 128, 128, 3) w_shape = (5, 5, 3, 8) @@ -303,8 +286,6 @@ def test_aot_executor(hexagon_launcher, aot_target_kind): relay_mod = tvm.IRModule.from_expr(f) relay_mod = relay.transform.InferType()(relay_mod) - target_hexagon, session_name = get_target_and_session(aot_target_kind) - weight_data = np.random.rand(w_shape[0], w_shape[1], w_shape[2], w_shape[3]).astype(dtype=dtype) input_data = np.random.rand( input_shape[0], input_shape[1], input_shape[2], input_shape[3] @@ -317,12 +298,12 @@ def test_aot_executor(hexagon_launcher, aot_target_kind): lowered = tvm.relay.build( relay_mod, params=params, - target=tvm.target.Target(target_hexagon, host=aot_target_kind), + target=tvm.target.Target(aot_target, host=aot_host_target), runtime=Runtime("cpp"), executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) - with hexagon_launcher.start_session(session_name=session_name) as hexagon_session: + with hexagon_launcher.start_session(session_name=rpc_session_name) as hexagon_session: aot_mod = hexagon_session.get_executor_from_factory(lowered) aot_mod.set_input(**inputs) aot_mod.run() @@ -346,7 +327,9 @@ def test_aot_executor(hexagon_launcher, aot_target_kind): @requires_hexagon_toolchain -def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind): +def test_aot_executor_multiple_conv2d( + hexagon_launcher, aot_host_target, aot_target, rpc_session_name +): dtype = "float32" input_shape = (1, 8, 8, 3) w1_shape = (5, 5, 3, 1) @@ -376,8 +359,6 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind): relay_mod = tvm.IRModule.from_expr(f) relay_mod = relay.transform.InferType()(relay_mod) - target_hexagon, session_name = get_target_and_session(aot_target_kind) - weight1_data = np.random.rand(w1_shape[0], w1_shape[1], w1_shape[2], w1_shape[3]).astype( dtype=dtype ) @@ -395,12 +376,12 @@ def test_aot_executor_multiple_conv2d(hexagon_launcher, aot_target_kind): lowered = tvm.relay.build( relay_mod, params=params, - target=tvm.target.Target(target_hexagon, host=aot_target_kind), + target=tvm.target.Target(aot_target, host=aot_host_target), runtime=Runtime("cpp"), executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) - with hexagon_launcher.start_session(session_name=session_name) as hexagon_session: + with hexagon_launcher.start_session(session_name=rpc_session_name) as hexagon_session: aot_mod = hexagon_session.get_executor_from_factory(lowered) aot_mod.set_input(**inputs) aot_mod.run() From 31a7e80b9bdcca2aee134544b235fcc327de532e Mon Sep 17 00:00:00 2001 From: Mehrdad Hessar Date: Thu, 21 Apr 2022 09:29:50 -0700 Subject: [PATCH 4/6] Address Eric suggesions --- python/tvm/contrib/hexagon/session.py | 74 +++++++++++++------ tests/python/contrib/test_hexagon/conftest.py | 12 +-- .../contrib/test_hexagon/test_launcher.py | 28 +++---- 3 files changed, 68 insertions(+), 46 deletions(-) diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index 4697b5dc7fe0..7aedd8027e03 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/python/tvm/contrib/hexagon/session.py @@ -44,9 +44,7 @@ class Session: Remote configs for RPC tracker. session_name : str - Hexagon RPC session name. Options are [hexagon-rpc, cpu-rpc]. - `hexagon-rpc` is used with hexagon target and `cpu-rpc` used with - hexagon as a sub-target of LLVM. + Hexagon RPC session name. remote_stack_size_bytes : int The stack size of the remote device, to be passed to @@ -67,10 +65,10 @@ def __init__( self._rpc_receive_buffer_size_bytes: int = rpc_receive_buffer_size_bytes self._remote_kw: dict = remote_kw self._rpc = None - self.device = None + self._device = None def __enter__(self): - if self.device: + if self._rpc: # Already initialized return self @@ -88,15 +86,6 @@ def __enter__(self): self._rpc_receive_buffer_size_bytes, ], ) - if self._session_name == "cpu-rpc": - self.device = self._rpc.cpu(0) - elif self._session_name == "hexagon-rpc": - self.device = self._rpc.hexagon(0) - else: - raise RuntimeError( - f"Incorrect session name: {self._session_name}.\n" - f"Options for session name are [hexagon-rpc, cpu-rpc]." - ) return self except RuntimeError as exception: @@ -105,6 +94,25 @@ def __enter__(self): def __exit__(self, exc_type, exc_value, exc_traceback): pass + @property + def device(self): + """Session device.""" + + if hasattr(self, "_device") and self._device is not None: + return self._device + + if not hasattr(self, "_requires_cpu_device"): + assert ( + False + ), "Device type is not set. 'set_device_type' should be called before accessing device." + + if self._requires_cpu_device: + self._device = self._rpc.cpu(0) + else: + self._device = self._rpc.hexagon(0) + + return self._device + def upload(self, local_path: Union[str, pathlib.Path], remote_filename: str): """Upload a local file to the remote workspace. @@ -143,9 +151,7 @@ def load_module(self, module: Union[str, pathlib.Path, tvm.runtime.Module]): TVM module object. """ - assert ( - self.device is not None - ), "Hexagon session must be started using __enter__ prior to use" + assert self._rpc is not None, "Hexagon session must be started using __enter__ prior to use" if isinstance(module, tvm.runtime.Module): with tempfile.TemporaryDirectory() as temp_dir: @@ -189,6 +195,7 @@ def get_graph_executor( """ graph_mod = self.load_module(module_name) + self.set_device_type(graph_mod) return tvm.contrib.graph_executor.create(graph_json, graph_mod, self.device) def get_aot_executor( @@ -216,6 +223,7 @@ def get_aot_executor( """ aot_mod = self.load_module(module_name) + self.set_device_type(aot_mod) return tvm.runtime.executor.AotModule(aot_mod["default"](self.device)) def get_executor_from_factory(self, module: ExecutorFactoryModule): @@ -236,6 +244,28 @@ def get_executor_from_factory(self, module: ExecutorFactoryModule): raise TypeError(f"Unsupported executor type: {type(module)}") + def set_device_type(self, module: Union[str, pathlib.Path, GraphExecutorFactoryModule]): + """Set session device type(hexagon, cpu) based on target in module. + + Parameters + ---------- + + module: TVMModule + TVM module object. + """ + # for cases when module is a single schedule without target attribute. + if not hasattr(module, "target"): + self._requires_cpu_device = False + else: + assert len(module.target.values()) == 1 + for target in module.target.values(): + target_type = str(target).split()[0] + + if target_type == "llvm": + self._requires_cpu_device = True + else: + self._requires_cpu_device = False + def _graph_executor_from_factory( self, module: Union[str, pathlib.Path, GraphExecutorFactoryModule], @@ -263,6 +293,7 @@ def _graph_executor_from_factory( graph_json = module.get_graph_json() graph_mod = self.load_module(module.get_lib()) + self.set_device_type(module) return tvm.contrib.graph_executor.create(graph_json, graph_mod, self.device) @@ -296,10 +327,11 @@ def _aot_executor_from_factory( for target in module.target.values() if "hexagon" in target.keys ) - assert len(module.target.values()) == 1 + + self.set_device_type(module) for target in module.target.values(): - target_kind = str(target).split()[0] + target_type = str(target).split()[0] assert hexagon_arch, "No hexagon target architecture found" assert len(hexagon_arch) == 1, f"Inconsistent hexagon architecture found, {hexagon_arch}" @@ -310,13 +342,13 @@ def _aot_executor_from_factory( binary_name = "test_binary.so" binary_path = temp_dir / binary_name - if target_kind == "hexagon": + if target_type == "hexagon": module.export_library( str(binary_path), fcompile=hexagon.create_aot_shared, hexagon_arch=hexagon_arch, ) - elif target_kind == "llvm": + elif target_type == "llvm": module.export_library( str(binary_path), cc=hexagon.hexagon_clang_plus(), diff --git a/tests/python/contrib/test_hexagon/conftest.py b/tests/python/contrib/test_hexagon/conftest.py index 1a917724b521..7a90317d5506 100644 --- a/tests/python/contrib/test_hexagon/conftest.py +++ b/tests/python/contrib/test_hexagon/conftest.py @@ -210,7 +210,7 @@ def terminate_rpc_servers(): ) -@pytest.fixture() +@tvm.testing.fixture def aot_target(aot_host_target): if aot_host_target == "c": yield tvm.target.hexagon("v68") @@ -218,13 +218,3 @@ def aot_target(aot_host_target): yield aot_host_target else: assert False, "Incorrect AoT host target: {aot_host_target}. Options are [c, llvm]." - - -@pytest.fixture() -def rpc_session_name(aot_host_target): - if aot_host_target == "c": - yield "hexagon-rpc" - elif aot_host_target.startswith("llvm"): - yield "cpu-rpc" - else: - assert False, "Incorrect AoT host target: {aot_host_target}. Options are [c, llvm]." diff --git a/tests/python/contrib/test_hexagon/test_launcher.py b/tests/python/contrib/test_hexagon/test_launcher.py index a6624969e7e9..c12f3f936fbf 100644 --- a/tests/python/contrib/test_hexagon/test_launcher.py +++ b/tests/python/contrib/test_hexagon/test_launcher.py @@ -43,6 +43,7 @@ def test_add(hexagon_session): ) mod = hexagon_session.load_module(func) + hexagon_session.set_device_type(mod) A_data = tvm.nd.array(np.array([2, 3], dtype=dtype), device=hexagon_session.device) assert (A_data.numpy() == np.array([2, 3])).all() @@ -68,6 +69,8 @@ def test_add_vtcm(hexagon_session): ) mod = hexagon_session.load_module(func) + hexagon_session.set_device_type(mod) + A_data = tvm.nd.empty(A.shape, A.dtype, hexagon_session.device, "global.vtcm") A_data.copyfrom(np.array([2, 3])) @@ -101,6 +104,7 @@ def test_matmul(self, hexagon_session, M, N, K): ) mod = hexagon_session.load_module(func) + hexagon_session.set_device_type(mod) x = np.random.uniform(size=[i.value for i in X.shape]).astype(X.dtype) y = np.random.uniform(size=[i.value for i in Y.shape]).astype(Y.dtype) @@ -267,7 +271,7 @@ def _workaround_create_aot_shared(): @requires_hexagon_toolchain -def test_aot_executor(hexagon_launcher, aot_host_target, aot_target, rpc_session_name): +def test_aot_executor(hexagon_session, aot_host_target, aot_target): dtype = "float32" input_shape = (1, 128, 128, 3) w_shape = (5, 5, 3, 8) @@ -303,11 +307,10 @@ def test_aot_executor(hexagon_launcher, aot_host_target, aot_target, rpc_session executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) - with hexagon_launcher.start_session(session_name=rpc_session_name) as hexagon_session: - aot_mod = hexagon_session.get_executor_from_factory(lowered) - aot_mod.set_input(**inputs) - aot_mod.run() - hexagon_output = aot_mod.get_output(0).numpy() + aot_mod = hexagon_session.get_executor_from_factory(lowered) + aot_mod.set_input(**inputs) + aot_mod.run() + hexagon_output = aot_mod.get_output(0).numpy() target_llvm = tvm.target.Target("llvm") with tvm.transform.PassContext(opt_level=3): @@ -327,9 +330,7 @@ def test_aot_executor(hexagon_launcher, aot_host_target, aot_target, rpc_session @requires_hexagon_toolchain -def test_aot_executor_multiple_conv2d( - hexagon_launcher, aot_host_target, aot_target, rpc_session_name -): +def test_aot_executor_multiple_conv2d(hexagon_session, aot_host_target, aot_target): dtype = "float32" input_shape = (1, 8, 8, 3) w1_shape = (5, 5, 3, 1) @@ -381,11 +382,10 @@ def test_aot_executor_multiple_conv2d( executor=Executor("aot", {"unpacked-api": False, "interface-api": "packed"}), ) - with hexagon_launcher.start_session(session_name=rpc_session_name) as hexagon_session: - aot_mod = hexagon_session.get_executor_from_factory(lowered) - aot_mod.set_input(**inputs) - aot_mod.run() - hexagon_output = aot_mod.get_output(0).numpy() + aot_mod = hexagon_session.get_executor_from_factory(lowered) + aot_mod.set_input(**inputs) + aot_mod.run() + hexagon_output = aot_mod.get_output(0).numpy() target_llvm = tvm.target.Target("llvm") with tvm.transform.PassContext(opt_level=3): From 4c77cc51c62a0621a06e99c4dfc657fdec9e09cd Mon Sep 17 00:00:00 2001 From: Mehrdad Hessar Date: Thu, 21 Apr 2022 13:30:41 -0700 Subject: [PATCH 5/6] address more comments --- python/tvm/contrib/hexagon/session.py | 18 ++++++------------ .../contrib/test_hexagon/test_launcher.py | 3 --- 2 files changed, 6 insertions(+), 15 deletions(-) diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index 7aedd8027e03..a7361cab6b3f 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/python/tvm/contrib/hexagon/session.py @@ -65,7 +65,7 @@ def __init__( self._rpc_receive_buffer_size_bytes: int = rpc_receive_buffer_size_bytes self._remote_kw: dict = remote_kw self._rpc = None - self._device = None + self._requires_cpu_device = False def __enter__(self): if self._rpc: @@ -98,14 +98,9 @@ def __exit__(self, exc_type, exc_value, exc_traceback): def device(self): """Session device.""" - if hasattr(self, "_device") and self._device is not None: + if hasattr(self, "_device"): return self._device - if not hasattr(self, "_requires_cpu_device"): - assert ( - False - ), "Device type is not set. 'set_device_type' should be called before accessing device." - if self._requires_cpu_device: self._device = self._rpc.cpu(0) else: @@ -195,7 +190,7 @@ def get_graph_executor( """ graph_mod = self.load_module(module_name) - self.set_device_type(graph_mod) + self._set_device_type(graph_mod) return tvm.contrib.graph_executor.create(graph_json, graph_mod, self.device) def get_aot_executor( @@ -223,7 +218,7 @@ def get_aot_executor( """ aot_mod = self.load_module(module_name) - self.set_device_type(aot_mod) + self._set_device_type(aot_mod) return tvm.runtime.executor.AotModule(aot_mod["default"](self.device)) def get_executor_from_factory(self, module: ExecutorFactoryModule): @@ -244,7 +239,7 @@ def get_executor_from_factory(self, module: ExecutorFactoryModule): raise TypeError(f"Unsupported executor type: {type(module)}") - def set_device_type(self, module: Union[str, pathlib.Path, GraphExecutorFactoryModule]): + def _set_device_type(self, module: Union[str, pathlib.Path, GraphExecutorFactoryModule]): """Set session device type(hexagon, cpu) based on target in module. Parameters @@ -293,7 +288,6 @@ def _graph_executor_from_factory( graph_json = module.get_graph_json() graph_mod = self.load_module(module.get_lib()) - self.set_device_type(module) return tvm.contrib.graph_executor.create(graph_json, graph_mod, self.device) @@ -328,7 +322,7 @@ def _aot_executor_from_factory( if "hexagon" in target.keys ) - self.set_device_type(module) + self._set_device_type(module) for target in module.target.values(): target_type = str(target).split()[0] diff --git a/tests/python/contrib/test_hexagon/test_launcher.py b/tests/python/contrib/test_hexagon/test_launcher.py index c12f3f936fbf..48b3dac2a2c9 100644 --- a/tests/python/contrib/test_hexagon/test_launcher.py +++ b/tests/python/contrib/test_hexagon/test_launcher.py @@ -43,7 +43,6 @@ def test_add(hexagon_session): ) mod = hexagon_session.load_module(func) - hexagon_session.set_device_type(mod) A_data = tvm.nd.array(np.array([2, 3], dtype=dtype), device=hexagon_session.device) assert (A_data.numpy() == np.array([2, 3])).all() @@ -69,7 +68,6 @@ def test_add_vtcm(hexagon_session): ) mod = hexagon_session.load_module(func) - hexagon_session.set_device_type(mod) A_data = tvm.nd.empty(A.shape, A.dtype, hexagon_session.device, "global.vtcm") A_data.copyfrom(np.array([2, 3])) @@ -104,7 +102,6 @@ def test_matmul(self, hexagon_session, M, N, K): ) mod = hexagon_session.load_module(func) - hexagon_session.set_device_type(mod) x = np.random.uniform(size=[i.value for i in X.shape]).astype(X.dtype) y = np.random.uniform(size=[i.value for i in Y.shape]).astype(Y.dtype) From 3d7d8c7c37df8979d36b42020e21f32ed4c0050a Mon Sep 17 00:00:00 2001 From: Mehrdad Hessar Date: Thu, 21 Apr 2022 13:42:14 -0700 Subject: [PATCH 6/6] fix lint --- python/tvm/contrib/hexagon/session.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/python/tvm/contrib/hexagon/session.py b/python/tvm/contrib/hexagon/session.py index a7361cab6b3f..7d2eecbc2c28 100644 --- a/python/tvm/contrib/hexagon/session.py +++ b/python/tvm/contrib/hexagon/session.py @@ -66,6 +66,7 @@ def __init__( self._remote_kw: dict = remote_kw self._rpc = None self._requires_cpu_device = False + self._device = None def __enter__(self): if self._rpc: @@ -98,7 +99,7 @@ def __exit__(self, exc_type, exc_value, exc_traceback): def device(self): """Session device.""" - if hasattr(self, "_device"): + if self._device is not None: return self._device if self._requires_cpu_device: