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_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_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_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_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_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_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_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_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_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_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, )