Revert "Qualcomm AI Engine Direct - ConvFormer Enablement (pytorch#6654)"

This reverts commit e8b5987.

Author: Zonglin Peng

backends/qualcomm/_passes/fuse_consecutive_transpose.py

@@ -15,18 +15,8 @@
 
 class FuseConsecutiveTranspose(ExportPass):
     """
-    This pass fuses consecutive transpose / permute into one or none to reduce runtime
-    overhead.
-    To simplify the fuse logic, we ensure each permute node's output has at most 1 permute node
-    by cloning transpose.
-    Example:
-    Before clone transpose:
-    relu -> permute1 ─> permute2
-               |──────> permute3
-
-    After clone transpose:
-    relu ─> permute1 ──────> permute2
-      |───> permute4(new) ─> permute3
+    This pass fuses consecutive transpose / permute into one to reduce runtime
+    overhead
     """
 
     def __init__(self):
@@ -37,81 +27,54 @@ def __init__(self):
         self.visited = set()
         self.nodes = []
 
-    def _clone_transpose(
-        self, graph_module: torch.fx.GraphModule
-    ) -> torch.fx.GraphModule:
-        graph = graph_module.graph
-        for n in graph_module.graph.nodes:
-            if n.target in self.op_map:
-                users = [user for user in list(n.users) if user.target in self.op_map]
-                if len(users) > 1:
-                    for i in range(1, len(users)):
-                        with graph.inserting_after(n):
-                            clone_permute_node = graph.create_node(
-                                "call_function",
-                                exir_ops.edge.aten.permute_copy.default,
-                                (n.args[0], n.args[1]),
-                            )
-                            clone_permute_node.meta = n.meta
-                            users[i].replace_input_with(n, clone_permute_node)
-
-    def _is_dispensable(self, axis_order):
-        for index, value in enumerate(axis_order):
-            if index != value:
-                return False
-        return True
-
     def _traverse(self, node):
         if node in self.visited or node.target not in self.op_map:
             return
 
         self.nodes.append(node)
         self.visited.add(node)
         next_users = [n for n in list(node.users) if n.target in self.op_map]
-
-        assert (
-            len(next_users) <= 1
-        ), "Each permute node should have at most 1 permute output node after _clone_transpose"
         if not next_users:
             return
-        else:
+
+        if len(next_users) == 1:
             self._traverse(list(node.users)[0])
+        else:
+            raise NotImplementedError(
+                f"Check the node {node}, wich encounter mutilple permute output case"
+            )
 
     def _fuse(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
         graph = graph_module.graph
         for n in graph_module.graph.nodes:
             self._traverse(n)
             if len(self.nodes) > 1:
+                permute_order = []
                 input_node, output_node = self.nodes[0].args[0], self.nodes[-1]
                 input_shape = input_node.meta["val"].shape
                 axis_order = torch.arange(len(input_shape)).tolist()
                 for node in self.nodes:
+                    permute_order.append(node.args[1])
                     axis_order = [axis_order[i] for i in node.args[1]]
-                # If axis order is just [0,1,2,3], we ignore permute node
-                if self._is_dispensable(axis_order):
-                    for user in output_node.users.copy():
-                        user.replace_input_with(output_node, n.args[0])
-                else:
-                    with graph.inserting_after(input_node):
-                        permute_op = exir_ops.edge.aten.permute_copy.default
-                        permute_node = graph.create_node(
-                            "call_function", permute_op, (input_node, axis_order)
-                        )
-                        users = output_node.users.copy()
-                        for user in users:
-                            user.replace_input_with(output_node, permute_node)
-
-                        # copy metadata
-                        permute_node.meta = output_node.meta
-                        # Without "qnn_permute", we might obtain wrong input shape
-                        if [pn.meta.get(QCOM_INSERTED_PERMUTE) for pn in self.nodes]:
-                            permute_node.meta[QCOM_INSERTED_PERMUTE] = True
+                with graph.inserting_after(input_node):
+                    permute_op = exir_ops.edge.aten.permute_copy.default
+                    permute_node = graph.create_node(
+                        "call_function", permute_op, (input_node, axis_order)
+                    )
+                    users = output_node.users.copy()
+                    for user in users:
+                        user.replace_input_with(output_node, permute_node)
+
+                    # copy metadata
+                    permute_node.meta = output_node.meta
+                    # Without "qnn_permute", we might obtain wrong input shape
+                    if [pn.meta.get(QCOM_INSERTED_PERMUTE) for pn in self.nodes]:
+                        permute_node.meta[QCOM_INSERTED_PERMUTE] = True
 
             # clear current stack
             self.nodes = []
 
     def call(self, graph_module: torch.fx.GraphModule):
-        self._clone_transpose(graph_module)
         self._fuse(graph_module)
         graph_module.recompile()
         dead_code_elimination_pass(graph_module)

backends/qualcomm/_passes/layout_transform.py

@@ -30,7 +30,6 @@ class LayoutTransform(ExportPass):
     """
 
     layout_sensitive_ops = {
-        exir_ops.edge.aten.adaptive_avg_pool2d.default,
         exir_ops.edge.aten.avg_pool2d.default,
         exir_ops.edge.aten.convolution.default,
         exir_ops.edge.aten.max_pool2d_with_indices.default,

backends/qualcomm/builders/__init__.py

@@ -7,7 +7,6 @@
 from . import (
     node_visitor,
     op_abs,
-    op_adaptive_avg_pool2d,
     op_add,
     op_arange,
     op_avg_pool2d,
@@ -79,7 +78,6 @@
 __all__ = [
     node_visitor,
     op_abs,
-    op_adaptive_avg_pool2d,
     op_add,
     op_arange,
     op_avg_pool2d,

backends/qualcomm/builders/op_adaptive_avg_pool2d.py

@@ -63,19 +63,15 @@ def define_node(
             nodes_to_wrappers,
         )
 
-        layer_norm_input_tensors = [input_tensor_wrapper, weight_tensor_wrapper]
-
         bias_node = node.args[3]
-        if bias_node is not None:
-            bias_tensor = get_parameter(bias_node, self.edge_program)
-            bias_tensor_wrapper = self.define_tensor(
-                bias_node,
-                node,
-                bias_tensor,
-                PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
-                nodes_to_wrappers,
-            )
-            layer_norm_input_tensors.append(bias_tensor_wrapper)
+        bias_tensor = get_parameter(bias_node, self.edge_program)
+        bias_tensor_wrapper = self.define_tensor(
+            bias_node,
+            node,
+            bias_tensor,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
+            nodes_to_wrappers,
+        )
 
         epsilon = node.args[4]
 
@@ -93,7 +89,9 @@ def define_node(
             QNN_OP_PACKAGE_NAME_QTI_AISW,
             OpLayerNorm.op_name,
         )
-        layer_norm_op.AddInputTensors(layer_norm_input_tensors)
+        layer_norm_op.AddInputTensors(
+            [input_tensor_wrapper, weight_tensor_wrapper, bias_tensor_wrapper]
+        )
         layer_norm_op.AddOutputTensors([output_tensor_wrapper])
         layer_norm_op.AddScalarParam(
             OpLayerNorm.param_epsilon,

backends/qualcomm/builders/op_layer_norm.py

@@ -66,7 +66,7 @@ def define_node(
             nodes_to_wrappers,
         )
 
-        # Fake node, nn module seems to be inconsistant with document
+        # Fake node, nn moudle seems to be inconsistant with document
         bias_tensor = torch.zeros(weight_tensor.shape)
         bias_node = torch.fx.Node(
             node.graph,

backends/qualcomm/builders/op_rms_norm.py

@@ -512,11 +512,6 @@ def annotate_sqrt(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_single_in_single_out(node, quantization_config)
 
 
-@register_annotator([torch.ops.aten.square.default])
-def annotate_square(node: Node, quantization_config: QuantizationConfig) -> None:
-    annotate_single_in_single_out(node, quantization_config)
-
-
 @register_annotator([torch.ops.aten.gelu.default])
 def annotate_gelu(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_single_in_single_out(node, quantization_config)

backends/qualcomm/quantizer/annotators.py

@@ -16,15 +16,6 @@ def forward(self, x):
         return torch.abs(x)
 
 
-class AdaptiveAvgPool2D(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-
-    def forward(self, x):
-        adaptive_avg_pool = torch.nn.AdaptiveAvgPool2d((1, 1))
-        return adaptive_avg_pool(x)
-
-
 class Add(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -694,24 +685,15 @@ def forward(self, x):
 
 
 class LayerNorm(torch.nn.Module):
-    def __init__(self, bias=True):
+    def __init__(self):
         super().__init__()
-        self.layer_norm = torch.nn.LayerNorm([768], eps=1e-6, bias=bias)
+        self.layer_norm = torch.nn.LayerNorm([768], eps=1e-6)
         self.linear = torch.nn.Linear(768, 196)
 
     def forward(self, x):
         return self.linear(self.layer_norm(x))
 
 
-class LayerNormAdd(torch.nn.Module):
-    def __init__(self):
-        super().__init__()
-        self.layer_norm = torch.nn.LayerNorm([512], eps=1e-6, bias=False)
-
-    def forward(self, x, y):
-        return self.layer_norm(x) + y
-
-
 class LeakyReLUDefault(torch.nn.Module):
     def __init__(self):
         super().__init__()