Qualcomm: backport remaining YOLO26 QNN fixes

backends/qualcomm/_passes/annotate_unbind.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.
 import torch
-from executorch.backends.qualcomm.builders.node_visitor import dq_ops
+from executorch.backends.qualcomm.builders.node_visitor import dq_ops, q_ops
 from executorch.backends.qualcomm.utils.constants import QCOM_QUANT_ATTRS
 from executorch.exir.pass_base import ExportPass, PassResult
 from torch.fx.passes.utils.source_matcher_utils import get_source_partitions
@@ -30,8 +30,12 @@ def _annotate_unbind(self, graph_module: torch.fx.GraphModule):
         )
         for src_partitions in partitions.values():
             for src_partition in src_partitions:
-                if src_partition.input_nodes[0].target in dq_ops:
-                    q_node = src_partition.input_nodes[0].args[0]
+                dq_node = src_partition.input_nodes[0]
+                if dq_node.target in dq_ops:
+                    # a folded quantized constant feeds its dequantize directly
+                    q_node = dq_node.args[0]
+                    if q_node.target not in q_ops:
+                        q_node = dq_node
                     quant_attrs = get_quant_attrs(self.edge_program, q_node)
                     for n in src_partition.nodes:
                         n.meta[QCOM_QUANT_ATTRS] = quant_attrs.copy()

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/builders/op_to.py

@@ -12,7 +12,7 @@
 
 from .node_visitor import NodeVisitor, QNN_TENSOR_TYPE_MAP
 from .node_visitor_manager import register_node_visitor
-from .qnn_constants import OpCast, OpConvert, QNN_OP_PACKAGE_NAME_QTI_AISW
+from .qnn_constants import OpCast, OpConvert, OpQuantize, QNN_OP_PACKAGE_NAME_QTI_AISW
 
 
 @register_node_visitor
@@ -138,6 +138,36 @@ def define_node(
             nodes_to_wrappers,
         )
 
+        # Cast copies an integer's raw value into a quantized tensor, and the HTP
+        # rejects Convert there: cast to float, then quantize.
+        if (
+            node.meta.get(QCOM_QUANT_ATTRS)
+            and not input_node.meta.get(QCOM_QUANT_ATTRS)
+            and not input_node.meta["val"].is_floating_point()
+        ):
+            float_tensor_wrapper = self.define_custom_tensor_wrapper(
+                node_name=node.name + "_float",
+                tensor_type=PyQnnManager.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+                dtype=PyQnnManager.Qnn_DataType_t.QNN_DATATYPE_FLOAT_32,
+                quant_encoding=PyQnnManager.Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_UNDEFINED,
+                quant_configs={},
+                dims=output_tensor.size(),
+                tensor=output_tensor,
+                is_fake_tensor=True,
+                nodes_to_wrappers=nodes_to_wrappers,
+            )
+            cast_op = PyQnnManager.PyQnnOpWrapper(
+                f"{node.name}_float", QNN_OP_PACKAGE_NAME_QTI_AISW, OpCast.op_name
+            )
+            cast_op.AddInputTensors(node_input_tensors)
+            cast_op.AddOutputTensors([float_tensor_wrapper])
+            quantize_op = PyQnnManager.PyQnnOpWrapper(
+                node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, OpQuantize.op_name
+            )
+            quantize_op.AddInputTensors([float_tensor_wrapper])
+            quantize_op.AddOutputTensors([output_tensor_wrapper])
+            return ops + [cast_op, quantize_op]
+
         qnn_op = OpCast if self.is_cast_node(node) else OpConvert
         op = PyQnnManager.PyQnnOpWrapper(
             node.name, QNN_OP_PACKAGE_NAME_QTI_AISW, qnn_op.op_name

backends/qualcomm/tests/models.py

@@ -378,6 +378,14 @@ def forward(self, x):
         return x.type(torch.IntTensor)
 
 
+class CastIntToFloat(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        return y + (x + 1).float()
+
+
 class CastMultiUsers(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -2580,6 +2588,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/test_passes.py

@@ -709,6 +709,41 @@ def forward(self, x):
             raise
         edge.to_executorch()
 
+    def test_unbind_of_quantized_constant_lowers(self):
+        """AnnotateUnbind read the quant attrs from the quantize node feeding unbind's
+        dequantize, but a folded quantized constant feeds the dequantize directly.
+        """
+
+        class UnbindConstant(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.table = torch.tensor([[1.0, 2.0, 3.0, 4.0], [5.0, 6.0, 7.0, 8.0]])
+
+            def forward(self, x):
+                a, b = torch.unbind(self.table, dim=0)
+                return x * a + b
+
+        sample_input = (torch.randn(4, 4),)
+        quantizer = QnnQuantizer()
+        quantizer.set_default_quant_config(quant_dtype=QuantDtype.use_16a8w)
+        prepared = prepare_pt2e(
+            torch.export.export(UnbindConstant(), sample_input).module(), quantizer
+        )
+        prepared(*sample_input)
+        compiler_specs = generate_qnn_executorch_compiler_spec(
+            soc_model=QcomChipset.SM8650,
+            backend_options=generate_htp_compiler_spec(use_fp16=False),
+        )
+        try:
+            edge = to_edge_transform_and_lower_to_qnn(
+                convert_pt2e(prepared), sample_input, compiler_specs
+            )
+        except RuntimeError as e:
+            if "QNN" in str(e) or "qnn" in str(e):
+                self.skipTest(f"QNN SDK not available: {e}")
+            raise
+        edge.to_executorch()
+
     def test_constant_only_partition_stays_on_cpu(self):
         """An int64 buffer gets an int32 cast from I64toI32. When the buffer's consumer is
         not delegated, that cast alone would form a partition with no graph input, which

backends/qualcomm/tests/test_qnn_delegate.py

@@ -445,10 +445,11 @@ def test_qnn_backend_bmm(self):
         self.lower_module_and_test_output(module, sample_input)
 
     def test_qnn_backend_cast(self):
-        modules = [Cast(), CastMultiUsers()]  # noqa: F405
+        modules = [Cast(), CastMultiUsers(), CastIntToFloat()]  # noqa: F405
         sample_inputs = [
             (10 * torch.rand((9, 4, 5, 3)),),
             (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3)),
+            (torch.randint(0, 80, size=(3, 3)), torch.randn(3, 3)),
         ]
         for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)):
             with self.subTest(i=i):
@@ -2246,6 +2247,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 = [
             {
@@ -3350,6 +3358,14 @@ def test_qnn_backend_cast(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_cast_int_to_float(self):
+        module = CastIntToFloat()  # noqa: F405
+        sample_input = (torch.randint(0, 80, size=(3, 3)), torch.randn(3, 3))
+        module = self.get_qdq_module(module, sample_input)
+        # absolute: the class rtol of 1 would accept an all-zero output
+        self.atol, self.rtol = 1, 0
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_cat(self):
         modules = [Cat2(), Cat3(), Cat4(), Cat5()]  # noqa: F405
         sample_input = (torch.randn(1, 1, 2, 2), torch.randn(1, 1, 4, 2))
@@ -5385,6 +5401,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 = [
             {