Fix weight_id collision when materializing multiple source functions

coremltools/converters/mil/mil/passes/defs/symbol_transform.py

@@ -307,9 +307,18 @@ def apply(self, prog: Program) -> None:
                                     ), "Only const may be absent from context"
                                     # The consts across function should share the same file value while lowering into milproto,
                                     # so we assign the weight_id, if not presented.
+                                    # A const name is only unique within a function, so the
+                                    # source function name has to be part of the weight_id:
+                                    # 2 source functions may own same-named consts that hold
+                                    # different values, and sharing a weight file value among
+                                    # those would corrupt the model. Note that this only affects
+                                    # the weight_id we invent here, so consts that got a weight_id
+                                    # elsewhere (e.g. by ``const_deduplication``) keep sharing
+                                    # their file value across functions.
                                     if source_input_var.op.weight_id is None:
                                         source_input_var.op.weight_id = (
-                                            f"const_{source_input_var.name}_weight_id"
+                                            f"const_{self.source_function_name}_"
+                                            f"{source_input_var.name}_weight_id"
                                         )
                                     context[source_input_var.name] = self._copy_construct_const_var(
                                         source_input_var

coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py

@@ -9,8 +9,11 @@
 import numpy as np
 import pytest
 
+import coremltools as ct
+from coremltools.converters.mil.converter import mil_convert
+from coremltools.converters.mil.frontend.milproto.load import load as _milproto_to_pymil
 from coremltools.converters.mil.mil import Builder as mb
-from coremltools.converters.mil.mil import get_new_symbol
+from coremltools.converters.mil.mil import get_new_symbol, types
 from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY
 from coremltools.converters.mil.testing_utils import (
     apply_pass_and_basic_check,
@@ -45,7 +48,91 @@ def prog(x):
         for block in prog.functions.values():
             const_ops = block.find_ops(op_type="const")
             assert len(const_ops) == 1
-            assert const_ops[0].weight_id == "const_y_weight_id"
+            assert const_ops[0].weight_id == "const_main_y_weight_id"
+
+    @staticmethod
+    def test_weight_id_no_collision_across_source_functions():
+        """
+        A const name is unique only within a function, so 2 source functions may each own
+        a same-named const that holds a different value. Since the auto assigned weight_id
+        determines which consts share a weight blob entry, deriving it from the const name
+        alone would make those 2 unrelated consts claim the same entry, and the resulting
+        model would be corrupted: one value gets written to the blob, while both functions
+        point to it, so the function whose shape disagrees with the blob content fails to load
+
+        Note that a const is only written to the weight blob if it is large enough
+        (small ones are stored immediately in the proto), so this test uses large consts
+        """
+        symbolic_shape = (get_new_symbol(),)
+        fixed_shape = (1,)
+        source_function_name_to_weight = {
+            "main": np.arange(4096, dtype=np.float32),
+            "func2": np.arange(8192, dtype=np.float32) + 0.5,
+        }
+
+        @mb.program(
+            input_specs=[mb.TensorSpec(shape=symbolic_shape, dtype=types.int32)],
+            opset_version=ct.target.iOS18,
+        )
+        def prog(x):
+            # gather from a runtime index, so the const cannot be folded away
+            weight = mb.const(val=source_function_name_to_weight["main"], name="shared_name")
+            return mb.gather(x=weight, indices=x, axis=0)
+
+        @mb.function(
+            input_specs=[mb.TensorSpec(shape=symbolic_shape, dtype=types.int32)],
+            opset_version=ct.target.iOS18,
+        )
+        def func2(x):
+            weight = mb.const(val=source_function_name_to_weight["func2"], name="shared_name")
+            return mb.gather(x=weight, indices=x, axis=0)
+
+        prog.add_function("func2", func2)
+
+        for source_function_name in source_function_name_to_weight:
+            graph_pass = PASS_REGISTRY["common::materialize_symbolic_shape_program"]
+            graph_pass.source_function_name = source_function_name
+            graph_pass.function_name_to_materialization_map = {
+                f"{source_function_name}_materialized": {"x": fixed_shape}
+            }
+            graph_pass.apply(prog)
+
+        # the same-named consts of the 2 source functions must not claim the same weight_id
+        weight_id_to_weight = {}
+        for function in prog.functions.values():
+            for const_op in function.find_ops(prefix="shared_name", op_type="const"):
+                assert const_op.weight_id is not None
+                weight = const_op.outputs[0].val
+                # a materialized const shares the weight blob entry of the source const
+                # it is cloned from, but consts of different source functions must not
+                np.testing.assert_array_equal(
+                    weight_id_to_weight.setdefault(const_op.weight_id, weight), weight
+                )
+        assert len(weight_id_to_weight) == len(source_function_name_to_weight)
+
+        # end-to-end: every materialized function has to read back its own weight
+        prog.export_as_multifunction = True
+        prog.skip_all_passes = True
+        mlmodel = mil_convert(
+            prog,
+            convert_from="milinternal",
+            convert_to="mlprogram",
+            specification_version=ct.target.iOS18,
+            compute_units=ct.ComputeUnit.CPU_ONLY,
+            skip_model_load=True,
+        )
+        spec = mlmodel.get_spec()
+        reloaded_prog = _milproto_to_pymil(
+            model_spec=spec,
+            specification_version=spec.specificationVersion,
+            file_weights_dir=mlmodel.weights_dir,
+        )
+        for source_function_name, weight in source_function_name_to_weight.items():
+            const_ops = reloaded_prog.functions[f"{source_function_name}_materialized"].find_ops(
+                prefix="shared_name", op_type="const"
+            )
+            assert len(const_ops) == 1
+            np.testing.assert_array_equal(const_ops[0].outputs[0].val, weight)
 
     @pytest.mark.parametrize("override_main_function", (True, False))
     def test_simple(self, override_main_function):