Add sym_float and trunc support for dynamic float-scale interpolate

coremltools/converters/mil/frontend/torch/ops.py

@@ -2463,6 +2463,28 @@ def size(context, node):
     context.add(size_node, node.name)
 
 
+@register_torch_op
+def sym_float(context, node):
+    """
+    torch.export emits ``sym_float`` nodes when a symbolic integer dimension
+    (from ``aten.sym_size``) is converted to float, e.g. in the output-size
+    computation of ``F.interpolate(..., scale_factor=..., recompute_scale_factor=True)``
+    over a dynamic input shape:
+
+        %h  = aten.sym_size.int(x, 2)   -> int32
+        %hf = sym_float(%h)             -> fp32
+
+    The conversion is an int32 -> fp32 cast; the value is a runtime dimension,
+    so it must remain dynamic.
+    """
+    inputs = _get_inputs(context, node, expected=1)
+    x = inputs[0]
+    if types.is_float(x.dtype):
+        context.add(x, node.name)
+    else:
+        context.add(mb.cast(x=x, dtype="fp32", name=node.name))
+
+
 @register_torch_op
 def _shape_as_tensor(context, node):
     inputs = _get_inputs(context, node, expected=1)
@@ -4697,9 +4719,12 @@ def _translate_torch_args(x, output_size, scales) -> Var:
         assert (
             isinstance(output_size, list) and len(output_size) == 1
         ), "for dynamic shape torch should give [output_size]"
+        output_height = output_size[0]
+        if output_height.dtype != types.int32:
+            output_height = mb.cast(x=output_height, dtype="int32")
         x = mb.torch_upsample_nearest_neighbor(
             x=x,
-            output_height=output_size[0],
+            output_height=output_height,
             output_width=1,
         )
     x = mb.squeeze(x=x, axes=[3], name=node.name)
@@ -4787,10 +4812,16 @@ def _translate_torch_args(x, output_size, scales_h, scales_w) -> Var:
         # the input shape is dynamic and recompute_scale_factor = True
         # need to trace the graph to find the scale factor
         # we define a torch front end op mb.torch_upsample_nearest_neighbor to resolve the const scaling factor
+        output_height = output_size[0]
+        output_width = output_size[1]
+        if output_height.dtype != types.int32:
+            output_height = mb.cast(x=output_height, dtype="int32")
+        if output_width.dtype != types.int32:
+            output_width = mb.cast(x=output_width, dtype="int32")
         upsample_nearest2d = mb.torch_upsample_nearest_neighbor(
             x=x,
-            output_height=output_size[0],
-            output_width=output_size[1],
+            output_height=output_height,
+            output_width=output_width,
             name=node.name,
         )
     context.add(upsample_nearest2d)
@@ -7866,6 +7897,33 @@ def floor(context, node):
     context.add(mb.floor(x=inputs[0], name=node.name))
 
 
+@register_torch_op
+def trunc(context, node):
+    """
+    torch.trunc rounds toward zero: trunc(x) = sign(x) * floor(|x|).
+    MIL has no native trunc op, so it is decomposed into sign/mul/abs/floor
+    (the same building blocks used by the ``frac`` lowering). On integer
+    tensors trunc is the identity.
+
+    torch.export also emits a ``trunc`` node when computing the output size of
+    ``F.interpolate(..., scale_factor=..., recompute_scale_factor=True)`` with
+    a float scale factor and a dynamic input shape; the SSA pass
+    ``torch_upsample_to_core_upsample`` traces this decomposition back to the
+    constant scale factor.
+    """
+    inputs = _get_inputs(context, node, expected=1)
+    x = inputs[0]
+    if types.is_int(x.dtype):
+        context.add(x, node.name)
+        return
+    floor_abs = mb.floor(
+        x=mb.abs(x=x, name=node.name + "_abs"), name=node.name + "_floor"
+    )
+    context.add(
+        mb.mul(x=floor_abs, y=mb.sign(x=x, name=node.name + "_sign"), name=node.name)
+    )
+
+
 @register_torch_op
 def reciprocal(context, node):
     inputs = _get_inputs(context, node, expected=1)

coremltools/converters/mil/frontend/torch/ssa_passes/torch_upsample_to_core_upsample.py

@@ -142,6 +142,42 @@ def _try_get_upsample_factor_pattern_1(output_size):
         return np.float32(op.y.val)
 
 
+def _try_get_upsample_factor_pattern_3(output_size):
+    """
+    Handles the torch.export decomposition for a FLOAT scale factor with a
+    dynamic input shape and ``recompute_scale_factor=True``:
+
+        %h    = aten.sym_size.int(x, 2)     -> gather(shape(x), 2)        (int32)
+        %hf   = sym_float(%h)               -> cast(fp32)
+        %hm   = mul(%hf, scale_factor)      -> mul  (y is const scale)
+        %ht   = trunc(%hm)                  -> mul(sign, floor(abs))      (fp32)
+        %out  = cast(%ht, int32)            -> cast(int32)   (output_size)
+
+    ``trunc`` is lowered as ``sign(x) * floor(abs(x))`` (see the torch
+    frontend ``trunc`` op), so we trace: cast(int32) -> mul -> floor -> abs
+    -> mul, and return the constant scale factor.
+    """
+    op = output_size
+    if op.op_type != "cast" or op.dtype.val != "int32":
+        return None
+
+    # trunc(x) = sign(x) * floor(abs(x)); locate the floor(abs(x)) input.
+    op = op.x.op
+    if op.op_type != "mul":
+        return None
+    floor_op = op.x.op if op.x.op.op_type == "floor" else op.y.op
+    if floor_op.op_type != "floor":
+        return None
+    abs_op = floor_op.x.op
+    if abs_op.op_type != "abs":
+        return None
+    mul_op = abs_op.x.op
+    if mul_op.op_type != "mul":
+        return None
+    assert mul_op.y.val is not None, "scale factor should be const"
+    return np.float32(mul_op.y.val)
+
+
 def _try_replace_with_core_upsample(op):
     """
     Inputs:
@@ -164,10 +200,19 @@ def _try_replace_with_core_upsample(op):
         scales_h = _try_get_upsample_factor_pattern_1(op.output_height.op)
         scales_w = _try_get_upsample_factor_pattern_1(op.output_width.op)
 
-        if scales_h is None or scales_w is None:
+        # Only fill in the scale factors that pattern 1 could not resolve, so a
+        # previously resolved value (e.g. the constant dummy width of a 1d
+        # upsample) is not overwritten.
+        if scales_h is None:
             scales_h = _try_get_upsample_factor_pattern_2(op.output_height.op, 2, op.x)
+        if scales_w is None:
             scales_w = _try_get_upsample_factor_pattern_2(op.output_width.op, 3, op.x)
 
+        if scales_h is None:
+            scales_h = _try_get_upsample_factor_pattern_3(op.output_height.op)
+        if scales_w is None:
+            scales_w = _try_get_upsample_factor_pattern_3(op.output_width.op)
+
         if scales_h is None or scales_w is None:
             return False
 

coremltools/converters/mil/frontend/torch/test/test_torch_ops.py

@@ -3310,6 +3310,100 @@ def forward(self, args):
                 if layer.WhichOneof("layer") == "upsample":
                     assert len(layer.upsample.fractionalScalingFactor) == 0
 
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend",
+        itertools.product(compute_units, backends, [TorchFrontend.TORCHEXPORT]),
+    )
+    def test_interpolate_nearest2d_with_float_scale_dynamic(
+        self, compute_unit, backend, frontend
+    ):
+        input_shape = (1, 3, 10, 10)
+
+        class Model(nn.Module):
+            def __init__(self, scale_factor):
+                super().__init__()
+                self.scale_factor = scale_factor
+
+            def forward(self, args):
+                return nn.functional.interpolate(
+                    args,
+                    scale_factor=self.scale_factor,
+                    mode="nearest",
+                    recompute_scale_factor=True,
+                )
+
+        model = Model((2.5, 1.5))
+
+        upper_bound_coreml = 20 if backend[0] == "mlprogram" else -1
+        upper_bound_torch = None if upper_bound_coreml == -1 else upper_bound_coreml
+        height = RangeDim(upper_bound=upper_bound_coreml)
+        width = RangeDim(upper_bound=upper_bound_coreml)
+        converter_input_type = [TensorType(shape=(1, 3, height, width), dtype=np.float32)]
+        torch_export_dynamic_shapes = {
+            "args": {
+                2: torch.export.Dim(name="height", max=upper_bound_torch),
+                3: torch.export.Dim(name="width", max=upper_bound_torch),
+            }
+        }
+
+        self.run_compare_torch(
+            input_shape,
+            model,
+            frontend=frontend,
+            backend=backend,
+            compute_unit=compute_unit,
+            converter_input_type=converter_input_type,
+            torch_export_dynamic_shapes=torch_export_dynamic_shapes,
+        )
+
+    @pytest.mark.parametrize(
+        "compute_unit, backend, frontend",
+        itertools.product(compute_units, backends, [TorchFrontend.TORCHEXPORT]),
+    )
+    def test_interpolate_bilinear2d_with_float_scale_dynamic(
+        self, compute_unit, backend, frontend
+    ):
+        input_shape = (1, 3, 9, 22)
+
+        class Model(nn.Module):
+            def __init__(self, scale_factor, align_corners):
+                super().__init__()
+                self.scale_factor = scale_factor
+                self.align_corners = align_corners
+
+            def forward(self, args):
+                return nn.functional.interpolate(
+                    args,
+                    scale_factor=self.scale_factor,
+                    mode="bilinear",
+                    align_corners=self.align_corners,
+                    recompute_scale_factor=True,
+                )
+
+        model = Model((2.5, 3.5), False)
+
+        upper_bound_coreml = 30 if backend[0] == "mlprogram" else -1
+        upper_bound_torch = None if upper_bound_coreml == -1 else upper_bound_coreml
+        height = RangeDim(upper_bound=upper_bound_coreml)
+        width = RangeDim(upper_bound=upper_bound_coreml)
+        converter_input_type = [TensorType(shape=(1, 3, height, width), dtype=np.float32)]
+        torch_export_dynamic_shapes = {
+            "args": {
+                2: torch.export.Dim(name="height", max=upper_bound_torch),
+                3: torch.export.Dim(name="width", max=upper_bound_torch),
+            }
+        }
+
+        self.run_compare_torch(
+            input_shape,
+            model,
+            frontend=frontend,
+            backend=backend,
+            compute_unit=compute_unit,
+            converter_input_type=converter_input_type,
+            torch_export_dynamic_shapes=torch_export_dynamic_shapes,
+        )
+
 
 class TestEmpty(TorchBaseTest):
     @pytest.mark.parametrize(