Add support to PropagateSlice for custom unary/binary ops

Summary: As titled.

Differential Revision: D119399660

Pull Request resolved: #22657

Author: DrJessop

backends/cadence/aot/reorder_ops.py

@@ -1179,17 +1179,23 @@ class PropagateSlice(RemoveOrReplacePassInterface):
     Handles any slice dim and any step size.
     """
 
-    def __init__(self) -> None:
+    def __init__(
+        self,
+        additional_unary_targets: Optional[list[EdgeOpOverload]] = None,
+        additional_binary_targets: Optional[list[EdgeOpOverload]] = None,
+    ) -> None:
         super().__init__()
-        elementwise_targets = [
+        unary_targets = [
             exir_ops.edge.quantized_decomposed.quantize_per_tensor.default,
             exir_ops.edge.cadence.quantize_per_tensor.default,
             exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default,
             exir_ops.edge.cadence.dequantize_per_tensor.default,
+            *(additional_unary_targets or []),
         ]
         binary_targets = [
             exir_ops.edge.aten.add.Tensor,
             exir_ops.edge.aten.mul.Tensor,
+            *(additional_binary_targets or []),
         ]
         self._dispatch: dict[
             EdgeOpOverload,
@@ -1198,7 +1204,7 @@ def __init__(self) -> None:
                 Callable[[torch.fx.Node, torch.fx.Node], bool],
             ],
         ] = {}
-        for t in elementwise_targets:
+        for t in unary_targets:
             self._dispatch[t] = (
                 self._should_swap_elementwise,
                 self._swap_elementwise_slice,
@@ -1224,7 +1230,8 @@ def _should_swap_elementwise(
     def _swap_elementwise_slice(
         self, op_node: torch.fx.Node, slice_node: torch.fx.Node
     ) -> bool:
-        op_input = get_arg(op_node, "input", torch.fx.Node)
+        op_input = op_node.args[0]
+        assert isinstance(op_input, torch.fx.Node)
         graph = slice_node.graph
 
         slice_dim = get_arg(slice_node, "dim", int)

backends/cadence/aot/tests/test_reorder_ops_passes.py

@@ -1358,6 +1358,73 @@ def test_unsupported_parent_not_swapped(self) -> None:
 
         self.assertFalse(result.modified)
 
+    def test_swap_additional_unary_target(self) -> None:
+        x_data = torch.randn(4, 60, 1, 1)
+        builder = GraphBuilder()
+        x = builder.placeholder("x", x_data)
+        relu = builder.call_operator(exir_ops.edge.aten.relu.default, args=(x,))
+        sliced = builder.call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor,
+            args=(relu, 0, 0, 4, 2),
+        )
+        builder.output([sliced])
+        gm = builder.get_graph_module()
+
+        result = transform_and_check_numerics(
+            gm,
+            (x_data,),
+            PropagateSlice(additional_unary_targets=[exir_ops.edge.aten.relu.default]),
+        )
+
+        self.assertTrue(result.modified)
+        slice_nodes = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+        )
+        self.assertEqual(len(slice_nodes), 1)
+        relu_nodes = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.relu.default
+        )
+        self.assertEqual(len(relu_nodes), 1)
+        self.assertIs(relu_nodes[0].args[0], slice_nodes[0])
+        self.assertEqual(list(relu_nodes[0].meta["val"].shape), [2, 60, 1, 1])
+
+    def test_swap_additional_binary_target(self) -> None:
+        lhs_data = torch.randn(1, 60, 1, 1)
+        rhs_data = torch.randn(4, 60, 1, 1)
+        builder = GraphBuilder()
+        lhs = builder.placeholder("lhs", lhs_data)
+        rhs = builder.placeholder("rhs", rhs_data)
+        sub = builder.call_operator(
+            exir_ops.edge.aten.sub.Tensor,
+            args=(lhs, rhs),
+        )
+        sliced = builder.call_operator(
+            exir_ops.edge.aten.slice_copy.Tensor,
+            args=(sub, 0, 0, 4, 2),
+        )
+        builder.output([sliced])
+        gm = builder.get_graph_module()
+
+        result = transform_and_check_numerics(
+            gm,
+            (lhs_data, rhs_data),
+            PropagateSlice(additional_binary_targets=[exir_ops.edge.aten.sub.Tensor]),
+        )
+
+        self.assertTrue(result.modified)
+        slice_nodes = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.slice_copy.Tensor
+        )
+        self.assertEqual(len(slice_nodes), 1)
+        self.assertEqual(slice_nodes[0].args[0].name, "rhs")
+        sub_nodes = gm.graph.find_nodes(
+            op="call_function", target=exir_ops.edge.aten.sub.Tensor
+        )
+        self.assertEqual(len(sub_nodes), 1)
+        self.assertIs(sub_nodes[0].args[0], lhs.node)
+        self.assertIs(sub_nodes[0].args[1], slice_nodes[0])
+        self.assertEqual(list(sub_nodes[0].meta["val"].shape), [2, 60, 1, 1])
+
     def test_swap_broadcast_mul_slice_on_broadcast_dim(self) -> None:
         """[1,60,1,1] * [4,1,1,1] → [4,60,1,1] → slice(dim=0, step=2)
         Only the [4,1,1,1] input should be sliced."""