diff --git a/python/tvm/contrib/utils.py b/python/tvm/contrib/utils.py index 4c5cb848feba..2c2baa849b40 100644 --- a/python/tvm/contrib/utils.py +++ b/python/tvm/contrib/utils.py @@ -136,6 +136,12 @@ def __truediv__(self, other): return self.path / other + def __enter__(self): + return self + + def __exit__(self, ptype, value, trace): + self.remove() + def __del__(self): temp_dirs = getattr(self, "TEMPDIRS", None) if temp_dirs is None: diff --git a/tests/python/codegen/test_target_codegen_llvm.py b/tests/python/codegen/test_target_codegen_llvm.py index b6d501b86baf..2105a2a2c31b 100644 --- a/tests/python/codegen/test_target_codegen_llvm.py +++ b/tests/python/codegen/test_target_codegen_llvm.py @@ -805,7 +805,8 @@ def test_llvm_crt_static_lib(): target=tvm.target.Target("llvm"), ) module.get_source() - module.save("test.o") + with utils.tempdir() as temp: + module.save(temp.relpath("test.o")) @tvm.testing.requires_llvm diff --git a/tests/python/relax/test_transform_codegen_pass.py b/tests/python/relax/test_transform_codegen_pass.py index 7560246b8af4..b997eb9c6bc0 100644 --- a/tests/python/relax/test_transform_codegen_pass.py +++ b/tests/python/relax/test_transform_codegen_pass.py @@ -24,6 +24,7 @@ import tvm import tvm.testing from tvm import relax, tir +from tvm.contrib import utils from tvm.relax.dpl import is_op, wildcard from tvm.relax.testing import transform from tvm.script import ir as I @@ -58,9 +59,9 @@ def check_executable(exec, dev, inputs, expected, entry_func_name): def check_roundtrip(exec0, dev, inputs, expected, entry_func_name="main"): - exec0.mod.export_library("exec.so") - exec1 = tvm.runtime.load_module("exec.so") - os.remove("exec.so") + with utils.tempdir() as temp: + exec0.mod.export_library(temp.relpath("exec.so")) + exec1 = tvm.runtime.load_module(temp.relpath("exec.so")) assert exec0.stats() == exec1["stats"]() assert exec0.as_text() == exec1["as_text"]() diff --git a/tests/python/relax/test_vm_build.py b/tests/python/relax/test_vm_build.py index 2078248880af..da8b905193fc 100644 --- a/tests/python/relax/test_vm_build.py +++ b/tests/python/relax/test_vm_build.py @@ -663,8 +663,10 @@ def te_func(A): target = tvm.target.Target("llvm", host="llvm") ex = relax.build(mod, target, exec_mode=exec_mode) - ex.export_library("exec.so") - vm = relax.VirtualMachine(tvm.runtime.load_module("exec.so"), tvm.cpu()) + with utils.tempdir() as temp: + ex.export_library(temp.relpath("exec.so")) + vm = relax.VirtualMachine(tvm.runtime.load_module(temp.relpath("exec.so")), tvm.cpu()) + inp = tvm.nd.array(np.random.rand(2).astype(np.float32)) inp2 = tvm.nd.array(np.random.rand(3).astype(np.float32)) @@ -956,7 +958,7 @@ def test_vm_mul(x: T.handle, y: T.handle, z: T.handle): # test returning a tuple @R.function def test_vm_tuple( - x: R.Tensor((), "int32") + x: R.Tensor((), "int32"), ) -> R.Tuple(R.Tensor((), "int32"), R.Tensor((), "int32")): return (x, x)