diff --git a/backends/qualcomm/_passes/i64_to_i32.py b/backends/qualcomm/_passes/i64_to_i32.py index a26956b5569..bcd0bb4cc25 100644 --- a/backends/qualcomm/_passes/i64_to_i32.py +++ b/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 diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py index 14aff8e7fbb..e4fec5cd0bd 100644 --- a/backends/qualcomm/tests/models.py +++ b/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__() diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 7b90fde61e2..28158e25956 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/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 diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 606fed01d72..8236028eaa2 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/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),) diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py index b5ae1d19d7d..2500b4066d9 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/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)