import tvm
from tvm import relay
import numpy as np
from tvm import rpc
from tvm.contrib import utils, ndk, graph_executor
import os
from tvm.relay import expr as _expr
from tvm.relay import function as _function
from tvm.relay import analysis
from tvm.ir import IRModule
target = "llvm -device=arm_cpu -model=android -mtriple=armv7a-linux-android -mattr=+neon"
os.environ["TVM_NDK_CC"] = \
'/Users/rqg/Library/Android/sdk/ndk/23.0.7530507/toolchains/llvm/prebuilt/darwin-x86_64/bin/armv7a-linux-androideabi19-clang++'
tracker_host = '192.168.31.60'
tracker_port = 9091
dev_key = 'a32'
if __name__ == '__main__':
input_name = "data"
data_shape = [1, 2]
expr_input = _expr.var(input_name, shape=data_shape)
expr_sm = relay.nn.softmax(expr_input)
func = _function.Function(analysis.free_vars(expr_sm), expr_sm)
params = {}
mod = IRModule.from_expr(func)
print(mod)
with tvm.transform.PassContext(opt_level=1):
lib = relay.build(mod, target=target, params=params)
# lib.get_lib().save('tt.ll')
tmp = utils.tempdir()
lib_fname = tmp.relpath("net.so")
fcompile = ndk.create_shared
lib.export_library(lib_fname, fcompile)
tracker = rpc.connect_tracker(tracker_host, tracker_port)
remote = tracker.request(dev_key, priority=0, session_timeout=60)
remote.upload(lib_fname)
rlib = remote.load_module("net.so")
dev = remote.cpu(0)
module = graph_executor.GraphModule(rlib["default"](dev))
data = np.array([[-0.55507267, 0.4956567]])
module.set_input(input_name, data)
module.run()
tvm_output = module.get_output(0).numpy()
# should be [[0.259085 0.740915]]
print(tvm_output)
Run relay.nn.softmax
on Android device with target:
-
armv7a-linux-android
. Wrong output -
aarch64-linux-android
. Correct output