alanmacd commented on code in PR #13046: URL: https://github.com/apache/tvm/pull/13046#discussion_r992893204
########## tests/python/relay/aot/test_cpp_aot.py: ########## @@ -203,5 +203,39 @@ def test_pass_wrong_device_arg(): # TODO write asserts for # and type of device. +@pytest.mark.parametrize("target_kind", ["c", "llvm"]) +@pytest.mark.parametrize("input_name", ["input:0", "input@0", "input_0"]) +def test_aot_input_name_with_special_character(target_kind: str, input_name: str): + dtype = "float32" + input = relay.var(input_name, shape=(10, 5), dtype=dtype) + weight = relay.var("weight", shape=(1, 5), dtype=dtype) + output = relay.add(input, weight) + func = relay.Function([input, weight], output) + + input_data = np.random.rand(10, 5).astype(dtype) + weight_data = np.random.rand(1, 5).astype(dtype) + expected_output = input_data + weight_data + params = {"weight": weight_data} + + with tvm.transform.PassContext(opt_level=3, config={"tir.disable_vectorize": True}): + mod = tvm.relay.build( + tvm.IRModule.from_expr(func), + target=target_kind, + params=params, + executor=tvm.relay.backend.Executor("aot", {"interface-api": "packed"}), + ) + temp_dir = tvm.contrib.utils.TempDirectory() + test_so_path = temp_dir / "test.so" + mod.export_library(test_so_path, cc="c++", options=["-std=gnu++17", "-g3", "-O0"]) + # test both original name and mangled name + for name in ["input_0", input_name]: + loaded_mod = tvm.runtime.load_module(test_so_path) + runner = tvm.runtime.executor.AotModule(loaded_mod["default"](tvm.cpu(0))) + inputs = {name: input_data} + runner.set_input(**inputs) + runner.run() + assert (runner.get_output(0).asnumpy() == expected_output).all() Review Comment: also add checks of `get_input `and `get_input_index `since they are separate code paths in aot_executor.cc? -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: commits-unsubscr...@tvm.apache.org For queries about this service, please contact Infrastructure at: us...@infra.apache.org