From 96417259f07cb753fcb6d1b0e30fdc6abe781303 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Mateusz=20S=C5=82uszniak?= <mateusz.sluszniak@swmansion.com>
Date: Wed, 21 Oct 2026 15:47:42 +0200
Subject: [PATCH] Qualcomm AI Engine Direct - Keep the selected element when
 removing a 0D select

Remove0DTensor replaced a 0D select with its whole source, which is only right for a
1-element source: y + x[1] became y + x, and torch.full(shape, t[i]) became a copy of
all of t (lowering then failed with "expand: ... length 3 -> 1"). A 1-element slice of
the selected element replaces it instead.
---
 backends/qualcomm/_passes/remove_0d_tensor.py | 19 ++++++++++++++++++-
 backends/qualcomm/tests/models.py             |  9 +++++++++
 backends/qualcomm/tests/rework/src/pattern.py |  2 ++
 backends/qualcomm/tests/test_qnn_delegate.py  | 15 +++++++++++++++
 4 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/backends/qualcomm/_passes/remove_0d_tensor.py b/backends/qualcomm/_passes/remove_0d_tensor.py
index 4056910f6ef..dc99dd598c5 100644
--- a/backends/qualcomm/_passes/remove_0d_tensor.py
+++ b/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)
diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py
index 14aff8e7fbb..fb96a09c8a6 100644
--- a/backends/qualcomm/tests/models.py
+++ b/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__()
diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py
index 7b90fde61e2..90f1f8fbc0e 100644
--- a/backends/qualcomm/tests/rework/src/pattern.py
+++ b/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:
diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py
index b5ae1d19d7d..2629280d3f7 100644
--- a/backends/qualcomm/tests/test_qnn_delegate.py
+++ b/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 = [
             {