Arm backend: Keep INT delegate boundaries quantized

backends/arm/test/misc/test_custom_partition.py

@@ -1,4 +1,4 @@
-# Copyright 2025 Arm Limited and/or its affiliates.
+# Copyright 2025-2026 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.
@@ -8,12 +8,17 @@
 
 import torch
 from executorch.backends.arm.test import common
-from executorch.backends.arm.test.tester.test_pipeline import TosaPipelineFP
+from executorch.backends.arm.test.tester.test_pipeline import (
+    EthosU55PipelineINT,
+    TosaPipelineFP,
+)
+from executorch.backends.test.harness.stages import StageType
 from executorch.exir.backend.operator_support import (
     DontPartition,
     DontPartitionModule,
     DontPartitionName,
 )
+from executorch.exir.delegate import executorch_call_delegate
 from executorch.exir.dialects._ops import ops as exir_ops
 
 input_t1 = Tuple[torch.Tensor, torch.Tensor]  # Input x, y
@@ -45,6 +50,62 @@ def forward(self, x: torch.Tensor, y: torch.Tensor):
         return self.nested(a, b)
 
 
+# Reproduce the shared partition boundary used by YOLO. One split output stays
+# in portable code while the other can enter another delegate. This exposes a
+# bug where the first delegate claims its shared output is FP32 but writes
+# INT16 data.
+#
+#     conv -> split -> reshape -> BMM (portable)
+#                 `-> relu            (next delegate)
+class ConvSplitBMM(torch.nn.Module):
+
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(2, 4, 1)
+
+    def forward(self, x, y):
+        lhs, other = self.conv(x).split([2, 2], dim=1)
+        lhs = lhs.reshape(1, 2, 4)
+        return torch.bmm(lhs, y), torch.relu(other)
+
+
+@common.XfailIfNoCorstone300
+def test_a16w8_u55_INT_split_boundary():
+    """Keep INT-only delegate outputs quantized across split boundaries."""
+    # Keep this partition boundary stable if U55 gains INT16 BMM support.
+    reject_bmm = DontPartition(exir_ops.edge.aten.bmm.default)
+    pipeline = EthosU55PipelineINT(
+        ConvSplitBMM(),
+        (torch.rand(1, 2, 2, 2), torch.rand(1, 4, 3)),
+        # The operator-support docs treat these expected-op lists as support
+        # evidence. Do not list BMM: this test deliberately runs it portably.
+        ["torch.ops.aten.conv2d.default"],
+        [],
+        a16w8_quantization=True,
+    )
+    pipeline.tester.use_portable_ops = True
+    pipeline.change_args("to_edge_transform_and_lower", additional_checks=[reject_bmm])
+    pipeline.change_args(
+        "check_count.exir",
+        {"torch.ops.higher_order.executorch_call_delegate": 2},
+    )
+
+    def check_delegate_output_dtypes():
+        artifact = pipeline.tester.get_artifact(StageType.TO_EDGE_TRANSFORM_AND_LOWER)
+        for node in artifact.exported_program().graph.nodes:
+            if node.target == executorch_call_delegate:
+                assert all(
+                    not output.dtype.is_floating_point for output in node.meta["val"]
+                )
+
+    pipeline.add_stage_after(
+        "check_count.exir", check_delegate_output_dtypes, suffix="output_dtypes"
+    )
+    pipeline.run()
+
+    assert reject_bmm.has_rejected_node()
+
+
 @common.parametrize("test_data", CustomPartitioning.inputs)
 def test_single_reject_tosa_FP(caplog, test_data: input_t1):
     caplog.set_level(logging.INFO)

backends/arm/test/setup_testing.sh

@@ -71,6 +71,8 @@ ops_list_u55=(
     aten::arange.start_out
     aten::eq.Tensor_out
     aten::logical_not.out
+    aten::split_with_sizes_copy.out
+    aten::bmm.out
     dim_order_ops::_clone_dim_order.out
     dim_order_ops::_to_dim_order_copy.out
     "${ops_list_quantized_decomposed[@]}"

backends/arm/tosa/partitioner.py

@@ -499,51 +499,76 @@ def _detag_boundary_nodes(
 
         Remove delegation tags from quantize nodes with inputs outside the
         partition and from dequantize nodes with outputs outside the partition.
+        This applies to all variants in ``Q_OPS`` and ``DQ_OPS``, independent
+        of their integer dtype.
 
-        For non Q/DQ nodes, remove the tag from the first node in the partition
-        if any input has floating-point dtype.
+        For INT-only partitions, also remove floating-point nodes at input and
+        output boundaries. Repeat until removing one node no longer exposes
+        another invalid boundary node.
 
         Args:
             tag: The delegation tag assigned to the partition.
             reporter: A reporter to log rejected nodes.
             module: The GraphModule containing the partition.
-            detag_first_fp_node: Whether to de-tag the first floating-point
-                node in a partition.
+            detag_first_fp_node: Whether to de-tag floating-point nodes at the
+                input and output boundaries of a partition.
 
         """
-        # De-tag outermost q-nodes upwards and dq-nodes downwards.
-        # De-tag if at least one input/output is not part of the partition.
-        for node in module.graph.nodes:
-            if not is_partitioned(node, tag):
-                continue
+        # Q_OPS and DQ_OPS cover the supported conversion variants. Their
+        # parameters determine the integer dtype, so this handles INT8 and
+        # INT16 graphs.
+        #
+        # Keep conversions and floating-point operations outside INT-only
+        # delegates:
+        #
+        #     integer delegate -> portable [DQ -> floating-point operations]
+        #
+        # Removing one boundary node can expose another Q/DQ or floating-point
+        # node, so repeat until the partition has a valid boundary.
+        modified = True
+        while modified:
+            modified = False
+            for node in module.graph.nodes:
+                if not is_partitioned(node, tag):
+                    continue
 
-            is_q_node = node.target in Q_OPS
-            is_dq_node = node.target in DQ_OPS
-            is_boundary_q_node = is_q_node and not is_partitioned(
-                node.all_input_nodes[0], tag
-            )
-            is_boundary_dq_node = is_dq_node and any(
-                not is_partitioned(user, tag) for user in node.users
-            )
+                is_q_node = node.target in Q_OPS
+                is_dq_node = node.target in DQ_OPS
+                is_boundary_q_node = is_q_node and not is_partitioned(
+                    node.all_input_nodes[0], tag
+                )
+                has_external_user = any(
+                    not is_partitioned(user, tag) for user in node.users
+                )
+                is_boundary_dq_node = is_dq_node and has_external_user
+                is_boundary_fp_output = (
+                    detag_first_fp_node
+                    and not is_q_node
+                    and not is_dq_node
+                    and has_external_user
+                    and get_first_fake_tensor(node).dtype.is_floating_point
+                )
 
-            if is_boundary_q_node or is_boundary_dq_node:
-                # Remove tag from quantize node with input outside partition,
-                # or dequantize node with any output outside partition
-                del node.meta["delegation_tag"]
-            elif detag_first_fp_node and not is_q_node and not is_dq_node:
-                # For non Q/DQ nodes, remove tag from first node in partition if any input has fp dtype
-                for input in node.all_input_nodes:
-                    if is_partitioned(input, tag) or isinstance(
-                        input.meta["val"], torch.SymInt
-                    ):
-                        continue
-                    if get_first_fake_tensor(input).dtype.is_floating_point:
-                        reporter.report_reject(
-                            node,
-                            f"Was first node in partition and input {input.name} had fp dtype.",
-                        )
-                        del node.meta["delegation_tag"]
-                        break
+                if is_boundary_q_node or is_boundary_dq_node or is_boundary_fp_output:
+                    del node.meta["delegation_tag"]
+                    modified = True
+                    continue
+
+                if detag_first_fp_node and not is_q_node and not is_dq_node:
+                    # Remove the first floating-point node at an input boundary.
+                    for input in node.all_input_nodes:
+                        if is_partitioned(input, tag) or isinstance(
+                            input.meta["val"], torch.SymInt
+                        ):
+                            continue
+                        if get_first_fake_tensor(input).dtype.is_floating_point:
+                            reporter.report_reject(
+                                node,
+                                f"Was first node in partition and input {input.name} had fp dtype.",
+                            )
+                            del node.meta["delegation_tag"]
+                            modified = True
+                            break
 
     def _preserve_io_quantization_enabled(self) -> bool:
         """Return True if compile specs preserve IO quantization."""