From dc2bffc8d8e8cc9b072fa49f2d81bb0787451f27 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 21:28:41 -0800
Subject: [PATCH 1/6] Qualcomm AI Engine Direct - Decompose kept-whole ops that
 HTP leaves on CPU

On HTP, a partition that reads nothing at runtime stays on the CPU,
because HTP refuses a graph without runtime inputs. That is decided after
decomposition. So a constant-only op that QNN keeps whole, such as
linear, could land in such a partition undecomposed, and lowering failed
with "was not decomposed or delegated".

The HTP filter of ops_to_not_decompose now finds, before the kept-whole
ops are decomposed, the ops that this rule would leave on the CPU, and
lets exir decompose the kept-whole ones among them. Ops in partitions
that stay delegated keep their whole form, and with it their
quantization encodings.

Running the partitioner here would take time that grows with the square
of the graph size: on a model with about 43,000 nodes, export sat in its
cycle search for more than 2 hours. So the estimate makes one pass over
the graph in order and one in reverse. Ops whose data types or shapes
rule QNN out, ops without a tensor output, such as the runtime asserts
export adds, and higher-order ops, which have no node visitor, count as
CPU ops. Every getitem counts as a QNN op, as in QNN's op check. A QNN
op that reads no runtime value is left on the CPU when:
- it reaches a QNN op through a CPU op before any op that reads a
  runtime value, since the partitioner may split it off to avoid a
  cycle;
- only getitems of CPU ops read a runtime value, since the partitioner
  moves those getitems out of their partitions;
- such a getitem follows a QNN op, and no chain of QNN ops links the op
  to an op that reads a runtime value.

Quantize and dequantize ops, and ops that read a dequantize, are never
decomposed, in case the estimate is wrong. Their pieces would have no
quantization encodings, so QNN would reject them where a partition would
take them, and export would then fail without the quantized CPU kernels.

The caller's skips are left out on purpose. Decomposing adds nodes, and
the new nodes take names that later nodes had before, so a skip that
changed what is decomposed could rename the node that a skip_node_id_set
entry means. Ops that QNN rejects or removes later are not known here
either. So a kept-whole op cut off only by those still fails as before,
as does one that exir has no decomposition for, such as elu. The
estimate errs toward decomposing: an op can be decomposed although a
partition that reads a runtime value takes it, in pieces.

exir calls ops_to_not_decompose several times on the ATen program and
uses only one of the filters, so the estimate runs once, on the filter's
first use. That is before exir renames the ops it keeps whole, which
would hide the dequantize ops. On the lowered program the filter stays
plain filter_fn, so exir still reports kept-whole ops left on the CPU.

The partitioner logs the names of the kept-whole ops the estimate
decomposes. Since decomposing renames later nodes, that is the trace to
look for when a skip_node_id_set entry stops matching, so the log is a
warning when skip_node_id_set is not empty.

New tests stub the QNN SDK setup, the QNN manager and the compile:
- test_htp_decomposes_kept_whole_ops_left_on_cpu: constant-only linears
  in each of the cases above are decomposed;
- the filter decomposes, and logs, a constant-only linear that reaches
  an op that reads the input only through a torch.cond, or only through
  an item;
- test_htp_keeps_linked_kept_whole_ops_whole: linears linked to the
  input, through a masked select, a mutated buffer or shared values,
  stay whole and delegated;
- test_htp_keeps_quantized_ops_whole: a quantized linear behind bfloat16
  casts stays whole and delegated, and weight-only and group-quantized
  linears keep their quantize and dequantize ops whole through exir's
  own filter loop;
- the filter is plain filter_fn on the lowered program and on GPU, and
  skips do not change it;
- on chains of 2 and 4 links, the estimate never runs the partitioner
  and reads the graph signature's input specs as often on both.
The constant-only partition test, which runs against the real SDK, now
also lowers constant-only linears behind bfloat16 casts, and one whose
pieces a delegated add still takes.
---
 .../qualcomm/partition/qnn_partitioner.py     | 174 +++++-
 backends/qualcomm/partition/utils.py          |  26 +-
 backends/qualcomm/tests/test_passes.py        | 508 +++++++++++++++++-
 3 files changed, 678 insertions(+), 30 deletions(-)

diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py
index 46c853331c6..b462c219423 100644
--- a/backends/qualcomm/partition/qnn_partitioner.py
+++ b/backends/qualcomm/partition/qnn_partitioner.py
@@ -4,9 +4,11 @@
 # 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 functools
 import logging
+import operator
 from collections import defaultdict
-from typing import Any, Callable, Dict, List, Optional, Tuple
+from typing import Any, Callable, Dict, List, Optional, Set, Tuple
 
 import torch
 from executorch.backends.qualcomm.builders import node_visitor_manager
@@ -14,10 +16,7 @@
     PER_CHANNEL_GROUP_ENCODING,
 )
 from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
-from executorch.backends.qualcomm.builders.utils import (
-    is_mutable_buffer_input,
-    is_parameter,
-)
+from executorch.backends.qualcomm.builders.utils import is_parameter
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
 from executorch.backends.qualcomm.serialization.qc_schema import (
     HtpArch,
@@ -47,8 +46,10 @@
 from executorch.exir.backend.utils import tag_constant_data, tag_mutated_buffer
 from executorch.exir.dialects._ops import ops as exir_ops
 from torch.export.exported_program import ExportedProgram
+from torch.export.graph_signature import InputKind
 from torch.fx.passes.infra.partitioner import Partition
 from torch.fx.passes.operator_support import OperatorSupportBase
+from torch.utils import _pytree as pytree
 
 from .common_defs import (
     allow_list_operator,
@@ -62,6 +63,8 @@
     get_skip_decomp_table,
     has_data_dependent_shape,
     has_unsupported_dtype,
+    is_quantized,
+    must_stay_on_cpu,
     warn_if_dilated_conv_may_not_fit_vtcm,
 )
 
@@ -327,6 +330,26 @@ def _drop_orphaned_parameter_dequants(
                 )
         return [partition for partition in partitions if partition.nodes]
 
+    @staticmethod
+    def _constant_inputs(edge_program: torch.export.ExportedProgram) -> Set[str]:
+        """
+        Names of the inputs that hold the same value on every run: parameters,
+        buffers that are never mutated, and lifted constants.
+        """
+        signature = edge_program.graph_signature
+        mutated = set(signature.buffers_to_mutate.values())
+        return {
+            spec.arg.name
+            for spec in signature.input_specs
+            if spec.kind
+            in (InputKind.PARAMETER, InputKind.BUFFER, InputKind.CONSTANT_TENSOR)
+            and spec.target not in mutated
+        }
+
+    @staticmethod
+    def _is_runtime_input(node: torch.fx.Node, constant_inputs: Set[str]) -> bool:
+        return node.op != "get_attr" and node.name not in constant_inputs
+
     @staticmethod
     def _drop_partitions_without_runtime_inputs(
         partitions: List[Partition], edge_program: torch.export.ExportedProgram
@@ -336,23 +359,112 @@ def _drop_partitions_without_runtime_inputs(
         only by constants once a CPU op splits it from the rest. Such a
         partition stays on CPU.
         """
-
-        def is_runtime_input(node: torch.fx.Node) -> bool:
-            return not is_parameter(node, edge_program) or is_mutable_buffer_input(
-                node, edge_program
-            )
-
+        constant_inputs = QnnPartitioner._constant_inputs(edge_program)
         return [
             partition
             for partition in partitions
             if any(
-                is_runtime_input(arg)
+                QnnPartitioner._is_runtime_input(arg, constant_inputs)
                 for node in partition.nodes
                 for arg in node.all_input_nodes
                 if arg not in partition.nodes
             )
         ]
 
+    @staticmethod
+    def _may_run_on_qnn(node: torch.fx.Node) -> bool:
+        """
+        Whether QNN's op check may accept the node, as far as can be told
+        before the kept-whole ops are decomposed. Ops whose dtypes or shapes
+        rule QNN out, ops without a tensor output, such as the runtime
+        asserts export adds, and higher-order ops, which have no node visitor,
+        stay on CPU. Every getitem is accepted, as in QNN's op check.
+        """
+        if node.op != "call_function":
+            return False
+        if node.target in allow_list_operator:
+            return True
+        if not isinstance(node.target, torch._ops.OpOverload):
+            return False
+        # Retracing can drop the "val" of an assert, but not its schema.
+        if not node.target._schema.returns:
+            return False
+        if "val" in node.meta and not any(
+            isinstance(val, torch.Tensor)
+            for val in pytree.tree_leaves(node.meta["val"])
+        ):
+            return False
+        return not must_stay_on_cpu(node)
+
+    @staticmethod
+    def _left_on_cpu(ep: ExportedProgram) -> Set[torch.fx.Node]:
+        """
+        Estimates, on the ATen program, which ops HTP leaves on CPU, so that
+        the kept-whole ones among them can be decomposed.
+        """
+
+        # Partitioning here would take time quadratic in the graph size, so
+        # only the ops that _may_run_on_qnn rules out count as CPU ops. A QNN
+        # op that reads no runtime value is left on CPU when:
+        # - it reaches a QNN op through a CPU op before any op that reads a
+        #   runtime value, since the partitioner may split it off to avoid a
+        #   cycle;
+        # - only getitems of CPU ops read a runtime value, since the
+        #   partitioner moves those getitems out of their partitions;
+        # - such a getitem follows a QNN op, and no chain of QNN ops links the
+        #   op to an op that reads a runtime value.
+        # Skips are left out, since a skip that changed what is decomposed
+        # would rename later nodes, so a skip_node_id_set entry could name a
+        # different node.
+        # Ops that QNN rejects or removes later are not known here either. So
+        # an op cut off only by them still fails as before.
+        # An op can also be decomposed although a partition that reads a
+        # runtime value takes it.
+        nodes = list(ep.graph.nodes)
+        qnn = set(filter(QnnPartitioner._may_run_on_qnn, nodes))
+        constant_inputs = QnnPartitioner._constant_inputs(ep)
+        runtime, after_qnn = set(), set()
+        for node in nodes:
+            if node in qnn or not after_qnn.isdisjoint(node.all_input_nodes):
+                after_qnn.add(node)
+            if node in qnn and any(
+                (
+                    arg in runtime
+                    if arg in qnn
+                    else QnnPartitioner._is_runtime_input(arg, constant_inputs)
+                )
+                for arg in node.all_input_nodes
+            ):
+                runtime.add(node)
+        to_qnn, before_cpu = set(), set()
+        for node in reversed(nodes):
+            if node in qnn or not to_qnn.isdisjoint(node.users):
+                to_qnn.add(node)
+            if any(
+                (user in before_cpu and user not in runtime)
+                or (user in to_qnn and user not in qnn)
+                for user in node.users
+            ):
+                before_cpu.add(node)
+        left_on_cpu = (qnn - runtime) & before_cpu
+        cpu_getitems = {
+            node
+            for node in runtime
+            if node.target is operator.getitem and node.args[0] not in qnn
+        }
+        if runtime <= cpu_getitems:
+            return qnn - runtime
+        if not any(node.args[0] in after_qnn for node in cpu_getitems):
+            return left_on_cpu
+        linked, stack = set(runtime), list(runtime)
+        while stack:
+            node = stack.pop()
+            for neighbor in (*node.all_input_nodes, *node.users):
+                if neighbor in qnn and neighbor not in linked:
+                    linked.add(neighbor)
+                    stack.append(neighbor)
+        return left_on_cpu | (qnn - runtime - linked)
+
     def tag_nodes(
         self, partitions: List[Partition], edge_program: torch.export.ExportedProgram
     ) -> None:
@@ -414,6 +526,42 @@ def ops_to_not_decompose(
         Determines which op should not be decomposed during partitioning.
         The list of operators is obtained from `get_skip_decomp_table()`.
         The filter function (`filter_fn`) can be used to further refine which nodes are not decomposed. (advanced use case)
+        On HTP, kept-whole ops that `_drop_partitions_without_runtime_inputs`
+        would leave on CPU, as far as can be told before they are decomposed,
+        are decomposed too. Quantized ops stay whole, and exir cannot
+        decompose some ops, such as elu and the scatter variants, so such an op
+        still fails as before.
         """
         do_not_decompose = get_skip_decomp_table()
-        return (do_not_decompose, filter_fn)
+        # On the lowered program exir reports every kept-whole op left there
+        # that the filter accepts, so it stays plain filter_fn.
+        if self.backend != QnnExecuTorchBackendType.kHtpBackend or ep.dialect == "EDGE":
+            return (do_not_decompose, filter_fn)
+
+        # exir calls this method several times on the ATen program but uses
+        # only one of the filters, so the estimate runs on first use. That is
+        # before exir renames the targets of the ops it keeps whole, which
+        # would hide the dequantize ops.
+        # Quantized ops stay whole, in case the estimate is wrong: their pieces
+        # would have no quantization encodings, so QNN would reject them where
+        # a partition would take them.
+        @functools.cache
+        def decomposed():
+            nodes = {node for node in self._left_on_cpu(ep) if not is_quantized(node)}
+            names = sorted(
+                node.name for node in nodes if node.target in do_not_decompose
+            )
+            message = f"[QnnPartitioner] Decomposing kept-whole ops HTP would leave on CPU: {names}"
+            if names and self.skip_node_id_set:
+                logger.warning(
+                    f"{message}. Decomposing renames later nodes, so check that "
+                    "skip_node_id_set still names the intended nodes."
+                )
+            elif names:
+                logger.info(message)
+            return nodes
+
+        return (
+            do_not_decompose,
+            lambda node: filter_fn(node) and node not in decomposed(),
+        )
diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py
index a608cc42960..015c69e16da 100644
--- a/backends/qualcomm/partition/utils.py
+++ b/backends/qualcomm/partition/utils.py
@@ -199,12 +199,32 @@ def has_data_dependent_shape(node: torch.fx.Node) -> bool:
     )
 
 
+def must_stay_on_cpu(node: torch.fx.Node) -> bool:
+    """
+    Whether the node's data types or shapes keep it off QNN. Both are already
+    known before decomposition.
+    """
+    return (
+        has_unsupported_dtype(node) and not _unsupported_dtypes_are_quantized(node)
+    ) or has_data_dependent_shape(node)
+
+
+def is_quantized(node: torch.fx.Node) -> bool:
+    """
+    Whether the node quantizes, dequantizes, or reads the output of a
+    dequantize op. Decomposing such a node drops its quantization encodings.
+    """
+    return (
+        node.target in _Q_OPS
+        or node.target in _DQ_OPS
+        or any(arg.target in _DQ_OPS for arg in node.all_input_nodes)
+    )
+
+
 # Logic to determine whether to skip decompose and has higher priority than get_skip_decomp_table()
 def filter_fn(node: torch.fx.Node) -> bool:
     # The node will stay on CPU, so it must not be kept from decomposition.
-    if has_unsupported_dtype(node) and not _unsupported_dtypes_are_quantized(node):
-        return False
-    if has_data_dependent_shape(node):
+    if must_stay_on_cpu(node):
         return False
     # QNN does not support int32/int64 IO for the following OPs.
     potential_i32_i64_io_ops = [
diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index a607bcb4d81..c93de755b84 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -1,7 +1,8 @@
+import contextlib
 import operator
 import unittest
 from collections import defaultdict
-from unittest.mock import MagicMock, patch
+from unittest.mock import MagicMock, patch, PropertyMock
 
 import torch
 from executorch.backends.qualcomm._passes import (
@@ -35,8 +36,14 @@
 )
 from executorch.backends.qualcomm.builders.qnn_constants import OpContextLoader
 from executorch.backends.qualcomm.builders.utils import is_parameter
-from executorch.backends.qualcomm.partition.qnn_partitioner import QnnOperatorSupport
-from executorch.backends.qualcomm.partition.utils import get_skip_decomp_table
+from executorch.backends.qualcomm.partition.qnn_partitioner import (
+    QnnOperatorSupport,
+    QnnPartitioner,
+)
+from executorch.backends.qualcomm.partition.utils import (
+    filter_fn,
+    get_skip_decomp_table,
+)
 from executorch.backends.qualcomm.qnn_preprocess import QnnBackend
 from executorch.backends.qualcomm.quantizer.quantizer import QnnQuantizer, QuantDtype
 from executorch.backends.qualcomm.serialization.qc_schema import (
@@ -58,20 +65,258 @@
     QCOM_SCALE,
     QCOM_ZERO_POINTS,
 )
+from executorch.backends.qualcomm.utils.qnn_manager_lifecycle import QnnManagerRegistry
 from executorch.backends.qualcomm.utils.utils import (
+    generate_gpu_compiler_spec,
     generate_htp_compiler_spec,
     generate_qnn_executorch_compiler_spec,
     to_edge_transform_and_lower_to_qnn,
 )
 from executorch.exir import EdgeCompileConfig, to_edge
+from executorch.exir.backend.backend_details import PreprocessResult
 from executorch.exir.debug_handle_utils import DEBUG_HANDLE_KEY
 from executorch.exir.dialects._ops import ops as exir_ops
+from executorch.exir.lowered_backend_module import get_lowered_backend_modules
+from executorch.exir.program._program import _replace_aten_ops_with_transformed_ops
 from torch.export.exported_program import OutputKind
+from torch.export.graph_signature import ExportGraphSignature
+from torch.fx.passes.infra.partitioner import CapabilityBasedPartitioner, Partition
 from torch.library import Library
 from torchao.quantization import IntxWeightOnlyConfig, PerAxis, quantize_
 from torchao.quantization.pt2e.quantize_pt2e import convert_pt2e, prepare_pt2e
 
 
+class ConstantBranches(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("table", torch.randn(4, 16))
+        self.register_buffer("weight", torch.randn(16, 16))
+        self.register_buffer("mask", torch.rand(4, 16) > 0.5)
+
+    def forward(self, x):
+        linear = torch.nn.functional.linear
+        # QNN has no bfloat16 type, and a masked select has no static shape,
+        # so these branches read nothing at runtime once those ops run on
+        # CPU. Their linears are decomposed.
+        cast = linear(self.table, self.weight).to(torch.bfloat16).float()
+        masked = linear(self.table, self.weight)[self.mask].sum()
+        # This linear feeds a delegated add both directly and through CPU
+        # casts, so the partitioner splits it off to avoid a cycle. It is
+        # decomposed.
+        detour = linear(self.table, self.weight)
+        detour = detour + detour.to(torch.bfloat16).float()
+        # This linear reads only constants too, but it shares a partition with
+        # the add that reads x, so it stays whole, even though that add
+        # reaches another QNN op through CPU casts.
+        delegated = linear(self.table, self.weight)
+        y = x + delegated + cast + masked + detour
+        return y.to(torch.bfloat16).float() + 1
+
+
+class MaskedReader(ConstantBranches):
+    def forward(self, x):
+        # Nothing after the masked select reaches QNN: the sum reads a
+        # data-dependent size, and the runtime asserts that export adds for
+        # that size have no tensor output. So the linear shares the add's
+        # partition and stays whole.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x + linear, linear[self.mask].sum()
+
+
+class SkippedRelu(ConstantBranches):
+    def forward(self, x):
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x + linear.relu()
+
+
+class QuantizedCast(ConstantBranches):
+    def forward(self, x):
+        # A quantized linear stays whole. Quantized casts get a QNN type, so
+        # here the whole model is delegated.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x + linear.to(torch.bfloat16).float()
+
+
+class NoRuntimeInput(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("table", torch.randn(4, 16))
+        self.register_buffer("weight", torch.randn(16, 16))
+
+    def forward(self, x):
+        # x only reaches a CPU cast, so no partition reads a runtime value, and
+        # the linear is decomposed.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x.to(torch.bfloat16), linear * 2
+
+
+class CastSplit(NoRuntimeInput):
+    def forward(self, x):
+        # QNN takes every getitem, so the getitems of this CPU split can join
+        # the add's partition. The linear would then close a cycle through the
+        # split, so it is left out and decomposed.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        low, high = linear.to(torch.bfloat16).split(8, dim=-1)
+        return x + linear, low, high
+
+
+class ConstantOutput(NoRuntimeInput):
+    def forward(self, x):
+        # No path links this linear to the ops that read x, but the partitioner
+        # also merges such ops, so the linear stays whole. The getitems of the
+        # split stay with it, since QNN takes it.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x + 1, *x.split(8, dim=-1), linear
+
+
+class CastSplitInput(NoRuntimeInput):
+    def forward(self, x):
+        # Only the getitems of a CPU split read x, and the partitioner moves
+        # them out of the linear's partition, so it reads no runtime value.
+        # The linear is decomposed.
+        low, high = x.to(torch.bfloat16).split(8, dim=-1)
+        return low, high, torch.nn.functional.linear(self.table, self.weight)
+
+
+class SplitAfterAdd(NoRuntimeInput):
+    def forward(self, x):
+        # These getitems come after a QNN op, so their partition cannot merge
+        # with the add's. The linear joins theirs, which reads no runtime value
+        # once they are moved out, so it is decomposed.
+        low, high = (x + 1).to(torch.bfloat16).split(8, dim=-1)
+        return low, high, torch.nn.functional.linear(self.table, self.weight)
+
+
+class CastSplitLinked(NoRuntimeInput):
+    def forward(self, x):
+        # As above, but the linear feeds the add, so it stays whole.
+        low, high = (x + 1).to(torch.bfloat16).split(8, dim=-1)
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return low, high, x + linear
+
+
+class LinearOfInput(NoRuntimeInput):
+    def forward(self, x):
+        # This linear reads x through the add, so it stays whole even though it
+        # reaches a QNN op through CPU casts.
+        linear = torch.nn.functional.linear(x + 1, self.weight)
+        return linear + linear.to(torch.bfloat16).float()
+
+
+class ReluBeforeCast(NoRuntimeInput):
+    def forward(self, x):
+        # The linear reaches the CPU cast through a relu that reads no runtime
+        # value either, so it is cut off and decomposed.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x + linear.relu().to(torch.bfloat16).float()
+
+
+class CondAfterLinear(NoRuntimeInput):
+    def forward(self, x):
+        # QNN has no node visitor for cond, so it runs on CPU and cuts the
+        # linear off from the add. The linear is decomposed.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        branch = torch.cond(x.sum() > 0, torch.neg, torch.abs, (linear,))
+        return x + branch
+
+
+class ItemAfterLinear(NoRuntimeInput):
+    def forward(self, x):
+        # item has no tensor output, so it runs on CPU, and the partition of
+        # the linear and the sum reads no runtime value. The linear is
+        # decomposed.
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x * linear.sum().item()
+
+
+class SharedConstant(NoRuntimeInput):
+    def forward(self, x):
+        # Like CastSplitLinked, but the linear reads a value that the add
+        # reads.
+        low, high = (x + 1).to(torch.bfloat16).split(8, dim=-1)
+        table = self.table * 2
+        linear = torch.nn.functional.linear(table, self.weight)
+        return low, high, x + table, linear
+
+
+class UnlinkedSplit(NoRuntimeInput):
+    def forward(self, x):
+        # No QNN op comes before this split, so its getitems cannot keep the
+        # linear out of the add's partition. It stays whole.
+        low, high = x.to(torch.bfloat16).split(8, dim=-1)
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x + 1, low, high, linear
+
+
+class MutatedTable(NoRuntimeInput):
+    def forward(self, x):
+        # The table changes on every run, so the linear reads a runtime value
+        # and stays whole, although only a CPU cast reads it.
+        self.table.mul_(2)
+        linear = torch.nn.functional.linear(self.table, self.weight)
+        return x + linear.to(torch.bfloat16).float()
+
+
+class QuantizedWeight(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("table", torch.randn(4, 16))
+        self.linear = torch.nn.Linear(16, 16, bias=False)
+
+    def forward(self, x):
+        return x + self.linear(self.table).to(torch.bfloat16).float()
+
+
+class GroupQuantizedWeight(NoRuntimeInput):
+    def __init__(self):
+        super().__init__()
+        self.register_buffer("scales", torch.full((16, 2), 0.1))
+        self.register_buffer("zero_points", torch.zeros(16, 2, dtype=torch.int8))
+
+    def forward(self, x):
+        group = (self.scales, self.zero_points, -8, 7, torch.int8, 8)
+        weight = torch.ops.quantized_decomposed.quantize_per_channel_group(
+            self.weight, *group
+        )
+        weight = torch.ops.quantized_decomposed.dequantize_per_channel_group(
+            weight, *group, torch.float32
+        )
+        linear = torch.nn.functional.linear(self.table, weight)
+        return x + linear.to(torch.bfloat16).float()
+
+
+@contextlib.contextmanager
+def stubbed_qnn_backend():
+    # Lowers without the QNN SDK: the backend accepts every op and compiles
+    # nothing.
+    with patch("executorch.backends.qualcomm.utils.utils.setup_qnn_sdk"), patch.object(
+        QnnManagerRegistry, "get_or_create_qnn_manager", return_value=MagicMock()
+    ), patch.object(
+        QnnBackend,
+        "preprocess_multimethod",
+        side_effect=lambda programs, _: {
+            method: [PreprocessResult(processed_bytes=b"qnn")] * len(partitions)
+            for method, partitions in programs.items()
+        },
+    ):
+        yield
+
+
+def stubbed_compiler_specs(backend_options):
+    with patch("executorch.backends.qualcomm.utils.utils.setup_qnn_sdk"):
+        return generate_qnn_executorch_compiler_spec(
+            soc_model=QcomChipset.SM8650, backend_options=backend_options
+        )
+
+
+def delegated_ops(edge):
+    return [
+        node.target
+        for lowered in get_lowered_backend_modules(edge.exported_program().graph_module)
+        for node in lowered.original_module.graph.nodes
+    ]
+
+
 class TestPasses(unittest.TestCase):
     def _build_context_loader_edge_program(self, op_name, check_ir_validity=True):
         graph_name = "forward"
@@ -1190,29 +1435,264 @@ class ConstantBranch(torch.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.register_buffer("table", torch.randn(4, 16))
+                self.register_buffer("weight", torch.randn(16, 16))
 
             def forward(self, x):
-                return x + torch.relu(self.table * 2 - 1)
+                # With relu on CPU, table * 2 - 1 reads nothing at runtime, and
+                # with the bfloat16 cast on CPU, neither does the linear. So the
+                # linear must be decomposed for this to lower.
+                linear = torch.nn.functional.linear(self.table, self.weight)
+                cast = linear.to(torch.bfloat16).float()
+                return x + torch.relu(self.table * 2 - 1) + cast
+
+        class DetourBranch(ConstantBranch):
+            def forward(self, x):
+                # The linear reaches the add only through CPU casts, so it is
+                # decomposed, but the partition of x + 1 takes its pieces.
+                linear = torch.nn.functional.linear(self.table, self.weight)
+                cast = (x + 1).to(torch.bfloat16).float()
+                return linear.to(torch.bfloat16).float() + cast
 
         sample_input = (torch.randn(4, 16),)
         compiler_specs = generate_qnn_executorch_compiler_spec(
             soc_model=QcomChipset.SM8650,
             backend_options=generate_htp_compiler_spec(use_fp16=True),
         )
-        # With relu on CPU, table * 2 - 1 reads nothing at runtime.
-        try:
+        # HTP accepts the pieces of a decomposed linear that a partition takes.
+        for model, takes_pieces in ((ConstantBranch(), False), (DetourBranch(), True)):
+            with self.subTest(model=type(model).__name__):
+                try:
+                    edge = to_edge_transform_and_lower_to_qnn(
+                        model.eval(),
+                        sample_input,
+                        compiler_specs,
+                        skip_node_op_set={"aten.relu.default"},
+                    )
+                except RuntimeError as e:
+                    if "Failed to initialize QNN backend" in str(e):
+                        self.skipTest(f"QNN SDK not available: {e}")
+                    raise
+                edge.to_executorch()
+                self.assertEqual(
+                    takes_pieces, exir_ops.edge.aten.mm.default in delegated_ops(edge)
+                )
+
+    def test_htp_decomposes_kept_whole_ops_left_on_cpu(self):
+        compiler_specs = stubbed_compiler_specs(
+            generate_htp_compiler_spec(use_fp16=True)
+        )
+        sample_input = (torch.randn(4, 16),)
+        with stubbed_qnn_backend():
             edge = to_edge_transform_and_lower_to_qnn(
-                ConstantBranch().eval(),
-                sample_input,
-                compiler_specs,
-                skip_node_op_set={"aten.relu.default"},
+                ConstantBranches().eval(), sample_input, compiler_specs
             )
-        except RuntimeError as e:
-            if "Failed to initialize QNN backend" in str(e):
-                self.skipTest(f"QNN SDK not available: {e}")
-            raise
+            for model in (
+                NoRuntimeInput(),
+                CastSplit(),
+                CastSplitInput(),
+                SplitAfterAdd(),
+                ReluBeforeCast(),
+            ):
+                with self.subTest(model=type(model).__name__):
+                    to_edge_transform_and_lower_to_qnn(
+                        model.eval(), sample_input, compiler_specs
+                    )
+        top_level_ops = [
+            node.target
+            for node in edge.exported_program().graph.nodes
+            if node.op == "call_function"
+        ]
+        self.assertNotIn(exir_ops.edge.aten.linear.default, top_level_ops)
+        self.assertIn(exir_ops.edge.aten.linear.default, delegated_ops(edge))
+        self.assertIn(exir_ops.edge.aten.add.Tensor, delegated_ops(edge))
         edge.to_executorch()
 
+    def test_htp_keeps_linked_kept_whole_ops_whole(self):
+        compiler_specs = stubbed_compiler_specs(
+            generate_htp_compiler_spec(use_fp16=True)
+        )
+        sample_input = (torch.randn(4, 16),)
+        for model in (
+            MaskedReader(),
+            ConstantOutput(),
+            CastSplitLinked(),
+            LinearOfInput(),
+            SharedConstant(),
+            UnlinkedSplit(),
+            MutatedTable(),
+        ):
+            with self.subTest(model=type(model).__name__), stubbed_qnn_backend():
+                edge = to_edge_transform_and_lower_to_qnn(
+                    model.eval(), sample_input, compiler_specs
+                )
+                self.assertIn(exir_ops.edge.aten.linear.default, delegated_ops(edge))
+
+    def test_htp_keeps_quantized_ops_whole(self):
+        compiler_specs = stubbed_compiler_specs(
+            generate_htp_compiler_spec(use_fp16=True)
+        )
+        sample_input = (torch.randn(4, 16),)
+        quantized = prepare_pt2e(
+            torch.export.export(QuantizedCast().eval(), sample_input).module(),
+            QnnQuantizer(),
+        )
+        quantized(*sample_input)
+        with stubbed_qnn_backend():
+            edge = to_edge_transform_and_lower_to_qnn(
+                convert_pt2e(quantized),
+                sample_input,
+                stubbed_compiler_specs(generate_htp_compiler_spec(use_fp16=False)),
+            )
+        self.assertIn(exir_ops.edge.aten.linear.default, delegated_ops(edge))
+
+        # These ops are left on CPU, but stay whole. exir renames each op it
+        # keeps whole while it filters, so the linear must still count as
+        # quantized once the dequantize of its weight is renamed.
+        weight_only = QuantizedWeight().eval()
+        quantize_(
+            weight_only,
+            IntxWeightOnlyConfig(weight_dtype=torch.int8, granularity=PerAxis(0)),
+        )
+        for model, quantized_ops in (
+            (weight_only, {torch.ops.torchao.dequantize_affine.default}),
+            (
+                GroupQuantizedWeight().eval(),
+                {
+                    torch.ops.quantized_decomposed.quantize_per_channel_group.default,
+                    torch.ops.quantized_decomposed.dequantize_per_channel_group.default,
+                },
+            ),
+        ):
+            with self.subTest(model=type(model).__name__):
+                kept_whole_ops = _replace_aten_ops_with_transformed_ops(
+                    "forward",
+                    torch.export.export(model, sample_input),
+                    {"forward": [QnnPartitioner(compiler_specs)]},
+                )
+                self.assertLessEqual(
+                    {torch.ops.aten.linear.default, *quantized_ops}, kept_whole_ops
+                )
+
+    def test_kept_whole_filter_is_filter_fn_on_edge_and_gpu(self):
+        sample_input = (torch.randn(4, 16),)
+        compiler_specs = stubbed_compiler_specs(
+            generate_htp_compiler_spec(use_fp16=True)
+        )
+        with stubbed_qnn_backend():
+            edge = to_edge_transform_and_lower_to_qnn(
+                ConstantBranches().eval(), sample_input, compiler_specs
+            )
+        # On the lowered program the filter must not let a kept-whole op left
+        # on CPU through, so exir still reports it.
+        _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(
+            edge.exported_program()
+        )
+        self.assertIs(keep_whole, filter_fn)
+        # GPU drops no partitions, so it keeps the same ops whole as before.
+        _, keep_whole = QnnPartitioner(
+            stubbed_compiler_specs(generate_gpu_compiler_spec())
+        ).ops_to_not_decompose(
+            torch.export.export(ConstantBranches().eval(), sample_input)
+        )
+        self.assertIs(keep_whole, filter_fn)
+
+    def test_htp_kept_whole_filter_ignores_skips(self):
+        # A skip of either kind that changed what is decomposed would rename
+        # later nodes, so a skip_node_id_set entry could name a different node.
+        # Skipping the relu cuts the linear off, so the model cannot be lowered.
+        compiler_specs = stubbed_compiler_specs(
+            generate_htp_compiler_spec(use_fp16=True)
+        )
+        ep = torch.export.export(SkippedRelu().eval(), (torch.randn(4, 16),))
+        _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(ep)
+        _, keep_whole_skipped = QnnPartitioner(
+            compiler_specs, skip_node_op_set={"aten.relu.default"}
+        ).ops_to_not_decompose(ep)
+        self.assertEqual(
+            [keep_whole(node) for node in ep.graph.nodes],
+            [keep_whole_skipped(node) for node in ep.graph.nodes],
+        )
+
+    def test_htp_kept_whole_filter_counts_cond_and_item_as_cpu(self):
+        # QNN's passes cannot lower either model yet, so the filter is checked
+        # alone.
+        compiler_specs = stubbed_compiler_specs(
+            generate_htp_compiler_spec(use_fp16=True)
+        )
+        # With a skip_node_id_set, the renaming is worth a warning.
+        for model, skip_node_id_set, level in (
+            (CondAfterLinear(), None, "INFO"),
+            (ItemAfterLinear(), {"linear"}, "WARNING"),
+        ):
+            with self.subTest(model=type(model).__name__):
+                ep = torch.export.export(model.eval(), (torch.randn(4, 16),))
+                _, keep_whole = QnnPartitioner(
+                    compiler_specs, skip_node_id_set=skip_node_id_set
+                ).ops_to_not_decompose(ep)
+                linear = next(
+                    node
+                    for node in ep.graph.nodes
+                    if node.target == torch.ops.aten.linear.default
+                )
+                with self.assertLogs(QnnPartitioner.__module__, "INFO") as logs:
+                    self.assertFalse(keep_whole(linear))
+                self.assertEqual(level, logs.records[-1].levelname)
+                self.assertIn(linear.name, logs.output[-1])
+
+    def test_htp_kept_whole_estimate_avoids_quadratic_steps(self):
+        class Chain(torch.nn.Module):
+            def __init__(self, length):
+                super().__init__()
+                self.scales = torch.nn.ParameterList(
+                    torch.rand(16) for _ in range(length)
+                )
+
+            def forward(self, x):
+                # Each scale * 2 reads only constants, like the dequantize of a
+                # quantized weight.
+                for scale in self.scales:
+                    x = torch.relu(x * (scale * 2) + x)
+                return x
+
+        with patch("executorch.backends.qualcomm.utils.utils.setup_qnn_sdk"):
+            compiler_specs = generate_qnn_executorch_compiler_spec(
+                soc_model=QcomChipset.SM8650,
+                backend_options=generate_htp_compiler_spec(use_fp16=True),
+            )
+
+        # Large models have tens of thousands of nodes and constants. The time
+        # to partition grows with the square of the graph size, and so does
+        # going over the input specs of the graph signature for each input, or
+        # running the estimate for each node exir filters.
+        def input_spec_reads(length):
+            ep = torch.export.export(Chain(length), (torch.randn(4, 16),))
+            input_specs = ep.graph_signature.input_specs
+            nodes = [node for node in ep.graph.nodes if node.op == "call_function"]
+            with patch.object(
+                CapabilityBasedPartitioner,
+                "propose_partitions",
+                autospec=True,
+                side_effect=CapabilityBasedPartitioner.propose_partitions,
+            ) as partition, patch.object(
+                ExportGraphSignature,
+                "input_specs",
+                new_callable=PropertyMock,
+                return_value=input_specs,
+                create=True,
+            ) as read_input_specs:
+                _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(ep)
+                for node in nodes:
+                    keep_whole(node)
+                QnnPartitioner._drop_partitions_without_runtime_inputs(
+                    [Partition(nodes=[node]) for node in nodes], ep
+                )
+            partition.assert_not_called()
+            return read_input_specs.call_count
+
+        reads = input_spec_reads(2)
+        self.assertGreater(reads, 0)
+        self.assertEqual(reads, input_spec_reads(4))
+
     def test_partitioner_falls_back_on_bfloat16(self):
         class Bf16Region(torch.nn.Module):
             def __init__(self):

From bc89e9e3039372c2b87a60a3d2e7752b18de66d6 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Sun, 1 Nov 2026 21:28:42 -0800
Subject: [PATCH 2/6] Document both calls to Partitioner.ops_to_not_decompose

to_edge_transform_and_lower calls ops_to_not_decompose during decomposition,
on the ATen program, to choose which nodes to keep whole, and after
lowering, on the EDGE program, to report kept-whole nodes that were neither
decomposed nor delegated. The docstring only described the first use.
QnnPartitioner tells the 2 calls apart by ep.dialect, and its HTP filter only
works on the nodes of the program it was given, so the docstring now names
both stages, their dialects, and that rule, which also covers the nodes of
the submodules of the program's top-level cond, map and scan ops. It also
says that the list
checked after lowering includes the lists of earlier partitioners for the
same method, and that some quantization ops are not reported, nor, when
that call returns a filter, ops that exir does not preserve. Docstring only.
---
 exir/backend/partitioner.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/exir/backend/partitioner.py b/exir/backend/partitioner.py
index c0edf0c71d8..c309b3e07a0 100644
--- a/exir/backend/partitioner.py
+++ b/exir/backend/partitioner.py
@@ -105,6 +105,19 @@ def ops_to_not_decompose(
         guaranteed that the program that the backend receives will not have any of these ops
         decomposed.
 
+        to_edge_transform_and_lower calls this method at two stages, possibly more than once
+        in each, and `ep.dialect` tells them apart. Each returned filter is only applied to
+        the nodes of the `ep` passed to that same call, including those of the submodules of
+        its top-level cond, map and scan ops.
+
+        * During decomposition, on the ATen program (`ep.dialect == "ATEN"`). The list and
+          the filter decide which nodes are kept whole.
+        * After lowering, on the EDGE program (`ep.dialect == "EDGE"`). A node still in the
+          graph whose op is in this list, or in the list of an earlier partitioner for the
+          same method, and that the filter accepts was neither decomposed nor delegated,
+          and raises an error, except for some quantization ops and, when this call
+          returns a filter, ops that exir does not preserve.
+
         Returns:
             List[torch._ops.OpOverload]: a list of operator names that should not be decomposed.
             Optional[Callable[[torch.fx.Node], bool]]]: an optional callable, acting as a filter, that users can provide

From 79ae02726cf53429f27ea05ad11fa248460960a5 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 14:02:00 -0800
Subject: [PATCH 3/6] Test reserve_node_names on its own

reserve_node_names was only tested through the Qualcomm pass manager,
and only in a test that needs the Qualcomm backend. The Samsung pass
manager calls it too. A torch-only test now exports a model whose input
takes the name "mul", so export renames the multiply to "mul_1". After
reserve_node_names, a new node asked to be named "mul_1" gets another
name, and the graph lints. Once export records the names it renames,
the test skips, like the QNN name collision test, since it then no
longer tests the helper.
---
 backends/transforms/targets.bzl               | 11 ++++++
 .../test/test_reserve_node_names.py           | 39 +++++++++++++++++++
 2 files changed, 50 insertions(+)
 create mode 100644 backends/transforms/test/test_reserve_node_names.py

diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl
index 4cd4db88fa6..49d06170065 100644
--- a/backends/transforms/targets.bzl
+++ b/backends/transforms/targets.bzl
@@ -607,6 +607,17 @@ def define_common_targets():
         ],
     )
 
+    runtime.python_test(
+        name = "test_reserve_node_names",
+        srcs = [
+            "test/test_reserve_node_names.py",
+        ],
+        deps = [
+            "//caffe2:torch",
+            ":utils",
+        ],
+    )
+
     runtime.python_test(
         name = "test_convert_dtype_pass",
         srcs = [
diff --git a/backends/transforms/test/test_reserve_node_names.py b/backends/transforms/test/test_reserve_node_names.py
new file mode 100644
index 00000000000..20c7b7a779a
--- /dev/null
+++ b/backends/transforms/test/test_reserve_node_names.py
@@ -0,0 +1,39 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+import unittest
+
+import torch
+from executorch.backends.transforms.utils import reserve_node_names
+
+
+class TestReserveNodeNames(unittest.TestCase):
+    def test_new_node_does_not_take_an_existing_name(self):
+        class Model(torch.nn.Module):
+            def forward(self, x, mul):
+                return x * 2 + mul
+
+        # The input takes the name "mul", so export renames the multiply to
+        # "mul_1" without recording that name in the graph's namespace.
+        ep = torch.export.export(Model(), (torch.randn(2), torch.randn(2)))
+        graph = ep.graph_module.graph
+        self.assertIn("mul_1", [node.name for node in graph.nodes])
+        # A skip, not a failure: a torch update that records the name is fine.
+        if "mul_1" in graph._graph_namespace._used_names:
+            self.skipTest(
+                "export now records the names it renames, so this no longer "
+                "tests reserve_node_names; see "
+                "https://github.com/pytorch/pytorch/issues/202927"
+            )
+
+        reserve_node_names(ep.graph_module)
+        (add,) = graph.find_nodes(op="call_function", target=torch.ops.aten.add.Tensor)
+        with graph.inserting_after(add):
+            node = graph.create_node(
+                "call_function", torch.ops.aten.neg.default, (add,), name="mul_1"
+            )
+        self.assertNotEqual(node.name, "mul_1")
+        graph.lint()

From b84f6410eb6ede570dc262355a4ed4a2e7cb388f Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 14:02:00 -0800
Subject: [PATCH 4/6] Keep the original error in the QNN name collision test

When the pass manager failed for a reason other than the name
collision, the test reported an assertion about the message, and the
real error showed up only as the context of that assertion. The test
now re-raises any error that is not the collision, with its own
traceback.
---
 backends/qualcomm/tests/test_passes.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index c93de755b84..a0f5023724e 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -502,7 +502,8 @@ def forward(self, input, k, v, mul):
                         ).graph.lint()
                     except Exception as e:
                         # The pass manager wraps the lint error in its own.
-                        self.assertIn("Node redefined name", str(e.__cause__ or e))
+                        if "Node redefined name" not in str(e.__cause__ or e):
+                            raise
                     else:
                         no_collision.append(type(model).__name__)
                     reserve.assert_called_once()

From c5b19058cc757d43ac01649d7547e92e81a85942 Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 14:02:10 -0800
Subject: [PATCH 5/6] Skip the QNN name collision test only when no case
 collides

The test skips when a torch update removes the name collision, since
the cases then no longer test reserve_node_names. It skipped as soon as
any case stopped colliding, which also hid the cases that still did.
It now skips only when none of them collide, and logs a warning that
names the cases that no longer collide.
---
 backends/qualcomm/tests/test_passes.py | 17 +++++++++++------
 1 file changed, 11 insertions(+), 6 deletions(-)

diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index a0f5023724e..c800c42bf2d 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -1,4 +1,5 @@
 import contextlib
+import logging
 import operator
 import unittest
 from collections import defaultdict
@@ -507,13 +508,17 @@ def forward(self, input, k, v, mul):
                     else:
                         no_collision.append(type(model).__name__)
                     reserve.assert_called_once()
-        # A skip, not a failure: a torch update that removes the collision is fine.
+        # A skip, not a failure: a torch update that removes the collision is
+        # fine. While any case still collides, the test still means something.
+        message = (
+            f"{no_collision} no longer collide without reserve_node_names, so "
+            "those cases no longer test it; see "
+            "https://github.com/pytorch/pytorch/issues/202927"
+        )
+        if len(no_collision) == len(cases):
+            self.skipTest(message)
         if no_collision:
-            self.skipTest(
-                f"{no_collision} no longer collide without reserve_node_names, so "
-                "those cases no longer test it; see "
-                "https://github.com/pytorch/pytorch/issues/202927"
-            )
+            logging.warning(message)
 
     def test_build_op_wrappers_returns_context_binary(self):
         op_name = "ctx_loader_build"

From 11862842d63d9043bc05a31395590f99bd0d79fd Mon Sep 17 00:00:00 2001
From: Anthony Shoumikhin <anthony@shoumikh.in>
Date: Mon, 2 Nov 2026 17:46:18 -0800
Subject: [PATCH 6/6] Qualcomm AI Engine Direct - Cast uint8 to float32 through
 int32 on HTP

On the HTP simulator (QNN SDK 2.51, SM8650, float16 graph), a cast from
uint8 to float32 gives 1 for every input above 0, as if the input were
bool. So a camera image passed as uint8 reached the model as a mask of
0s and 1s. The cast op is built with the right data types (UINT_8 in,
FLOAT_32 out). HTP casts uint8 to int32 correctly, and int32 to float32
correctly, but not uint8 to float32 in one step. Chains such as uint8 to
float16 to float32 gave the same result, because consecutive casts are
fused into one uint8 to float32 cast before the graph is built.

The cast builder already goes through int32 for casts to or from int64
that it cannot do in one step. On HTP, a uint8 to float32 cast with a
plain float output now takes the same 2 steps and gives 0 to 255. Casts
to float16, casts with a quantized output, and every cast on GPU are
built as before. To tell the backends apart, get_node_visitors takes a
backend_type and passes it to every builder. The partitioner and
preprocess pass the one from the compiler specs, and the test helpers
that rebuild the HTP op wrapper list pass HTP.

A new test builds the cast and checks that it takes 2 ops for float32 on
HTP, and 1 op for float32 on GPU and for float16 on HTP.
test_qnn_backend_cast in the floating point operator tests now also
scales a uint8 tensor holding 0 to 255 to float32 and to float16, the
way a camera input is prepared, and the quantized test scales it to
float32. The quantized test checks its outputs with a relative tolerance
of 0.1 instead of 1, since with 1 a cast that gives only 0 or 1 would
also pass. GPU cannot run on the x86 host, so only the HTP simulator
ran the operator tests.
---
 backends/qualcomm/builders/node_visitor.py    |  2 ++
 .../qualcomm/builders/node_visitor_manager.py |  8 ++++-
 backends/qualcomm/builders/op_to.py           | 26 ++++++++++++---
 .../qualcomm/partition/qnn_partitioner.py     |  1 +
 backends/qualcomm/qnn_preprocess.py           |  1 +
 backends/qualcomm/tests/models.py             |  9 +++++
 backends/qualcomm/tests/rework/src/utils.py   |  7 +++-
 backends/qualcomm/tests/test_passes.py        | 27 +++++++++++++++
 backends/qualcomm/tests/test_qnn_delegate.py  | 33 +++++++++++++++----
 9 files changed, 100 insertions(+), 14 deletions(-)

diff --git a/backends/qualcomm/builders/node_visitor.py b/backends/qualcomm/builders/node_visitor.py
index 02fbe08e8b0..a416f62c329 100644
--- a/backends/qualcomm/builders/node_visitor.py
+++ b/backends/qualcomm/builders/node_visitor.py
@@ -137,6 +137,7 @@ def __init__(
         edge_program: torch.export.ExportedProgram,
         enable_tensor_dump,
         is_qnn_partitioner=True,
+        backend_type=None,
     ) -> None:
         # is_qnn_partitioner: True for the real qnn_partitioner / qnn_preprocess
         # flows, where external_ids is built from the program being processed so
@@ -149,6 +150,7 @@ def __init__(
         self.edge_program = edge_program
         self.enable_tensor_dump = enable_tensor_dump
         self.is_qnn_partitioner = is_qnn_partitioner
+        self.backend_type = backend_type
 
     def get_node(self, node):
         """
diff --git a/backends/qualcomm/builders/node_visitor_manager.py b/backends/qualcomm/builders/node_visitor_manager.py
index f866c900a6c..f4c5a7c9f6b 100644
--- a/backends/qualcomm/builders/node_visitor_manager.py
+++ b/backends/qualcomm/builders/node_visitor_manager.py
@@ -54,6 +54,7 @@ def get_node_visitors(
     enable_tensor_dump=False,
     op_package_infos: List[QnnExecuTorchOpPackageInfo] = None,
     is_qnn_partitioner=True,
+    backend_type=None,
 ) -> Dict[str, NodeVisitor]:
     """Create a new class instance at runtime, and put them in a dict"""
     node_to_external_map = generate_node_to_external_map(edge_program)
@@ -63,7 +64,11 @@ def get_node_visitors(
             visitor
         ), f"Expecting a callable class, but got {visitor} of type {type(visitor)}"
         node_visitors[target] = visitor(
-            node_to_external_map, edge_program, enable_tensor_dump, is_qnn_partitioner
+            node_to_external_map,
+            edge_program,
+            enable_tensor_dump,
+            is_qnn_partitioner,
+            backend_type,
         )
     if op_package_infos:
         custom_ops = []
@@ -75,6 +80,7 @@ def get_node_visitors(
                     edge_program,
                     enable_tensor_dump,
                     is_qnn_partitioner,
+                    backend_type,
                 )
                 node_visitors[op_package_info.custom_op_name] = custom_op_builder
                 custom_ops.append(op_package_info.custom_op_name)
diff --git a/backends/qualcomm/builders/op_to.py b/backends/qualcomm/builders/op_to.py
index 9b024307934..76d22059406 100644
--- a/backends/qualcomm/builders/op_to.py
+++ b/backends/qualcomm/builders/op_to.py
@@ -8,6 +8,9 @@
 import executorch.backends.qualcomm.python.PyQnnManagerAdaptor as PyQnnManager
 
 import torch
+from executorch.backends.qualcomm.serialization.qc_schema import (
+    QnnExecuTorchBackendType,
+)
 from executorch.backends.qualcomm.utils.constants import QCOM_QUANT_ATTRS
 
 from .node_visitor import NodeVisitor, QNN_TENSOR_TYPE_MAP
@@ -97,13 +100,26 @@ def define_node(
         # since int32 is the only source that can be cast into int64
         # this is mainly for validation purpose, redundant cast ops will be fused
         # in preprocess stage.
+        # HTP casts uint8 to float32 as if the input were bool, while uint8 to
+        # int32 and int32 to float32 are right, so on HTP a plain uint8 to
+        # float32 cast goes through int32 too. Here the step is needed for the
+        # values, and nothing fuses it away.
         ops = []
         if (
-            node.meta["val"].dtype == torch.int64
-            and input_node.meta["val"].dtype != torch.int32
-        ) or (
-            input_node.meta["val"].dtype == torch.int64
-            and node.meta["val"].dtype != torch.int32
+            (
+                node.meta["val"].dtype == torch.int64
+                and input_node.meta["val"].dtype != torch.int32
+            )
+            or (
+                input_node.meta["val"].dtype == torch.int64
+                and node.meta["val"].dtype != torch.int32
+            )
+            or (
+                input_node.meta["val"].dtype == torch.uint8
+                and node.meta["val"].dtype == torch.float32
+                and not node.meta.get(QCOM_QUANT_ATTRS)
+                and self.backend_type == QnnExecuTorchBackendType.kHtpBackend
+            )
         ):
             input_quant_encoding, input_quant_configs = self.get_quant_encoding_conf(
                 input_node, node
diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py
index b462c219423..772512c17ee 100644
--- a/backends/qualcomm/partition/qnn_partitioner.py
+++ b/backends/qualcomm/partition/qnn_partitioner.py
@@ -95,6 +95,7 @@ def __init__(
             edge_program,
             op_package_infos=python_options.op_package_options.op_package_infos,
             is_qnn_partitioner=is_qnn_partitioner,
+            backend_type=python_options.backend_options.backend_type,
         )
 
         self.skip_node_op_set = skip_node_op_set
diff --git a/backends/qualcomm/qnn_preprocess.py b/backends/qualcomm/qnn_preprocess.py
index dda4687e59c..9e9bb112916 100644
--- a/backends/qualcomm/qnn_preprocess.py
+++ b/backends/qualcomm/qnn_preprocess.py
@@ -108,6 +108,7 @@ def _build_op_wrappers(
             edge_program,
             enable_tensor_dump=enable_tensor_dump,
             op_package_infos=op_package_infos,
+            backend_type=backend_type,
         )
         py_op_wrapper_list = []
         for node in graph_module.graph.nodes:
diff --git a/backends/qualcomm/tests/models.py b/backends/qualcomm/tests/models.py
index 926f5e1fb67..61882533093 100644
--- a/backends/qualcomm/tests/models.py
+++ b/backends/qualcomm/tests/models.py
@@ -402,6 +402,15 @@ def forward(self, x, y):
         return res + index.to(torch.int32)
 
 
+class CastToFloat(torch.nn.Module):
+    def __init__(self, dtype=torch.float32):
+        super().__init__()
+        self.dtype = dtype
+
+    def forward(self, x):
+        return x.to(self.dtype) / 255
+
+
 class Cat2(torch.nn.Module):
     def __init__(self):
         super().__init__()
diff --git a/backends/qualcomm/tests/rework/src/utils.py b/backends/qualcomm/tests/rework/src/utils.py
index 9b564f83a7c..2b1232e5fff 100644
--- a/backends/qualcomm/tests/rework/src/utils.py
+++ b/backends/qualcomm/tests/rework/src/utils.py
@@ -36,6 +36,9 @@
 from executorch.backends.qualcomm.partition.utils import (
     warn_if_dilated_conv_may_not_fit_vtcm,
 )
+from executorch.backends.qualcomm.serialization.qc_schema import (
+    QnnExecuTorchBackendType,
+)
 from executorch.backends.qualcomm.tests.rework.conftest import (
     calibrate,
     check_exception,
@@ -333,7 +336,9 @@ def _build_op_wrapper_list(module, inputs):
         )
         nodes_to_wrappers = defaultdict(dict)
         node_visitors = get_node_visitors(
-            delegated_program.exported_program, enable_tensor_dump=False
+            delegated_program.exported_program,
+            enable_tensor_dump=False,
+            backend_type=QnnExecuTorchBackendType.kHtpBackend,
         )
         py_op_wrapper_list = []
         for node in graph_module.graph.nodes:
diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py
index c800c42bf2d..53bd790d7bd 100644
--- a/backends/qualcomm/tests/test_passes.py
+++ b/backends/qualcomm/tests/test_passes.py
@@ -54,6 +54,7 @@
 )
 from executorch.backends.qualcomm.tests.models import (
     BroadcastAndMutate,
+    CastToFloat,
     ConvBackward,
     HardSigmoid,
     Reciprocal,
@@ -1018,6 +1019,32 @@ def forward(self, x, i):
                 self.skipTest(f"QNN SDK not available: {e}")
             raise
 
+    def test_uint8_to_float32_cast_goes_through_int32_only_on_htp(self):
+        cases = [
+            (torch.float32, QnnExecuTorchBackendType.kHtpBackend, 2),
+            (torch.float32, QnnExecuTorchBackendType.kGpuBackend, 1),
+            (torch.float16, QnnExecuTorchBackendType.kHtpBackend, 1),
+        ]
+        for dtype, backend_type, expected in cases:
+            with self.subTest(dtype=dtype, backend_type=backend_type):
+                edge_program = to_edge(
+                    torch.export.export(
+                        CastToFloat(dtype),
+                        (torch.arange(4, dtype=torch.uint8),),
+                        strict=True,
+                    )
+                ).exported_program()
+                visitors = get_node_visitors(edge_program, backend_type=backend_type)
+                cast = next(
+                    node
+                    for node in edge_program.graph.nodes
+                    if node.op == "call_function" and node.args[0].op == "placeholder"
+                )
+                ops = visitors[cast.target.__name__].define_node(
+                    cast, defaultdict(dict)
+                )
+                self.assertEqual(expected, len(ops))
+
     def test_mha_to_sha(self):
         from executorch.backends.qualcomm.utils.utils import convert_linear_to_conv2d
         from executorch.examples.models.llama.model_args import ModelArgs
diff --git a/backends/qualcomm/tests/test_qnn_delegate.py b/backends/qualcomm/tests/test_qnn_delegate.py
index 0c64c690d47..33aec5a5f2e 100644
--- a/backends/qualcomm/tests/test_qnn_delegate.py
+++ b/backends/qualcomm/tests/test_qnn_delegate.py
@@ -529,10 +529,17 @@ 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(),  # noqa: F405
+            CastMultiUsers(),  # noqa: F405
+            CastToFloat(),  # noqa: F405
+            CastToFloat(torch.float16),  # noqa: F405
+        ]
         sample_inputs = [
             (10 * torch.rand((9, 4, 5, 3)),),
             (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3)),
+            (torch.arange(256, dtype=torch.uint8).reshape(16, 16),),
+            (torch.arange(256, dtype=torch.uint8).reshape(16, 16),),
         ]
         for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)):
             with self.subTest(i=i):
@@ -3726,10 +3733,18 @@ def test_qnn_backend_bmm(self):
                 self.lower_module_and_test_output(qdq_module, sample_input)
 
     def test_qnn_backend_cast(self):
-        module = CastMultiUsers()  # noqa: F405
-        sample_input = (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3))
-        module = self.get_qdq_module(module, sample_input)
-        self.lower_module_and_test_output(module, sample_input)
+        # With rtol 1, any output between 0 and twice the reference passes, so
+        # a uint8 cast that gives only 0 or 1 would pass too.
+        self.rtol = 1e-1
+        modules = [CastMultiUsers(), CastToFloat()]  # noqa: F405
+        sample_inputs = [
+            (torch.randint(0, 3, size=(3, 3)), torch.randn(3, 3)),
+            (torch.arange(256, dtype=torch.uint8).reshape(16, 16),),
+        ]
+        for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)):
+            with self.subTest(i=i):
+                module = self.get_qdq_module(module, sample_input)
+                self.lower_module_and_test_output(module, sample_input)
 
     def test_qnn_backend_cat(self):
         modules = [Cat2(), Cat3(), Cat4(), Cat5()]  # noqa: F405
@@ -7754,7 +7769,9 @@ def test_qnn_backend_draw_graph(self):
         )
         nodes_to_wrappers = defaultdict(dict)
         node_visitors = get_node_visitors(
-            delegated_program.exported_program, enable_tensor_dump=False
+            delegated_program.exported_program,
+            enable_tensor_dump=False,
+            backend_type=QnnExecuTorchBackendType.kHtpBackend,
         )
 
         py_op_wrapper_list = []
@@ -8916,7 +8933,9 @@ def test_qnn_backend_draw_graph(self):
         )
         nodes_to_wrappers = defaultdict(dict)
         node_visitors = get_node_visitors(
-            delegated_program.exported_program, enable_tensor_dump=False
+            delegated_program.exported_program,
+            enable_tensor_dump=False,
+            backend_type=QnnExecuTorchBackendType.kHtpBackend,
         )
 
         py_op_wrapper_list = []