diff --git a/backends/qualcomm/_passes/__init__.py b/backends/qualcomm/_passes/__init__.py index 2a059b47471..d7d916c4b6a 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 @@ -81,6 +82,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 2bf5e3f28b3..54b74c4c268 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, @@ -128,6 +129,7 @@ def get_default_pass_activations(cls): (AnnotateStack, True), (AnnotateUnbind, True), (CanonicalizeConv, True), + (ConstantFolding, True), (ConvertBmmToMatmul, False), (ConvertLinearToConv2d, False), (DecomposeAcos, True), @@ -288,6 +290,7 @@ def get_passes_dependency_for_capture_program(cls): AnnotateStack: [RemoveRedundancy], AnnotateUnbind: [RemoveRedundancy], CanonicalizeConv: [FoldQDQ], + ConstantFolding: [LayoutTransform], ConvertBmmToMatmul: [RecomposePixelUnshuffle], ConvertLinearToConv2d: [FoldQDQ], DecomposeAcos: [RemoveRedundancy], diff --git a/backends/qualcomm/builders/utils.py b/backends/qualcomm/builders/utils.py index 1c5dde9a630..6cbdd986457 100755 --- a/backends/qualcomm/builders/utils.py +++ b/backends/qualcomm/builders/utils.py @@ -85,15 +85,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 73e02c0a457..b08420e00c2 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 edacd370067..53f81ae741f 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 @@ -485,6 +486,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):