This is an automated email from the ASF dual-hosted git repository.
tlopex pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/tvm.git
The following commit(s) were added to refs/heads/main by this push:
new 67bd1ea1a1 [Metal] Let compile callback declare payload format via
(payload, fmt) (#19924)
67bd1ea1a1 is described below
commit 67bd1ea1a16390db0b6760fea57f55a893fd9375
Author: Egor Churaev <[email protected]>
AuthorDate: Fri Jul 10 09:24:17 2026 +0300
[Metal] Let compile callback declare payload format via (payload, fmt)
(#19924)
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 +++++++-
tests/python/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 af914f6821..1469423f09 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 <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>
@@ -465,7 +466,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";
@@ -489,7 +492,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)));
}
diff --git a/tests/python/codegen/test_target_codegen_metal.py
b/tests/python/codegen/test_target_codegen_metal.py
index db4b0bd095..794a861c3a 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
@@ -230,6 +231,97 @@ def test_func_with_trailing_pod_params():
assert occurrences == 1, occurrences
[email protected]
[email protected](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)
+
+
[email protected]
[email protected](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):
