diff --git a/coremltools/converters/mil/mil/passes/defs/symbol_transform.py b/coremltools/converters/mil/mil/passes/defs/symbol_transform.py index c81c76f79..a4f5e8314 100644 --- a/coremltools/converters/mil/mil/passes/defs/symbol_transform.py +++ b/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 diff --git a/coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py b/coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py index df0b73382..a512a0900 100644 --- a/coremltools/converters/mil/mil/passes/tests/test_symbol_transform.py +++ b/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):