-
Notifications
You must be signed in to change notification settings - Fork 107
/
tvm_compiler_utils.py
69 lines (56 loc) · 2.92 KB
/
tvm_compiler_utils.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
#!/usr/bin/env python3
import os
import sys
import tvm
from tvm import relay
from tvm.contrib import ndk
sys.setrecursionlimit(1031)
print("sys.getrecursionlimit()", sys.getrecursionlimit())
def tvm_compile(func, params, arch, dlr_model_name):
###arch x86_64
if arch == 'x86_64':
target = "llvm -model=N3350 -target=x86_64-linux-android -mattr=+ssse3,+sse4.2"
sysroot="/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/sysroot"
toolchain="/opt/android-ndk/toolchains/x86_64-4.9/prebuilt/linux-x86_64"
os.environ['TVM_NDK_CC'] = "/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/bin/x86_64-linux-android28-clang++"
###arch x86 i686
elif arch == 'x86':
target = "llvm -model=x5-Z8350 -target=i686-linux-android -mattr=+ssse3"
sysroot="/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/sysroot"
toolchain="/opt/android-ndk/toolchains/x86-4.9/prebuilt/linux-x86_64"
os.environ['TVM_NDK_CC'] = "/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/bin/i686-linux-android21-clang++"
###arch arm64 aarch64
elif arch == 'arm64-v8a':
target = "llvm -device=arm_cpu -model=SM8150 -target=aarch64-linux-android"
sysroot="/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/sysroot"
toolchain="/opt/android-ndk/toolchains/aarch64-linux-android-4.9/prebuilt/linux-x86_64"
os.environ['TVM_NDK_CC'] = "/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/bin/aarch64-linux-android28-clang++"
###arch armv7
## More info on armv7 hard/soft abi for Android https://android.googlesource.com/platform/ndk/+/master/docs/HardFloatAbi.md
elif arch == 'armeabi-v7a':
target = "llvm -device=arm_cpu -model=MSM8940 -target=armv7a-linux-androideabi -mfloat-abi=soft -mattr=+neon,+thumb-mode"
sysroot="/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/sysroot"
toolchain="/opt/android-ndk/toolchains/arm-linux-androideabi-4.9/prebuilt/linux-x86_64"
os.environ['TVM_NDK_CC'] = "/opt/android-ndk/toolchains/llvm/prebuilt/linux-x86_64/bin/armv7a-linux-androideabi21-clang++"
else:
print("Valid arch: arm64-v8a, armeabi-v7a, x86_64, x86")
return
print('target:', target)
print("Compiling...")
with relay.build_config(opt_level=3):
graph, lib, params = relay.build(func, target, params=params)
print("Compilation done")
print("lib type_key: ", lib.type_key)
print("Saving files")
out_folder = arch + "/" + dlr_model_name + "/"
os.makedirs(out_folder, exist_ok=True)
# save the graph, lib and params into separate files
path_lib = out_folder + "model.so"
options=["-shared", "-fPIC", "--sysroot", sysroot, "--gcc-toolchain="+toolchain, "-static-libstdc++"]
lib.export_library(path_lib, ndk.create_shared, options=options)
print("export_library done")
with open(out_folder + "model.json", "w") as fo:
fo.write(graph)
with open(out_folder + "model.params", "wb") as fo:
fo.write(relay.save_param_dict(params))
print("Files saved to", out_folder)