Merge branch 'main' into dynamic_unbound_kv_cache

Author: psiddh

backends/arm/_passes/arm_pass_manager.py

@@ -481,16 +481,14 @@ def _tosa_pipeline(
                 ConvertFullLikeToFullPass(),
                 MatchArgDtypePass(),
                 UnsqueezeScalarPlaceholdersPass(exported_program),
-                # TODO: Move DecomposeNotEqualPass to before or after this block of
-                # passes. Ticket: MLETORCH-1540
-                DecomposeNotEqualPass(),
                 MatchArgRanksPass(exported_program),
             ]
         )
 
         # Node transformation passes (post scalar-removal)
         self.add_passes(
             [
+                DecomposeNotEqualPass(),
                 NormalizeIndexPutNoneIndicesPass(),
                 NormalizeIndexPutBoolIndexTensorPass(),
                 RewriteIndexPutPass(),

backends/arm/_passes/arm_pass_utils.py

@@ -9,7 +9,7 @@
 import operator
 import traceback
 from inspect import isclass
-from typing import cast, List, Optional, Sequence, Tuple
+from typing import cast, Optional, Sequence
 
 import torch
 import torch.fx
@@ -19,10 +19,6 @@
 from executorch.exir import ExportedProgram
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.dialects.edge._ops import EdgeOpOverload
-from executorch.exir.graph_module import (
-    _get_control_flow_submodules,
-    get_control_flow_submodules,
-)
 from executorch.exir.pass_base import NodeMetadata
 
 from torch._export.utils import (
@@ -36,7 +32,6 @@
 from torch._ops import OpOverload
 from torch._subclasses.fake_tensor import FakeTensor
 from torch.export.graph_signature import InputKind
-from torch.fx import GraphModule, Node
 
 
 def is_submodule_node(node: torch.fx.Node):
@@ -364,48 +359,6 @@ def set_node_arg(node: torch.fx.Node, i: int | str, value):
         raise RuntimeError("Invalid type")
 
 
-def is_nested_control_flow_graph(graph_module: GraphModule) -> bool:
-    """Returns True if graph_module is a nested control-flow graph."""
-
-    # Find all top-level control-flow submodules
-    top_cf = get_control_flow_submodules(graph_module)
-    # For each submodule, see if it itself has control-flow inside
-    for _, submod, _ in top_cf:
-        if get_control_flow_submodules(submod):
-            return True
-    return False
-
-
-def get_cond_while_submodules_nested(
-    graph_module: GraphModule,
-    apply_quantization: bool = False,
-) -> List[Tuple[str, GraphModule, Node]]:
-    """Recursively find cond/while_loop submodules in an GraphModule.
-
-    In nested control flow graphs, FX records the submodule functions
-    (true/false or cond/body) in reverse order compared to top-level graphs. We
-    must swap the indices when nested so that cond (first) and body/true_fn
-    (second) are consistently identified across all nesting levels.
-
-    """
-
-    # Determine arg indices based on nesting and whether only cond branch is needed
-    nested = is_nested_control_flow_graph(graph_module)
-    # cond: [true_fn, false_fn] or swapped if nested
-    cond_indices = [2, 1] if nested else [1, 2]
-    # while_loop: [cond_fn, body_fn] or swapped if nested
-    while_indices = [1, 0] if nested else [0, 1]
-    if apply_quantization:
-        # only keep the cond_fn for while_loop (first index) when quantizing.
-        while_indices = [while_indices[0]]
-    mapping = {
-        torch.ops.higher_order.cond: cond_indices,
-        torch.ops.higher_order.while_loop: while_indices,
-    }
-    # collect cond/while submodules (using mapping indices)
-    return _get_control_flow_submodules(graph_module, mapping)
-
-
 def to_2tuple(value):
     """Normalizes scalars, and 1-element sequences to a tuple of length 2."""
     if isinstance(value, int):

backends/arm/_passes/control_flow_const_inline.py

@@ -7,12 +7,10 @@
 
 import torch
 from executorch.backends.arm._passes.arm_pass import ArmPass
-from executorch.backends.arm._passes.arm_pass_utils import (
-    get_cond_while_submodules_nested,
-    is_submodule_node,
-)
+from executorch.backends.arm._passes.arm_pass_utils import is_submodule_node
 from executorch.backends.transforms.utils import is_get_attr_node
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.graph_module import get_cond_while_submodules
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx import GraphModule
 
@@ -37,7 +35,7 @@ class ControlFlowConstInlinePass(ArmPass):
 
     def _convert_getattr(self, graph_module):
         modified = False
-        for _, submodule, _ in get_cond_while_submodules_nested(graph_module):
+        for _, submodule, _ in get_cond_while_submodules(graph_module):
             for submodule_node in submodule.graph.nodes:
                 if submodule_node.target in self._targeted_ops:
                     self._convert_getattr(submodule)

backends/arm/_passes/insert_const_shapes.py

@@ -26,6 +26,10 @@ class InsertConstShapesPass(ArmPass):
         exir_ops.edge.aten.repeat.default,
     }
 
+    def __init__(self) -> None:
+        super().__init__()
+        self._const_shape_cache: dict[tuple[int, ...], Any] = {}
+
     @staticmethod
     def _is_shape_arg(arg: Any) -> bool:
         """Return True when `arg` looks like a literal shape list/tuple."""
@@ -46,13 +50,17 @@ def call_operator(self, op, args, kwargs, meta, updated: Optional[bool] = False)
                     # Insert a const node for the shape argument
                     if op == exir_ops.edge.aten.view_copy.default:
                         arg = meta.data["val"].shape
-                    const_node = super().call_shape_operator(
-                        exir_ops.backend.tosa.CONST_SHAPE.default,
-                        (arg,),
-                        {},
-                        meta,
-                        True,
-                    )
+                    shape = tuple(arg)
+                    const_node = self._const_shape_cache.get(shape)
+                    if const_node is None:
+                        const_node = super().call_shape_operator(
+                            exir_ops.backend.tosa.CONST_SHAPE.default,
+                            (arg,),
+                            {},
+                            meta,
+                            True,
+                        )
+                        self._const_shape_cache[shape] = const_node
                     new_args.append(const_node)
                     updated = True
                 else:

backends/arm/_passes/insert_rescales_pass.py

@@ -509,7 +509,13 @@ def _rescale_submodule_inputs(
             input_node = input_nodes[qargs_index]
             if len(input_node.users) == 0:
                 continue
-            if len(out_qparams_map := input_node.meta.get("output_qparams", {})) != 1:
+            out_qparams_map = input_node.meta.get("output_qparams", {})
+            if len(out_qparams_map) == 0:
+                # Nested control-flow submodules may also expose frozen captured
+                # values as placeholders. Those are not control-flow boundary
+                # inputs, so there is no qparam pair to bridge with a RESCALE.
+                continue
+            if len(out_qparams_map) != 1:
                 raise ValueError(
                     f"Expected submodule input {input_node} to have exactly one output qparam, got {out_qparams_map}"
                 )

backends/arm/_passes/match_arg_ranks_pass.py

@@ -57,6 +57,7 @@ def __init__(self, exported_program: ExportedProgram, *args, **kwargs) -> None:
         exir_ops.edge.aten.ge.Tensor,
         exir_ops.edge.aten.lt.Tensor,
         exir_ops.edge.aten.le.Tensor,
+        exir_ops.edge.aten.ne.Tensor,
         exir_ops.edge.aten.pow.Tensor_Tensor,
         exir_ops.edge.aten.remainder.Tensor,
         exir_ops.edge.aten.where.self,

backends/arm/_passes/scalars_to_attribute_pass.py

@@ -8,11 +8,9 @@
 
 import torch
 from executorch.backends.arm._passes import ArmPass
-from executorch.backends.arm._passes.arm_pass_utils import (
-    get_cond_while_submodules_nested,
-    get_first_fake_tensor,
-)
+from executorch.backends.arm._passes.arm_pass_utils import get_first_fake_tensor
 from executorch.backends.arm._passes.match_arg_ranks_pass import MatchArgRanksPass
+from executorch.exir.graph_module import get_cond_while_submodules
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx import GraphModule, Node
 from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix
@@ -98,7 +96,7 @@ def handle_control_nodes(self, graph_module: GraphModule) -> None:
         """Apply scalar argument conversion on subgraphs of control-flow
         nodes.
         """
-        for _, submodule, _ in get_cond_while_submodules_nested(graph_module):
+        for _, submodule, _ in get_cond_while_submodules(graph_module):
             for submodule_node in submodule.graph.nodes:
                 self._convert_scalar_args(submodule, submodule_node)
 

backends/arm/operator_support/TARGETS

@@ -6,6 +6,7 @@ runtime.python_library(
     deps = [
         "//executorch/backends/arm:constants",
         "//executorch/backends/arm/_passes:passes",
+        "//executorch/backends/arm/tosa:resize_utils",
         "//executorch/backends/arm/tosa:tosa",
         "//executorch/backends/transforms:remove_getitem_op",
         "//executorch/backends/xnnpack/_passes:xnnpack_passes",

backends/arm/operator_support/control_flow_support.py

@@ -19,6 +19,13 @@
 from torch.fx.passes.operator_support import OperatorSupportBase
 
 
+def _owning_graph_module(node: fx.Node) -> fx.GraphModule:
+    graph_module = getattr(node.graph, "owning_module", None)
+    if not isinstance(graph_module, fx.GraphModule):
+        raise RuntimeError(f"Could not resolve owning GraphModule for node {node}")
+    return graph_module
+
+
 def _fully_partitioned(submodule: fx.GraphModule) -> bool:
     """Check that all nested control-flow ops within this submodule are also
     fully partitioned.
@@ -27,8 +34,8 @@ def _fully_partitioned(submodule: fx.GraphModule) -> bool:
 
     for submodule_node in submodule.graph.nodes:
         if submodule_node.target in ControlFlowOpSupported._targeted_ops:
-            if _submodules_fully_partitioned(submodule_node, submodule):
-                return True
+            if not _submodules_fully_partitioned(submodule_node, submodule):
+                return False
 
         if submodule_node.op != "call_function":
             continue
@@ -56,13 +63,18 @@ def _fully_partitioned(submodule: fx.GraphModule) -> bool:
     return True
 
 
-def _submodules_fully_partitioned(node: fx.Node, graph_module: fx.GraphModule) -> bool:
+def _submodules_fully_partitioned(
+    node: fx.Node, graph_module: fx.GraphModule | None = None
+) -> bool:
     """Returns whether the submodule arguments to a cond node were fully
     partitioned.
 
     Updates "val" meta of the submodules if they are.
 
     """
+    if graph_module is None:
+        graph_module = _owning_graph_module(node)
+
     match node.target:
         case torch.ops.higher_order.cond:
             submodule_args = node.args[1:3]
@@ -129,9 +141,7 @@ def is_node_supported(
                         node, f"Submodule had unsupported user {user}"
                     )
                     return False
-                if not _submodules_fully_partitioned(
-                    user, self.exported_program.graph_module
-                ):
+                if not _submodules_fully_partitioned(user):
                     self.reporter.report_reject(
                         node, "One submodule was not fully partitioned"
                     )
@@ -174,9 +184,7 @@ def is_node_supported(
                 )
                 return False
 
-            if not _submodules_fully_partitioned(
-                node, self.exported_program.graph_module
-            ):
+            if not _submodules_fully_partitioned(node):
                 self.reporter.report_reject(
                     node, "Submodule was not fully partitioned."
                 )

backends/arm/operator_support/upsample_support.py

@@ -13,19 +13,65 @@
     SupportedTOSAOperatorCheck,
 )
 from executorch.backends.arm.tosa import TosaSpecification
+from executorch.backends.arm.tosa.resize_utils import get_tosa_resize_validation_error
 from executorch.exir.dialects._ops import ops as exir_ops
 
 
+def _is_upsample_node_tosa_supported(
+    support_check: SupportedTOSAOperatorCheck,
+    node: fx.Node,
+    tosa_spec: TosaSpecification,
+    *,
+    align_corners: bool,
+) -> bool:
+    input_node = ensure_type(fx.Node, node.args[0])
+    input_size_yx = get_first_fake_tensor(input_node).shape[2:]
+    output_size_yx = get_first_fake_tensor(node).shape[2:]
+
+    try:
+        scale_y_n, scale_y_d, offset_y, border_y = (
+            RewriteUpsamplePass.get_resize_parameters_1d(
+                input_size_yx[0], output_size_yx[0], align_corners
+            )
+        )
+        scale_x_n, scale_x_d, offset_x, border_x = (
+            RewriteUpsamplePass.get_resize_parameters_1d(
+                input_size_yx[1], output_size_yx[1], align_corners
+            )
+        )
+    except RuntimeError as err:
+        support_check.reporter.report_reject(node, str(err))
+        return False
+
+    # Validate the exact TOSA RESIZE parameters that RewriteUpsamplePass will
+    # emit so support checks and fake-op validation reject the same cases.
+    validation_error = get_tosa_resize_validation_error(
+        input_hw=input_size_yx,
+        output_hw=output_size_yx,
+        scale=[scale_y_n, scale_y_d, scale_x_n, scale_x_d],
+        offset=[offset_y, offset_x],
+        border=[border_y, border_x],
+        tosa_spec=tosa_spec,
+    )
+    if validation_error is not None:
+        support_check.reporter.report_reject(node, validation_error)
+        return False
+
+    return True
+
+
 @register_tosa_support_check
 class UpsampleNearest2dSupported(SupportedTOSAOperatorCheck):
     """Provide the explicit TOSA support gate for nearest upsample."""
 
     targets = [exir_ops.edge.aten.upsample_nearest2d.vec]
 
     def is_node_tosa_supported(
-        self, _node: fx.Node, _tosa_spec: TosaSpecification
+        self, node: fx.Node, tosa_spec: TosaSpecification
     ) -> bool:  # type: ignore[override, misc]
-        return True
+        return _is_upsample_node_tosa_supported(
+            self, node, tosa_spec, align_corners=False
+        )
 
 
 @register_tosa_support_check
@@ -37,33 +83,9 @@ class UpsampleBilinear2dSupported(SupportedTOSAOperatorCheck):
     targets = [exir_ops.edge.aten.upsample_bilinear2d.vec]
 
     def is_node_tosa_supported(
-        self, node: fx.Node, _tosa_spec: TosaSpecification
+        self, node: fx.Node, tosa_spec: TosaSpecification
     ) -> bool:  # type: ignore[override, misc]
-        input_node = ensure_type(fx.Node, node.args[0])
         align_corners = ensure_type(bool, node.args[2])
-        input_size_yx = get_first_fake_tensor(input_node).shape[2:]
-        output_size_yx = get_first_fake_tensor(node).shape[2:]
-
-        try:
-            scale_y_n, scale_y_d, _, _ = RewriteUpsamplePass.get_resize_parameters_1d(
-                input_size_yx[0], output_size_yx[0], align_corners
-            )
-            scale_x_n, scale_x_d, _, _ = RewriteUpsamplePass.get_resize_parameters_1d(
-                input_size_yx[1], output_size_yx[1], align_corners
-            )
-        except RuntimeError as err:
-            self.reporter.report_reject(node, str(err))
-            return False
-
-        # get_resize_parameters_1d() returns the TOSA RESIZE scale fraction for
-        # each spatial dimension. For align_corners=False, this is the effective
-        # output_size / input_size ratio, so the 1/16 boundary is checked
-        # directly in the same representation that RESIZE lowering will use.
-        if scale_y_d >= 16 * scale_y_n or scale_x_d >= 16 * scale_x_n:
-            self.reporter.report_reject(
-                node,
-                "Bilinear RESIZE downscale must be strictly greater than 1/16",
-            )
-            return False
-
-        return True
+        return _is_upsample_node_tosa_supported(
+            self, node, tosa_spec, align_corners=align_corners
+        )