def test_load_params_with_constants_in_ext_codegen():
# After binding params and partitioning graph_module.get_params()
# might contain parameters that are not an graph executor input but
# for example constants in external function.
y_in = np.ones((1, )).astype("float32")
params = {"y": y_in}
mod = tvm.IRModule()
x = relay.var("x", shape=(1, 10))
y = relay.var("y", shape=(1, ))
xcb = compiler_begin(x, "ccompiler")
ycb = compiler_begin(y, "ccompiler")
z = relay.add(xcb, ycb)
zce = compiler_end(z, "ccompiler")
mod["main"] = relay.Function([x, y], zce)
mod["main"] = bind_params_by_name(mod["main"], params)
mod = relay.transform.PartitionGraph()(mod)
graph_module = relay.build(mod, target="llvm", params=params)
# Params will be stored in metadata module.
assert len(graph_module.get_params()) == 0
lib = update_lib(graph_module.get_lib())
rt_mod = tvm.contrib.graph_executor.create(graph_module.get_graph_json(),
lib, tvm.cpu(0))
rt_mod.load_params(runtime.save_param_dict(graph_module.get_params()))
def test_byoc_microtvm(merge_compiler_regions):
"""This is a simple test to check BYOC capabilities of AOT - with and without merging compiler regions to test for https://github.com/apache/tvm/issues/9036"""
use_unpacked_api = False
interface_api = "packed"
test_runner = AOTTestRunner(pass_config={"tir.usmp.enable": True})
x = relay.var("x", shape=(10, 10))
w0 = relay.var("w0", shape=(10, 10))
w1 = relay.var("w1", shape=(10, 10))
# z0 = x + w0
x_ = compiler_begin(x, "ccompiler")
w0_ = compiler_begin(w0, "ccompiler")
z0_ = relay.add(x_, w0_)
z0 = compiler_end(z0_, "ccompiler")
# z1 = z0 + w1
z0__ = compiler_begin(z0, "ccompiler")
w1_ = compiler_begin(w1, "ccompiler")
z1_ = relay.add(z0__, w1_)
z1 = compiler_end(z1_, "ccompiler")
# z2 = z0 + z1
z2 = relay.add(z0, z1)
f = relay.Function([x, w0, w1], z2)
mod = tvm.IRModule()
mod["main"] = f
if merge_compiler_regions:
mod = transform.MergeCompilerRegions()(mod)
mod = transform.PartitionGraph("mod_name")(mod)
mod = transform.InferType()(mod)
x_data = [("x", np.random.rand(10, 10).astype("float32"))]
w_data = [("w{}".format(i), np.random.rand(10, 10).astype("float32"))
for i in range(2)]
map_inputs = OrderedDict(x_data + w_data)
output_list = generate_ref_data(mod, map_inputs)
compiled_test_mods = compile_models(
AOTTestModel(name="my_mod",
module=mod,
inputs=map_inputs,
outputs=output_list),
interface_api=interface_api,
use_unpacked_api=use_unpacked_api,
pass_config=test_runner.pass_config,
)
for compiled_model in compiled_test_mods:
check_for_no_tvm_backendallocworkspace_calls(
compiled_model.executor_factory.lib)
run_and_check(
models=compiled_test_mods,
runner=test_runner,
interface_api=interface_api,
)
def annotated():
in_1 = relay.var("in_1", shape=(10, 10), dtype="float32")
in_2 = relay.var("in_2", shape=(10, 10), dtype="float32")
in_3 = relay.var("in_3", shape=(10, 10), dtype="float32")
in_4 = relay.var("in_4", shape=(10, 10), dtype="float32")
in_5 = relay.var("in_5", shape=(10, 10), dtype="float32")
in_6 = relay.var("in_6", shape=(10, 10), dtype="float32")
in_7 = relay.var("in_7", shape=(10, 10), dtype="float32")
in_8 = relay.var("in_8", shape=(10, 10), dtype="float32")
in_9 = relay.var("in_9", shape=(10, 10), dtype="float32")
in_10 = relay.var("in_10", shape=(10, 10), dtype="float32")
begin0 = compiler_begin(in_1, "test")
begin1 = compiler_begin(in_2, "test")
begin2 = compiler_begin(in_3, "test")
begin3 = compiler_begin(in_4, "test")
node0 = relay.add(begin0, begin1)
node1 = relay.add(begin2, begin3)
end0 = compiler_end(node0, "test")
end1 = compiler_end(node1, "test")
begin4 = compiler_begin(end0, "test")
begin5 = compiler_begin(end1, "test")
node2 = relay.add(begin4, begin5)
end2 = compiler_end(node2, "test")
dbegin0 = compiler_begin(in_5, "default")
dbegin1 = compiler_begin(in_6, "default")
node3 = relay.subtract(dbegin0, dbegin1)
dbegin2 = compiler_begin(in_7, "default")
dend1 = compiler_end(node3, "default")
dbegin3 = compiler_begin(dend1, "default")
node4 = relay.subtract(dbegin2, dbegin3)
dend2 = compiler_end(node4, "default")
begin6 = compiler_begin(end2, "test")
begin7 = compiler_begin(dend2, "test")
node5 = relay.add(begin6, begin7)
end3 = compiler_end(node5, "test")
end4 = compiler_end(node5, "test")
dbegin4 = compiler_begin(in_8, "default")
dbegin5 = compiler_begin(end3, "default")
node6 = relay.subtract(dbegin4, dbegin5)
begin8 = compiler_begin(in_9, "test")
begin9 = compiler_begin(end4, "test")
node7 = relay.add(begin8, begin9)
end5 = compiler_end(node7, "test")
dend3 = compiler_end(node6, "default")
begin10 = compiler_begin(dend3, "test")
begin11 = compiler_begin(end5, "test")
node8 = relay.add(begin10, begin11)
end6 = compiler_end(node8, "test")
begin12 = compiler_begin(in_10, "test")
begin13 = compiler_begin(end6, "test")
node9 = relay.add(begin12, begin13)
end7 = compiler_end(node9, "test")
f = relay.Function(
[in_1, in_2, in_3, in_4, in_5, in_6, in_7, in_8, in_9, in_10],
end7)
mod = tvm.IRModule.from_expr(f)
return mod
def test_partition():
in_1 = relay.var("in_1", shape=(10, 10), dtype="float32")
in_2 = relay.var("in_2", shape=(10, 10), dtype="float32")
in_3 = relay.var("in_3", shape=(10, 10), dtype="float32")
in_4 = relay.var("in_4", shape=(10, 10), dtype="float32")
in_5 = relay.var("in_5", shape=(10, 10), dtype="float32")
in_6 = relay.var("in_6", shape=(10, 10), dtype="float32")
in_7 = relay.var("in_7", shape=(10, 10), dtype="float32")
in_8 = relay.var("in_8", shape=(10, 10), dtype="float32")
in_9 = relay.var("in_9", shape=(10, 10), dtype="float32")
in_10 = relay.var("in_10", shape=(10, 10), dtype="float32")
begin0 = compiler_begin(in_1, "onnx")
begin1 = compiler_begin(in_2, "onnx")
begin2 = compiler_begin(in_3, "onnx")
begin3 = compiler_begin(in_4, "onnx")
node0 = relay.add(begin0, begin1)
node1 = relay.add(begin2, begin3)
end0 = compiler_end(node0, "onnx")
end1 = compiler_end(node1, "onnx")
begin4 = compiler_begin(end0, "onnx")
begin5 = compiler_begin(end1, "onnx")
node2 = relay.add(begin4, begin5)
end2 = compiler_end(node2, "onnx")
dbegin0 = compiler_begin(in_5, "default")
dbegin1 = compiler_begin(in_6, "default")
node3 = relay.subtract(dbegin0, dbegin1)
dbegin2 = compiler_begin(in_7, "default")
dend1 = compiler_end(node3, "default")
dbegin3 = compiler_begin(dend1, "default")
node4 = relay.subtract(dbegin2, dbegin3)
dend2 = compiler_end(node4, "default")
begin6 = compiler_begin(end2, "onnx")
begin7 = compiler_begin(dend2, "onnx")
node5 = relay.add(begin6, begin7)
end3 = compiler_end(node5, "onnx")
end4 = compiler_end(node5, "onnx")
dbegin4 = compiler_begin(in_8, "default")
dbegin5 = compiler_begin(end3, "default")
node6 = relay.subtract(dbegin4, dbegin5)
begin8 = compiler_begin(in_9, "onnx")
begin9 = compiler_begin(end4, "onnx")
node7 = relay.multiply(begin8, begin9)
end5 = compiler_end(node7, "onnx")
dend3 = compiler_end(node6, "default")
begin10 = compiler_begin(dend3, "onnx")
begin11 = compiler_begin(end5, "onnx")
node8 = relay.add(begin10, begin11)
end6 = compiler_end(node8, "onnx")
begin12 = compiler_begin(in_10, "onnx")
begin13 = compiler_begin(end6, "onnx")
node9 = relay.add(begin12, begin13)
end7 = compiler_end(node9, "onnx")
func = relay.Function(
[in_1, in_2, in_3, in_4, in_5, in_6, in_7, in_8, in_9, in_10], end7)
target = "llvm"
mod = IRModule.from_expr(func)
mod = transform.PartitionGraph()(mod)
with tvm.transform.PassContext(opt_level=3, disabled_pass=["FuseOps"]):
graph_json, mod1, params = relay.build(mod, target)
assert mod1.type_key == "metadata"
assert mod1.imported_modules[0].type_key == "llvm"
assert mod1.imported_modules[0].get_source()
assert mod1.imported_modules[1].type_key == "onnx"
assert mod1.imported_modules[1].get_source()
