Qualcomm AI Engine Direct - Keep the selected element when removing a 0D select

backends/qualcomm/_passes/remove_0d_tensor.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import copy
+
 import torch
 from executorch.exir.dialects._ops import ops as exir_ops
 from executorch.exir.pass_base import ExportPass, PassResult
@@ -28,8 +30,23 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
         for node in graph.nodes:
             if node.target in self.remove_ops and len(node.meta["val"].shape) == 0:
+                source, dim, index = node.args
+                replacement = source
+                size = source.meta["val"].shape[dim]
+                if size != 1:
+                    # Keep the selected element as a 1-element slice; the source
+                    # itself is only equivalent when it holds a single element.
+                    index %= size
+                    with graph.inserting_before(node):
+                        replacement = graph.create_node(
+                            "call_function",
+                            exir_ops.edge.aten.slice_copy.Tensor,
+                            (source, dim, index, index + 1),
+                        )
+                    replacement.meta = copy.copy(node.meta)
+                    replacement.meta["val"] = source.meta["val"].narrow(dim, index, 1)
                 for user_n in list(node.users.keys()):
-                    user_n.replace_input_with(node, node.args[0])
+                    user_n.replace_input_with(node, replacement)
                 graph.erase_node(node)
 
         dead_code_elimination_pass(graph_module)

backends/qualcomm/tests/models.py

@@ -2763,6 +2763,15 @@ def forward(self, x, y):
         )
 
 
+class SelectElement(torch.nn.Module):
+    def __init__(self, index):
+        super().__init__()
+        self.index = index
+
+    def forward(self, x, y):
+        return y + x[self.index]
+
+
 class SelectScatter(torch.nn.Module):
     def __init__(self, dim, index):
         super().__init__()

backends/qualcomm/tests/rework/src/pattern.py

@@ -5012,6 +5012,8 @@ def test(
 
         assertions.assert_no_target(gm, exir_ops.edge.aten.select_copy.int)
         assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1)
+        # x has 4 elements, so the selected one must be kept, not x itself
+        assertions.assert_target_count(gm, exir_ops.edge.aten.slice_copy.Tensor, 1)
 
 
 class SeqMSE:

backends/qualcomm/tests/test_qnn_delegate.py

@@ -2493,6 +2493,13 @@ def test_qnn_backend_select_copy(self):
         sample_input = (torch.randn([1, 3, 3, 3]),)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_select_element(self):
+        sample_input = (torch.randn(4), torch.randn(3, 4))
+        for i, index in enumerate([1, -1]):
+            with self.subTest(i=i):
+                module = SelectElement(index)  # noqa: F405
+                self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_select_scatter(self):
         test_comb = [
             {
@@ -6110,6 +6117,14 @@ def test_qnn_backend_select_copy(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_select_element(self):
+        sample_input = (torch.randn(4), torch.randn(3, 4))
+        for i, index in enumerate([1, -1]):
+            with self.subTest(i=i):
+                module = SelectElement(index)  # noqa: F405
+                module = self.get_qdq_module(module, sample_input)
+                self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_select_scatter(self):
         test_comb = [
             {