Skip to content

Commit c4d9c77

Browse files
mehrdadhmehrdadh
committed
Hexagon AOT tests work
1 parent 9a4a50b commit c4d9c77

8 files changed

Lines changed: 373 additions & 25 deletions

File tree

‎python/tvm/contrib/hexagon/build.py‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -254,6 +254,24 @@ def get_graph_executor(self, libmod, remote_libmod_filename: str):
254254
libmod.get_graph_json(), hexagon_mod, self.session.device
255255
)
256256

257+
def get_aot_executor(self, libmod, remote_libmod_filename: str):
258+
"""Create a local GraphModule which consumes a remote libmod.
259+
Parameters
260+
----------
261+
libmod : tvm.runtime.Module
262+
The module of the corresponding function.
263+
This library module is for remote hexagon runtime.
264+
remote_libmod_filename : str
265+
Module filename on remote. It is assumed this file lives under self._workspace path.
266+
Returns
267+
-------
268+
graph_module : GraphModule
269+
Runtime graph module that can be used to execute the graph.
270+
"""
271+
self.session.__enter__()
272+
hexagon_mod = self.get_module(remote_libmod_filename)
273+
return tvm.runtime.executor.AotModule(hexagon_mod["default"](self.session.device))
274+
257275
def close(self):
258276
"""Close RPC server on Android"""
259277
# Kill process childs

‎python/tvm/contrib/hexagon/hexagon.py‎

Lines changed: 47 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import functools as ft
2121
import os
22+
import pathlib
2223
import tvm
2324
import tvm.ir
2425
import tvm.contrib.cc as cc
@@ -39,10 +40,20 @@
3940
#
4041
# Subsequent calls to 'link_shared' will use the newly registered linker.
4142

42-
hexagon_toolchain_root = os.environ.get("HEXAGON_TOOLCHAIN") or "" # pylint: disable=invalid-name
43-
hexagon_link_main = os.path.join( # pylint: disable=invalid-name
44-
hexagon_toolchain_root, "bin", "hexagon-link"
43+
HEXAGON_TOOLCHAIN = os.environ.get(
44+
"HEXAGON_TOOLCHAIN", default=None
45+
) # pylint: disable=invalid-name
46+
HEXAGON_SDK_PATH = os.environ.get("HEXAGON_SDK_PATH", default=None) # pylint: disable=invalid-name
47+
HEXAGON_LINK_MAIN = os.path.join( # pylint: disable=invalid-name
48+
HEXAGON_TOOLCHAIN, "bin", "hexagon-link"
4549
)
50+
HEXAGON_CLANG_PLUS = os.path.join( # pylint: disable=invalid-name
51+
HEXAGON_TOOLCHAIN, "bin", "hexagon-clang++"
52+
)
53+
HEXAGON_SDK_INCLUDE_DIRS = [ # pylint: disable=invalid-name
54+
pathlib.Path(HEXAGON_SDK_PATH) / "incs",
55+
pathlib.Path(HEXAGON_SDK_PATH) / "incs" / "stddef",
56+
]
4657

4758

4859
def register_linker(f):
@@ -53,7 +64,13 @@ def register_linker(f):
5364
@register_func("tvm.contrib.hexagon.hexagon.hexagon_link")
5465
def hexagon_link():
5566
"""Return path to the Hexagon linker."""
56-
return hexagon_link_main
67+
return HEXAGON_LINK_MAIN
68+
69+
70+
@register_func("tvm.contrib.hexagon.hexagon.hexagon_clang_plus")
71+
def hexagon_clang_plus():
72+
"""Return path to the Hexagon clang++."""
73+
return HEXAGON_CLANG_PLUS
5774

5875

5976
@register_func("tvm.contrib.hexagon.hexagon.link_shared")
@@ -101,12 +118,12 @@ def to_str(s):
101118
message += (
102119
" Please verify the value of the HEXAGON_LINKER environment variable "
103120
+ '(currently set to "'
104-
+ hexagon_toolchain_root
121+
+ HEXAGON_TOOLCHAIN
105122
+ '").'
106123
)
107124
raise Exception(message)
108125

109-
libpath = os.path.join(hexagon_toolchain_root, "target", "hexagon", "lib", "v66", "G0")
126+
libpath = os.path.join(HEXAGON_TOOLCHAIN, "target", "hexagon", "lib", "v66", "G0")
110127
cc.create_shared(
111128
so_name,
112129
objs,
@@ -248,3 +265,27 @@ def transform(func, mod, ctx):
248265

249266
def ir_lower_vtcm_pass():
250267
return [(3, ir_lower_vtcm())]
268+
269+
270+
@register_func("tvm.contrib.hexagon.hexagon.aot_export")
271+
def aot_export(so_name, files, **kwargs):
272+
tvm_dir = pathlib.Path(os.path.dirname(os.path.realpath(__file__))) / ".." / ".." / ".." / ".."
273+
options = [
274+
f"-I{tvm_dir / 'include'}",
275+
f"-I{tvm_dir / '3rdparty' / 'dlpack' / 'include'}",
276+
f"-I{tvm_dir / '3rdparty' / 'dmlc-core' / 'include'}",
277+
f"-I{tvm_dir / 'src' / 'runtime' / 'hexagon' / 'android' / 'sim' / 'driver'}",
278+
f"-DDMLC_USE_LOGGING_LIBRARY=<tvm/runtime/logging.h>",
279+
f"-D_MACH_I32=int",
280+
]
281+
282+
# For debugging
283+
for path in HEXAGON_SDK_INCLUDE_DIRS:
284+
options.append(f"-I{str(path)}")
285+
286+
cross_compile = cc.cross_compiler(
287+
compile_func=tvm.get_global_func("tvm.contrib.hexagon.hexagon.hexagon_clang_plus")()
288+
)
289+
cross_compile.output_format = "o"
290+
c_files = [str(file) for file in files]
291+
cross_compile(so_name, c_files, options=options)

‎python/tvm/runtime/module.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -472,14 +472,13 @@ def export_library(self, file_name, fcompile=None, addons=None, workspace_dir=No
472472
# The imports could contain a c module but the object format could be tar
473473
# Thus, it would not recognize the following include paths as options
474474
# which are there assuming a c compiler is the fcompile.
475-
if has_c_module and not file_name.endswith(".tar"):
475+
if has_c_module and not file_name.endswith(".tar") and not file_name.endswith(".so"):
476476
options = []
477477
if "options" in kwargs:
478478
opts = kwargs["options"]
479479
options = opts if isinstance(opts, (list, tuple)) else [opts]
480480
opts = options + ["-I" + path for path in find_include_path()]
481481
kwargs.update({"options": opts})
482-
483482
return fcompile(file_name, files, **kwargs)
484483

485484

‎src/relay/backend/aot_executor_codegen.cc‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1172,6 +1172,9 @@ class AOTExecutorCodegenModule : public runtime::ModuleNode {
11721172
if (!target_host.defined() && it.second->kind->device_type == kDLCPU) {
11731173
target_host = it.second;
11741174
}
1175+
if (!target_host.defined() && it.second->kind->device_type == kDLHexagon) {
1176+
target_host = *(new Target("c"));
1177+
}
11751178
ICHECK(dev_type);
11761179
targets[static_cast<DLDeviceType>(dev_type->value)] = it.second;
11771180
}

‎src/runtime/aot_executor/aot_executor.cc‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -125,8 +125,8 @@ PackedFunc AotExecutor::GetFunction(const std::string& name,
125125
}
126126

127127
void AotExecutor::Run() {
128-
LOG(INFO) << "Get entrypoint " << metadata_->mod_name() << "_run_model";
129-
auto pf = module_.GetFunction(metadata_->mod_name() + "_run_model", true /* query_imports */);
128+
LOG(INFO) << "Get entrypoint " << metadata_->mod_name() << "___tvm_main__";
129+
auto pf = module_.GetFunction(metadata_->mod_name() + "___tvm_main__", true /* query_imports */);
130130
ICHECK(pf != nullptr) << "Module entrypoint is not defined";
131131

132132
const int num_args = args_.size();

‎src/runtime/hexagon/hexagon/hexagon_device_api_v2.cc‎

Lines changed: 9 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -56,7 +56,7 @@ void HexagonDeviceAPIv2::GetAttr(Device dev, DeviceAttrKind kind, TVMRetValue* r
5656
// DataSpace: static allocations for Hexagon
5757
void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, int ndim, const int64_t* shape,
5858
DLDataType dtype, Optional<String> mem_scope) {
59-
CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
59+
CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
6060

6161
// Forcing contiguous allocation, for now
6262
// TODO(Straw): Enable discontiguous allocation after RFC 39 lands
@@ -84,7 +84,12 @@ void* HexagonDeviceAPIv2::AllocDataSpace(Device dev, size_t nbytes, size_t align
8484
}
8585

8686
void HexagonDeviceAPIv2::FreeDataSpace(Device dev, void* ptr) {
87-
CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
87+
bool device = false;
88+
if ((TVMDeviceExtType(dev.device_type) == kDLHexagon) ||
89+
(DLDeviceType(dev.device_type) == kDLCPU)) {
90+
device = true;
91+
}
92+
CHECK(device) << "dev.device_type: " << dev.device_type;
8893
auto* hexbuf = static_cast<HexagonBuffer*>(ptr);
8994
CHECK(hexbuf != nullptr);
9095
delete hexbuf;
@@ -97,7 +102,7 @@ struct HexagonWorkspacePool : public WorkspacePool {
97102
};
98103

99104
void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType type_hint) {
100-
CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
105+
CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
101106
auto* hexbuf = static_cast<HexagonBuffer*>(
102107
dmlc::ThreadLocalStore<HexagonWorkspacePool>::Get()->AllocWorkspace(dev, size));
103108

@@ -109,7 +114,7 @@ void* HexagonDeviceAPIv2::AllocWorkspace(Device dev, size_t size, DLDataType typ
109114
}
110115

111116
void HexagonDeviceAPIv2::FreeWorkspace(Device dev, void* data) {
112-
CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon);
117+
CHECK(TVMDeviceExtType(dev.device_type) == kDLHexagon) << "dev.device_type: " << dev.device_type;
113118
auto it = workspace_allocations_.find(data);
114119
CHECK(it != workspace_allocations_.end())
115120
<< "Attempt made to free unknown or already freed workspace allocation";

0 commit comments

Comments
 (0)