Qualcomm AI Engine Direct - XR model enablement pipe_clean

summary
- support linalg_vector_norm, instance_norm
- expand coverage of quantization annotator
- test cases
- small refactor for _pass importing

Author: haowhsu-quic

backends/qualcomm/_passes/__init__.py

@@ -1,34 +1,52 @@
 from .annotate_and_quant_scalar import AnnotateAndQuantScalar
 from .annotate_decomposed import AnnotateDecomposed
 from .annotate_quant_attrs import AnnotateQuantAttrs
+from .convert_binary_op_with_scalar import ConvertBinaryOpsWithScalar
 from .convert_bmm_to_matmul import ConvertBmmToMatmul
 from .convert_interpolate_with_upsample2d import ConvertInterpolateWithUpsample2D
 from .convert_prelu import ConvertPReLU
 from .convert_to_linear import ConvertToLinear
+from .decompose_einsum import DecomposeEinsum
+from .decompose_linalg_vector_norm import DecomposeLinalgVectorNorm
+from .decompose_silu import DecomposeSilu
 from .expand_broadcast_tensor_shape import ExpandBroadcastTensorShape
 from .fold_qdq import FoldQDQ
+from .fuse_consecutive_transpose import FuseConsecutiveTranspose
 from .i64_to_i32 import I64toI32
+from .insert_io_qdq import InsertIOQDQ
+from .insert_requantize import InsertRequantize
 from .layout_transform import LayoutTransform
 from .recompose_pixel_unshuffle import RecomposePixelUnshuffle
 from .recompose_rms_norm import RecomposeRmsNorm
+from .reduce_dynamic_range import ReduceDynamicRange
 from .remove_redundancy import RemoveRedundancy
 from .replace_index_put_input import ReplaceIndexPutInput
+from .replace_inf_buffer import ReplaceInfBuffer
 
 
 __all__ = [
     AnnotateAndQuantScalar,
     AnnotateDecomposed,
     AnnotateQuantAttrs,
     ConvertBmmToMatmul,
+    ConvertBinaryOpsWithScalar,
     ConvertInterpolateWithUpsample2D,
     ConvertPReLU,
     ConvertToLinear,
+    DecomposeEinsum,
+    DecomposeLinalgVectorNorm,
+    DecomposeSilu,
     ExpandBroadcastTensorShape,
     FoldQDQ,
+    FuseConsecutiveTranspose,
     I64toI32,
+    InsertIOQDQ,
+    InsertRequantize,
     LayoutTransform,
     RecomposePixelUnshuffle,
     RecomposeRmsNorm,
+    ReduceDynamicRange,
     RemoveRedundancy,
     ReplaceIndexPutInput,
+    ReplaceInfBuffer,
 ]

backends/qualcomm/_passes/convert_to_linear.py

@@ -39,6 +39,7 @@ class ConvertToLinear(ExportPass):
     mm = exir_ops.edge.aten.mm.default
 
     addmm_patterns = [
+        {view_copy: 1, permute_copy: 1, addmm: 1},
         {view_copy: 2, permute_copy: 1, addmm: 1},
         {permute_copy: 1, addmm: 1},
     ]

backends/qualcomm/_passes/decompose_linalg_vector_norm.py

@@ -0,0 +1,85 @@
+# 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 torch
+from executorch.exir import to_edge
+from executorch.exir.pass_base import ExportPass, PassResult
+
+
+class LinalgVectorNorm(torch.nn.Module):
+    def __init__(self, exp, dim, keepdim):
+        super().__init__()
+        self.exp = exp
+        self.dim = tuple(dim) if dim is not None else None
+        self.keepdim = keepdim
+
+    def forward(self, x):
+        if self.dim is None:
+            x = torch.flatten(x)
+            self.dim = 0
+
+        x = torch.abs(x)
+        x = torch.pow(x, self.exp)
+        x = torch.sum(x, dim=self.dim, keepdim=self.keepdim)
+        return torch.pow(x, 1.0 / self.exp)
+
+
+class DecomposeLinalgVectorNorm(ExportPass):
+    """
+    Decompose for math equivalent op.
+    """
+
+    def __init__(self, quantization_capture=False) -> None:
+        super().__init__()
+        self.quantization_capture = quantization_capture
+
+    def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
+        graph = graph_module.graph
+        for node in graph.nodes:
+            if "linalg_vector_norm" in str(node.target):
+                ord = node.args[1] if len(node.args) > 1 else 2.0
+                dim = node.args[2] if len(node.args) > 2 else None
+                keepdim = node.args[3] if len(node.args) > 3 else False
+                model = LinalgVectorNorm(ord, dim, keepdim)
+                if self.quantization_capture:
+                    decomposed_module = torch.export.export(
+                        model, (node.args[0].meta["val"],)
+                    ).module()
+                else:
+                    edge_mgr = to_edge(
+                        torch.export.export(model, (node.args[0].meta["val"],))
+                    )
+                    decomposed_module = edge_mgr.exported_program()
+
+                with graph.inserting_before(node):
+                    # remap is used to map original node values to new node values,
+                    # which ensures that reference to nodes are correclty updated in the new graph
+                    remap = {"x": node.args[0]}
+
+                    for decomposed_node in decomposed_module.graph.nodes:
+                        # no need to copy existent 'output'
+                        if decomposed_node.op == "output":
+                            for user in node.users.copy():
+                                # remap
+                                user.replace_input_with(
+                                    node,
+                                    remap[decomposed_node.args[0][0]],
+                                )
+                        # no need to copy existent placeholders
+                        elif decomposed_node.op == "placeholder":
+                            # replace node map from string to graph node
+                            remap[decomposed_node] = remap.pop(decomposed_node.name)
+                        else:
+                            remap[decomposed_node] = graph.node_copy(
+                                decomposed_node,
+                                arg_transform=lambda x, remap=remap: remap[x],
+                            )
+
+                    graph.erase_node(node)
+
+        graph.eliminate_dead_code()
+        graph_module.recompile()
+        return PassResult(graph_module, True)

backends/qualcomm/_passes/layout_transform.py

@@ -33,8 +33,10 @@ class LayoutTransform(ExportPass):
         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.instance_norm.default,
         exir_ops.edge.aten.max_pool2d_with_indices.default,
         exir_ops.edge.aten._native_batch_norm_legit_no_training.default,
+        exir_ops.edge.aten._native_batch_norm_legit.no_stats,
         exir_ops.edge.aten.native_group_norm.default,
         exir_ops.edge.aten.pixel_shuffle.default,
         exir_ops.edge.aten.pixel_unshuffle.default,

backends/qualcomm/_passes/utils.py

@@ -64,6 +64,7 @@ def get_passes_dependency_for_capture_program():
         ConvertInterpolateWithUpsample2D,
         ConvertPReLU,
         ConvertToLinear,
+        DecomposeLinalgVectorNorm,
         ExpandBroadcastTensorShape,
         FoldQDQ,
         I64toI32,
@@ -81,6 +82,7 @@ def get_passes_dependency_for_capture_program():
         ConvertPReLU: [RemoveRedundancy],
         ConvertBmmToMatmul: [ConvertToLinear],
         ConvertInterpolateWithUpsample2D: [RemoveRedundancy],
+        DecomposeLinalgVectorNorm: [RemoveRedundancy],
         I64toI32: [RemoveRedundancy],
         AnnotateQuantAttrs: [
             RecomposePixelUnshuffle,

backends/qualcomm/builders/__init__.py

@@ -35,6 +35,7 @@
     op_hardtanh,
     op_index,
     op_index_put,
+    op_instance_norm,
     op_layer_norm,
     op_le,
     op_linear,
@@ -109,6 +110,7 @@
     op_hardsigmoid,
     op_index,
     op_index_put,
+    op_instance_norm,
     op_layer_norm,
     op_le,
     op_linear,

backends/qualcomm/builders/op_batch_norm.py

@@ -9,10 +9,14 @@
 
 import torch
 from executorch.backends.qualcomm.utils.constants import (
+    QCOM_AXIS_ORDER,
     QCOM_QUANT_ATTRS,
     QCOM_QUANT_MAX,
+    QCOM_QUANT_MIN,
     QCOM_SCALE,
+    QCOM_ZERO_POINT,
 )
+from executorch.exir.dialects._ops import ops as exir_ops
 
 from .node_visitor import NodeVisitor, register_node_visitor
 from .qnn_constants import OpBatchnorm, QNN_OP_PACKAGE_NAME_QTI_AISW
@@ -21,7 +25,10 @@
 
 @register_node_visitor
 class BatchNorm(NodeVisitor):
-    target = ["aten._native_batch_norm_legit_no_training.default"]
+    target = [
+        "aten._native_batch_norm_legit_no_training.default",
+        "aten._native_batch_norm_legit.no_stats",
+    ]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)
@@ -43,9 +50,13 @@ def define_node(
         input_node = node.args[0]
         input_tensor = self.get_tensor(input_node, node)
 
-        mean_node, var_node, eps = node.args[3], node.args[4], 1e-9
-        mean_tensor = get_parameter(mean_node, self.edge_program)
-        var_tensor = get_parameter(var_node, self.edge_program)
+        eps = 1e-9
+        if "no_stats" in str(node.target):
+            mean_tensor = torch.Tensor([node.args[4]])
+            var_tensor = torch.Tensor([node.args[5]])
+        else:
+            mean_tensor = get_parameter(node.args[3], self.edge_program)
+            var_tensor = get_parameter(node.args[4], self.edge_program)
 
         input_tensor_wrapper = self.define_tensor(
             input_node,
@@ -54,22 +65,43 @@ def define_node(
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
         )
+        batch_norm_input_tensors = [input_tensor_wrapper]
 
-        bias_node = node.args[2]
-        bias_tensor = get_parameter(bias_node, self.edge_program)
-        filter_node = node.args[1]
-        filter_tensor = get_parameter(filter_node, self.edge_program)
-
-        amount = (filter_tensor * mean_tensor) / torch.sqrt(var_tensor + eps)
-        bias_tensor = bias_tensor - amount
-        self.update_encoding(bias_node, bias_tensor, eps)
-        bias_tensor_wrapper = self.define_tensor(
-            bias_node,
+        output_tensor = self.get_tensor(node, node, 0)
+        output_tensor_wrapper = self.define_tensor(
             node,
-            bias_tensor,
-            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
+            node,
+            output_tensor,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
             nodes_to_wrappers,
         )
+        batch_norm_output_tensors = [output_tensor_wrapper]
+
+        n_feature = output_tensor.shape[-1 if QCOM_AXIS_ORDER in node.meta else 1]
+        bias_node = node.args[2]
+        if bias_node is not None:
+            bias_tensor = get_parameter(bias_node, self.edge_program)
+
+        filter_node = node.args[1]
+        if filter_node is not None:
+            filter_tensor = get_parameter(filter_node, self.edge_program)
+        else:
+            # 'graph', 'name', 'op', 'target', 'args', and 'kwargs'
+            filter_node = torch.fx.Node(
+                node.graph,
+                node.name + "_filter",
+                "call_function",
+                exir_ops.edge.aten.scalar_tensor.default,
+                (),  # args
+                {},  # kwargs
+            )
+            filter_tensor = torch.ones(n_feature)
+            if quant_attrs := node.meta.get(QCOM_QUANT_ATTRS):
+                quant_attrs = quant_attrs.copy()
+                quant_range = quant_attrs[QCOM_QUANT_MAX] - quant_attrs[QCOM_QUANT_MIN]
+                quant_attrs[QCOM_ZERO_POINT] = 0
+                quant_attrs[QCOM_SCALE] = 1.0 / quant_range
+                filter_node.meta[QCOM_QUANT_ATTRS] = quant_attrs
 
         filter_tensor = filter_tensor / torch.sqrt(var_tensor + eps)
         self.update_encoding(filter_node, filter_tensor, eps)
@@ -80,22 +112,20 @@ def define_node(
             PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
             nodes_to_wrappers,
         )
-
-        batch_norm_input_tensors = [
-            input_tensor_wrapper,
-            filter_tensor_wrapper,
-            bias_tensor_wrapper,
-        ]
-
-        output_tensor = self.get_tensor(node, node, 0)
-        output_tensor_wrapper = self.define_tensor(
-            node,
-            node,
-            output_tensor,
-            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
-            nodes_to_wrappers,
-        )
-        batch_norm_output_tensors = [output_tensor_wrapper]
+        batch_norm_input_tensors.append(filter_tensor_wrapper)
+
+        if bias_node is not None:
+            amount = (filter_tensor * mean_tensor) / torch.sqrt(var_tensor + eps)
+            bias_tensor = bias_tensor - amount
+            self.update_encoding(bias_node, bias_tensor, eps)
+            bias_tensor_wrapper = self.define_tensor(
+                bias_node,
+                node,
+                bias_tensor,
+                PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_STATIC,
+                nodes_to_wrappers,
+            )
+            batch_norm_input_tensors.append(bias_tensor_wrapper)
 
         batch_norm_op = PyQnnWrapper.PyQnnOpWrapper(
             node.name,