Keep floor_div(x, 1) for float inputs in noop_elimination

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)

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):