Update base for Update on "[ET-VK] Add symint infrastructure to VulkanBackend and ComputeGraph"

Extend the Vulkan backend runtime infrastructure to better support symbolic
integer (symint) arguments. This is a prerequisite for operators that need
to handle dynamic shapes via symint values.

Changes:
- VulkanBackend.cpp: Compute output offset from end of args instead of
  assuming outputs follow inputs directly. Add scalar-to-tensor input
  handling so that Int/Bool EValues can populate tensor inputs. Support
  symint inputs provided as raw Int EValues (not just scalar tensors).
  Add symint output handling to write values back as tensor or Int EValue.
- ComputeGraph.h: Add SymInt case to extract_scalar<T>() so operators
  can transparently read symint values as scalars.
- ComputeGraph.cpp: Add Int fallback in read_symint() so values stored
  as plain Int (rather than SymInt objects) can be read uniformly.

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

cc manuelcandales digantdesai cbilgin

[ghstack-poisoned]

Author: ssjia

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.