def build_module(mod, target, params=None, enable_bnns=True, tvm_ops=0):
"""Build module with option to build for BNNS."""
if isinstance(mod, tvm.relay.expr.Call):
mod = tvm.IRModule.from_expr(mod)
with tvm.transform.PassContext(opt_level=3):
if enable_bnns:
mod = partition_for_bnns(mod)
relay.backend.te_compiler.get().clear()
return relay.build(mod, target=target, target_host=target, params=params)
def run(mod, target, simplify=True, with_bnns=False):
with tvm.transform.PassContext(opt_level=3):
if simplify:
mod = simplify_model(mod)
if with_bnns:
mod = partition_for_bnns(mod)
graph_module = relay.build(mod, target=target, target_host=target, params=params)
lib_name = "deploy.tar"
path_dso = temp.relpath(lib_name)
graph_module.export_library(path_dso)
dev = tvm.cpu(0)
loaded_lib = tvm.runtime.load_module(path_dso)
module = graph_runtime.GraphModule(loaded_lib["default"](dev))
module.run()
return module.get_output(0).asnumpy()
def partition(exp):
"""Apply BNNS specific partitioning transformation"""
mod = tvm.IRModule.from_expr(exp)
with tvm.transform.PassContext(opt_level=3):
mod = partition_for_bnns(mod)
return mod