diff --git a/backends/qualcomm/builders/op_to.py b/backends/qualcomm/builders/op_to.py index 9b024307934..192cbd25c61 100644 --- a/backends/qualcomm/builders/op_to.py +++ b/backends/qualcomm/builders/op_to.py @@ -97,13 +97,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 a plain uint8 to float32 + # cast goes through int32 too. Builders do not know the backend, so + # GPU gets the step as well. 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 diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 46c853331c6..7ca6147ac3e 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, + must_stay_on_cpu, + reads_dequantized, 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,123 @@ 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. Such + # an op is decomposed even when a partition that reads one still + # takes it. + # The getitems of a CPU op join partitions while the partitioner + # merges them, but are moved out afterwards. So when only such + # getitems read a runtime value, every partition is dropped. And when + # such a getitem comes after a QNN op, its partition cannot merge with + # that op's, so it can keep other ops out of every partition that reads + # a runtime value. Then 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, 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 +537,26 @@ 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. Ops that read dequantized values stay whole, so + in a quantized program 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) + + 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. + # An op that reads dequantized values stays whole, in case the estimate + # is wrong: its pieces would have no quantization encodings, so QNN + # would reject them where a partition would take them. + return ( + do_not_decompose, + lambda node: filter_fn(node) + and (node not in left_on_cpu() or reads_dequantized(node)), + ) diff --git a/backends/qualcomm/partition/utils.py b/backends/qualcomm/partition/utils.py index a608cc42960..eb00751235e 100644 --- a/backends/qualcomm/partition/utils.py +++ b/backends/qualcomm/partition/utils.py @@ -199,12 +199,28 @@ 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 reads_dequantized(node: torch.fx.Node) -> bool: + """ + Whether the node reads the output of a dequantize op, which makes it part + of a quantized region of the model. + """ + return 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/models.py b/backends/qualcomm/tests/models.py index 926f5e1fb67..6a8722b5eff 100644 --- a/backends/qualcomm/tests/models.py +++ b/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__() diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index a607bcb4d81..cc97e42b2cc 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -1,7 +1,7 @@ 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 +35,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,15 +64,21 @@ 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 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 @@ -257,12 +269,14 @@ 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. - if no_collision: + # A skip, not a failure: a torch update that removes the collision is + # fine. While any case still collides, the test still means something. + if len(no_collision) == len(cases): self.skipTest( f"{no_collision} no longer collide without reserve_node_names, so " "those cases no longer test it; see " @@ -1213,6 +1227,304 @@ def forward(self, x): raise edge.to_executorch() + def test_htp_decomposes_kept_whole_ops_left_on_cpu(self): # noqa: C901 + 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 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 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 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() + + sample_input = (torch.randn(4, 16),) + quantized = prepare_pt2e( + torch.export.export(QuantizedCast().eval(), sample_input).module(), + QnnQuantizer(), + ) + quantized(*sample_input) + # Runs 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() + }, + ): + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + gpu_compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_gpu_compiler_spec(), + ) + quantized_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=False), + ) + edge = to_edge_transform_and_lower_to_qnn( + ConstantBranches().eval(), sample_input, compiler_specs + ) + 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 + ) + kept_whole = {} + for model in ( + MaskedReader(), + ConstantOutput(), + CastSplitLinked(), + LinearOfInput(), + SharedConstant(), + UnlinkedSplit(), + ): + name = type(model).__name__ + with self.subTest(model=name): + kept_whole[name] = to_edge_transform_and_lower_to_qnn( + model.eval(), sample_input, compiler_specs + ) + with self.subTest(model="QuantizedCast"): + kept_whole["QuantizedCast"] = to_edge_transform_and_lower_to_qnn( + convert_pt2e(quantized), sample_input, quantized_specs + ) + program = edge.exported_program() + top_level_ops = { + node.target for node in program.graph.nodes if node.op == "call_function" + } + self.assertNotIn(exir_ops.edge.aten.linear.default, top_level_ops) + delegated_ops = { + node.target + for lowered in get_lowered_backend_modules(program.graph_module) + for node in lowered.original_module.graph.nodes + if node.op == "call_function" + } + self.assertIn(exir_ops.edge.aten.linear.default, delegated_ops) + self.assertIn(exir_ops.edge.aten.add.Tensor, delegated_ops) + # 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(program) + self.assertIs(keep_whole, filter_fn) + edge.to_executorch() + # GPU drops no partitions, so it keeps the same ops whole as before. + _, keep_whole = QnnPartitioner(gpu_compiler_specs).ops_to_not_decompose( + torch.export.export(ConstantBranches().eval(), sample_input) + ) + self.assertIs(keep_whole, filter_fn) + for name, edge in kept_whole.items(): + graph_module = edge.exported_program().graph_module + with self.subTest(model=name): + self.assertIn( + exir_ops.edge.aten.linear.default, + [ + node.target + for lowered in get_lowered_backend_modules(graph_module) + for node in lowered.original_module.graph.nodes + ], + ) + + # Skips must not change what is decomposed: decomposing adds nodes, + # which would rename the nodes that skip_node_id_set names. Skipping + # the relu cuts the linear off, so the model cannot be lowered. + ep = torch.export.export(SkippedRelu().eval(), sample_input) + _, 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_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..1d0e8d3b237 100644 --- a/backends/qualcomm/tests/test_qnn_delegate.py +++ b/backends/qualcomm/tests/test_qnn_delegate.py @@ -529,10 +529,11 @@ def test_qnn_backend_bmm(self): self.lower_module_and_test_output(module, sample_input) def test_qnn_backend_cast(self): - modules = [Cast(), CastMultiUsers()] # noqa: F405 + modules = [Cast(), CastMultiUsers(), CastToFloat()] # 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),), ] for i, (module, sample_input) in enumerate(zip(modules, sample_inputs)): with self.subTest(i=i): @@ -3726,10 +3727,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 diff --git a/backends/transforms/targets.bzl b/backends/transforms/targets.bzl index 4dd160aa44d..3722955f1b2 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..a3bf97b7f9e 100644 --- a/exir/backend/partitioner.py +++ b/exir/backend/partitioner.py @@ -105,6 +105,18 @@ 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. + + * 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 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