[Metal] Let compile callback declare payload format via (payload, fmt)

src/backend/metal/codegen/codegen_metal.cc

@@ -23,6 +23,7 @@
 #include "codegen_metal.h"
 
 #include <tvm/ffi/cast.h>
+#include <tvm/ffi/container/array.h>
 #include <tvm/ffi/container/map.h>
 #include <tvm/ffi/reflection/registry.h>
 #include <tvm/runtime/logging.h>
@@ -471,7 +472,9 @@ ffi::Module BuildMetal(IRModule mod, Target target) {
   // Map<String, Bytes> across all multi-shader backends.
   ffi::Map<ffi::String, ffi::Bytes> 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<PrimFuncNode>()) << "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<std::string>();
+      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<ffi::Array<ffi::Any>>()) {
+        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<std::string>();
+        kernel_fmt = (*ret_tuple)[1].cast<std::string>();
+        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<std::string>();
+        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)));
   }

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):