diff --git a/coremltools/converters/mil/mil/passes/defs/cleanup/noop_elimination.py b/coremltools/converters/mil/mil/passes/defs/cleanup/noop_elimination.py index 6accb300e..b18b0f1a4 100644 --- a/coremltools/converters/mil/mil/passes/defs/cleanup/noop_elimination.py +++ b/coremltools/converters/mil/mil/passes/defs/cleanup/noop_elimination.py @@ -6,6 +6,7 @@ import numpy as np +from coremltools.converters.mil.mil import types from coremltools.converters.mil.mil.passes.graph_pass import AbstractGraphPass from coremltools.converters.mil.mil.passes.helper import block_context_manager from coremltools.converters.mil.mil.passes.pass_registry import register_pass @@ -111,7 +112,14 @@ def remove_elementwise(op): return _remove_elementwise_binary(op, 0, 0) elif op.op_type in {"mul"}: return _remove_elementwise_binary(op, 1, 1) - elif op.op_type in {"floor_div", "pow", "real_div"}: + elif op.op_type in {"pow", "real_div"}: + return _remove_elementwise_binary(op, None, 1) + elif op.op_type in {"floor_div"}: + # floor_div(x, 1) == floor(x), which only equals x when x is already + # integral; removing it for a float input drops the floor and changes + # the result, so restrict this no-op to integer inputs. + if not types.is_int(op.x.dtype): + return False return _remove_elementwise_binary(op, None, 1) elif op.op_type in {"sub"}: return _remove_elementwise_binary(op, None, 0) diff --git a/coremltools/converters/mil/mil/passes/tests/test_cleanup_passes.py b/coremltools/converters/mil/mil/passes/tests/test_cleanup_passes.py index 074be3939..bee3fc754 100644 --- a/coremltools/converters/mil/mil/passes/tests/test_cleanup_passes.py +++ b/coremltools/converters/mil/mil/passes/tests/test_cleanup_passes.py @@ -15,11 +15,14 @@ from coremltools.converters.mil import mil from coremltools.converters.mil.mil import Builder as mb from coremltools.converters.mil.mil import Function, Symbol, get_new_symbol, types +from coremltools.converters.mil.mil.ops.tests.iOS14 import backends +from coremltools.converters.mil.mil.ops.tests.testing_utils import run_compare_builder from coremltools.converters.mil.mil.passes.defs.cleanup import topological_reorder from coremltools.converters.mil.mil.passes.defs.cleanup.remove_redundant_ops import ( remove_redundant_ops, ) from coremltools.converters.mil.mil.passes.pass_registry import PASS_REGISTRY +from coremltools.converters.mil.testing_reqs import compute_units from coremltools.converters.mil.testing_utils import ( apply_pass_and_basic_check, assert_model_is_valid, @@ -1355,9 +1358,11 @@ def prog(x): elif op_type in {"real_div"}: if pos == "y" and (val == 1.0 or val == [1.0, 1.0, 1.0, 1.0]): new_program = ["relu"] - elif op_type in {"pow", "floor_div"}: + elif op_type in {"pow"}: if pos == "y" and (val == 1.0 or val == [1.0, 1.0, 1.0, 1.0]): new_program = ["relu"] + # floor_div(x, 1) == floor(x) is a no-op only for integer x; the fp32 input here + # keeps it (integer removal is covered by test_floor_div_elimination_by_dtype). elif op_type in {"sub"}: if pos == "y" and (val == 0.0 or val == [0.0, 0.0, 0.0, 0.0]): new_program = ["relu"] @@ -1370,6 +1375,57 @@ def prog(x): expected_output_shapes={block.outputs[0].name: (2, 4)}, ) + @pytest.mark.parametrize( + "dtype, one, expected", + [ + (types.int32, 1, ["add"]), + (types.fp32, 1.0, ["floor_div", "add"]), + ], + ) + def test_floor_div_elimination_by_dtype(self, dtype, one, expected): + # floor_div(x, 1) == floor(x) is a no-op only when x is integral. It must be + # removed for integer inputs but kept for float inputs -- dropping it there + # would remove the floor and silently change the result. (add x, 1 is a + # non-eliminable consumer that accepts both int and float and keeps floor_div + # interior so the pass is allowed to remove it.) + @mb.program(input_specs=[mb.TensorSpec(shape=(2, 4), dtype=dtype)]) + def prog(x): + r1 = mb.floor_div(x=x, y=one) + return mb.add(x=r1, y=one) + + prev_prog, _, _ = apply_pass_and_basic_check(prog, "common::noop_elimination") + assert get_op_types_in_program(prev_prog) == ["floor_div", "add"] + assert get_op_types_in_program(prog) == expected + + @pytest.mark.parametrize( + "compute_unit, backend", + itertools.product(compute_units, backends), + ) + def test_floor_div_float_prediction_is_floored(self, compute_unit, backend): + # End-to-end guard that the fp32 floor_div is not silently dropped by + # noop_elimination. floor_div(x, 1) == floor(x) for a float input, so removing + # it as a "no-op" changes the numeric output. Here add(_, 1) is a non-eliminable + # consumer that keeps the floor_div interior, so the pass is free to remove it: + # with the fix the floor survives and the model returns floor(x) + 1; without it + # the floor_div is eliminated and the model would return x + 1 instead. + x = np.array([[2.7, -1.3, 0.5, 4.9], [1.1, 3.6, -2.8, 0.0]], dtype=np.float32) + expected_outputs = np.floor(x) + 1.0 + + def build(x): + return mb.add(x=mb.floor_div(x=x, y=1.0), y=1.0) + + input_placeholders = {"x": mb.placeholder(shape=x.shape)} + input_values = {"x": x} + run_compare_builder( + build, + input_placeholders, + input_values, + expected_output_types=(2, 4, types.fp32), + expected_outputs=expected_outputs, + compute_unit=compute_unit, + backend=backend, + ) + def test_elementwise_broadcast(self): @mb.program(input_specs=[mb.TensorSpec(shape=[4])]) def prog(x):