Qualcomm AI Engine Direct - Cast TopK int64 indices to int32

backends/qualcomm/_passes/i64_to_i32.py

@@ -34,6 +34,10 @@ class I64toI32(ExportPass):
         exir_ops.edge.aten.scalar_tensor.default,
         exir_ops.edge.dim_order_ops._to_dim_order_copy.default,
     }
+    # Multi-output ops whose int64 outputs (indices) QNN produces as int32
+    I64_MULTI_OUTPUT_OPS = {
+        exir_ops.edge.aten.topk.default,
+    }
     # This dict is to ensure that the input of the OPs are int64 due to Pytorch restrictions.
     # For example, scatter op can only accept args[2], the index, as int64.
     # Key: Ops to cast input to i64
@@ -68,6 +72,16 @@ def call_operator(self, op, args, kwargs, meta):
             )
         return super().call_operator(op, args, kwargs, meta)
 
+    def call_getitem(self, value, key, meta):
+        res = super().call_getitem(value, key, meta)
+        if value.node.target in self.I64_MULTI_OUTPUT_OPS and self._is_tensor_of_dtype(
+            meta["val"], torch.int64
+        ):
+            return super().call_operator(
+                self.copy_op, (res,), {"dtype": torch.int32}, meta
+            )
+        return res
+
     def _record_original_output_dtype(self, graph_module: torch.fx.GraphModule):
         for n in graph_module.graph.nodes:
             # Keep track of original output dtype so we ensure the dtype of the graph is consistent with nn.Module

backends/qualcomm/tests/models.py

@@ -2997,6 +2997,11 @@ def forward(self, x):
         return a + self.idx_source[b]
 
 
+class TopKIndices(torch.nn.Module):
+    def forward(self, x):
+        return torch.topk(x, 3)[1]
+
+
 class Triu(torch.nn.Module):
     def __init__(self, diagonal: Optional[int] = None):
         super().__init__()

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

@@ -4958,8 +4958,9 @@ def test(
             "getitem": 1 + offset,
             "getitem_1": 1 + offset,
             "aten_view_copy_default": 2 + offset,
-            "aten_index_tensor": 3 + offset,
-            "aten_add_tensor": 4 + offset,
+            "aten__to_copy_default": 3 + offset,
+            "aten_index_tensor": 4 + offset,
+            "aten_add_tensor": 5 + offset,
         }
         if quantizer is not None:
             name_handle_map["quantized_decomposed_dequantize_per_tensor_default"] = 1

backends/qualcomm/tests/test_passes.py

@@ -620,8 +620,9 @@ def test_resolve_debug_handle(self):
             "getitem": 1,
             "getitem_1": 1,
             "aten_view_copy_default": 2,
-            "aten_index_tensor": 3,
-            "aten_add_tensor": 4,
+            "aten__to_copy_default": 3,
+            "aten_index_tensor": 4,
+            "aten_add_tensor": 5,
         }
         module = TopKandIndex()  # noqa: F405
         sample_input = (torch.randn(3, 10),)

backends/qualcomm/tests/test_qnn_delegate.py

@@ -3068,6 +3068,19 @@ def test_qnn_backend_topk_and_index(self):
         sample_input = (torch.randn(3, 10),)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_topk_indices(self):
+        module = TopKIndices()  # noqa: F405
+        sample_input = (torch.randn(3, 10),)
+        self.lower_module_and_test_output(module, sample_input)
+
+    def test_qnn_backend_topk_on_cpu_and_index(self):
+        # int64 indices from a TopK left on CPU feed the delegated index
+        module = TopKandIndex()  # noqa: F405
+        sample_input = (torch.randn(3, 10),)
+        self.lower_module_and_test_output(
+            module, sample_input, skip_node_op_set={"aten.topk.default"}
+        )
+
     def test_qnn_backend_view_permute_matmul(self):
         module = ViewPermuteMatMul()  # noqa: F405
         torch.manual_seed(8)