Update base for Update on "[ET Device Support] TensorImpl carries device info"

This diff extends `TensorImpl` to carry device information, enabling the runtime tensor to track which device its data resides on (CPU, CUDA, etc.). This is a prerequisite for parsing device info from the schema and allocating device memory.

Differential Revision: [D93635655](https://our.internmc.facebook.com/intern/diff/D93635655/)

[ghstack-poisoned]

Author: Gasoonjia

backends/arm/_passes/TARGETS

@@ -1,13 +1,55 @@
 load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
 
 runtime.python_library(
-    name = "passes",
-    srcs = glob(["*.py"]),
+    name = "core",
+    srcs = [
+        "arm_pass.py",
+        "arm_pass_utils.py",
+        "quant_args.py",
+    ],
     deps = [
         "//executorch/backends/arm:common",
         "//executorch/backends/arm:constants",
         "//executorch/backends/arm/tosa:utils",
         "//executorch/backends/arm/tosa/dialect:lib",
+        "//executorch/backends/transforms:utils",
+        "//executorch/exir:lib",
+    ],
+)
+
+runtime.python_library(
+    name = "arm_pass_manager_base",
+    srcs = ["arm_pass_manager.py"],
+    deps = [
+        ":core",
+        "//executorch/backends/arm:common",
+        "//executorch/backends/arm/tosa:utils",
+        "//executorch/backends/arm/tosa/dialect:lib",
+        "//executorch/exir:lib",
+    ],
+)
+
+runtime.python_library(
+    name = "arm_pass_manager_fb",
+    srcs = [],
+    deps = [
+        ":arm_pass_manager_base",
+        # @oss-disable[end= ]: "//executorch/backends/arm/_passes/fb:fb",
+    ],
+)
+
+runtime.python_library(
+    name = "passes",
+    srcs = glob(["*.py"], exclude = [
+        "arm_pass.py",
+        "arm_pass_utils.py",
+        "quant_args.py",
+    ]),
+    deps = [
+        ":core",
+        ":arm_pass_manager_base" if runtime.is_oss else ":arm_pass_manager_fb",
+        "//executorch/backends/arm/tosa:utils",
+        "//executorch/backends/arm/tosa/dialect:lib",
         "//executorch/backends/transforms:fuse_view_copy",
         "//executorch/backends/transforms:remove_getitem_op",
         "//executorch/backends/transforms:replace_scalar_with_tensor",

backends/arm/_passes/__init__.py

@@ -138,6 +138,9 @@
     RewriteBoolToFp32CastViaInt8Pass,
 )
 from .rewrite_conv_pass import RewriteConvPass  # noqa
+from .rewrite_high_rank_singleton_permute_pass import (  # noqa
+    RewriteHighRankSingletonPermutePass,
+)
 from .rewrite_index_put_pass import RewriteIndexPutPass  # noqa
 from .rewrite_le_lt_to_ge_gt_pass import RewriteLeLtToGeGtPass  # noqa
 from .rewrite_matmul import RewriteMatmulPass  # noqa
@@ -155,4 +158,13 @@
 from .control_flow_const_inline import (  # noqa  # usort: skip
     ControlFlowConstInlinePass,
 )
-from .arm_pass_manager import ArmPassManager  # noqa  # usort: skip
+
+# Import all subpackages to allow extensions to patch classes
+import importlib  # noqa: E402
+import pkgutil  # noqa: E402
+
+for _, _modname, _ispkg in pkgutil.iter_modules(__path__, __name__ + "."):
+    if _ispkg:
+        importlib.import_module(_modname)
+
+from .arm_pass_manager import ArmPassManager  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -8,6 +8,7 @@
 import logging
 from collections import defaultdict
 from collections.abc import Sequence
+from dataclasses import dataclass, field
 
 import executorch.backends.arm.tosa.dialect  # noqa: unused
 from executorch.backends.arm._passes import (
@@ -121,6 +122,7 @@
     RewriteBoolBitwiseToLogicalPass,
     RewriteBoolToFp32CastViaInt8Pass,
     RewriteConvPass,
+    RewriteHighRankSingletonPermutePass,
     RewriteIndexPutPass,
     RewriteLeLtToGeGtPass,
     RewriteMatmulPass,
@@ -156,11 +158,21 @@
 logger = logging.getLogger(__name__)
 
 
+@dataclass
+class PassInsertions:
+    """Holds lists of passes to be inserted before and after a target pass."""
+
+    before_passes: list = field(default_factory=list)
+    after_passes: list = field(default_factory=list)
+
+
 class ArmPassManager(PassManager):
     def __init__(self, compile_spec: ArmCompileSpec) -> None:
         self.compile_spec = compile_spec
         self.tosa_spec = compile_spec.tosa_spec
         self._skip_pass_types: tuple[type, ...] = ()
+        self._pass_insertions: dict[type, PassInsertions] = {}
+        self._insertions_applied = False
         super().__init__()
         self.configure_skip_passes()
 
@@ -222,6 +234,98 @@ def validate_constraints_mandatory(self):
 
             raise RuntimeError(error_msg)
 
+    def insert_passes_before(
+        self, target_pass_type: type, passes: list[ExportPass]
+    ) -> None:
+        """Register passes to be inserted before instances of target_pass_type.
+        Insertions are deferred and applied via _apply_pass_insertions().
+
+        Args:
+            target_pass_type: The pass class to insert before (e.g., InsertTableOpsPass)
+            passes: List of pass instances to insert
+
+        """
+        self._pass_insertions.setdefault(
+            target_pass_type, PassInsertions()
+        ).before_passes.extend(passes)
+
+    def insert_passes_after(
+        self, target_pass_type: type, passes: list[ExportPass]
+    ) -> None:
+        """Register passes to be inserted after instances of target_pass_type.
+        Insertions are deferred and applied via _apply_pass_insertions().
+
+        Args:
+            target_pass_type: The pass class to insert after
+            passes: List of pass instances to insert
+
+        """
+        self._pass_insertions.setdefault(
+            target_pass_type, PassInsertions()
+        ).after_passes.extend(passes)
+
+    def _apply_pass_insertions(self) -> None:
+        """Apply all registered pass insertions to the collected passes.
+
+        Called ONCE after all add_passes() calls are complete, before execution.
+
+        Raises:
+            ValueError: If any registered target pass type is not found in the pipeline.
+
+        """
+        if self._insertions_applied or not self._pass_insertions:
+            return
+
+        # Fail fast if any target pass type is missing from the pipeline
+        existing_pass_types = {type(p) for p in self.passes}
+        for target_type in self._pass_insertions:
+            if target_type not in existing_pass_types:
+                available = [type(p).__name__ for p in self.passes]
+                raise ValueError(
+                    f"Target pass {target_type.__name__} not found in the pass "
+                    f"pipeline. Available passes: {available}"
+                )
+
+        # Build new pass list with insertions applied
+        new_passes = []
+        for pass_obj in self.passes:
+            pass_type = type(pass_obj)
+
+            # Insert passes BEFORE this pass
+            if pass_type in self._pass_insertions:
+                insertions = self._pass_insertions[pass_type]
+                for before_pass in insertions.before_passes:
+                    # Check if we should skip this inserted pass
+                    if type(before_pass) not in self._skip_pass_types:
+                        new_passes.append(before_pass)
+
+            # Add the original pass
+            new_passes.append(pass_obj)
+
+            # Insert passes AFTER this pass
+            if pass_type in self._pass_insertions:
+                insertions = self._pass_insertions[pass_type]
+                for after_pass in insertions.after_passes:
+                    # Check if we should skip this inserted pass
+                    if type(after_pass) not in self._skip_pass_types:
+                        new_passes.append(after_pass)
+
+        # Replace the passes list
+        self.passes = new_passes
+        self._insertions_applied = True
+
+    def _configure_pass_insertions(self, exported_program: ExportedProgram) -> None:
+        """Hook for subclasses to configure pass insertions. Called at the START
+        of pipeline construction, before any passes are added.
+
+        Subclasses should override this to call insert_passes_before/after.
+
+        Args:
+            exported_program: The exported program being transformed
+
+        """
+        pass
+
     def add_passes(self, passes: Sequence[ExportPass | None]):
         for p in passes:
             if p is not None:
@@ -240,6 +344,9 @@ def add_pass(self, pipeline_pass):
     def _tosa_pipeline(
         self, exported_program: ExportedProgram, graph_module: GraphModule
     ) -> GraphModule:
+        # Allow subclasses to configure pass insertions before building pipeline
+        self._configure_pass_insertions(exported_program)
+
         # Preprocessing passes
         self.add_pass(AnnotateOutputDimOrderPass())
 
@@ -366,6 +473,7 @@ def _tosa_pipeline(
                 CastToInt32Pass(),
                 BroadcastArgsPass(),
                 ConvertPermuteSingletonToViewPass(),
+                RewriteHighRankSingletonPermutePass(),
                 FuseViewCopyTransformPass(),
                 DecomposeConvWithInt16ActivationPass(),
                 DecomposeSumPass(),
@@ -396,6 +504,9 @@ def _tosa_pipeline(
             ]
         )
 
+        # Apply all pass insertions once after all passes are collected
+        self._apply_pass_insertions()
+
         self.validate_constraints_mandatory()
         return self._transform(graph_module)
 
@@ -468,6 +579,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
                 DecomposeLeakyReLUPass(tfa_pass=True),
                 DecomposeLinalgVectorNormPass(tfa_pass=True),
                 DecomposeSqrtPass(tfa_pass=True),
+                DecomposeAdaptiveAvgPool2dPass(tfa_pass=True),
                 DecomposeAvgPool2dPass(tfa_pass=True),
                 DecomposeSoftmaxUnstablePass(tfa_pass=True),
                 DecomposeSoftmaxPass(

backends/arm/_passes/decompose_adaptive_avg_pool2d_pass.py

@@ -49,7 +49,7 @@ class DecomposeAdaptiveAvgPool2dPass(ArmPass):
     _passes_required_after: Set[Type[ExportPass]] = {DecomposeAvgPool2dPass}
 
     def call_operator(self, op, args, kwargs, meta, updated=False):
-        if op not in (edge_ops + aten_ops):
+        if op not in (edge_ops + aten_ops) or not self.allowed_to_transform(meta):
             return super().call_operator(op, args, kwargs, meta, updated)
 
         avg_pool2d_op, slice_op, cat_op = _get_decomposition(op)

backends/arm/_passes/decompose_slice_scatter_pass.py

@@ -41,12 +41,16 @@ def _get_slice_scatter_decomposition(op) -> tuple:
 
 def _fixup_start(start, dim_size: int) -> int:
     s = 0 if start is None else int(start)
-    return max(0, min(s % dim_size if s < 0 else s, dim_size))
+    if s < 0:
+        s += dim_size
+    return max(0, min(s, dim_size))
 
 
 def _fixup_end(end, dim_size: int) -> int:
     e = dim_size if end is None else int(end)
-    return max(0, min(e % dim_size if e < 0 else e, dim_size))
+    if e < 0:
+        e += dim_size
+    return max(0, min(e, dim_size))
 
 
 class DecomposeSliceScatterPass(ArmPass):
@@ -131,31 +135,34 @@ def call_operator(self, op, args, kwargs, meta):
 
         # ---- fast path: contiguous update (step == 1) ----
         if step == 1:
-            # prefix = input[..., :start_i, ...] along dim_norm
-            prefix = super().call_operator(
-                slice_copy_op,
-                (input, dim_norm, 0, start_i, 1),
-                {},
-                meta,
-                updated=True,
+            prefix = None
+            suffix = None
+
+            # prefix = input[..., :start_i, ...] along dim_norm (only if non-empty)
+            if start_i > 0:
+                prefix = super().call_operator(
+                    slice_copy_op,
+                    (input, dim_norm, 0, start_i, 1),
+                    {},
+                    meta,
+                    updated=True,
+                )
+
+            # suffix = input[..., end_i:, ...] along dim_norm (only if non-empty)
+            if end_i < dim_size:
+                suffix = super().call_operator(
+                    slice_copy_op,
+                    (input, dim_norm, end_i, dim_size, 1),
+                    {},
+                    meta,
+                    updated=True,
+                )
+
+            parts = [x for x in (prefix, src, suffix) if x is not None]
+
+            return super().call_operator(
+                cat_op, (parts, dim_norm), {}, meta, updated=True
             )
-            # suffix = input[..., end_i:, ...] along dim_norm
-            suffix = super().call_operator(
-                slice_copy_op,
-                (input, dim_norm, end_i, dim_size, 1),
-                {},
-                meta,
-                updated=True,
-            )
-            # concat(prefix, src, suffix) along dim_norm
-            updated = super().call_operator(
-                cat_op,
-                ([prefix, src, suffix], dim_norm),
-                {},
-                meta,
-                updated=True,
-            )
-            return updated
 
         # ---- general path: strided update (step > 1) ----
         # Move updated dim to front to use a single index tensor.