From 95c1004ba2e8dc5df7048cea37ff6f31a58f74c4 Mon Sep 17 00:00:00 2001 From: Egor Churaev Date: Wed, 1 Jul 2026 23:41:14 +0300 Subject: [PATCH] [Metal] Let compile callback declare payload format via (payload, fmt) A tvm_callback_metal_compile used purely for debugging/inspection may return the MSL source unchanged. Previously, merely registering the callback forced the module format to "metallib", so the runtime tried to load the text source as a binary metallib and failed with "Invalid library file" (issue #18798). The callback may now return a (payload, format) pair to declare the payload format. A bare str/bytes return keeps the legacy metallib behavior. All kernels of a module share a single declared format, so a callback that mixes formats across kernels (including a legacy metallib return alongside a (payload, "metal") return) is rejected at codegen time instead of producing a module that fails to load. --- src/backend/metal/codegen/codegen_metal.cc | 30 +++++- .../codegen/test_target_codegen_metal.py | 92 +++++++++++++++++++ 2 files changed, 121 insertions(+), 1 deletion(-) diff --git a/src/backend/metal/codegen/codegen_metal.cc b/src/backend/metal/codegen/codegen_metal.cc index 83ee2722c248..2317ef88390f 100644 --- a/src/backend/metal/codegen/codegen_metal.cc +++ b/src/backend/metal/codegen/codegen_metal.cc @@ -23,6 +23,7 @@ #include "codegen_metal.h" #include +#include #include #include #include @@ -471,7 +472,9 @@ ffi::Module BuildMetal(IRModule mod, Target target) { // Map across all multi-shader backends. ffi::Map smap; const auto fmetal_compile = tvm::ffi::Function::GetGlobal("tvm_callback_metal_compile"); + // Default payload format. A callback may override it per the contract below. std::string fmt = fmetal_compile ? "metallib" : "metal"; + bool fmt_locked = false; for (auto kv : mod->functions) { TVM_FFI_ICHECK(kv.second->IsInstance()) << "CodeGenMetal: Can only take PrimFunc"; @@ -495,7 +498,32 @@ ffi::Module BuildMetal(IRModule mod, Target target) { std::string fsource = cg.Finish(); source_maker << fsource << "\n"; if (fmetal_compile) { - fsource = (*fmetal_compile)(fsource, target).cast(); + ffi::Any ret = (*fmetal_compile)(fsource, target); + // Backward-compatible contract for tvm_callback_metal_compile: + // * returning a str/bytes -> treated as a compiled metallib payload + // * returning (payload, fmt) -> the callback declares the payload format + std::string kernel_fmt; + if (auto ret_tuple = ret.try_cast>()) { + TVM_FFI_ICHECK_EQ(ret_tuple->size(), 2) + << "tvm_callback_metal_compile must return either a payload or a " + "(payload, format) pair, but got a tuple of size " + << ret_tuple->size(); + fsource = (*ret_tuple)[0].cast(); + kernel_fmt = (*ret_tuple)[1].cast(); + TVM_FFI_ICHECK(kernel_fmt == "metal" || kernel_fmt == "metallib") + << "tvm_callback_metal_compile returned unsupported format \"" << kernel_fmt + << "\"; expected \"metal\" or \"metallib\""; + } else { + // Backward-compatible behavior + fsource = ret.cast(); + kernel_fmt = "metallib"; + } + // All kernels of a module share a single declared format + TVM_FFI_ICHECK(!fmt_locked || fmt == kernel_fmt) + << "tvm_callback_metal_compile returned inconsistent formats across kernels: \"" << fmt + << "\" vs \"" << kernel_fmt << "\""; + fmt = kernel_fmt; + fmt_locked = true; } smap.Set(func_name, ffi::Bytes(std::move(fsource))); } diff --git a/tests/python/codegen/test_target_codegen_metal.py b/tests/python/codegen/test_target_codegen_metal.py index 0237367a57dd..2d06ce77f7e5 100644 --- a/tests/python/codegen/test_target_codegen_metal.py +++ b/tests/python/codegen/test_target_codegen_metal.py @@ -16,6 +16,7 @@ # under the License. import numpy as np import pytest +import tvm_ffi import tvm import tvm.testing @@ -211,6 +212,97 @@ def compile_metal(src, target): assert occurrences == 1, occurrences +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_metal(), reason="need metal") +def test_metal_compile_callback_source_passthrough(): + n = 1024 + + @I.ir_module(s_tir=True) + class Module: + @T.prim_func(s_tir=True) + def main(A: T.Buffer((n,), "float32"), B: T.Buffer((n,), "float32")): + T.func_attr({"tirx.noalias": True}) + for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"): + for i_1 in T.thread_binding(32, thread="threadIdx.x"): + with T.sblock("B"): + v_i = T.axis.spatial(n, i_0 * 32 + i_1) + T.reads(A[v_i]) + T.writes(B[v_i]) + B[v_i] = A[v_i] + 1.0 + + seen = {} + + def inspect_callback(src, target): + # Pure inspection callback: capture the source, return it untouched and + # declare it is still textual MSL so it is compiled at load time. + seen["src"] = src + return (src, "metal") + + tvm.register_global_func("tvm_callback_metal_compile", inspect_callback, override=True) + try: + f = tvm.compile(Module, target="metal") + dev = tvm.metal() + a = np.random.rand(n).astype("float32") + a_nd = tvm.runtime.tensor(a, dev) + b_nd = tvm.runtime.empty((n,), "float32", dev) + f(a_nd, b_nd) + dev.sync() + finally: + tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile") + + assert "src" in seen and len(seen["src"]) > 0 + tvm.testing.assert_allclose(b_nd.numpy(), a + 1.0, atol=1e-5, rtol=1e-5) + + +@pytest.mark.gpu +@pytest.mark.skipif(not env.has_metal(), reason="need metal") +def test_metal_compile_callback_mixed_formats_rejected(): + n = 1024 + + @I.ir_module(s_tir=True) + class Module: + @T.prim_func(s_tir=True) + def main( + A: T.Buffer((n,), "float32"), + B: T.Buffer((n,), "float32"), + C: T.Buffer((n,), "float32"), + ): + T.func_attr({"tirx.noalias": True}) + # Two independent thread-bound regions -> two device kernels, so the + # compile callback is invoked twice within one module. + for i_0 in T.thread_binding(n // 32, thread="blockIdx.x"): + for i_1 in T.thread_binding(32, thread="threadIdx.x"): + with T.sblock("B"): + v_i = T.axis.spatial(n, i_0 * 32 + i_1) + T.reads(A[v_i]) + T.writes(B[v_i]) + B[v_i] = A[v_i] + 1.0 + for j_0 in T.thread_binding(n // 32, thread="blockIdx.x"): + for j_1 in T.thread_binding(32, thread="threadIdx.x"): + with T.sblock("C"): + v_j = T.axis.spatial(n, j_0 * 32 + j_1) + T.reads(A[v_j]) + T.writes(C[v_j]) + C[v_j] = A[v_j] + 2.0 + + calls = {"n": 0} + + def mixed_callback(src, target): + calls["n"] += 1 + if calls["n"] == 1: + # Treated as a compiled metallib payload. + return src + # Second kernel declares textual MSL, contradicting the metallib above. + return (src, "metal") + + tvm.register_global_func("tvm_callback_metal_compile", mixed_callback, override=True) + try: + with pytest.raises(Exception, match="inconsistent formats"): + tvm.compile(Module, target="metal") + finally: + tvm_ffi.registry.remove_global_func("tvm_callback_metal_compile") + + @pytest.mark.gpu @pytest.mark.skipif(not env.has_metal(), reason="need metal") def test_export_load_with_fallback(monkeypatch, tmp_path):