Update on "Arm backend: Add INT16 support to rescale operation"

Add INT16 support for RequantizeNode rescale operations in ExecutorTorch ARM backend.

This follows the pattern established for linear, mul, sigmoid, tanh, slice, view/transpose, cat, and FCNode operations, extending int16 support to RequantizeNode rescale operations.

Changes:
- Add INT16 dtype validation support in op_rescale.py
- Enable rescale operations for 16A8W quantization configuration

The 16A8W configuration uses 16-bit activations with 8-bit weights, enabling higher precision for activations while maintaining weight efficiency. RequantizeNode rescale operations are essential for proper quantization scaling in the 16A8W pipeline.

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

cc digantdesai freddan80 per zingo oscarandersson8218

[ghstack-poisoned]

Author: Ninja91

.github/workflows/pull.yml

@@ -971,6 +971,13 @@ jobs:
         ./cmake-out/backends/vulkan/test/custom_ops/q4gsw_linear
         ./cmake-out/backends/vulkan/test/custom_ops/choose_qparams_per_row
 
+        # "Classic" Operator tests
+        PYTHON_EXECUTABLE=python bash backends/vulkan/test/scripts/test_op.sh --build
+        # TODO(ssjia): figure out how to run custom op tests in CI. Currently, they are
+        # failing due to to the libstdc++.so.6 installed with conda not supporting
+        # GLIBCXX_3.4.30. These tests are still run in Meta internal CI.
+        # ./cmake-out/backends/vulkan/test/op_tests/vulkan_sdpa_test
+
         # Run e2e testing for selected operators. More operators will be tested via this
         # route in the future.
         python -m unittest backends/vulkan/test/test_vulkan_delegate.py -k "*pt2e*"

backends/arm/README.md

@@ -206,14 +206,6 @@ The current TOSA version does not support int64. However, int64 is commonly used
 - For quantized models, these transformations will be automatically handled during annotation before the export stage.
 
 List of model specific and optional passes:
-- InsertCastForOpsWithInt64InputPass
-    - Functionality:
-        - For LLMs such as LLama, some opeartors like aten.embedding have int64 input. In order to lower these operators to TOSA, this pass will insert a casting node that converts the input from int64 to int32.
-    - Supported Ops:
-        - aten.embedding.default, aten.slice_copy.Tensor
-    - Example usage:
-        - backends/arm/test/models/test_llama.py
-
 - ConvertInt64ConstOpsToInt32Pass
     - Functionalities:
       - Rewrites constant-producing ops that output int64 to instead output int32, when values are within int32 bounds.
@@ -244,3 +236,16 @@ List of model specific and optional passes:
     - Example usage:
       - (Functionality 1) backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py
       - (Functionality 2) backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py
+
+- InsertInt32CastsAfterInt64PlaceholdersPass
+    - Functionalities:
+      - Inserts an int64 -> int32 cast immediately after each int64 placeholder (graph input).
+      - Redirects all uses of each int64 placeholder to its int32 cast output.
+      - Inserts local int32 -> int64 casts at call sites where an operator requires int64 inputs, e.g. `torch.nn.functional.one_hot`
+    - Pass ordering:
+      - When used with `ConvertInt64ConstOpsToInt32Pass` and `ConvertInt64OutputOpsToInt32Pass`, run this pass last.
+      - Rationale: Those passes may cause retracing to re-infer some int64 placeholders as int32. Running this pass last casts only inputs that remain int64, minimizing inserted casts.
+    - Example usage:
+      - backends/arm/test/models/test_llama.py
+      - backends/arm/test/models/stable_diffusion/test_CLIPTextModelWithProjection.py
+      - backends/arm/test/models/stable_diffusion/test_T5EncoderModel.py

backends/arm/_passes/__init__.py

@@ -75,8 +75,8 @@
 from .fuse_constant_ops_pass import ComputeConstantOpsAOT, FuseConstantArgsPass  # noqa
 from .fuse_equal_placeholders_pass import FuseEqualPlaceholdersPass  # noqa
 from .fuse_quantized_activation_pass import FuseQuantizedActivationPass  # noqa
-from .insert_int64_input_cast_pass import (  # noqa  # noqa
-    InsertCastForOpsWithInt64InputPass,
+from .insert_int32_casts_after_int64_placeholders import (  # noqa
+    InsertInt32CastsAfterInt64PlaceholdersPass,
 )
 from .insert_rescales_pass import InsertRescalePass  # noqa
 from .insert_table_ops import InsertTableOpsPass  # noqa

backends/arm/_passes/arm_pass_manager.py

@@ -76,7 +76,7 @@
     FuseConstantArgsPass,
     FuseEqualPlaceholdersPass,
     FuseQuantizedActivationPass,
-    InsertCastForOpsWithInt64InputPass,
+    InsertInt32CastsAfterInt64PlaceholdersPass,
     InsertRescalePass,
     InsertTableOpsPass,
     MatchArgDtypePass,
@@ -277,7 +277,7 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         )  # ConvertInt64ConstOpsToInt32Pass requires this pass to remove the assertation in Graph
         self.add_pass(ConvertInt64ConstOpsToInt32Pass())
         self.add_pass(ConvertInt64OutputOpsToInt32Pass())
-        self.add_pass(InsertCastForOpsWithInt64InputPass())
+        self.add_pass(InsertInt32CastsAfterInt64PlaceholdersPass())
         self.add_pass(DecomposeEmbeddingPass())
         self.add_pass(DecomposeScaledDotProductAttention())
         self.add_pass(DecomposeRoundPass())

backends/arm/_passes/insert_int32_casts_after_int64_placeholders.py

@@ -0,0 +1,122 @@
+# Copyright 2025 Arm Limited and/or its affiliates.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+# pyre-unsafe
+
+
+import logging
+
+import torch
+from executorch.backends.arm._passes.arm_pass_utils import create_node
+from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.pass_base import EdgeOpOverload, ExportPass, PassResult
+from torch._subclasses.fake_tensor import FakeTensor
+
+
+logger = logging.getLogger(__name__)
+
+
+class InsertInt32CastsAfterInt64PlaceholdersPass(ExportPass):
+    """
+    Insert an int64->int32 cast after each int64 placeholder.
+
+    Note: Overflow checks are not applied in this pass. It is the user's responsibility to ensure that values fit within
+    the int32 range.
+    """
+
+    # Ops that require i64 inputs → positions of args to upcast.
+    # Key: op overload; Value: zero-based indices of positional args that must be i64.
+    I64_INPUT_ARG_POSITIONS = {
+        torch.ops.aten.one_hot.default: (0,),
+    }
+
+    def _insert_callsite_i32_to_i64_casts(self, graph_module: torch.fx.GraphModule):
+        """
+        If an operator requires int64 inputs but dtype propagation (via call_operator)
+        produced int32, insert a local int32→int64 cast at the call site to satisfy
+        PyTorch's operator input validation.
+        """
+        modified = False
+        graph = graph_module.graph
+        for node in graph.nodes:
+            if node.op != "call_function":
+                continue
+            if node.target not in self.I64_INPUT_ARG_POSITIONS:
+                continue
+
+            with graph.inserting_before(node):
+                arg_positions = self.I64_INPUT_ARG_POSITIONS.get(node.target)
+                args_list = list(node.args)
+                for pos in arg_positions:  # type: ignore[union-attr]
+                    input_arg = args_list[pos]
+                    to_copy_op = self._get_decomposition(graph)
+                    cast_node = graph_module.graph.create_node(
+                        "call_function",
+                        to_copy_op,
+                        (input_arg,),
+                        {"dtype": torch.int64},
+                    )
+                    cast_node.meta["val"] = node.meta["val"].to(torch.int64)
+                    args_list[pos] = cast_node
+                node.args = tuple(args_list)
+                modified = True
+        return modified
+
+    def _graph_uses_edge_ops(self, graph: torch.fx.Graph) -> bool:
+        for n in graph.nodes:
+            if n.op == "call_function":
+                if isinstance(n.target, EdgeOpOverload):
+                    return True
+        return False
+
+    def _get_decomposition(self, graph: torch.fx.Graph):
+        if self._graph_uses_edge_ops(graph):
+            return exir_ops.edge.dim_order_ops._to_dim_order_copy.default
+        else:
+            return torch.ops.dim_order_ops._to_dim_order_copy.default
+
+    def _is_tensor_of_dtype(self, node_val, dtype: torch.dtype) -> bool:
+        return isinstance(node_val, FakeTensor) and node_val.dtype == dtype
+
+    def _insert_placeholder_i64_to_i32_casts(self, graph_module: torch.fx.GraphModule):
+        modified = False
+        graph = graph_module.graph
+        for node in graph.nodes:
+            if node.op != "placeholder":
+                continue
+            node_val = node.meta["val"]
+            if not self._is_tensor_of_dtype(node_val, torch.int64):
+                continue
+
+            to_copy_op = self._get_decomposition(graph)
+            with graph.inserting_after(node):
+                cast_after = create_node(
+                    graph,
+                    to_copy_op,
+                    args=(node,),
+                    kwargs={
+                        "dtype": torch.int32,
+                    },
+                )
+                users = [user for user in node.users if user != cast_after]
+                for user in users:
+                    user.replace_input_with(node, cast_after)
+                logger.warning(
+                    f"Inserting a casting node {cast_after.name} after {node.name} to cast int64 placeholder"
+                    f" to int32 for {node.name} defined in {node.meta.get('stack_trace','[no stack trace found]')}"
+                )
+                modified = True
+        return modified
+
+    def call(self, graph_module: torch.fx.GraphModule):
+        modified = False
+        modified |= self._insert_placeholder_i64_to_i32_casts(graph_module)
+        modified |= self._insert_callsite_i32_to_i64_casts(graph_module)
+
+        if modified:
+            graph_module.graph.eliminate_dead_code()
+            graph_module.recompile()
+            graph_module = super().call(graph_module).graph_module
+        return PassResult(graph_module, modified)

backends/arm/_passes/insert_int64_input_cast_pass.py

@@ -73,7 +73,9 @@ def define_node(
             abs_output = output
 
         # Do the INT32 Abs
-        tosa_graph.addOperator(
+        self._serialize_operator(
+            node,
+            tosa_graph,
             ts.TosaOp.Op().ABS,
             [
                 rescaled_inputs[0].name,

backends/arm/operators/op_abs.py

@@ -67,7 +67,9 @@ def define_node(
             dtype=ts.DType.INT32,
         )
 
-        tosa_graph.addOperator(
+        self._serialize_operator(
+            node,
+            tosa_graph,
             ts.TosaOp.Op().REDUCE_SUM,
             [rescaled_inputs[0].name],
             [intermediate.name],
@@ -111,7 +113,9 @@ def define_node(
         attr = ts.TosaSerializerAttribute()
         attr.ReduceSumAttribute(tensor.dim_order.index(dim))
 
-        tosa_graph.addOperator(
+        self._serialize_operator(
+            node,
+            tosa_graph,
             ts.TosaOp.Op().REDUCE_SUM,
             [tensor.name],
             [output.name],

backends/arm/operators/op_sum.py

@@ -44,6 +44,7 @@ python_library(
         "//executorch/backends/arm:ethosu_partitioner",
         "//executorch/backends/arm/quantizer:lib",
         "//executorch/backends/arm/tosa:mapping",
+        "//executorch/backends/arm:vgf_partitioner",
         "//executorch/devtools/backend_debug:delegation_info",
         "//executorch/exir/backend:operator_support",
         "fbsource//third-party/pypi/tabulate:tabulate",

backends/arm/test/TARGETS

@@ -11,7 +11,7 @@
 from executorch.backends.arm._passes import (
     ConvertInt64ConstOpsToInt32Pass,
     ConvertInt64OutputOpsToInt32Pass,
-    InsertCastForOpsWithInt64InputPass,
+    InsertInt32CastsAfterInt64PlaceholdersPass,
 )
 
 from executorch.backends.arm.test import common
@@ -33,10 +33,9 @@ class TestCLIPTextModelWithProjection(unittest.TestCase):
     # for that is some assert ops are removed by passes in the
     # .to_executorch step, i.e. after Arm partitioner.
     ops_after_partitioner = {
-        "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 3,
-        "executorch_exir_dialects_edge__ops_aten_view_copy_default": 1,
         "executorch_exir_dialects_edge__ops_aten_argmax_default": 1,
-        "torch.ops.higher_order.executorch_call_delegate": 1,
+        "executorch_exir_dialects_edge__ops_dim_order_ops__to_dim_order_copy_default": 2,
+        "torch.ops.higher_order.executorch_call_delegate": 2,
     }
 
     def _prepare_inputs(
@@ -71,9 +70,9 @@ def test_CLIPTextModelWithProjection_tosa_FP(self):
                     example_inputs=text_encoder_model_inputs,
                     compile_spec=common.get_tosa_compile_spec(tosa_spec="TOSA-1.0+FP"),
                     transform_passes=[
-                        InsertCastForOpsWithInt64InputPass(),
                         ConvertInt64ConstOpsToInt32Pass(),
                         ConvertInt64OutputOpsToInt32Pass(),
+                        InsertInt32CastsAfterInt64PlaceholdersPass(),
                     ],
                 )
                 .export()