diff --git a/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py b/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py index 82da49e41096..bc3ae64b46c1 100644 --- a/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py +++ b/tests/python/all-platform-minimal-test/test_minimal_target_codegen_llvm.py @@ -20,7 +20,7 @@ import tvm.testing from tvm import te from tvm import topi -from tvm.contrib import utils, clang +from tvm.contrib import utils import numpy as np import ctypes import math @@ -65,43 +65,3 @@ def check_llvm(): tvm.testing.assert_allclose(c.numpy(), a.numpy() + b.numpy()) check_llvm() - - -@tvm.testing.requires_llvm -def test_llvm_import(): - """all-platform-minimal-test: check shell dependent clang behavior.""" - # extern "C" is necessary to get the correct signature - cc_code = """ - extern "C" float my_add(float x, float y) { - return x + y; - } - """ - n = 10 - A = te.placeholder((n,), name="A") - B = te.compute( - (n,), lambda *i: tvm.tir.call_pure_extern("float32", "my_add", A(*i), 1.0), name="B" - ) - - def check_llvm(use_file): - if not clang.find_clang(required=False): - print("skip because clang is not available") - return - temp = utils.tempdir() - ll_path = temp.relpath("temp.ll") - ll_code = clang.create_llvm(cc_code, output=ll_path) - s = te.create_schedule(B.op) - if use_file: - s[B].pragma(s[B].op.axis[0], "import_llvm", ll_path) - else: - s[B].pragma(s[B].op.axis[0], "import_llvm", ll_code) - # BUILD and invoke the kernel. - f = tvm.build(s, [A, B], "llvm") - dev = tvm.cpu(0) - # launch the kernel. - a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev) - b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev) - f(a, b) - tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1.0) - - check_llvm(use_file=True) - check_llvm(use_file=False) diff --git a/tests/python/unittest/test_target_codegen_llvm.py b/tests/python/unittest/test_target_codegen_llvm.py index e5e93ed2c940..8c8d601672ac 100644 --- a/tests/python/unittest/test_target_codegen_llvm.py +++ b/tests/python/unittest/test_target_codegen_llvm.py @@ -23,7 +23,7 @@ import tvm.testing from tvm import te from tvm import topi -from tvm.contrib import utils +from tvm.contrib import utils, clang import numpy as np import ctypes import math @@ -845,5 +845,45 @@ def make_call_extern(caller, callee): assert matches == sorted(matches) +@tvm.testing.requires_llvm +def test_llvm_import(): + """all-platform-minimal-test: check shell dependent clang behavior.""" + # extern "C" is necessary to get the correct signature + cc_code = """ + extern "C" float my_add(float x, float y) { + return x + y; + } + """ + n = 10 + A = te.placeholder((n,), name="A") + B = te.compute( + (n,), lambda *i: tvm.tir.call_pure_extern("float32", "my_add", A(*i), 1.0), name="B" + ) + + def check_llvm(use_file): + if not clang.find_clang(required=False): + print("skip because clang is not available") + return + temp = utils.tempdir() + ll_path = temp.relpath("temp.ll") + ll_code = clang.create_llvm(cc_code, output=ll_path) + s = te.create_schedule(B.op) + if use_file: + s[B].pragma(s[B].op.axis[0], "import_llvm", ll_path) + else: + s[B].pragma(s[B].op.axis[0], "import_llvm", ll_code) + # BUILD and invoke the kernel. + f = tvm.build(s, [A, B], "llvm") + dev = tvm.cpu(0) + # launch the kernel. + a = tvm.nd.array(np.random.uniform(size=n).astype(A.dtype), dev) + b = tvm.nd.array(np.random.uniform(size=n).astype(B.dtype), dev) + f(a, b) + tvm.testing.assert_allclose(b.numpy(), a.numpy() + 1.0) + + check_llvm(use_file=True) + check_llvm(use_file=False) + + if __name__ == "__main__": sys.exit(pytest.main([__file__] + sys.argv[1:]))