def test_transpose(interface_api, use_unpacked_api, test_runner): """Test that non-inpleaceable operations (e.g., transpose) do not happen in-place.""" dtype = "float32" x = relay.var("x", shape=(10, 5), dtype=dtype) y = relay.var("y", shape=(10, 5), dtype=dtype) t = relay.var("z", shape=(), dtype=dtype) a = relay.add(x, y) b = relay.transpose(a) z = relay.add(b, t) # Check result. func = relay.Function([x, y, t], z) x_data = np.random.rand(10, 5).astype(dtype) y_data = np.random.rand(10, 5).astype(dtype) t_data = np.random.uniform(size=()).astype(dtype) inputs = {"x": x_data, "y": y_data, "z": t_data} output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, enable_op_fusion=False, )
def test_aot_codegen_checks_returns(): """This test checks whether AoT lowering creates calls that check the return value correctly""" x = relay.var("x", shape=(1, 10)) y = relay.var("y", shape=(1, 10)) z = relay.add(x, y) func = relay.Function([x, y], z) compiled_test_mods = compile_models( models=AOTTestModel(module=IRModule.from_expr(func), inputs=None, outputs=None), interface_api="c", use_unpacked_api=True, ) source = compiled_test_mods[0].executor_factory.lib.imported_modules[ 0].get_source() main_ir_module = compiled_test_mods[ 0].executor_factory.lowered_ir_mods.items()[0][1] main_func = main_ir_module["__tvm_main__"] # Check operator call is wrapped properly assert (str( main_func.body[1]) == "tir.tvm_check_return(0, -1, tir.call_extern(" + '"tvmgen_default_fused_add",' + " x_buffer_var, y_buffer_var, output_buffer_var))\n") # TODO(Mousius) - Create a better place for C codegen tests assert ( "if (tvmgen_default_fused_add(x_buffer_var, y_buffer_var, output_buffer_var) != 0 ) return -1;" in source)
def test_name_sanitiser_name_clash(): """Test that 2 input tensors with names that clash once sanitized, generates an error""" interface_api = "c" use_unpacked_api = True test_runner = AOT_DEFAULT_RUNNER dtype = "float32" x = relay.var("input::-1", shape=(10, 5), dtype=dtype) # Next 2 input tensor names will clash once sanitized. y = relay.var("input::-2", shape=(10, 5), dtype=dtype) t = relay.var("input:--2", shape=(), dtype=dtype) a = relay.add(x, y) b = relay.transpose(a) z = relay.add(b, t) # Check result. func = relay.Function([x, y, t], z) x_data = np.random.rand(10, 5).astype(dtype) y_data = np.random.rand(10, 5).astype(dtype) t_data = np.random.uniform(size=()).astype(dtype) inputs = {"input::-1": x_data, "input::-2": y_data, "input:--2": t_data} output_list = generate_ref_data(func, inputs) with pytest.raises(TVMError, match="Sanitized input tensor name clash"): compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, enable_op_fusion=False, )
def test_deprecated_target_arguments(capsys): """Tests we can still use relay.build with -executor, -runtime and -link-params""" interface_api = "c" use_unpacked_api = True test_runner = AOT_DEFAULT_RUNNER x = relay.var("x", shape=(1, 10)) y = relay.var("y", shape=(1, 10)) z = relay.add(x, y) func = relay.Function([x, y], z) x_in = np.ones((1, 10)).astype("float32") y_in = np.random.uniform(size=(1, 10)).astype("float32") params = {"x": x_in} inputs = {"y": y_in} output_list = generate_ref_data(func, inputs, params) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list, params=params), test_runner, interface_api, use_unpacked_api, use_runtime_executor=False, target= "c -executor=aot --link-params -runtime=c -interface-api=c --unpacked-api", )
def test_add_with_params(interface_api, use_unpacked_api, test_runner): """Tests compilation of add with parameters""" input_x = relay.var("x", shape=(1, 10)) input_y = relay.var("y", shape=(1, 10)) input_z = relay.add(input_x, input_y) func = relay.Function([input_x, input_y], input_z) input_x_data = np.ones((1, 10)).astype("float32") input_y_data = np.random.uniform(size=(1, 10)).astype("float32") params = {"x": input_x_data} inputs = {"y": input_y_data} output_list = generate_ref_data(func, inputs, params) compile_and_run( AOTTestModel( module=IRModule.from_expr(func), inputs=inputs, outputs=output_list, params=params, ), test_runner, interface_api, use_unpacked_api, )
def test_condition_mutually_exclusive(): mod = IRModule.from_expr(concat_func_3) with tvm.transform.PassContext(config={"tir.LoopPartition": {"partition_const_loop": True}}): mod = tvm.tir.transform.FlattenBuffer()(mod) mod = tvm.tir.transform.LoopPartition()(mod) mod = tvm.tir.transform.Simplify()(mod) mod = tvm.tir.transform.RemoveNoOp()(mod) assert tvm.ir.structural_equal(mod["main"], partitioned_concat_3)
def test_tuple_getitem(interface_api, use_unpacked_api, test_runner): func = relay.Function([], relay.TupleGetItem(relay.Tuple([relay.const(1), relay.const(2)]), 0)) output_list = generate_ref_data(func, {}) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs={}, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_lower_build_tir_module(): func = matmul.with_attr("global_symbol", "main") func = func.with_attr("tir.noalias", True) ir_mod = IRModule({"main": func}) # check lowering lowered_mod = tvm.lower(ir_mod) tvm.ir.assert_structural_equal(lowered_mod, LoweredModule()) # check building mod = tvm.build(ir_mod, target="llvm") _check_module_with_numpy(mod)
def test_add_const(interface_api, use_unpacked_api, test_runner): two = relay.add(relay.const(1), relay.const(1)) func = relay.Function([], two) output_list = generate_ref_data(func, {}) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs={}, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_meta_schedule_integration_apply_history_best(): @derived_object class DummyDatabase(PyDatabase): def __init__(self): super().__init__() self.records = [] self.workload_reg = [] def has_workload(self, mod: IRModule) -> Workload: for workload in self.workload_reg: if tvm.ir.structural_equal(workload.mod, mod): return True return False def commit_tuning_record(self, record: TuningRecord) -> None: self.records.append(record) def commit_workload(self, mod: IRModule) -> Workload: for workload in self.workload_reg: if tvm.ir.structural_equal(workload.mod, mod): return workload workload = Workload(mod) self.workload_reg.append(workload) return workload def get_top_k(self, workload: Workload, top_k: int) -> List[TuningRecord]: return list( filter( lambda x: x.workload == workload, sorted(self.records, key=lambda x: sum(x.run_secs) / len(x.run_secs)), ))[:int(top_k)] def __len__(self) -> int: return len(self.records) def print_results(self) -> None: print("\n".join([str(r) for r in self.records])) mod, _, _ = get_network(name="resnet_18", input_shape=[1, 3, 224, 224]) database = DummyDatabase() env = ApplyHistoryBest(database) target = Target("llvm") workload = database.commit_workload(MockModule) database.commit_tuning_record( TuningRecord(Schedule(MockModule).trace, [1.0], workload, target, [])) mod = env.query(task_name="mock-task", mod=mod, target=target, dispatched=[MockModule]) mod = IRModule({"main": mod}) assert tvm.ir.structural_equal(mod, workload.mod)
def test_lower_build_tir_module(): func = matmul.with_attr("global_symbol", "main") func = func.with_attr("tir.noalias", True) ir_mod = IRModule({"main": func}) # check lowering with the CSE pass disabled as otherwise it would do some commoning with tvm.transform.PassContext(opt_level=3, disabled_pass=["tir.CommonSubexprElimTIR"]): lowered_mod = tvm.lower(ir_mod) tvm.ir.assert_structural_equal(lowered_mod, LoweredTIRModule) # check building mod = tvm.build(ir_mod, target="llvm") _check_module_with_numpy(mod)
def test_subtract(interface_api, use_unpacked_api, test_runner): i = relay.var("i", shape=[], dtype="int32") sub = relay.subtract(i, relay.const(1, dtype="int32")) func = relay.Function([i], sub, ret_type=relay.TensorType([], "int32")) i_data = np.array(1, dtype="int32") inputs = {"i": i_data} output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_id(interface_api, use_unpacked_api, test_runner): x = relay.var("x", "float32") ident = relay.Function([x], x) one = np.array(1.0, "float32") inputs = {"x": one} output_list = generate_ref_data(ident, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(ident), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_aot_uses_anf(): """Checks that A-Normal Form is being used in the AOT lowering pipeline.""" x = relay.var("x", shape=(1, 10, 10, 10)) y = relay.var("y", shape=(1, 10, 10, 10)) z = relay.add(x, y) func = relay.Function([x, y], z) @pass_instrument class CheckANFRuns: def __init__(self): self.did_run_anf = False def run_before_pass(self, _, info): if info.name == "ToANormalForm": self.did_run_anf = True if info.name == "LowerTE": assert self.did_run_anf, "ToANormalForm pass should run before LowerTE." check_run_anf = CheckANFRuns() model = AOTTestModel(module=IRModule.from_expr(func), inputs=None, outputs=None) runtime = Runtime("crt") executor = Executor( "aot", { "workspace-byte-alignment": 8, "interface-api": "c", "unpacked-api": True, }, ) config = {"tir.disable_vectorize": True} with tvm.transform.PassContext(opt_level=3, config=config, instruments=[check_run_anf]): tvm.relay.build( model.module, tvm.target.Target("c"), executor=executor, runtime=runtime, workspace_memory_pools=None, params=model.params, mod_name=model.name, ) assert check_run_anf.did_run_anf, "Expected ToANormalForm pass to have run."
def test_tuple_output(interface_api, use_unpacked_api, test_runner): x = relay.var("x", shape=(6, 9)) y = relay.split(x, 3).astuple() a = relay.TupleGetItem(y, 0) b = relay.TupleGetItem(y, 1) out = relay.Tuple([a, b]) func = relay.Function([x], out) x_data = np.random.rand(6, 9).astype("float32") inputs = {"x": x_data} output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_mul_param(interface_api, use_unpacked_api, test_runner): x = relay.var("x", shape=(10, 10)) y = relay.var("y", shape=(1, 10)) func = relay.Function([x, y], relay.multiply(x, y)) x_data = np.random.rand(10, 10).astype("float32") y_data = np.random.rand(1, 10).astype("float32") inputs = OrderedDict([("x", x_data), ("y", y_data)]) output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def compile_to_main_func(interface_api="c", use_unpacked_api=True): test_runner = AOT_DEFAULT_RUNNER compiled_models = compile_models( models=AOTTestModel( module=IRModule.from_expr(func), inputs=inputs, outputs=output_list, ), interface_api=interface_api, use_unpacked_api=use_unpacked_api, workspace_byte_alignment=16, pass_config=test_runner.pass_config, ) main_ir_module = list(compiled_models[0].executor_factory.lowered_ir_mods.values())[0] main_func = main_ir_module["run_model"] return main_func
def test_error_c_interface_with_packed_api(): interface_api = "c" use_unpacked_api = False test_runner = AOT_DEFAULT_RUNNER two = relay.add(relay.const(1), relay.const(1)) func = relay.Function([], two) with pytest.raises(tvm.TVMError, match="Packed interface required for packed operators"): compile_and_run( AOTTestModel( module=IRModule.from_expr(func), inputs={}, outputs=generate_ref_data(func, {}) ), test_runner, interface_api, use_unpacked_api, )
def test_nested_tuples(interface_api, use_unpacked_api, test_runner): x = relay.var("x", shape=(10,)) x1 = x + relay.const(1.0) x2 = x1 + relay.const(1.0) x3 = x2 + relay.const(1.0) x4 = x3 + relay.const(1.0) out = relay.Tuple([x1, relay.Tuple([relay.Tuple([x2, x3]), x4])]) func = relay.Function([x], out) x_data = np.random.uniform(size=(10,)).astype(np.float32) inputs = {"x": x_data} output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_name_sanitiser(): """Test that input tensors with special characters in the name don't break compilation""" interface_api = "c" use_unpacked_api = True test_runner = AOT_DEFAULT_RUNNER func = relay.var("input-x::2", "float32") ident = relay.Function([func], func) one = np.array(1.0, "float32") inputs = {"input-x::2": one} output_list = generate_ref_data(ident, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, enable_op_fusion=False, )
def test_add_with_params(interface_api, use_unpacked_api, test_runner): x = relay.var("x", shape=(1, 10)) y = relay.var("y", shape=(1, 10)) z = relay.add(x, y) func = relay.Function([x, y], z) x_in = np.ones((1, 10)).astype("float32") y_in = np.random.uniform(size=(1, 10)).astype("float32") params = {"x": x_in} inputs = {"y": y_in} output_list = generate_ref_data(func, inputs, params) compile_and_run( AOTTestModel( module=IRModule.from_expr(func), inputs=inputs, outputs=output_list, params=params ), test_runner, interface_api, use_unpacked_api, )
def test_concatenate(interface_api, use_unpacked_api, test_runner): dtype = "float32" x = relay.var("x", shape=(10, 5), dtype=dtype) y = relay.var("y", shape=(10, 5), dtype=dtype) t = relay.var("z", shape=(), dtype=dtype) z = relay.concatenate((x, y), axis=1) z = relay.add(z, t) # Check result. func = relay.Function([x, y, t], z) x_data = np.random.rand(10, 5).astype(dtype) y_data = np.random.rand(10, 5).astype(dtype) t_data = np.random.uniform(size=()).astype(dtype) inputs = OrderedDict([("x", x_data), ("y", y_data), ("z", t_data)]) output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_error_c_interface_with_packed_api(): interface_api = "c" use_unpacked_api = False test_runner = AOT_DEFAULT_RUNNER two = relay.add(relay.const(1), relay.const(1)) func = relay.Function([], two) with pytest.raises( tvm.TVMError, match=re.escape( 'Either need interface_api == "packed" (got: c) or ' "unpacked-api == true (got: 0) when targeting " "c runtime"), ): compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs={}, outputs=generate_ref_data(func, {})), test_runner, interface_api, use_unpacked_api, )
def test_concatenate(interface_api, use_unpacked_api, test_runner): """Tests compilation of concatenate""" dtype = "float32" input_x = relay.var("x", shape=(10, 5), dtype=dtype) input_y = relay.var("y", shape=(10, 5), dtype=dtype) input_z = relay.var("z", shape=(), dtype=dtype) concat_inputs = relay.concatenate((input_x, input_y), axis=1) func_output = relay.add(input_z, concat_inputs) # Check result. func = relay.Function([input_x, input_y, input_z], func_output) x_data = np.random.rand(10, 5).astype(dtype) y_data = np.random.rand(10, 5).astype(dtype) t_data = np.random.uniform(size=()).astype(dtype) inputs = OrderedDict([("x", x_data), ("y", y_data), ("z", t_data)]) output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
def test_nested_tuples(interface_api, use_unpacked_api, test_runner): """Tests compilation of functions with nested tuple outputs""" input_x = relay.var("x", shape=(10, )) output_1 = input_x + relay.const(1.0) output_2 = output_1 + relay.const(1.0) output_3 = output_2 + relay.const(1.0) output_4 = output_3 + relay.const(1.0) full_output = relay.Tuple( [output_1, relay.Tuple([relay.Tuple([output_2, output_3]), output_4])]) func = relay.Function([input_x], full_output) x_data = np.random.uniform(size=(10, )).astype(np.float32) inputs = {"x": x_data} output_list = generate_ref_data(func, inputs) compile_and_run( AOTTestModel(module=IRModule.from_expr(func), inputs=inputs, outputs=output_list), test_runner, interface_api, use_unpacked_api, )
ir_module = MyModule print(type(ir_module)) print(ir_module.script()) ################################################################################################ # Besides, we can also use tensor expression DSL to write simple operators, and convert them # to an IRModule. # from tvm import te A = te.placeholder((8,), dtype="float32", name="A") B = te.compute((8,), lambda *i: A(*i) + 1.0, name="B") func = te.create_prim_func([A, B]) ir_module_from_te = IRModule({"main": func}) print(ir_module_from_te.script()) ################################################################################################ # Build and Run an IRModule # ------------------------- # We can build the IRModule into a runnable module with specific target backends. # mod = tvm.build(ir_module, target="llvm") # The module for CPU backends. print(type(mod)) ################################################################################################ # Prepare the input array and output array, then run the module. #