This is an automated email from the ASF dual-hosted git repository. tqchen pushed a commit to branch tirx/standalone-package-cleanup in repository https://gitbox.apache.org/repos/asf/tvm.git
commit 1d8feb685786d01f1114c2b9825c23f439fb6637 Author: Tianqi Chen <[email protected]> AuthorDate: Thu Oct 15 05:25:18 2026 +0000 [FIX][Metal] Restore the compiler callback after its owning test --- .../python/backend/metal/codegen/test_metal.py | 25 ++++++++++++++++------ 1 file changed, 18 insertions(+), 7 deletions(-) diff --git a/tirx/tests/python/backend/metal/codegen/test_metal.py b/tirx/tests/python/backend/metal/codegen/test_metal.py index 1fb69ddc2a..d582ef1cc0 100644 --- a/tirx/tests/python/backend/metal/codegen/test_metal.py +++ b/tirx/tests/python/backend/metal/codegen/test_metal.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import pytest +import tvm_ffi import tvm_tirx import tvm_tirx.testing @@ -202,16 +203,26 @@ def test_func_with_trailing_pod_params(): with T.launch_thread("threadIdx.x", 16) as i: B[i] = A[i] + x - @tvm_tirx.register_global_func("tvm_callback_metal_compile") + previous_callback = tvm_tirx.get_global_func("tvm_callback_metal_compile", allow_missing=True) + + @tvm_tirx.register_global_func("tvm_callback_metal_compile", override=True) def compile_metal(src, target): return xcode.compile_metal(src) - mod = tvm_tirx.IRModule({"main": func}) - - f = tvm_tirx.tirx.compile(mod, target="metal").mod - src: str = f.imports[0].inspect_source() - occurrences = src.count("struct func_kernel_args_t") - assert occurrences == 1, occurrences + try: + mod = tvm_tirx.IRModule({"main": func}) + + f = tvm_tirx.tirx.compile(mod, target="metal").mod + src: str = f.imports[0].inspect_source() + occurrences = src.count("struct func_kernel_args_t") + assert occurrences == 1, occurrences + finally: + if previous_callback is None: + tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile") + else: + tvm_tirx.register_global_func( + "tvm_callback_metal_compile", previous_callback, override=True + ) def test_codegen_simdgroup_buffer_data():
