From 770b70829423f208b7c1e4afac77c4f766b72be2 Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Wed, 2 Sep 2026 14:58:21 +0800 Subject: [PATCH 1/2] Qualcomm AI Engine Direct - Support Constant Propogation --- backends/qualcomm/_passes/__init__.py | 2 + backends/qualcomm/_passes/constant_folding.py | 255 ++++++++++++++ backends/qualcomm/_passes/qnn_pass_manager.py | 3 + backends/qualcomm/builders/utils.py | 8 +- backends/qualcomm/tests/rework/passes/test.py | 12 + backends/qualcomm/tests/rework/src/pattern.py | 324 ++++++++++++++++++ 6 files changed, 600 insertions(+), 4 deletions(-) create mode 100644 backends/qualcomm/_passes/constant_folding.py diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index 81dc1770856..723f4850bd4 100644 --- a/backends/qualcomm/_passes/__init__.py +++ b/backends/qualcomm/_passes/__init__.py @@ -11,6 +11,7 @@ from .annotate_unbind import AnnotateUnbind from .build_quant_io import BuildQuantIo from .canonicalize_conv import CanonicalizeConv +from .constant_folding import ConstantFolding from .convert_bmm_to_matmul import ConvertBmmToMatmul from .convert_linear_to_conv2d import ConvertLinearToConv2d from .convert_mha_to_sha import ConvertMhaToSha @@ -84,6 +85,7 @@ AnnotateUnbind, BuildQuantIo, CanonicalizeConv, + ConstantFolding, ConvertBmmToMatmul, ConvertLinearToConv2d, ConvertMhaToSha, diff --git a/backends/qualcomm/_passes/constant_folding.py b/backends/qualcomm/_passes/constant_folding.py new file mode 100644 index 00000000000..3702517e5c2 --- /dev/null +++ b/backends/qualcomm/_passes/constant_folding.py @@ -0,0 +1,255 @@ +# Copyright (c) Qualcomm Innovation Center, Inc. +# 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 logging + +import torch +from executorch.backends.qualcomm._passes.utils import copy_meta +from executorch.backends.qualcomm.builders.utils import ( + get_parameter, + is_mutable_buffer_input, + is_parameter, +) +from executorch.exir.operator.util import _QUANT_PRIMITIVES +from executorch.exir.pass_base import ExportPass, PassResult +from executorch.exir.passes import dead_code_elimination_pass +from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge +from torch._guards import detect_fake_mode +from torch.utils import _pytree as pytree +from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix + + +# These are qdq edge targets +_EDGE_QDQ_TARGETS = set(_QUANT_PRIMITIVES) | { + aten_to_edge(op) for op in _QUANT_PRIMITIVES +} + +# copied from executorch/exir/passes/const_prop_pass.py +_PRIMITIVE_TYPES = ( + float, + int, + bool, + str, + torch.Tensor, + torch.device, + torch.dtype, + torch.layout, +) + + +class ConstantFolding(ExportPass): + """ + executorch/exir/passes/const_prop_pass.py runs at to_executorch stage, and it won't + work if we run the pass at edge stage. This pass is to address the issue. + One of biggest reason for this pass is that some ops are not supported by QNN and will + fail during op validation. If the op can be constant folded, then there won't be partitions. + + + Folds subgraphs whose leaves are parameters, lifted tensor constants, or + *non-mutated* buffers. Mutated buffers are intentionally excluded — view ops + over them would alias mutable storage if folded, and downstream passes / + `run_decompositions` mis-handle that aliasing. + """ + + _TENSOR_CONSTANT_PREFIX = "_prop_tensor_constant_" + + def __init__(self, edge_program: torch.export.ExportedProgram): + super().__init__() + # Run decomposition is required so graph_signature stores information about mutable buffer. + self.decomposed_aten_program = edge_program.run_decompositions({}) + + def _get_const_placeholders( + self, graph_module: torch.fx.GraphModule + ) -> dict[torch.fx.Node, torch.Tensor]: + """ + Find all constant_tensor and store it in a dict {node : tensor} + The tensor in dict has actual tensor with values, not fake tensor. + """ + + node_to_tensor = {} + for node in graph_module.graph.nodes: + if node.op != "placeholder": + continue + # Don't fold mutable buffer + if is_mutable_buffer_input(node, self.decomposed_aten_program): + continue + if not is_parameter(node, self.decomposed_aten_program): + continue + node_to_tensor[node] = get_parameter(node, self.decomposed_aten_program) + return node_to_tensor + + def _propagate( # noqa: C901 + self, + graph_module: torch.fx.GraphModule, + node_to_tensor: dict[torch.fx.Node, torch.Tensor], + ) -> None: + """ + Iterate the call_function node. compute the output of that node if it can be constant folded, + and save it in node_to_tensor. Based on torch/fx/graph.py, graph.nodes should always return + nodes in topological order. + Example: + constant_value_1 -> upsample_bicubic2d ----\ + add_1 --------------------------\ + constant_value_2 ----/ add_2 ----> output + user_input(non-constant)------/ + + With graph above, when first enter this method, node_to_tensor dict looks like: + node_to_tensor = { + constant_value_1_node : constant_value_1_tensor, + constant_value_2_node : constant_value_2_tensor, + } + + + At the end of the method, node_to_tensor dict looks like following: + node_to_tensor = { + constant_value_1_node : constant_value_1_tensor, + constant_value_2_node : constant_value_2_tensor, + upsample_bicubic2d_node : upsample_bicubic2d_output_tensor, + add_1_node : add_1_output_tensor, + } + + Tensor in dict will be actual values instead of fake tensor. + For example, add_1_output_tensor value would be the result of adding upsample_bicubic2d_output_tensor and constant_value_2. + """ + + # If previous node is visited, this function won't recursive search all the way back to input node. + # It should fall into the base case. + # This prevents recursive search back to input source node everytime. + def _is_const(arg, node_to_tensor): + + # For args case. + if isinstance(arg, (tuple, list)): + return all(_is_const(x, node_to_tensor) for x in arg) + + # For kwargs case + if isinstance(arg, dict): + return all(_is_const(x, node_to_tensor) for x in arg.values()) + + # If a tensor is optional and not provided, it will be None. + # Primitive_types is for constants like integers. These should be able to be folded. + if arg is None or isinstance(arg, _PRIMITIVE_TYPES): + return True + + if isinstance(arg, torch.fx.Node): + # Base case + return arg in node_to_tensor + else: + # If there are some unexpected args that doesn't know how to handle, just return False to be safe. + return False + + # `pytree.tree_map` flattens containers and treats Node as a leaf, so this + # only ever receives leaves. + def _get_data(arg, node_to_tensor): + if arg is None or isinstance(arg, _PRIMITIVE_TYPES): + return arg + if isinstance(arg, torch.fx.Node): + return node_to_tensor.get(arg) + return None + + for node in graph_module.graph.nodes: + if node.op != "call_function": + continue + if not _is_const(node.args, node_to_tensor): + continue + if not _is_const(node.kwargs, node_to_tensor): + continue + + # Copied from executorch/exir/passes/const_prop_pass.py + # Retrieves args and kwargs required for the node to perform inference. + args_data, kwargs_data = pytree.tree_map( + lambda x: _get_data(x, node_to_tensor), + (node.args, node.kwargs), + ) + # Perform node inference + with torch.no_grad(): + try: + result = node.target(*args_data, **kwargs_data) + except Exception: + logging.warning( + f"Unable to fold the node {node.name}. Skip folding.", + exc_info=True, + ) + continue + + if isinstance(result, torch.Tensor): + result = result.detach().clone(memory_format=torch.contiguous_format) + + # Save the node's result to the map. + node_to_tensor[node] = result + + def _materialize_as_buffer( + self, + graph_module: torch.fx.GraphModule, + node: torch.fx.Node, + tensor: torch.Tensor, + ) -> None: + buffer_name = get_new_attr_name_with_prefix(self._TENSOR_CONSTANT_PREFIX)( + graph_module + ) + graph_module.register_buffer(buffer_name, tensor) + val = node.meta.get("val") + fake_mode = detect_fake_mode(val) if val is not None else None + with graph_module.graph.inserting_before(node): + get_attr_node = graph_module.graph.get_attr(buffer_name) + get_attr_node.meta = copy_meta( + node.meta, + lambda m: ( + { + **m, + "val": fake_mode.fake_tensor_converter.from_real_tensor( + fake_mode, tensor + ), + } + if fake_mode is not None + else m + ), + ) + # Replace node's user with const node as input + node.replace_all_uses_with(get_attr_node) + + def _materialize( + self, + graph_module: torch.fx.GraphModule, + node_to_tensor: dict[torch.fx.Node, torch.Tensor], + ) -> None: + """ + The term "boundary" here refers to where node can no longer be folded. + Boundry will be before add for this graph is const -> relu1 -> sqrt --- > add -> output + input _| + + Rules: + 1) When creating buffer, start with reverse order, so just create the buffer before boundary. + 2) Only the boundary buffer will be created, won't recursive trace args and create unused buffer. + 3) For nodes like conv2d with quantizer, preserve dq node right afer weight and bias. + """ + + # Reverse order: process later (more-derived) constants first, so a + # chain collapses to a single buffer at its boundary. + # Align with rule 1. + for node, tensor in reversed(list(node_to_tensor.items())): + if node.op == "placeholder": + continue + + # If all users can be constant folded, then don't need to fold at this level. + # This behavior aligns with rule 2. + if all(user in node_to_tensor for user in node.users): + continue + + # Guarding cases like weight -> dq -> conv2d. Don't fold dq here. + # Aligns with rule 3 + if node.target in _EDGE_QDQ_TARGETS: + continue + + self._materialize_as_buffer(graph_module, node, tensor) + + def call(self, graph_module: torch.fx.GraphModule): + node_to_tensor = self._get_const_placeholders(graph_module) + if len(node_to_tensor) == 0: + return PassResult(graph_module, False) + self._propagate(graph_module, node_to_tensor) + self._materialize(graph_module, node_to_tensor) + dead_code_elimination_pass(graph_module) + graph_module.recompile() + return PassResult(graph_module, True) diff --git a/backends/qualcomm/_passes/qnn_pass_manager.py b/backends/qualcomm/_passes/qnn_pass_manager.py index 40af0946366..18819a1ea4f 100644 --- a/backends/qualcomm/_passes/qnn_pass_manager.py +++ b/backends/qualcomm/_passes/qnn_pass_manager.py @@ -16,6 +16,7 @@ AnnotateStack, AnnotateUnbind, CanonicalizeConv, + ConstantFolding, ConvertBmmToMatmul, ConvertLinearToConv2d, ConvertMhaToSha, @@ -131,6 +132,7 @@ def get_default_pass_activations(cls): (AnnotateStack, True), (AnnotateUnbind, True), (CanonicalizeConv, True), + (ConstantFolding, True), (ConvertBmmToMatmul, False), (ConvertLinearToConv2d, False), (DecomposeAcos, True), @@ -300,6 +302,7 @@ def get_passes_dependency_for_capture_program(cls): AnnotateStack: [RemoveRedundancy], AnnotateUnbind: [RemoveRedundancy], CanonicalizeConv: [FoldQDQ, FuseBatchNormWithConv], + ConstantFolding: [LayoutTransform], ConvertBmmToMatmul: [RecomposePixelUnshuffle], ConvertLinearToConv2d: [FoldQDQ], DecomposeAcos: [RemoveRedundancy], diff --git a/backends/qualcomm/builders/utils.py b/backends/qualcomm/builders/utils.py index c691ef7ffd3..037a75c4438 100755 --- a/backends/qualcomm/builders/utils.py +++ b/backends/qualcomm/builders/utils.py @@ -104,15 +104,15 @@ def is_graph_input( def is_mutable_buffer_input( - tensor: torch.fx.Node, edge_program: torch.export.ExportedProgram + node: torch.fx.Node, edge_program: torch.export.ExportedProgram ) -> bool: """ Check if the given tensor is a mutable buffer input Args: - tensor: EdgeIR Tensor that is being checked for mutable buffer input + node: EdgeIR node that is being checked for mutable buffer input """ - if tensor.op == "placeholder" and is_buffer(edge_program, tensor): - fqn = edge_program.graph_signature.inputs_to_buffers[tensor.target] + if node.op == "placeholder" and is_buffer(edge_program, node): + fqn = edge_program.graph_signature.inputs_to_buffers[node.target] # if the buffer is mutated then record that return fqn in edge_program.graph_signature.buffers_to_mutate.values() diff --git a/backends/qualcomm/tests/rework/passes/test.py b/backends/qualcomm/tests/rework/passes/test.py index 1bf2c76d28c..7d863fffa68 100644 --- a/backends/qualcomm/tests/rework/passes/test.py +++ b/backends/qualcomm/tests/rework/passes/test.py @@ -59,6 +59,18 @@ def test_canonicalize_conv(request, kwargs): CanonicalizeConv.test(request, kwargs) # noqa: F405 +@enumerate_backends() +@repack_pass_fixtures +def test_constant_folding(request, kwargs): + ConstantFolding.test(request, kwargs) # noqa: F405 + + +@enumerate_backends() +@repack_pass_fixtures +def test_constant_folding_idempotent(request, kwargs): + ConstantFolding.test_idempotent(request, kwargs) # noqa: F405 + + @enumerate_backends() @repack_pass_fixtures def test_convert_bmm_to_matmul(request, kwargs): diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 7b90fde61e2..8a9b6ff331f 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -9,6 +9,7 @@ import inspect import math import operator +from functools import partial from typing import TYPE_CHECKING import pytest @@ -489,6 +490,329 @@ def test( CanonicalizeConv._assert_no_dilation(gm) +class ConstantFolding: + # --- Group 1: foldable const chains --- + class _ConstChain(torch.nn.Module): + """relu(a*b) -> permute -> reshape: whole chain collapses to one const.""" + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(1, 3, 4, 4)) + self.register_buffer("b", torch.randn(1, 3, 4, 4)) + + def forward(self, x): + t = torch.relu(self.a * self.b).permute(0, 1, 3, 2).reshape(1, 3, 4, 4) + return x + t + + class _ComplicatedChain(torch.nn.Module): + """Ensure the pass handles the case where mutli user and a user can be folded while another cannot.""" + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(1, 3, 4, 4)) + self.register_buffer("b", torch.randn(1, 3, 4, 4)) + self.register_buffer("c", torch.randn(1, 3, 4, 4)) + + def forward(self, x): + const1 = self.a + self.b + const2 = const1 + self.c + non_const1 = x + const1 + return const2 + non_const1 + + class _MultiUserConst(torch.nn.Module): + """One folded const feeding two distinct real ops: fold once, share it.""" + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(1, 3, 4, 4)) + self.register_buffer("b", torch.randn(1, 3, 4, 4)) + + def forward(self, x, y): + t = self.a * self.b + return x + t, y - t + + class _Bicubic(torch.nn.Module): + """Motivating case: bicubic upsample of a position-embedding buffer.""" + + def __init__(self): + super().__init__() + self.register_buffer("pos_emb", torch.randn(1, 3, 4, 4)) + + def forward(self, x): + up = torch.nn.functional.interpolate( + self.pos_emb, scale_factor=2.0, mode="bicubic" + ) + return x + up + + class _ConstReduction(torch.nn.Module): + """Reduction over a const view: shape args are ints, result is a tensor.""" + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(2, 6)) + + def forward(self, x): + return x + self.a.reshape(2, 3, 2).sum(-1) + + class _ConstDtypeChange(torch.nn.Module): + """Const compare + cast: the folded value changes dtype (bool) mid-chain, + so the materialized constant must carry the final dtype, not the leaves'.""" + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(4, 4)) + self.register_buffer("b", torch.randn(4, 4)) + + def forward(self, x): + return x + (self.a > self.b).to(torch.float32) + + class _ConstGraphOutput(torch.nn.Module): + """A folded const that *is* a graph output must stay materialized.""" + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(4, 4)) + self.register_buffer("b", torch.randn(4, 4)) + + def forward(self, x): + return self.a * self.b, x + x + + # --- Group 2: chains that must NOT fold --- + class _MutableBuffer(torch.nn.Module): + """Ensure Mutable Buffer is not folded""" + + def __init__(self): + super().__init__() + self.register_buffer("cache", torch.zeros(1, 3, 4, 4)) + self.register_buffer("w", torch.randn(1, 3, 4, 4)) + + def forward(self, x, idx): + scaled = self.cache * self.w + self.cache.index_put_((idx,), x) + return x + scaled + + class _UserInputChain(torch.nn.Module): + """Const mixed with a user input: the op depends on x, so it cannot fold.""" + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(1, 3, 4, 4)) + + def forward(self, x): + return (self.a * x) + self.a + + class _Conv2D(torch.nn.Module): + # Ensure weight and bias dq is not folded + + def __init__(self): + super().__init__() + self.conv = torch.nn.Conv2d(3, 4, 3, padding=1, bias=True) + + def forward(self, x): + return self.conv(x) + + @staticmethod + def _lower(pass_pipeline, module, inputs, backend_type, compile_spec, quantizer): + return pass_pipeline.lower_edge_ep( + module=module, + sample_input=inputs, + target_pass=_passes.ConstantFolding, + backend_type=backend_type, + compile_spec=compile_spec, + quantizer=quantizer, + ) + + # TODO: Make this a general helper test and verify in all passes. + @staticmethod + def _assert_numeric_match(edge_ep, module, inputs, is_fp): + """ + Bigger tolerance for quantized since weights got quantized. + """ + got = edge_ep.module()(*inputs) + got = got if isinstance(got, (tuple, list)) else (got,) + ref = module(*inputs) + ref = ref if isinstance(ref, (tuple, list)) else (ref,) + tolerance = 1e-5 if is_fp else 3e-1 + for i, (g, r) in enumerate(zip(got, ref)): + assert torch.allclose(g, r, atol=tolerance, rtol=tolerance), ( + f"ConstantFolding changed output {i}: " + f"max abs diff = {(g - r).abs().max().item()}" + ) + + @staticmethod + @unpack_pass_fixtures + def test( + subtests, + quantizer, + compile_spec, + backend_type: QnnExecuTorchBackendType, + assertions: Assertions, + pass_pipeline: PassPipeline, + ): + lower = partial( + ConstantFolding._lower, + pass_pipeline, + backend_type=backend_type, + compile_spec=compile_spec, + quantizer=quantizer, + ) + is_fp = quantizer is None + + # --- Group 1: const chains that fold on an unquantized graph --- + with subtests.test(msg="const_chain_collapses"): + inputs = (torch.randn(1, 3, 4, 4),) + module = ConstantFolding._ConstChain() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + for target in ( + exir_ops.edge.aten.mul.Tensor, + exir_ops.edge.aten.relu.default, + exir_ops.edge.aten.permute_copy.default, + exir_ops.edge.aten.view_copy.default, + ): + assertions.assert_no_target(gm, target) + assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + with subtests.test(msg="complicated_chain"): + inputs = (torch.randn(1, 3, 4, 4),) + module = ConstantFolding._ComplicatedChain() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 2) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + with subtests.test(msg="upsample_bicubic2d"): + inputs = (torch.randn(1, 3, 8, 8),) + module = ConstantFolding._Bicubic() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_no_target(gm, exir_ops.edge.aten.upsample_bicubic2d.vec) + assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + with subtests.test(msg="const_shared_by_two_users"): + inputs = (torch.randn(1, 3, 4, 4), torch.randn(1, 3, 4, 4)) + module = ConstantFolding._MultiUserConst() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + # One fold serves both users; neither real op is duplicated. + assertions.assert_no_target(gm, exir_ops.edge.aten.mul.Tensor) + assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) + assertions.assert_target_count(gm, exir_ops.edge.aten.sub.Tensor, 1) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + with subtests.test(msg="const_reduction"): + inputs = (torch.randn(2, 3),) + module = ConstantFolding._ConstReduction() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_no_target(gm, exir_ops.edge.aten.sum.dim_IntList) + assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + with subtests.test(msg="const_dtype_change"): + inputs = (torch.randn(4, 4),) + module = ConstantFolding._ConstDtypeChange() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_no_target(gm, exir_ops.edge.aten.gt.Tensor) + assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + with subtests.test(msg="folded_const_as_graph_output"): + inputs = (torch.randn(4, 4),) + module = ConstantFolding._ConstGraphOutput() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_no_target(gm, exir_ops.edge.aten.mul.Tensor) + out_args = gm.graph.output_node().args[0] + assert len(out_args) == 2, f"expected 2 graph outputs, got {len(out_args)}" + const_out = out_args[0] + assert const_out.op in ("placeholder", "get_attr"), ( + f"folded const output should be a materialized constant, " + f"got {const_out.op} ({const_out.target})" + ) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + # --- Group 2: chains that must never fold, quantized or not --- + with subtests.test(msg="mutable_buffer_not_folded"): + inputs = (torch.randn(1, 3, 4, 4), torch.tensor([0])) + module = ConstantFolding._MutableBuffer() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.mul.Tensor, 1) + cache_nodes = [ + n for n in gm.graph.nodes if n.op == "placeholder" and "cache" in n.name + ] + assert ( + len(cache_nodes) == 1 + ), "mutable buffer placeholder disappeared entirely" + mul_nodes = [ + n + for n in gm.graph.nodes + if n.op == "call_function" and n.target == exir_ops.edge.aten.mul.Tensor + ] + assert len(mul_nodes) == 1, "mul node disappeared entirely" + cache_node = cache_nodes[0] + mul_node = mul_nodes[0] + assert ( + len(cache_node.users) == 2 + ), "mutable buffer has no users left: should have multiply and index_put" + assert cache_node in mul_node.args, ( + f"mul no longer reads the mutable buffer directly; " + f"args={[getattr(a, 'name', a) for a in mul_node.args]}" + ) + + with subtests.test(msg="user_input_blocks_fold"): + inputs = (torch.randn(1, 3, 4, 4),) + module = ConstantFolding._UserInputChain() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + # `a * x` depends on x, so neither it nor the following add is_fp. + assertions.assert_target_count(gm, exir_ops.edge.aten.mul.Tensor, 1) + assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + # --- Group 3: quantized-only invariant --- + if quantizer is not None: + with subtests.test(msg="preserve_conv_dq"): + inputs = (torch.randn(1, 3, 8, 8),) + module = ConstantFolding._Conv2D() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_target_count(gm, dq_ops, 2) + assertions.assert_target_count( + gm, exir_ops.edge.aten.convolution.default, 1 + ) + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + + @staticmethod + @unpack_pass_fixtures + def test_idempotent( + quantizer, + compile_spec, + backend_type: QnnExecuTorchBackendType, + pass_pipeline: PassPipeline, + ): + """Ensuring that running this pass 1 time V.S multiple times returns the same graph.""" + inputs = (torch.randn(1, 3, 4, 4),) + edge_ep = ConstantFolding._lower( + pass_pipeline, + ConstantFolding._ConstChain(), + inputs, + backend_type=backend_type, + compile_spec=compile_spec, + quantizer=quantizer, + ) + before = [n.name for n in edge_ep.graph.nodes] + gm = _passes.ConstantFolding(edge_ep)(edge_ep.graph_module).graph_module + after = [n.name for n in gm.graph.nodes] + assert ( + before == after + ), f"ConstantFolding is not idempotent:\nbefore={before}\nafter ={after}" + + class ConvertBmmToMatmul: class _Basic(torch.nn.Module): def forward(self, x, y): From 63f7ca47d02f3dbd8520feb9add56e3029673b1e Mon Sep 17 00:00:00 2001 From: Winston Kuo Date: Wed, 7 Oct 2026 14:50:21 +0800 Subject: [PATCH 2/2] Address Claude Code Review --- backends/qualcomm/_passes/constant_folding.py | 47 ++++++------- backends/qualcomm/tests/rework/src/pattern.py | 70 ++++++++++++++++++- 2 files changed, 90 insertions(+), 27 deletions(-) diff --git a/backends/qualcomm/_passes/constant_folding.py b/backends/qualcomm/_passes/constant_folding.py index 3702517e5c2..cddf40b13f2 100644 --- a/backends/qualcomm/_passes/constant_folding.py +++ b/backends/qualcomm/_passes/constant_folding.py @@ -4,6 +4,7 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. import logging +from typing import Any import torch from executorch.backends.qualcomm._passes.utils import copy_meta @@ -15,6 +16,7 @@ from executorch.exir.operator.util import _QUANT_PRIMITIVES from executorch.exir.pass_base import ExportPass, PassResult from executorch.exir.passes import dead_code_elimination_pass +from executorch.exir.passes.constant_prop_pass import _PRIMITIVE_TYPES from executorch.exir.passes.replace_aten_with_edge_pass import aten_to_edge from torch._guards import detect_fake_mode from torch.utils import _pytree as pytree @@ -26,18 +28,6 @@ aten_to_edge(op) for op in _QUANT_PRIMITIVES } -# copied from executorch/exir/passes/const_prop_pass.py -_PRIMITIVE_TYPES = ( - float, - int, - bool, - str, - torch.Tensor, - torch.device, - torch.dtype, - torch.layout, -) - class ConstantFolding(ExportPass): """ @@ -83,7 +73,7 @@ def _get_const_placeholders( def _propagate( # noqa: C901 self, graph_module: torch.fx.GraphModule, - node_to_tensor: dict[torch.fx.Node, torch.Tensor], + node_to_tensor: dict[torch.fx.Node, Any], ) -> None: """ Iterate the call_function node. compute the output of that node if it can be constant folded, @@ -93,15 +83,14 @@ def _propagate( # noqa: C901 constant_value_1 -> upsample_bicubic2d ----\ add_1 --------------------------\ constant_value_2 ----/ add_2 ----> output - user_input(non-constant)------/ - + user_input(non-constant)------/ + With graph above, when first enter this method, node_to_tensor dict looks like: node_to_tensor = { constant_value_1_node : constant_value_1_tensor, constant_value_2_node : constant_value_2_tensor, } - - + At the end of the method, node_to_tensor dict looks like following: node_to_tensor = { constant_value_1_node : constant_value_1_tensor, @@ -109,7 +98,7 @@ def _propagate( # noqa: C901 upsample_bicubic2d_node : upsample_bicubic2d_output_tensor, add_1_node : add_1_output_tensor, } - + Tensor in dict will be actual values instead of fake tensor. For example, add_1_output_tensor value would be the result of adding upsample_bicubic2d_output_tensor and constant_value_2. """ @@ -151,12 +140,15 @@ def _get_data(arg, node_to_tensor): for node in graph_module.graph.nodes: if node.op != "call_function": continue + # Avoid folding nodes like torch.rand since we shouldn't precompute. + if node.is_impure(): + continue if not _is_const(node.args, node_to_tensor): continue if not _is_const(node.kwargs, node_to_tensor): continue - # Copied from executorch/exir/passes/const_prop_pass.py + # Copied from executorch/exir/passes/constant_prop_pass.py # Retrieves args and kwargs required for the node to perform inference. args_data, kwargs_data = pytree.tree_map( lambda x: _get_data(x, node_to_tensor), @@ -173,9 +165,6 @@ def _get_data(arg, node_to_tensor): ) continue - if isinstance(result, torch.Tensor): - result = result.detach().clone(memory_format=torch.contiguous_format) - # Save the node's result to the map. node_to_tensor[node] = result @@ -188,6 +177,7 @@ def _materialize_as_buffer( buffer_name = get_new_attr_name_with_prefix(self._TENSOR_CONSTANT_PREFIX)( graph_module ) + tensor = tensor.detach().clone(memory_format=torch.contiguous_format) graph_module.register_buffer(buffer_name, tensor) val = node.meta.get("val") fake_mode = detect_fake_mode(val) if val is not None else None @@ -212,17 +202,17 @@ def _materialize_as_buffer( def _materialize( self, graph_module: torch.fx.GraphModule, - node_to_tensor: dict[torch.fx.Node, torch.Tensor], + node_to_tensor: dict[torch.fx.Node, Any], ) -> None: """ The term "boundary" here refers to where node can no longer be folded. - Boundry will be before add for this graph is const -> relu1 -> sqrt --- > add -> output - input _| + Boundary will be before add for this graph is const -> relu1 -> sqrt --- > add -> output + input _| Rules: 1) When creating buffer, start with reverse order, so just create the buffer before boundary. 2) Only the boundary buffer will be created, won't recursive trace args and create unused buffer. - 3) For nodes like conv2d with quantizer, preserve dq node right afer weight and bias. + 3) For nodes like conv2d with quantizer, preserve dq node right after weight and bias. """ # Reverse order: process later (more-derived) constants first, so a @@ -232,6 +222,11 @@ def _materialize( if node.op == "placeholder": continue + # A folded value can be a scalar, SymInt or tuple (`_local_scalar_dense`, + # `sym_size`, multi-output ops). Only a tensor can become a buffer. + if not isinstance(tensor, torch.Tensor): + continue + # If all users can be constant folded, then don't need to fold at this level. # This behavior aligns with rule 2. if all(user in node_to_tensor for user in node.users): diff --git a/backends/qualcomm/tests/rework/src/pattern.py b/backends/qualcomm/tests/rework/src/pattern.py index 8a9b6ff331f..946eccdb36a 100644 --- a/backends/qualcomm/tests/rework/src/pattern.py +++ b/backends/qualcomm/tests/rework/src/pattern.py @@ -611,6 +611,25 @@ def __init__(self): def forward(self, x): return self.conv(x) + class _ImpureOp(torch.nn.Module): + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(1, 3, 4, 4)) + self.register_buffer("b", torch.randn(1, 3, 4, 4)) + + def forward(self, x): + return x + (self.a * self.b) + torch.rand(1, 3, 4, 4) + + class _NonTensorBoundary(torch.nn.Module): + + def __init__(self): + super().__init__() + self.register_buffer("a", torch.randn(1, 3, 4, 4)) + + def forward(self, x): + return torch.clamp(x, min=self.a.min().item()) + @staticmethod def _lower(pass_pipeline, module, inputs, backend_type, compile_spec, quantizer): return pass_pipeline.lower_edge_ep( @@ -769,11 +788,31 @@ def test( module = ConstantFolding._UserInputChain() edge_ep = lower(module, inputs) gm = edge_ep.graph_module - # `a * x` depends on x, so neither it nor the following add is_fp. + # `a * x` depends on x, so neither it nor the following add is folded. assertions.assert_target_count(gm, exir_ops.edge.aten.mul.Tensor, 1) assertions.assert_target_count(gm, exir_ops.edge.aten.add.Tensor, 1) ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + with subtests.test(msg="impure_op_not_folded"): + inputs = (torch.randn(1, 3, 4, 4),) + module = ConstantFolding._ImpureOp() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assertions.assert_target_count(gm, exir_ops.edge.aten.rand.default, 1) + assertions.assert_no_target(gm, exir_ops.edge.aten.mul.Tensor) + + with subtests.test(msg="non_tensor_boundary_not_materialized"): + inputs = (torch.randn(1, 3, 4, 4),) + module = ConstantFolding._NonTensorBoundary() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + assert not [ + n + for n in gm.graph.nodes + if _passes.ConstantFolding._TENSOR_CONSTANT_PREFIX in n.name + ], "the scalar boundary was materialized as a buffer" + ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + # --- Group 3: quantized-only invariant --- if quantizer is not None: with subtests.test(msg="preserve_conv_dq"): @@ -787,6 +826,35 @@ def test( ) ConstantFolding._assert_numeric_match(edge_ep, module, inputs, is_fp) + with subtests.test(msg="quantized_fold_boundary_keeps_quant_attrs"): + inputs = (torch.randn(1, 3, 8, 8),) + module = ConstantFolding._Bicubic() + edge_ep = lower(module, inputs) + gm = edge_ep.graph_module + # The boundary here is upsample_bicubic2d, not a dq, so the folded + # constant only stays quantized if it inherits the boundary's attrs. + user_inputs = set(edge_ep.graph_signature.user_inputs) + add_node = [ + n + for n in gm.graph.nodes + if n.op == "call_function" + and n.target == exir_ops.edge.aten.add.Tensor + ][0] + folded = [ + a + for a in add_node.args + if isinstance(a, torch.fx.Node) + and a.op == "placeholder" + and a.name not in user_inputs + ] + assert ( + len(folded) == 1 + ), f"expected one folded constant feeding add, got {folded}" + assert QCOM_QUANT_ATTRS in folded[0].meta, ( + f"folded constant {folded[0].name} lost its quant attrs; a float " + f"value would reach add in a quantized graph" + ) + @staticmethod @unpack_pass_fixtures def test_idempotent(