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 1f7c99d3f4 [Relax][CoreML] Fix CoreML partition pass (#19711)
1f7c99d3f4 is described below
commit 1f7c99d3f45b33a99c06061702da0da911fc60ec
Author: Shushi Hong <[email protected]>
AuthorDate: Tue Jun 9 22:38:28 2026 -0400
[Relax][CoreML] Fix CoreML partition pass (#19711)
This fixes CoreML Relax partitioning after `FoldDataflowBlockOutput` was
folded into `CanonicalizeBindings`.
`partition_for_coreml` still called
`relax.transform.FoldDataflowBlockOutput()`, but that pass is no longer
exported. This caused CoreML partition tests to fail with:
```text
AttributeError: module 'tvm.relax.transform' has no attribute
'FoldDataflowBlockOutput'
```
The old pass behavior is now covered by `CanonicalizeBindings`, so this
updates CoreML partitioning to call `CanonicalizeBindings()` instead.
The existing CoreML test file now also includes a partition-only
regression test. The end-to-end CoreML tests remain guarded by the
CoreML runtime requirements, while the partition test can run without
coremltools, Xcode, or a CoreML runtime.
---
python/tvm/relax/backend/metal/coreml.py | 2 +-
tests/python/relax/test_codegen_coreml.py | 44 ++++++++++++++++++++++++++++---
2 files changed, 41 insertions(+), 5 deletions(-)
diff --git a/python/tvm/relax/backend/metal/coreml.py
b/python/tvm/relax/backend/metal/coreml.py
index 7dd8ea3958..e719d2f1c7 100644
--- a/python/tvm/relax/backend/metal/coreml.py
+++ b/python/tvm/relax/backend/metal/coreml.py
@@ -169,7 +169,7 @@ def partition_for_coreml(mod):
"""
patterns = get_patterns_with_prefix("coreml")
- mod = transform.FoldDataflowBlockOutput()(mod)
+ mod = transform.CanonicalizeBindings()(mod)
mod = transform.FuseOpsByPattern(patterns, bind_constants=True,
annotate_codegen=False)(mod)
mod = transform.MergeCompositeFunctions()(mod)
return mod
diff --git a/tests/python/relax/test_codegen_coreml.py
b/tests/python/relax/test_codegen_coreml.py
index e9b9bcb09c..2068044bbe 100644
--- a/tests/python/relax/test_codegen_coreml.py
+++ b/tests/python/relax/test_codegen_coreml.py
@@ -14,6 +14,8 @@
# KIND, either express or implied. See the License for the
# specific language governing permissions and limitations
# under the License.
+import importlib.util
+
import numpy as np
import pytest
@@ -21,7 +23,6 @@ import tvm
import tvm.testing
from tvm import relax
-requires_coremltools = tvm.testing.requires_package("coremltools")
target, dev = "llvm", tvm.cpu()
@@ -36,12 +37,35 @@ def _has_xcode():
return False
-pytestmark = pytest.mark.skipif(
- not (requires_coremltools and _has_xcode()),
+requires_coreml_runtime = pytest.mark.skipif(
+ not (importlib.util.find_spec("coremltools") and _has_xcode()),
reason="coreml is not enabled.",
)
+def test_partition_for_coreml_uses_current_relax_passes():
+ from tvm.relax.backend.metal.coreml import partition_for_coreml
+
+ x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
+ y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
+ bb = relax.BlockBuilder()
+ with bb.function("main", [x, y]):
+ with bb.dataflow():
+ lv0 = bb.emit(relax.op.add(x, y))
+ gv = bb.emit_output(lv0)
+ bb.emit_func_output(gv)
+
+ partitioned = partition_for_coreml(bb.get())
+
+ assert relax.analysis.well_formed(partitioned)
+ assert any(
+ getattr(func, "attrs", None) is not None
+ and "Codegen" in func.attrs
+ and str(func.attrs["Codegen"]) == "coreml"
+ for func in partitioned.functions.values()
+ )
+
+
def verify(mod, inputs):
from tvm.relax.backend.metal.coreml import partition_for_coreml
@@ -65,6 +89,7 @@ def verify(mod, inputs):
tvm.testing.assert_allclose(out1.numpy(), out2.numpy(), rtol=1e-3,
atol=1e-3)
+@requires_coreml_runtime
def test_add():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
@@ -80,6 +105,7 @@ def test_add():
verify(mod, [x_data, y_data])
+@requires_coreml_runtime
def test_add_const():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
y = relax.const(np.ones([10, 10]), "float32")
@@ -94,6 +120,7 @@ def test_add_const():
verify(mod, [x_data])
+@requires_coreml_runtime
def test_multiply():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
@@ -110,6 +137,7 @@ def test_multiply():
verify(mod, [x_data, y_data])
+@requires_coreml_runtime
def test_matmul():
x = relax.Var("x", relax.TensorStructInfo([8, 10], "float32"))
y = relax.Constant(tvm.runtime.tensor(np.random.rand(10,
8).astype("float32"), dev))
@@ -139,6 +167,7 @@ def test_matmul():
verify(mod, [x_data, y_data])
+@requires_coreml_runtime
def test_clip():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
bb = relax.BlockBuilder()
@@ -168,6 +197,7 @@ def test_clip():
verify(mod, [x_data])
+@requires_coreml_runtime
def test_expand_dims():
def get_mod(axis):
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
@@ -184,6 +214,7 @@ def test_expand_dims():
verify(get_mod(axis=1), [x_data])
+@requires_coreml_runtime
def test_relu():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
bb = relax.BlockBuilder()
@@ -198,6 +229,7 @@ def test_relu():
verify(mod, [x_data])
+@requires_coreml_runtime
def test_batch_flatten():
x = relax.Var("x", relax.TensorStructInfo([10, 10, 10], "float32"))
bb = relax.BlockBuilder()
@@ -212,7 +244,7 @@ def test_batch_flatten():
verify(mod, [x_data])
-@requires_coremltools
+@requires_coreml_runtime
def test_softmax():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
bb = relax.BlockBuilder()
@@ -227,6 +259,7 @@ def test_softmax():
verify(mod, [x_data])
+@requires_coreml_runtime
def test_conv2d():
x = relax.Var("x", relax.TensorStructInfo([1, 3, 224, 224], "float32"))
w = relax.const(np.zeros((16, 3, 3, 3), dtype="float32"))
@@ -241,6 +274,7 @@ def test_conv2d():
verify(mod, [x_data])
+@requires_coreml_runtime
def test_global_avg_pool2d():
x = relax.Var("x", relax.TensorStructInfo([1, 1, 10, 10], "float32"))
bb = relax.BlockBuilder()
@@ -254,6 +288,7 @@ def test_global_avg_pool2d():
verify(mod, [x_data])
+@requires_coreml_runtime
def test_subgraph1():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
@@ -270,6 +305,7 @@ def test_subgraph1():
verify(mod, [x_data, y_data])
+@requires_coreml_runtime
def test_subgraph2():
x = relax.Var("x", relax.TensorStructInfo([10, 10], "float32"))
y = relax.Var("y", relax.TensorStructInfo([10, 10], "float32"))
