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()