Qualcomm: decide graph-output publishing per output, not per node

backends/qualcomm/builders/node_visitor.py

@@ -4,7 +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.
 
-from typing import Any, Dict, Tuple
+from typing import Any, Dict, Optional, Tuple
 
 import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager
 
@@ -382,11 +382,12 @@ def get_tensor_type(
         self,
         node: torch.fx.Node,
         tensor_type: PyQnnManager.Qnn_TensorType_t,
+        wrapper_idx: Optional[int] = None,
     ) -> PyQnnManager.Qnn_TensorType_t:
         is_input = is_graph_input(node, self.edge_program) or is_mutable_buffer_input(
             node, self.edge_program
         )
-        is_output = is_graph_output(node)
+        is_output = is_graph_output(node, wrapper_idx)
         # handle logic for input/output tensors
         if is_input or is_output:
             # For more info about existence of self.is_qnn_partitioner, check constructor for explanation.
@@ -458,7 +459,7 @@ def get_tensor_name(
                     self.edge_program.graph_signature.buffers_to_mutate.keys()
                 ).index(node.name)
                 tensor_name = f"output_mutbuf_{position_index}_{tensor_name}"
-            elif is_graph_output(node):
+            elif is_graph_output(node, wrapper_idx):
                 tensor_name = f"output_{tensor_name}"
 
         # Only add qcom_tensor_name when enable tensor dump.
@@ -533,7 +534,7 @@ def define_tensor(
         tensor_name = self.get_tensor_name(tensor_source_node, wrapper_idx)
         dims = torch.Size([1]) if len(tensor.size()) == 0 else tensor.size()
         dynamic_dims, nominal_dims = self.get_dynamic_dimension(dims)
-        tensor_type = self.get_tensor_type(tensor_source_node, tensor_type)
+        tensor_type = self.get_tensor_type(tensor_source_node, tensor_type, wrapper_idx)
         quant_encoding, quant_configs = self.get_quant_encoding_conf(
             tensor_source_node, target_build_node
         )

backends/qualcomm/builders/utils.py

@@ -98,19 +98,27 @@ def is_mutable_buffer_input(
         return fqn in edge_program.graph_signature.buffers_to_mutate.values()
 
 
-def is_graph_output(node: torch.fx.Node) -> bool:
+def is_graph_output(node: torch.fx.Node, output_index: Optional[int] = None) -> bool:
     """
     Check if the given tensor is used as a graph output
 
     Args:
-        tensor: EdgeIR Tensor that is being checked for graph input
+        node: EdgeIR Tensor that is being checked for graph output
+        output_index: for a multi-output node, restrict the check to this
+            output. Without it a single escaping output would mark every
+            output of the node as a graph output, publishing values that
+            nothing consumes.
     """
     for user in node.users.keys():
-        # getitem node is skipped, check the op_skip_ops.py
-        if user.op == "output" or (
-            user.target.__name__ == "getitem" and is_graph_output(user)
-        ):
+        if user.op == "output":
             return True
+        # getitem node is skipped, check the op_skip_ops.py
+        # call_module targets are plain strings and have no __name__
+        if getattr(user.target, "__name__", "") == "getitem":
+            if output_index is not None and user.args[1] != output_index:
+                continue
+            if is_graph_output(user):
+                return True
     return False
 
 

backends/qualcomm/tests/test_passes.py

@@ -847,6 +847,94 @@ def test_backend_bundle_cache_survives_an_expired_entry(self):
         self.assertEqual(2, created, f"expected 2 backend creations, got {created}")
         self.assertGreaterEqual(reused, 1, "third manager must reuse the live bundle")
 
+    def test_is_graph_output_tolerates_users_without_a_name(self):
+        """call_module targets are plain strings; reading __name__ raised."""
+        from executorch.backends.qualcomm.builders.utils import is_graph_output
+
+        graph = torch.fx.Graph()
+        placeholder = graph.placeholder("x")
+        # a call_module user, whose target is a str and so has no __name__
+        graph.create_node("call_module", "_guards_fn", (placeholder,))
+        graph.output((placeholder,))
+
+        self.assertTrue(any(u.op == "call_module" for u in placeholder.users))
+        self.assertTrue(is_graph_output(placeholder))
+        self.assertTrue(is_graph_output(placeholder, 0))
+
+    def test_is_graph_output_is_per_output_not_per_node(self):
+        """Only the outputs a graph actually consumes may be published.
+
+        A multi-output node used to be treated as a whole: one escaping output
+        marked every output of that node as a graph output, so QNN published
+        values nothing reads. The published set then no longer matched the one
+        ExecuTorch bound, and the two are paired by position.
+        """
+        from executorch.backends.qualcomm.builders.utils import is_graph_output
+
+        class OnlyValues(torch.nn.Module):
+            def forward(self, x):
+                # max.dim returns (values, indices); only values is returned.
+                return torch.max(x, dim=1).values
+
+        gm = torch.export.export(
+            OnlyValues().eval(), (torch.randn(2, 3),), strict=True
+        ).module()
+        node = next(
+            n
+            for n in gm.graph.nodes
+            if n.op == "call_function" and "max" in str(n.target)
+        )
+        getitems = {
+            u.args[1]: u
+            for u in node.users
+            if getattr(u.target, "__name__", "") == "getitem"
+        }
+        self.assertEqual({0, 1}, set(getitems), "both outputs should be unpacked")
+        self.assertTrue(getitems[0].users, "values reaches the graph output")
+        self.assertFalse(getitems[1].users, "indices is dead")
+
+        self.assertTrue(is_graph_output(node), "node does reach a graph output")
+        self.assertTrue(is_graph_output(node, 0), "values is consumed")
+        self.assertFalse(
+            is_graph_output(node, 1),
+            "indices has no consumer and must not be published",
+        )
+
+    def test_unconsumed_output_is_native_not_app_read(self):
+        """The tensor type is what actually reaches QNN, so assert on it.
+
+        An unconsumed output published as APP_READ makes the QNN graph declare
+        more outputs than ExecuTorch binds, and the two are paired by position.
+        """
+        import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager
+        from executorch.backends.qualcomm.builders.node_visitor import NodeVisitor
+
+        class OnlyValues(torch.nn.Module):
+            def forward(self, x):
+                return torch.max(x, dim=1).values
+
+        edge_program = torch.export.export(
+            OnlyValues().eval(), (torch.randn(2, 3),), strict=True
+        )
+        node = next(
+            n
+            for n in edge_program.graph_module.graph.nodes
+            if n.op == "call_function" and "max" in str(n.target)
+        )
+        visitor = NodeVisitor({node: 0}, edge_program, enable_tensor_dump=False)
+
+        native = PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE
+        self.assertEqual(
+            PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_APP_READ,
+            visitor.get_tensor_type(node, native, 0),
+            "values is consumed and must be published",
+        )
+        self.assertEqual(
+            native,
+            visitor.get_tensor_type(node, native, 1),
+            "indices has no consumer and must stay internal",
+        )
+
 
 if __name__ == "__main__":
     unittest.main()