Claim getitem when its producer is claimed

Summary:
NativeSupportedOperators now claims operator.getitem, so a supported
multi-output op and the getitems that unpack it stay in one partition. Without
this every getitem is a partition boundary, and fused-projection models (chunked
gate_up / qkv) fragment into many delegates.

The claim is conditional on the producer also being claimed. partition() passes
allows_single_node_partition=True, which disables the capability partitioner's
filter that would otherwise drop getitem-only partitions -- getitem is in its
default non-compute set, so such a partition normally scores zero compute nodes
and is deleted (see torch/fx/passes/infra/partitioner.py). With that filter off,
claiming getitem unconditionally can emit a delegate that holds no computation
and is fed by a tuple crossing the delegate boundary, where boundaries carry
tensors. Backends that claim getitem unconditionally (Qualcomm, OpenVINO) leave
the filter enabled and rely on it.

Reviewed By: digantdesai

Differential Revision: D117907800

Pull Request resolved: pytorch#22974

Author: metascroy

backends/native/partitioner.py

@@ -13,6 +13,7 @@
 ExecuTorch forbids a partitioner from mutating the graph module.
 """
 
+import operator
 from typing import Callable, final, List, Mapping, Optional, Tuple
 
 # Registers torch.ops.torchao.dequantize_gguf, referenced in _SUPPORTED_NON_CORE_OPS.
@@ -71,6 +72,23 @@ def is_node_supported(
             return False
         if node.op != "call_function":
             return False
+        # getitem only unpacks a multi-output node's result (e.g. split/chunk); it
+        # is not a real op. Claim it so a supported multi-output op and its
+        # unpackers stay in one partition -- otherwise every getitem is a partition
+        # boundary and fused-projection models (chunked gate_up / qkv) fragment into
+        # many delegates.
+        #
+        # Only claim it when the producer is claimed too. partition() passes
+        # allows_single_node_partition=True, which disables the capability
+        # partitioner's filter that would otherwise drop getitem-only partitions
+        # (getitem counts as non-compute there). Claiming unconditionally can
+        # therefore emit a delegate holding no computation, fed by a tuple crossing
+        # the delegate boundary.
+        if node.target is operator.getitem:
+            producer = node.args[0] if node.args else None
+            return isinstance(producer, Node) and self.is_node_supported(
+                submodules, producer
+            )
         if isinstance(node.target, torch._ops.HigherOrderOperator):
             return False
 

backends/native/test/test_partitioner.py

@@ -4,6 +4,7 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import operator
 import unittest
 from unittest.mock import MagicMock
 
@@ -32,10 +33,13 @@
 )
 
 
-def _make_node(op, target):
-    node = MagicMock()
+def _make_node(op, target, args=()):
+    # spec=Node so isinstance(node, Node) holds; the getitem support check
+    # resolves its producer through node.args and requires a real Node.
+    node = MagicMock(spec=torch.fx.Node)
     node.op = op
     node.target = target
+    node.args = args
     return node
 
 
@@ -88,6 +92,34 @@ def test_accepts_view_copy_tagged_op(self):
         self.assertIn(torch.Tag.view_copy, op.tags)
         self.assertTrue(self.sup.is_node_supported({}, _make_node("call_function", op)))
 
+    def test_accepts_getitem_on_supported_producer(self):
+        # getitem unpacks a multi-output result; it must be claimed alongside a
+        # producer the backend claims, so the two stay in one partition.
+        producer = _make_node("call_function", torch.ops.aten.split_copy.Tensor)
+        self.assertTrue(self.sup.is_node_supported({}, producer))
+        getitem = _make_node("call_function", operator.getitem, args=(producer, 0))
+        self.assertTrue(self.sup.is_node_supported({}, getitem))
+
+    def test_rejects_getitem_on_unsupported_producer(self):
+        # partition() sets allows_single_node_partition=True, which disables the
+        # capability partitioner's getitem-only partition filter. Claiming a
+        # getitem whose producer is not claimed would strand it in a delegate
+        # containing no computation, fed by a tuple crossing the boundary.
+        producer = _make_node(
+            "call_function", torch.ops.aten.linalg_solve_triangular.default
+        )
+        self.assertFalse(self.sup.is_node_supported({}, producer))
+        getitem = _make_node("call_function", operator.getitem, args=(producer, 0))
+        self.assertFalse(self.sup.is_node_supported({}, getitem))
+
+    def test_rejects_getitem_without_node_producer(self):
+        # A getitem whose first arg is not a Node (or which has no args) has no
+        # producer to validate, so it is not claimed.
+        for args in ((), (None, 0), ("not_a_node", 0)):
+            with self.subTest(args=args):
+                node = _make_node("call_function", operator.getitem, args=args)
+                self.assertFalse(self.sup.is_node_supported({}, node))
+
     def test_rejects_non_core_unsupported_op(self):
         op = torch.ops.aten.linalg_solve_triangular.default
         self.assertNotIn(torch.Tag.core, op.tags)