Migrate to Graph.materialize_symints; stop passing raw SymInts (pytorch#21919)

Differential Revision: D109610966

Pull Request resolved: pytorch#21919

Author: laithsakka

backends/arm/_passes/convert_full_like_to_full_pass.py

@@ -42,7 +42,9 @@ def call_operator(self, op, args, kwargs, meta):
             return super().call_operator(op, args, kwargs, meta)
 
         tensor = args[0].data
-        full_args = (list(tensor.shape), args[1])
+        # Each entry is an int (static dim) or an aten.sym_size.int ProxyValue (dynamic dim).
+        size_args = self.call_size_operator_all(args[0], meta, edge_dialect=True)
+        full_args = (size_args, args[1])
         full_kwargs = {"dtype": tensor.dtype}
         return super().call_operator(
             exir_ops.edge.aten.full.default, full_args, full_kwargs, meta

backends/arm/_passes/decompose_linear_pass.py

@@ -56,7 +56,23 @@ def call(self, graph_module):
             # weights have shape (Co, Ci)
             weights_reshaped_shape = [weights_shape[0], weights_shape[1], 1, 1]
 
+            # Materialize SymInts into FX nodes; one call shares
+            # the symbol subgraph across all three shape lists.
             with graph_module.graph.inserting_before(node):
+                n_in, n_w = (
+                    len(input_reshaped_shape),
+                    len(weights_reshaped_shape),
+                )
+                all_sizes = graph_module.graph.materialize_symints(
+                    [
+                        *input_reshaped_shape,
+                        *weights_reshaped_shape,
+                        *output_shape,
+                    ]
+                )
+                input_reshaped_shape = all_sizes[:n_in]
+                weights_reshaped_shape = all_sizes[n_in : n_in + n_w]
+                output_size = all_sizes[n_in + n_w :]
                 # Reshape input to 4D with shape (N, Ci, 1, 1)
                 input_reshaped = create_node(
                     graph=graph_module.graph,
@@ -103,7 +119,7 @@ def call(self, graph_module):
                 output = create_node(
                     graph=graph_module.graph,
                     op_target=exir_ops.edge.aten.view_copy.default,
-                    args=(conv, list(output_shape)),
+                    args=(conv, output_size),
                     kwargs={},
                     from_node=node,
                     inherit_qparams=False,

backends/arm/_passes/decompose_select.py

@@ -56,10 +56,18 @@ def call(self, graph_module: torch.fx.GraphModule):
                 index = size_at_dim - abs(index)
 
             with graph_module.graph.inserting_before(node):
+                # ``Graph.create_node`` rejects raw SymInts in call_function
+                # args; ``materialize_symints`` lifts symbolic ``index`` /
+                # ``index + 1`` (from negative-index wrap-around against a
+                # dynamic shape) into graph nodes in a single call (shares
+                # producer-discovery + hash-cons). Static ints pass through.
+                start_arg, end_arg = graph_module.graph.materialize_symints(
+                    [index, index + 1]
+                )
                 slice_node = create_node(
                     graph_module.graph,
                     slice_op,
-                    (input_node, dim, index, index + 1),
+                    (input_node, dim, start_arg, end_arg),
                     from_node=node,
                     inherit_qparams=False,
                 )

backends/arm/_passes/insert_dynamic_padding.py

@@ -5,8 +5,6 @@
 
 from typing import Set, Type
 
-import torch
-
 from executorch.backends.arm._passes import ArmOpTargetedPass
 from executorch.backends.arm.tosa.dialect.shape import is_shape_op_node
 
@@ -41,7 +39,7 @@ def _is_dynamic_padding(
         return (isinstance(padding, ProxyValue) and is_shape_op_node(padding.node)) or (
             (
                 isinstance(padding, (list, tuple))
-                and any(isinstance(p, torch.SymInt) for p in padding)
+                and any(isinstance(p, ProxyValue) for p in padding)
             )
         )
 

backends/arm/_passes/rewrite_conv_pass.py

@@ -817,6 +817,29 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                     dilation,
                 )
 
+            # Compute fake tensor BEFORE materializing SymInts into FX nodes,
+            # since the underlying op expects ints/SymInts (not FX Nodes).
+            bias_fake_tensor = get_first_fake_tensor(bias) if bias else None
+            tosa_node_fake_tensor = target_op(
+                input_tensor_for_tosa_fake,
+                weight_fake_tensor,
+                bias_fake_tensor,
+                *conv_args[3:],
+            )
+
+            # ``Graph.create_node`` rejects raw SymInts in call_function args.
+            # If ``pad`` contains symbolic entries, materialize them into FX
+            # nodes so the TOSA conv node references the producing graph
+            # subgraph instead of holding raw SymInts.
+            if isinstance(pad, (list, tuple)) and any(
+                isinstance(p, torch.SymInt) for p in pad
+            ):
+                with graph_module.graph.inserting_before(node):
+                    materialized_pad = graph_module.graph.materialize_symints(pad)
+                new_conv_args = list(conv_args)
+                new_conv_args[4] = materialized_pad
+                conv_args = tuple(new_conv_args)
+
             with graph_module.graph.inserting_after(node):
                 tosa_op = create_node(
                     graph=graph_module.graph,
@@ -825,13 +848,6 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:  # noqa: C901
                     from_node=node,
                     inherit_qparams=True,
                 )
-            bias_fake_tensor = get_first_fake_tensor(bias) if bias else None
-            tosa_node_fake_tensor = target_op(
-                input_tensor_for_tosa_fake,
-                weight_fake_tensor,
-                bias_fake_tensor,
-                *conv_args[3:],
-            )
             tosa_op.meta["val"] = tosa_node_fake_tensor
 
             node_replacement, node_replacement_fake_tensor = (

backends/arm/_passes/size_adjust_input_pass.py

@@ -83,13 +83,13 @@ def _get_slice_adjustment(
     if exact_values is not None:
         adjustments = {max(value - pad, 0) for value in exact_values}
         if len(adjustments) == 1:
-            adjustment = next(iter(adjustments))
-            return adjustment if adjustment > 0 else None
+            exact_adjustment = next(iter(adjustments))
+            return exact_adjustment if exact_adjustment > 0 else None
 
     if pad >= stride - 1:
         return None
 
-    adjustment: SymIntLike | None = None  # type: ignore[no-redef]
+    adjustment: SymIntLike | None = None
     for threshold in range(pad + 1, stride):
         term = (remainder + stride - threshold) // stride
         adjustment = term if adjustment is None else adjustment + term
@@ -272,10 +272,29 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
 
             parent_node = node.args[0]
             with graph_module.graph.inserting_before(node):
+                # ``Graph.create_node`` rejects raw SymInts in
+                # call_function args. Aggregate every slice arg across
+                # all entries and lift via a single ``materialize_symints``
+                # call -- the helper walks the graph once for producer
+                # discovery, so a single call amortises that cost and lets
+                # symints with shared sub-expressions get hash-consed into
+                # one subgraph. Plain ints pass through unchanged.
+                flat = [a for args in slice_args for a in args]
+                materialized = iter(graph.materialize_symints(flat))
+                # Pop exactly len(args) values from the materialized iterator
+                # and pack them back into a tuple -- regroups the flat list
+                # of lifted values into the original (dim, start, end) shape.
+                lifted_slice_args = [
+                    tuple(next(materialized) for _ in args) for args in slice_args
+                ]
+
                 last_node = cast(torch.fx.Node, parent_node)
-                for args in slice_args:
+                for args in lifted_slice_args:
                     slice_node = create_node(
-                        graph, slice_op, (last_node,) + args, from_node=node
+                        graph,
+                        slice_op,
+                        (last_node,) + args,
+                        from_node=node,
                     )
                     last_node = slice_node
                 node.replace_input_with(cast(torch.fx.Node, parent_node), last_node)

backends/arm/test/passes/test_decompose_select_pass.py

@@ -63,11 +63,18 @@ def test_decompose_select_negative_symbolic_index_uses_symbolic_sub() -> None:
 
     slice_node = slice_nodes[0]
     assert slice_node.args[1] == 1
-    assert slice_node.args[2] != -1
-    assert isinstance(slice_node.args[2], torch.SymInt)
-    assert isinstance(slice_node.args[3], torch.SymInt)
-    assert str(slice_node.args[2]).endswith(" - 1")
-    assert str(slice_node.args[3]) in str(slice_node.args[2])
+    # Start/end are now FX nodes (materialized from SymInts) rather than raw
+    # SymInts, since Graph.create_node rejects raw symbolic leaves in
+    # call_function args. The original SymInt is preserved in meta['val'].
+    start_arg, end_arg = slice_node.args[2], slice_node.args[3]
+    assert isinstance(start_arg, torch.fx.Node)
+    assert isinstance(end_arg, torch.fx.Node)
+    start_val = start_arg.meta["val"]
+    end_val = end_arg.meta["val"]
+    assert isinstance(start_val, torch.SymInt)
+    assert isinstance(end_val, torch.SymInt)
+    assert str(start_val).endswith(" - 1")
+    assert str(end_val) in str(start_val)
     assert squeeze_nodes[0].args == (slice_node, [1])
 
     result.graph_module.graph.lint()