diff --git a/.github/workflows/pull.yml b/.github/workflows/pull.yml index 8a7494f02e3..33bfbebc24e 100644 --- a/.github/workflows/pull.yml +++ b/.github/workflows/pull.yml @@ -1325,6 +1325,10 @@ jobs: # test_mha_to_sha imports oss_scripts/llama which transitively needs torchtune pip install -r requirements-examples.txt + # The QNN manager loads the SDK libraries by name. Without this, every test + # that needs the real backend skips instead of running. + export LD_LIBRARY_PATH="$QNN_SDK_ROOT/lib/x86_64-linux-clang/:${LD_LIBRARY_PATH:-}" + # Run QNN pass unit tests pytest -xvs backends/qualcomm/tests/test_passes.py 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 46c853331c6..772512c17ee 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, ) @@ -92,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 @@ -327,6 +331,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 +360,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 +527,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/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..1188183ab84 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, @@ -177,6 +180,14 @@ class DrawGraph: dims: [1, 4, 8, 8] quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET > color=black fillcolor=transparent shape=box style=rounded] + "aten_permute_copy_default@0" [label=< + + + + + + +
name: aten_permute_copy_default@0
data_type: Qnn_DataType_t.QNN_DATATYPE_UFIXED_POINT_8
tensor_type: Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE
dims: [1, 8, 8, 4]
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET
> color=black fillcolor=transparent shape=box style=rounded] "input_1_y@0" [label=< @@ -193,6 +204,14 @@ class DrawGraph:
name: input_1_y@0
dims: [1, 4, 8, 8]
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET
> color=black fillcolor=transparent shape=box style=rounded] + "aten_permute_copy_default_1@0" [label=< + + + + + + +
name: aten_permute_copy_default_1@0
data_type: Qnn_DataType_t.QNN_DATATYPE_UFIXED_POINT_8
tensor_type: Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE
dims: [1, 8, 8, 4]
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET
> color=black fillcolor=transparent shape=box style=rounded] "b__frozen_param0@0" [label=< @@ -225,6 +244,14 @@ class DrawGraph:
name: b__frozen_param0@0
dims: [1, 8, 8, 4]
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET
> color=black fillcolor=transparent shape=box style=rounded] + "aten_permute_copy_default_3@0" [label=< + + + + + + +
name: aten_permute_copy_default_3@0
data_type: Qnn_DataType_t.QNN_DATATYPE_UFIXED_POINT_8
tensor_type: Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE
dims: [1, 4, 8, 8]
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET
> color=black fillcolor=transparent shape=box style=rounded] "output_quantized_decomposed_dequantize_per_tensor_default@0" [label=< @@ -297,6 +324,14 @@ class DrawGraph:
name: output_quantized_decomposed_dequantize_per_tensor_default@0
dims: [1, 8, 8, 4]
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET
> color=black fillcolor=transparent shape=box style=rounded] + "aten_permute_copy_default_2@0" [label=< + + + + + + +
name: aten_permute_copy_default_2@0
data_type: Qnn_DataType_t.QNN_DATATYPE_UFIXED_POINT_8
tensor_type: Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE
dims: [1, 4, 8, 8]
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_SCALE_OFFSET
> color=black fillcolor=transparent shape=box style=rounded] "output_quantized_decomposed_dequantize_per_tensor_default_1@0" [label=< @@ -306,13 +341,16 @@ class DrawGraph:
name: output_quantized_decomposed_dequantize_per_tensor_default_1@0
quantization_encoding: Qnn_QuantizationEncoding_t.QNN_QUANTIZATION_ENCODING_UNDEFINED
> color=black fillcolor=transparent shape=box style=rounded] "input_0_x@0" -> "quantized_decomposed_quantize_per_tensor_default@0" + "quantized_decomposed_quantize_per_tensor_default@0" -> "aten_permute_copy_default@0" "input_1_y@0" -> "quantized_decomposed_quantize_per_tensor_default_1@0" - "quantized_decomposed_quantize_per_tensor_default@0" -> "aten_convolution_default@0" + "quantized_decomposed_quantize_per_tensor_default_1@0" -> "aten_permute_copy_default_1@0" + "aten_permute_copy_default@0" -> "aten_convolution_default@0" "b__frozen_param0@0" -> "aten_convolution_default@0" "b__frozen_param1@0" -> "aten_convolution_default@0" "aten_convolution_default@0" -> "aten_relu_default@0" - "aten_relu_default@0" -> "output_quantized_decomposed_dequantize_per_tensor_default@0" - "quantized_decomposed_quantize_per_tensor_default_1@0" -> "aten_convolution_default_1@0" + "aten_relu_default@0" -> "aten_permute_copy_default_3@0" + "aten_permute_copy_default_3@0" -> "output_quantized_decomposed_dequantize_per_tensor_default@0" + "aten_permute_copy_default_1@0" -> "aten_convolution_default_1@0" "b__frozen_param2@0" -> "aten_convolution_default_1@0" "b__frozen_param3@0" -> "aten_convolution_default_1@0" "aten_convolution_default_1@0" -> "aten_relu_default_1@0" @@ -321,7 +359,8 @@ class DrawGraph: "aten_add_tensor@0" -> "aten_convolution_default_2@0" "b__frozen_param4@0" -> "aten_convolution_default_2@0" "b__frozen_param5@0" -> "aten_convolution_default_2@0" - "aten_convolution_default_2@0" -> "output_quantized_decomposed_dequantize_per_tensor_default_1@0" + "aten_convolution_default_2@0" -> "aten_permute_copy_default_2@0" + "aten_permute_copy_default_2@0" -> "output_quantized_decomposed_dequantize_per_tensor_default_1@0" } """ @@ -333,7 +372,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 a607bcb4d81..53bd790d7bd 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -1,7 +1,9 @@ +import contextlib +import logging 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 +37,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 ( @@ -46,6 +54,7 @@ ) from executorch.backends.qualcomm.tests.models import ( BroadcastAndMutate, + CastToFloat, ConvBackward, HardSigmoid, Reciprocal, @@ -58,20 +67,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" @@ -257,17 +504,22 @@ 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() - # 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" @@ -767,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 @@ -1190,29 +1468,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): 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 = [] 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() 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