Qualcomm AI Engine Direct - Lift scalar arguments

Differential Revision: D70069871

Pull Request resolved: #8642

Author: chunit-quic

backends/qualcomm/_passes/__init__.py

@@ -1,11 +1,8 @@
-from .annotate_and_quant_scalar import AnnotateAndQuantScalar
 from .annotate_decomposed import AnnotateDecomposed
 from .annotate_quant_attrs import AnnotateQuantAttrs
 from .constant_i64_to_i32 import ConstantI64toI32
-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_any import DecomposeAny
 from .decompose_einsum import DecomposeEinsum
@@ -17,7 +14,9 @@
 from .insert_io_qdq import InsertIOQDQ
 from .insert_requantize import InsertRequantize
 from .layout_transform import LayoutTransform
+from .lift_constant_scalar_operands import LiftConstantScalarOperands
 from .recompose_pixel_unshuffle import RecomposePixelUnshuffle
+from .recompose_prelu import RecomposePReLU
 from .recompose_rms_norm import RecomposeRmsNorm
 from .reduce_dynamic_range import ReduceDynamicRange
 from .remove_redundancy import RemoveRedundancy
@@ -27,14 +26,12 @@
 
 
 __all__ = [
-    AnnotateAndQuantScalar,
     AnnotateDecomposed,
     AnnotateQuantAttrs,
     ConstantI64toI32,
     ConvertBmmToMatmul,
-    ConvertBinaryOpsWithScalar,
     ConvertInterpolateWithUpsample2D,
-    ConvertPReLU,
+    RecomposePReLU,
     ConvertToLinear,
     DecomposeAny,
     DecomposeEinsum,
@@ -46,6 +43,7 @@
     InsertIOQDQ,
     InsertRequantize,
     LayoutTransform,
+    LiftConstantScalarOperands,
     RecomposePixelUnshuffle,
     RecomposeRmsNorm,
     ReduceDynamicRange,

backends/qualcomm/_passes/annotate_and_quant_scalar.py

@@ -32,9 +32,9 @@ class DecomposeLinalgVectorNorm(ExportPass):
     Decompose for math equivalent op.
     """
 
-    def __init__(self, quantization_capture=False) -> None:
+    def __init__(self, aten_dialect_capture=False) -> None:
         super().__init__()
-        self.quantization_capture = quantization_capture
+        self.aten_dialect_capture = aten_dialect_capture
 
     def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         graph = graph_module.graph
@@ -44,7 +44,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
                 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:
+                if self.aten_dialect_capture:
                     decomposed_module = torch.export.export(
                         model, (node.args[0].meta["val"],)
                     ).module()

backends/qualcomm/_passes/convert_binary_op_with_scalar.py

@@ -30,13 +30,15 @@ def call(self, graph_module: torch.fx.GraphModule):
                 silu_node_input = node.args[0]
                 with graph_module.graph.inserting_after(silu_node_input):
                     sigmoid_node = graph.create_node(
-                        "call_function", torch.ops.aten.sigmoid, (silu_node_input,)
+                        "call_function",
+                        torch.ops.aten.sigmoid.default,
+                        (silu_node_input,),
                     )
                     sigmoid_node.meta = self._copy_meta(silu_node.meta)
                     with graph_module.graph.inserting_after(sigmoid_node):
                         mul_node = graph.create_node(
                             "call_function",
-                            torch.ops.aten.mul,
+                            torch.ops.aten.mul.Tensor,
                             (silu_node_input, sigmoid_node),
                         )
                         mul_node.meta = self._copy_meta(silu_node.meta)

backends/qualcomm/_passes/decompose_linalg_vector_norm.py

@@ -53,20 +53,15 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.clamp.default,
         exir_ops.edge.aten.constant_pad_nd.default,
         exir_ops.edge.aten.div.Tensor,
-        exir_ops.edge.aten.eq.Scalar,
         exir_ops.edge.aten.eq.Tensor,
         exir_ops.edge.aten.full.default,
         exir_ops.edge.aten.full_like.default,
-        exir_ops.edge.aten.ge.Scalar,
         exir_ops.edge.aten.ge.Tensor,
         exir_ops.edge.aten.gelu.default,
-        exir_ops.edge.aten.gt.Scalar,
         exir_ops.edge.aten.gt.Tensor,
         exir_ops.edge.aten.hardswish.default,
         exir_ops.edge.aten.hardsigmoid.default,
         exir_ops.edge.aten.hardtanh.default,
-        exir_ops.edge.aten.leaky_relu.default,
-        exir_ops.edge.aten.le.Scalar,
         exir_ops.edge.aten.le.Tensor,
         exir_ops.edge.aten.linear.default,
         exir_ops.edge.aten.log.default,