Arm backend: Support dyn. shapes in rewrite_slice (pytorch#23053)

Preserve symbolic and proxy slice bounds when normalizing start and end
indices.

Build dynamic slice sizes from shape operators and relax TOSA SLICE
validation for non-static start and size values.


cc @digantdesai @freddan80 @per @zingo @mansnils @Sebastian-Larsson
@robell @rascani

Signed-off-by: Oscar Andersson <oscar.andersson@arm.com>

Author: oscarandersson8218

backends/arm/_passes/rewrite_slice.py

@@ -2,6 +2,7 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+import operator
 from typing import Set, Type
 
 from executorch.backends.arm._passes import ArmOpTargetedPass
@@ -18,16 +19,29 @@ class RewriteSlicePass(ArmOpTargetedPass):
     _passes_required_after: Set[Type[ExportPass]] = set()
     target_ops = (exir_ops.edge.aten.slice_copy.Tensor,)
 
-    def _fixup_start(self, start, input_shape, dim) -> int:
-        """Convert negative and out-of-bounds start indices to valid positive
-        indices.
-        """
-        idx = start or 0
-        if idx < 0:
-            idx = start % input_shape[dim]
-        if idx >= input_shape[dim]:
-            idx = input_shape[dim] - 1
-        return idx
+    def _normalize_slice_bound(
+        self,
+        bound: int | None,
+        dim_size: int | SymInt | ProxyValue,
+        *,
+        is_start: bool,
+        meta,
+    ) -> int | SymInt | ProxyValue:
+        if isinstance(dim_size, int):
+            return (
+                slice(bound, None).indices(dim_size)[0]
+                if is_start
+                else slice(None, bound).indices(dim_size)[1]
+            )
+        if bound is None:
+            return 0 if is_start else dim_size
+        if bound < 0:
+            return super().call_sym(
+                max,
+                (super().call_sym(operator.add, (dim_size, bound), meta), 0),
+                meta,
+            )
+        return super().call_sym(min, (bound, dim_size), meta)
 
     def call_operator(self, op, args, kwargs, meta, updated=False) -> ProxyValue:
         if op not in self.target_ops:
@@ -38,14 +52,32 @@ def call_operator(self, op, args, kwargs, meta, updated=False) -> ProxyValue:
                 f"Only slice with 4 arguments and step of 1 is supported, got {len(args)} arguments and step {args[4]}"
             )
         input, dim, start, end = args[:4]
+        if isinstance(start, (SymInt, ProxyValue)) or isinstance(
+            end, (SymInt, ProxyValue)
+        ):
+            return super().call_operator(op, args, kwargs, meta, updated)
+
         input_shape = input.data.shape
-        start_index = self._fixup_start(start, input_shape, dim)
+        size_list = list(meta.data["val"].shape)
+        # TOSA SLICE rejects statically known non-positive sizes, including
+        # dimensions that are not being sliced.
+        if any(isinstance(size, int) and size <= 0 for size in size_list):
+            return super().call_operator(op, args, kwargs, meta, updated)
 
-        start_list = [0] * len(input_shape)
+        dim_size = input_shape[dim]
+        if isinstance(dim_size, SymInt):
+            dim_size = super().call_size_operator(input, dim, meta)
+        start_index = self._normalize_slice_bound(
+            start, dim_size, is_start=True, meta=meta
+        )
+        end_index = self._normalize_slice_bound(
+            end, dim_size, is_start=False, meta=meta
+        )
+
+        start_list: list[int | SymInt | ProxyValue] = [0] * len(input_shape)
         start_list[dim] = start_index
-        size_list = list(meta.data["val"].shape)
 
-        if any(isinstance(dim, SymInt) for dim in start_list):
+        if any(isinstance(dim, (SymInt, ProxyValue)) for dim in start_list):
             starts = start_list
         else:
             starts = super().call_shape_operator(  # type: ignore[assignment]
@@ -55,8 +87,15 @@ def call_operator(self, op, args, kwargs, meta, updated=False) -> ProxyValue:
                 meta,
                 True,
             )
-        if any(isinstance(dim, SymInt) for dim in size_list):
-            sizes = size_list
+        if any(isinstance(d, SymInt) for d in size_list):
+            # Express sizes as ProxyValues to allow for symbolic shapes
+            input_shape_proxy = super().call_size_operator_all(input, meta)
+            sizes = input_shape_proxy
+            if isinstance(size_list[dim], SymInt):
+                size = super().call_sym(operator.sub, (end_index, start_index), meta)
+                sizes[dim] = super().call_sym(max, (size, 0), meta)
+            else:
+                sizes[dim] = size_list[dim]
         else:
             sizes = super().call_shape_operator(  # type: ignore[assignment]
                 exir_ops.backend.tosa.CONST_SHAPE.default,

backends/arm/_passes/symbolic_to_tosa_shape_pass.py

@@ -28,6 +28,8 @@
     operator.mul: exir_ops.backend.tosa.MUL_SHAPE.default,
     operator.mod: exir_ops.backend.tosa.MOD_SHAPE.default,
     operator.floordiv: exir_ops.backend.tosa.DIV_FLOOR_SHAPE.default,
+    max: exir_ops.backend.tosa.MAX_SHAPE.default,
+    min: exir_ops.backend.tosa.MIN_SHAPE.default,
 }
 
 

backends/arm/operator_support/slice_copy_support.py

@@ -31,24 +31,53 @@ class SliceCopySupported(SupportedTOSAOperatorCheck):
 
     targets = [exir_ops.edge.aten.slice_copy.Tensor]
 
-    def is_node_tosa_supported(
-        self, node: fx.Node, tosa_spec: TosaSpecification
-    ) -> bool:  # type: ignore[override, misc]
+    @staticmethod
+    def _has_symbolic_bound(node: fx.Node) -> bool:
+        return any(
+            isinstance(bound, torch.SymInt)
+            or isinstance(bound, fx.Node)
+            and isinstance(bound.meta.get("val"), torch.SymInt)
+            for bound in node.args[2:4]
+        )
+
+    @staticmethod
+    def _has_no_non_positive_static_output_size(node: fx.Node) -> bool:
+        return all(
+            not isinstance(output_size, int) or output_size > 0
+            for output_size in node.meta["val"].shape
+        )
+
+    def _has_valid_slice_arguments(self, node: fx.Node) -> bool:
         if len(node.args) not in (4, 5):
             self.reporter.report_reject(
                 node,
                 f"{node.target}: expected 4 or 5 args, got {len(node.args)}.",
             )
             return False
+        if len(node.args) == 5 and node.args[4] <= 0:  # type: ignore[operator]
+            self.reporter.report_reject(
+                node, f"{node.target}: step must be > 0, got {node.args[4]}."
+            )
+            return False
+        return True
 
-        if len(node.args) == 5:
-            step = node.args[4]
-            if step <= 0:  # type: ignore[operator]
-                self.reporter.report_reject(
-                    node,
-                    f"{node.target}: step must be > 0, got {step}.",
-                )
-                return False
+    def is_node_tosa_supported(
+        self, node: fx.Node, tosa_spec: TosaSpecification
+    ) -> bool:  # type: ignore[override, misc]
+        if not self._has_valid_slice_arguments(node):
+            return False
+
+        if self._has_symbolic_bound(node):
+            self.reporter.report_reject(
+                node, "Symbolic slice bounds cannot be lowered to TOSA SLICE."
+            )
+            return False
+
+        if not self._has_no_non_positive_static_output_size(node):
+            self.reporter.report_reject(
+                node, "TOSA SLICE requires a guaranteed positive output size."
+            )
+            return False
 
         values_dtype = node.args[0].meta["val"].dtype  # type: ignore[union-attr]
 

backends/arm/test/misc/tosa_dialect/test_tosa_shape_support.py

@@ -12,10 +12,12 @@
 import executorch.backends.arm.operator_support.sym_size_int_support  # noqa: F401
 import pytest
 import torch
+from executorch.backends.arm.operator_support.slice_copy_support import (
+    SliceCopySupported,
+)
 from executorch.backends.arm.operator_support.symint_arithmetic_support import (
     SymIntArithmeticSupport,
 )
-
 from executorch.backends.arm.operator_support.tosa_supported_operators import (
     tosa_support_factory,
 )
@@ -604,3 +606,51 @@ def test_shape_extension_rejects_non_symint_arithmetic():
     )
 
     assert support.is_node_supported({}, arithmetic_node) is False
+
+
+def test_shape_extension_rejects_slice_with_symbolic_bound():
+    class SymbolicBoundSlice(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.slice.Tensor(x, 1, 0, x.shape[0], 1)
+
+    inputs = (torch.randn(4, 5),)
+    exported_program = _exported_program(
+        SymbolicBoundSlice(),
+        inputs,
+        dynamic_shapes=({0: Dim("batch", min=2, max=5)},),
+    )
+    support, reporter = _support("TOSA-1.1+FP+shape", exported_program)
+    slice_node = _find_node(exported_program, exir_ops.edge.aten.slice_copy.Tensor)
+
+    assert support.is_node_supported(exported_program.graph_module, slice_node) is False
+    assert "Symbolic slice bounds" in reporter.get_table_report()
+
+
+def test_shape_extension_rejects_slice_with_empty_unsliced_dimension():
+    class SliceSecondDimension(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.slice.Tensor(x, 1, 0, 3, 1)
+
+    exported_program = _exported_program(SliceSecondDimension(), (torch.randn(0, 5),))
+    reporter = WhyNoPartitionReporter()
+    support = SliceCopySupported(
+        TosaSpecification.create_from_string("TOSA-1.1+FP+shape"), reporter
+    )
+    slice_node = _find_node(exported_program, exir_ops.edge.aten.slice_copy.Tensor)
+
+    assert support.is_node_supported(exported_program.graph_module, slice_node) is False
+
+
+def test_shape_extension_rejects_empty_slice():
+    class EmptySlice(torch.nn.Module):
+        def forward(self, x):
+            return torch.ops.aten.slice.Tensor(x, 1, 0, -6, 1)
+
+    exported_program = _exported_program(EmptySlice(), (torch.randn(2, 5),))
+    reporter = WhyNoPartitionReporter()
+    support = SliceCopySupported(
+        TosaSpecification.create_from_string("TOSA-1.1+FP+shape"), reporter
+    )
+    slice_node = _find_node(exported_program, exir_ops.edge.aten.slice_copy.Tensor)
+
+    assert support.is_node_supported(exported_program.graph_module, slice_node) is False