Qualcomm AI Engine Direct - Decompose kept-whole ops that HTP leaves on CPU

backends/qualcomm/builders/op_to.py

@@ -97,13 +97,25 @@ 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 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)
+            )
         ):
             input_quant_encoding, input_quant_configs = self.get_quant_encoding_conf(
                 input_node, node

backends/qualcomm/partition/qnn_partitioner.py

@@ -4,20 +4,19 @@
 # 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
 from executorch.backends.qualcomm.builders.node_visitor import (
     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,7 @@
     get_skip_decomp_table,
     has_data_dependent_shape,
     has_unsupported_dtype,
+    must_stay_on_cpu,
     warn_if_dilated_conv_may_not_fit_vtcm,
 )
 
@@ -310,6 +312,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
@@ -319,23 +341,114 @@ 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, and ops without a tensor output, such as the runtime
+        asserts export adds, 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
+        # Retracing can drop the "val" of an assert, but not its schema.
+        if (
+            isinstance(node.target, torch._ops.OpOverload)
+            and 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.
+        """
+
+        # The real partitions are only known once the kept-whole ops are
+        # decomposed, and partitioning would be too slow here, since its time
+        # grows with the square of the graph size. So only the ops that
+        # _may_run_on_qnn rules out count as CPU ops here.
+        # Only a partition of ops that read no runtime value is dropped. The
+        # partitioner also merges ops that no path links, so an op that reads
+        # no runtime value joins a partition that reads one, unless it reaches
+        # a QNN op through a CPU op without passing an op that reads one,
+        # since the partitioner may then split it off to avoid a cycle.
+        # In 2 cases this is not enough. When no op reads a runtime value,
+        # every partition is dropped. And the getitems of a CPU op join
+        # partitions while the partitioner merges them, but are moved out
+        # afterwards, so they can keep an op out of every partition that reads
+        # a runtime value. In these cases an op is also left on CPU when no
+        # chain of QNN ops links it to an op that reads one. Such an op can
+        # still end up delegated, in pieces.
+        # The caller's skips are left out, since decomposing adds nodes, and a
+        # skip must not rename the nodes that skip_node_id_set names. Ops that
+        # QNN rejects while partitioning, or that its passes remove before
+        # then, are not known yet either. So when such an op is all that cuts
+        # a kept-whole op off, or all that reads a runtime value, the op is
+        # not decomposed and still fails as before. And when such an op reads
+        # only constants and feeds a kept-whole op, the kept-whole op can be
+        # decomposed although it ends up delegated.
+        nodes = list(ep.graph.nodes)
+        qnn = set(filter(QnnPartitioner._may_run_on_qnn, nodes))
+        constant_inputs = QnnPartitioner._constant_inputs(ep)
+        runtime = set()
+        for node in nodes:
+            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
+        if runtime and not any(
+            node.target is operator.getitem and node.args[0] not in qnn
+            for node in runtime
+        ):
+            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:
@@ -397,6 +510,21 @@ 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.
         """
         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)
+
+        left_on_cpu = functools.cache(lambda: self._left_on_cpu(ep))
+
+        # exir calls this method several times on the ATen program but uses
+        # only one of the filters, so the estimate runs on first use.
+        return (
+            do_not_decompose,
+            lambda node: node not in left_on_cpu() and filter_fn(node),
+        )

backends/qualcomm/partition/utils.py

@@ -199,12 +199,20 @@ 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)
+
+
 # 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 = [

backends/qualcomm/tests/models.py

@@ -402,6 +402,14 @@ def forward(self, x, y):
         return res + index.to(torch.int32)
 
 
+class CastToFloat(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return x.to(torch.float32) / 255
+
+
 class Cat2(torch.nn.Module):
     def __init__(self):
         super().__init__()