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.
#
