Qualcomm AI Engine Direct - Add QNN support for to_edge_transform_and_lower

summary:

- Support `to_edge_transform_and_lower`
  - Replace capture_program with new API `to_edge_transform_and_lower_to_qnn`
  - Replace capture_program with to_edge_transform_and_lower_to_qnn for unit_test
  - Replace capture_program with to_edge_transform_and_lower_to_qnn for examples
  - Replace capture_program with to_edge_transform_and_lower_to_qnn for llama
- Add QnnPassManager to manage all passes in different stage
  - Deprecated _transform in export_llama_lib with qnn_pass_manager
  - Add transform_for_export_pipeline for LiftConstantScalarOperands to avoid creating temporary tensors in the operation builder.
    However, this pass will create a get_attr node, which should be converted into a lifted tensor constant by the lift_constant_tensor_pass.
    If placed in the to_edge_transform_passes, it will be executed after the lift_constant_tensor_pass,
    causing the operation builder to fail to correctly retrieve the parameter by the get_parameter for get_attr node.
- Refactor the passes
  - Fix the output dtype doesn't match in runtime after build quant io
  - Combine constant_i64_to_i32 and tensor_i64_to_i32 into i64_to_i32
  - Replace convert_to_linear pass with fixed_linear_keep_dim pass
    - Since QNN has no keep dims for linear op, we will need to add squeeze and unsqueeze around linear node
  - Add TagQuantIO pass to tag io nodes to avoid inserting q/dq in qnn_preprocess
  - Add prelu, leaky_relu, linear, rms_norm into decompose_table
    - Remove recompose_prelu.py
  - Remove unused variable in insert_requantize.py, and replace_index_put_input.py

- Support aten.split_with_sizes_copy.default
- Support leaky_relu with inplace=True

Author: shewu-quic

backends/qualcomm/_passes/__init__.py

@@ -1,43 +1,41 @@
 from .annotate_decomposed import AnnotateDecomposed
 from .annotate_quant_attrs import AnnotateQuantAttrs
-from .constant_i64_to_i32 import ConstantI64toI32
 from .convert_bmm_to_matmul import ConvertBmmToMatmul
-from .convert_to_linear import ConvertToLinear
 from .decompose_any import DecomposeAny
 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 .fixed_linear_keep_dim import FixedLinearKeepDim
 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 .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
 from .replace_index_put_input import ReplaceIndexPutInput
 from .replace_inf_buffer import ReplaceInfBuffer
-from .tensor_i64_to_i32 import TensorI64toI32
+from .tag_quant_io import TagQuantIO
 
 
 __all__ = [
     AnnotateDecomposed,
     AnnotateQuantAttrs,
-    ConstantI64toI32,
     ConvertBmmToMatmul,
-    RecomposePReLU,
-    ConvertToLinear,
     DecomposeAny,
     DecomposeEinsum,
     DecomposeLinalgVectorNorm,
     DecomposeSilu,
     ExpandBroadcastTensorShape,
+    FixedLinearKeepDim,
     FoldQDQ,
     FuseConsecutiveTranspose,
+    I64toI32,
     InsertIOQDQ,
     InsertRequantize,
     LayoutTransform,
@@ -48,5 +46,5 @@
     RemoveRedundancy,
     ReplaceIndexPutInput,
     ReplaceInfBuffer,
-    TensorI64toI32,
+    TagQuantIO,
 ]

backends/qualcomm/_passes/build_quant_io.py

@@ -27,25 +27,17 @@ def _make_spec(self, x):
             return None
 
     def _build(self, graph_module: torch.fx.GraphModule) -> torch.fx.GraphModule:
-        # forcely update delegate node's meta['spec'] to get correct output
+        # Forcedly update delegate node's meta['spec'] to get correct output
         # tensor size in runtime
         call_delegate = [
             node
             for node in graph_module.graph.nodes
             if node.op == "call_function" and node.name == "executorch_call_delegate"
         ]
         assert len(call_delegate) == 1
-        spec = []
         for n in graph_module.graph.nodes:
             if QCOM_QUANTIZED_IO in n.meta:
                 n.meta["val"] = n.meta["val"].to(dtype=n.meta[QCOM_QUANTIZED_IO])
-            if n.op == "call_function" and "getitem" in n.name:
-                fake_tensor = n.meta["val"]
-                if QCOM_QUANTIZED_IO in n.meta:
-                    fake_tensor = fake_tensor.to(dtype=n.meta[QCOM_QUANTIZED_IO])
-                spec.append(self._make_spec(fake_tensor))
-
-        call_delegate[0].meta["spec"] = tuple(spec)
 
     def call(self, graph_module: torch.fx.GraphModule):
         self._build(graph_module)