Pytest Refactor - Passes

Author: winskuo-quic

backends/qualcomm/_passes/__init__.py

@@ -8,6 +8,7 @@
 from .annotate_quant_attrs import AnnotateQuantAttrs
 from .annotate_stack import AnnotateStack
 from .annotate_unbind import AnnotateUnbind
+from .build_quant_io import BuildQuantIo
 from .canonicalize_conv import CanonicalizeConv
 from .convert_bmm_to_matmul import ConvertBmmToMatmul
 from .convert_linear_to_conv2d import ConvertLinearToConv2d
@@ -73,6 +74,7 @@
     AnnotateQuantAttrs,
     AnnotateStack,
     AnnotateUnbind,
+    BuildQuantIo,
     CanonicalizeConv,
     ConvertBmmToMatmul,
     ConvertLinearToConv2d,

backends/qualcomm/_passes/annotate_avg_pool1d.py

@@ -27,6 +27,7 @@ class AnnotateAvgPool1D(ExportPass):
         torch.avg_pool1d,
         torch.ops.aten.adaptive_avg_pool1d.default,
         torch.adaptive_avg_pool1d,
+        torch.nn.modules.pooling.AvgPool1d,
     ]
 
     def __init__(self, edge_program: torch.export.ExportedProgram):

backends/qualcomm/_passes/annotate_stack.py

@@ -18,16 +18,14 @@ class AnnotateStack(ExportPass):
     generated after quantization process.
     """
 
-    decomp_ops = [torch.ops.aten.stack.default]
+    _SOURCE_OPS = [torch.stack, torch.ops.aten.stack.default, "stack"]
 
     def __init__(self, edge_program: torch.export.ExportedProgram):
         super(AnnotateStack, self).__init__()
         self.edge_program = edge_program
 
     def _annotate_stack(self, graph_module: torch.fx.GraphModule):
-        partitions = get_source_partitions(
-            graph_module.graph, [torch.stack, torch.ops.aten.stack.default, "stack"]
-        )
+        partitions = get_source_partitions(graph_module.graph, self._SOURCE_OPS)
         for src_partitions in partitions.values():
             for src_partition in src_partitions:
                 output = src_partition.output_nodes[0]

backends/qualcomm/_passes/annotate_unbind.py

@@ -18,16 +18,14 @@ class AnnotateUnbind(ExportPass):
     generated after quantization process.
     """
 
-    decomp_ops = [torch.ops.aten.unbind.int]
+    _SOURCE_OPS = [torch.unbind, torch.ops.aten.unbind.int, "unbind"]
 
     def __init__(self, edge_program: torch.export.ExportedProgram):
         super(AnnotateUnbind, self).__init__()
         self.edge_program = edge_program
 
     def _annotate_unbind(self, graph_module: torch.fx.GraphModule):
-        partitions = get_source_partitions(
-            graph_module.graph, [torch.unbind, torch.ops.aten.unbind.int, "unbind"]
-        )
+        partitions = get_source_partitions(graph_module.graph, self._SOURCE_OPS)
         for src_partitions in partitions.values():
             for src_partition in src_partitions:
                 if src_partition.input_nodes[0].target in dq_ops:

backends/qualcomm/_passes/decompose_any.py

@@ -4,6 +4,8 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
+import re
+
 import torch
 from executorch.exir import to_edge
 from executorch.exir.pass_base import ExportPass, PassResult
@@ -33,13 +35,13 @@ class DecomposeAny(ExportPass):
     Decompose for math equivalent op.
     """
 
-    def __init__(self, quantization_capture=False) -> None:
+    def __init__(self) -> None:
         super().__init__()
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
         for node in graph.nodes:
-            if "any.dim" in str(node.target):
+            if re.search(r"any\.(dim|default)", str(node.target)):
                 dim = node.args[1] if len(node.args) > 1 else None
                 keepdim = node.args[2] if len(node.args) > 2 else False
                 model = Any(dim, keepdim)

backends/qualcomm/_passes/decompose_expm1.py

@@ -22,7 +22,10 @@ def __init__(self) -> None:
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
         for node in graph.nodes:
-            if node.target == torch.ops.aten.special_expm1.default:
+            if node.target in {
+                torch.ops.aten.special_expm1.default,
+                torch.ops.aten.expm1.default,
+            }:
                 input_node = node.args[0]
                 with graph_module.graph.inserting_after(input_node):
                     exp_op = torch.ops.aten.exp.default

backends/qualcomm/_passes/decompose_maxpool3d.py

@@ -94,13 +94,19 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                         dilation *= 3
 
                 ceil_mode = node.args[5] if num_args > 5 else False
-                return_indices = node.args[6] if num_args > 6 else False
+                # since the indices output might not be connected
+                # only traverse getitem user and check if any of them referring to the indices
+                return_indices = (
+                    any(user.args[1] == 1 for user in node.users)
+                    if len(node.meta["val"]) > 1
+                    else False
+                )
                 if return_indices:
                     warnings.warn(
-                        "[QNN Delegate Op Builder]: The case return_indices=True is not be support, fallback",
+                        "[QNN Delegate Op Builder]: The case return_indices=True is not supported, fallback",
                         stacklevel=1,
                     )
-                    return
+                    continue
 
                 model = ModelMaxPool3D(
                     filter_size, stride, padding, dilation, return_indices, ceil_mode

backends/qualcomm/_passes/decompose_pad.py

@@ -29,7 +29,6 @@ class DecomposePad(ExportPass):
 
     _PAD_TARGETS = {
         torch.ops.aten.pad.default,
-        exir_ops.edge.aten.pad.default,
     }
 
     _PAD_OPS = {

backends/qualcomm/_passes/decompose_var.py

@@ -5,12 +5,9 @@
 # LICENSE file in the root directory of this source tree.
 
 import torch
-from executorch.exir.dialects._ops import ops as exir_ops
-from executorch.exir.dialects.edge._ops import EdgeOpOverload
 from executorch.exir.pass_base import ExportPass, PassResult
-from torchao.quantization.pt2e.utils import get_new_attr_name_with_prefix
 
-from .utils import copy_meta, create_const_node
+from .utils import copy_meta
 
 
 class DecomposeVar(ExportPass):
@@ -29,17 +26,12 @@ def __init__(self):
         self.var_targets = {
             torch.ops.aten.var.correction,
             torch.ops.aten.var.dim,
-            exir_ops.edge.aten.var.correction,
-            exir_ops.edge.aten.var.dim,
         }
 
     def _get_correction(self, node):
         """Extract the correction factor from node args based on op variant."""
         target = node.target
-        if target in (
-            torch.ops.aten.var.correction,
-            exir_ops.edge.aten.var.correction,
-        ):
+        if target in {torch.ops.aten.var.correction}:
             # var.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False)
             # correction is a kwarg, but in the graph it may appear in kwargs
             correction = node.kwargs.get("correction", None)
@@ -54,10 +46,7 @@ def _get_correction(self, node):
     def _get_dim_and_keepdim(self, node):
         """Extract dim and keepdim from node args based on op variant."""
         target = node.target
-        if target in (
-            torch.ops.aten.var.correction,
-            exir_ops.edge.aten.var.correction,
-        ):
+        if target in {torch.ops.aten.var.correction}:
             # var.correction(Tensor self, int[1]? dim=None, *, Scalar? correction=None, bool keepdim=False)
             dim = node.args[1] if len(node.args) > 1 else None
             keepdim = node.kwargs.get("keepdim", False)
@@ -70,31 +59,18 @@ def _get_dim_and_keepdim(self, node):
 
     def call(self, graph_module: torch.fx.GraphModule):
         graph = graph_module.graph
-        const_cache = {}
 
         for node in list(graph.nodes):
             if node.op == "call_function" and node.target in self.var_targets:
                 x_node = node.args[0]
-                is_edge = isinstance(node.target, EdgeOpOverload)
                 meta = node.meta
 
                 correction = self._get_correction(node)
                 dim, keepdim = self._get_dim_and_keepdim(node)
 
-                mean_op = (
-                    exir_ops.edge.aten.mean.dim if is_edge else torch.ops.aten.mean.dim
-                )
-                sub_op = (
-                    exir_ops.edge.aten.sub.Tensor
-                    if is_edge
-                    else torch.ops.aten.sub.Tensor
-                )
-                mul_op = (
-                    exir_ops.edge.aten.mul.Tensor
-                    if is_edge
-                    else torch.ops.aten.mul.Tensor
-                )
-
+                mean_op = torch.ops.aten.mean.dim
+                sub_op = torch.ops.aten.sub.Tensor
+                mul_op = torch.ops.aten.mul.Tensor
                 # Handle dim=None: reduce over all dimensions
                 input_shape = node.args[0].meta["val"].shape
                 if dim is None:
@@ -148,22 +124,8 @@ def call(self, graph_module: torch.fx.GraphModule):
                         # Guard against division by zero (e.g. single-element dim with correction=1).
                         # Using inf matches the native PyTorch behavior where 0 * inf → nan.
                         scale = float("inf") if denom == 0 else float(n) / denom
-
-                        if is_edge:
-                            cache_key = ("_var_scale_", scale)
-                            if cache_key not in const_cache:
-                                attr_name = get_new_attr_name_with_prefix(
-                                    "_var_scale_const_"
-                                )(graph_module)
-                                const_cache[cache_key] = create_const_node(
-                                    graph, graph_module, attr_name, scale, node
-                                )
-                            scale_node = const_cache[cache_key]
-                        else:
-                            scale_node = scale
-
                         result_node = graph.create_node(
-                            "call_function", mul_op, (var_node, scale_node)
+                            "call_function", mul_op, (var_node, scale)
                         )
                         result_node.meta = copy_meta(meta)
                     else:

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -129,10 +129,8 @@ def get_default_pass_activations(cls):
             (DecomposeAny, True),
             (DecomposeAtan2, True),
             (DecomposeColIm, True),
-            (DecomposeCDist, True),
             (DecomposeDiagonal, True),
             (DecomposeDivMode, True),
-            (DecomposeFill, True),
             (DecomposeHyperbolicVariants, True),
             (DecomposeLogVariants, True),
             (DecomposeMaxPool3d, True),
@@ -141,7 +139,6 @@ def get_default_pass_activations(cls):
             (DecomposeRemainder, True),
             (DecomposeTan, True),
             (DecomposeTrunc, True),
-            (DecomposeVar, True),
             (ExpandBroadcastTensorShape, True),
             (FixedLinearKeepDim, True),
             (FoldQDQ, True),
@@ -205,7 +202,6 @@ def get_export_passes(
         passes = [
             DecomposeBinaryAlpha,
             DecomposeCDist,
-            DecomposePad,
             DecomposeScaledDotProductAttention,
             DecomposeRoll,
             DecomposeSelectScatter,
@@ -288,19 +284,16 @@ def get_passes_dependency_for_capture_program(cls):
             DecomposeAny: [RemoveRedundancy],
             DecomposeAtan2: [RemoveRedundancy],
             DecomposeColIm: [FoldQDQ],
-            DecomposeCDist: [RemoveRedundancy],
             DecomposeDiagonal: [RemoveRedundancy],
             DecomposeDivMode: [RemoveRedundancy],
             DecomposeFill: [RemoveRedundancy],
             DecomposeHyperbolicVariants: [RemoveRedundancy],
-            DecomposeLinalgVectorNorm: [RemoveRedundancy],
             DecomposeLogVariants: [RemoveRedundancy],
             DecomposeMaxPool3d: [RemoveRedundancy],
             DecomposePad: [RemoveRedundancy],
             DecomposeRemainder: [RemoveRedundancy],
             DecomposeTan: [RemoveRedundancy],
             DecomposeTrunc: [RemoveRedundancy],
-            DecomposeVar: [RemoveRedundancy],
             ExpandBroadcastTensorShape: [FoldQDQ],
             FixedLinearKeepDim: [FoldQDQ],
             FoldQDQ: [AnnotateQuantAttrs, AnnotateStack, AnnotateUnbind],
@@ -474,8 +467,6 @@ def transform_for_preprocess_pipeline(
             insert_permute=True,
         )
         self._transform(exported_program.graph_module)
-        # Update inputs_to_buffers and buffers_to_mutate in graph signature for mutable buffer
-        # Since I/O will be inserted Q/DQ, it results in failed to mapping output node names and buffer
         exported_program._graph_signature = _get_updated_graph_signature(
             exported_program.graph_signature,
             exported_program.graph_module,