diff --git a/backends/qualcomm/_passes/BUCK b/backends/qualcomm/_passes/BUCK index f5e635c739a..6d46f1972e1 100644 --- a/backends/qualcomm/_passes/BUCK +++ b/backends/qualcomm/_passes/BUCK @@ -14,6 +14,7 @@ fbcode_target(_kind = runtime.python_library, "//executorch/backends/transforms:addmm_mm_to_linear", "//executorch/backends/transforms:decompose_sdpa", "//executorch/backends/transforms:fuse_batch_norm_with_conv", + "//executorch/backends/transforms:utils", "//executorch/exir/backend:backend_details", "//executorch/exir/backend:compile_spec_schema", "//executorch/backends/qualcomm/quantizer:quantizer", diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index e8745524512..4c8d7e0cd8e 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -87,6 +87,7 @@ from executorch.backends.transforms.decompose_sdpa import ( DecomposeScaledDotProductAttention, ) +from executorch.backends.transforms.utils import reserve_node_names from executorch.exir import ExportedProgram from executorch.exir.pass_manager import PassManager from executorch.exir.program._program import ( @@ -100,6 +101,7 @@ class QnnPassManager(PassManager): def _transform(self, graph_module: GraphModule): + reserve_node_names(graph_module) return self(graph_module).graph_module def _reset(self): diff --git a/backends/qualcomm/partition/qnn_partitioner.py b/backends/qualcomm/partition/qnn_partitioner.py index 640d6fbf478..6cb4b1ca864 100644 --- a/backends/qualcomm/partition/qnn_partitioner.py +++ b/backends/qualcomm/partition/qnn_partitioner.py @@ -68,6 +68,15 @@ logger = logging.getLogger(__name__) logger.setLevel(logging.DEBUG) + +def _is_runtime_input( + node: torch.fx.Node, edge_program: torch.export.ExportedProgram +) -> bool: + return not is_parameter(node, edge_program) or is_mutable_buffer_input( + node, edge_program + ) + + # dequantize ops with a CPU out variant; others cannot fall back CPU_DQ_OPS = { exir_ops.edge.quantized_decomposed.dequantize_per_tensor.default, @@ -319,17 +328,11 @@ 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 - ) - return [ partition for partition in partitions if any( - is_runtime_input(arg) + _is_runtime_input(arg, edge_program) for node in partition.nodes for arg in node.all_input_nodes if arg not in partition.nodes @@ -396,7 +399,37 @@ 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) + The filter is `filter_fn`. On HTP, ops that read only constants are also + decomposed, except quantize and dequantize ops. """ do_not_decompose = get_skip_decomp_table() - return (do_not_decompose, filter_fn) + # On the lowered program this filter tells exir which kept-whole ops + # may stay on CPU, and none may, so it stays plain filter_fn there. + if self.backend != QnnExecuTorchBackendType.kHtpBackend or ep.dialect == "EDGE": + return (do_not_decompose, filter_fn) + # An op that reads only constants can end up in a partition that + # _drop_partitions_without_runtime_inputs keeps on CPU after + # decomposition, so every such op is decomposed. Partitions are not + # known yet, so this also decomposes ones whose partition stays + # delegated; folding constant-only subgraphs first would avoid that. + constant_only = set() + for node in ep.graph.nodes: + if node.op == "placeholder": + if not _is_runtime_input(node, ep): + constant_only.add(node) + elif node.op == "call_function" and all( + arg in constant_only for arg in node.all_input_nodes + ): + constant_only.add(node) + # Quantize and dequantize ops stay whole: QNN builds weight encodings + # from them. + quant_ops = { + op + for op in do_not_decompose + if op.namespace in ("quantized_decomposed", "torchao") + } + return ( + do_not_decompose, + lambda node: (node not in constant_only or node.target in quant_ops) + and filter_fn(node), + ) diff --git a/backends/qualcomm/tests/test_passes.py b/backends/qualcomm/tests/test_passes.py index 58ce89a6e30..59542c8d267 100644 --- a/backends/qualcomm/tests/test_passes.py +++ b/backends/qualcomm/tests/test_passes.py @@ -31,7 +31,10 @@ ) 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.qnn_partitioner import ( + QnnOperatorSupport, + QnnPartitioner, +) 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 ( @@ -188,6 +191,74 @@ def forward(self, x): ) torch.testing.assert_close(gm(x), model(x)) + def test_export_pipeline_keeps_node_names_unique(self): + # torch.export leaves the input name t_1 unrecorded, and renames q * 2 + # to mul_1 to make room for the input mul without recording that + # either. The decompositions would otherwise reuse those names. + class Autocast(torch.nn.Module): + def __init__(self): + super().__init__() + self.w = torch.nn.Parameter(torch.randn(8, 8)) + + def forward(self, x, t_1): + with torch.autocast("cpu", dtype=torch.bfloat16): + # The second transpose is named t_1 inside the region, and + # the decomposition copies it next to the input t_1. + y = x @ self.w.t() @ self.w.t() + return y.float() + t_1 + + class Attention(torch.nn.Module): + def forward(self, q, k, v, mul): + return ( + torch.nn.functional.scaled_dot_product_attention(q * 2, k, v) + mul + ) + + cases = ( + ( + Autocast(), + (torch.randn(4, 8), torch.randn(4, 8)), + torch._higher_order_ops.wrap.wrap_with_autocast, + ), + ( + Attention(), + tuple(torch.randn(1, 2, 4, 8) for _ in range(4)), + torch.ops.aten.scaled_dot_product_attention.default, + ), + ) + no_collision = [] + for model, inputs, decomposed_op in cases: + with self.subTest(model=type(model).__name__): + ep = torch.export.export(model.eval(), inputs) + self.assertIn(decomposed_op, [n.target for n in ep.graph.nodes]) + ep = get_qnn_pass_manager_cls()().transform_for_export_pipeline(ep) + self.assertNotIn(decomposed_op, [n.target for n in ep.graph.nodes]) + # ep.module() copies the graph and renames duplicates, so lint + # the graph the passes produced. + ep.graph.lint() + torch.testing.assert_close(ep.module()(*inputs), model(*inputs)) + ep = torch.export.export(model.eval(), inputs) + with patch( + "executorch.backends.qualcomm._passes.qnn_pass_manager." + "reserve_node_names" + ) as reserve: + try: + get_qnn_pass_manager_cls()().transform_for_export_pipeline( + ep + ).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)) + 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: + 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" + ) + def test_build_op_wrappers_returns_context_binary(self): op_name = "ctx_loader_build" ctx_bin = b"qnn_context_binary" @@ -946,16 +1017,19 @@ 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) + table = torch.nn.functional.linear(self.table, self.weight) + return x + torch.relu(table * 2 - 1) 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. + # With relu on CPU, the branch before it reads nothing at runtime. It + # also holds linear, which QNN otherwise keeps from decomposing. try: edge = to_edge_transform_and_lower_to_qnn( ConstantBranch().eval(), @@ -967,8 +1041,86 @@ def forward(self, x): if "Failed to initialize QNN backend" in str(e): self.skipTest(f"QNN SDK not available: {e}") raise + top_level_ops = { + node.target + for node in edge.exported_program().graph.nodes + if node.op == "call_function" + } + self.assertIn(exir_ops.edge.aten.mul.Tensor, top_level_ops) + self.assertIn(exir_ops.edge.aten.sub.Tensor, top_level_ops) + self.assertNotIn(exir_ops.edge.aten.add.Tensor, top_level_ops) + self.assertNotIn(exir_ops.edge.aten.linear.default, top_level_ops) + self.assertIn(torch.ops.higher_order.executorch_call_delegate, top_level_ops) edge.to_executorch() + def test_htp_keeps_whole_only_ops_that_read_a_runtime_input(self): + class Model(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( + "qweight", torch.randint(-8, 8, (16, 32), dtype=torch.int8) + ) + self.register_buffer("scales", torch.rand(16, 2) + 0.1) + self.register_buffer("zero_points", torch.zeros(16, 2).long()) + + def forward(self, x, y): + weight = torch.ops.quantized_decomposed.dequantize_per_channel_group( + self.qweight, + self.scales, + self.zero_points, + -8, + 7, + torch.int8, + 16, + torch.float32, + ) + table = torch.nn.functional.linear(self.table, self.weight) + return torch.nn.functional.linear(x, self.weight) + table, ( + torch.nn.functional.linear(y, weight) + ) + + ep = torch.export.export(Model(), (torch.randn(4, 16), torch.randn(4, 32))) + compiler_specs = generate_qnn_executorch_compiler_spec( + soc_model=QcomChipset.SM8650, + backend_options=generate_htp_compiler_spec(use_fp16=True), + ) + _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(ep) + linears = { + node.args[0].name: node + for node in ep.graph.nodes + if node.target == torch.ops.aten.linear.default + } + placeholders = {b: p for p, b in ep.graph_signature.inputs_to_buffers.items()} + self.assertFalse(keep_whole(linears[placeholders["table"]])) + self.assertTrue(keep_whole(linears["x"])) + (dequantize,) = ( + node + for node in ep.graph.nodes + if node.target + == torch.ops.quantized_decomposed.dequantize_per_channel_group.default + ) + self.assertTrue(keep_whole(dequantize)) + # After lowering, exir uses the filter to report a kept-whole op left + # on CPU, so there it must not let a constant-only one through. + edge = to_edge( + ep, + compile_config=EdgeCompileConfig( + _check_ir_validity=False, + preserve_ops=[torch.ops.aten.linear.default], + ), + ).exported_program() + _, keep_whole = QnnPartitioner(compiler_specs).ops_to_not_decompose(edge) + placeholders = {b: p for p, b in edge.graph_signature.inputs_to_buffers.items()} + (constant_linear,) = ( + node + for node in edge.graph.nodes + if node.target == exir_ops.edge.aten.linear.default + and node.args[0].name == placeholders["table"] + ) + self.assertTrue(keep_whole(constant_linear)) + def test_partitioner_falls_back_on_bfloat16(self): class Bf16Region(torch.nn.Module): def __init__(self): diff --git a/backends/samsung/_passes/enn_pass_manager.py b/backends/samsung/_passes/enn_pass_manager.py index 874aa7115c9..d148c0a57d5 100644 --- a/backends/samsung/_passes/enn_pass_manager.py +++ b/backends/samsung/_passes/enn_pass_manager.py @@ -33,6 +33,7 @@ ) from executorch.backends.transforms.remove_clone_ops import RemoveCloneOpsTransform from executorch.backends.transforms.remove_getitem_op import RemoveGetItemPass +from executorch.backends.transforms.utils import reserve_node_names from executorch.exir import ExportedProgram from executorch.exir.pass_manager import PassManager @@ -44,6 +45,7 @@ def __init__(self) -> None: super().__init__() def _transform(self, graph_module: GraphModule): + reserve_node_names(graph_module) return self(graph_module).graph_module # before annotator, for quant models diff --git a/backends/transforms/utils.py b/backends/transforms/utils.py index 92794364cb0..314dd8100d5 100644 --- a/backends/transforms/utils.py +++ b/backends/transforms/utils.py @@ -59,6 +59,19 @@ def _get_fake_tensor_mode(graph: torch.fx.Graph, data: torch.Tensor) -> torch.Te return FakeTensorConverter().from_real_tensor(example_fake_tensor.fake_mode, t=data) +def reserve_node_names(graph_module: torch.fx.GraphModule) -> None: + """ + Records every node name of the top-level graph in its namespace, so that + nodes a pass adds there later get names that are not already taken. + """ + # TODO: Remove once torch.export records the names it gives inputs and the + # ops it renames for them (https://github.com/pytorch/pytorch/issues/202927). + # torch.fx has no public way to reserve a name. + namespace = graph_module.graph._graph_namespace + for node in graph_module.graph.nodes: + namespace._used_names.add(node.name) + + def is_get_attr_node(node: torch.fx.Node) -> bool: """ Returns true if the given node is a get attr node for a tensor of the model