Qualcomm AI Engine Direct - test framework refactor (pytorch#22072)

Co-author: @winskuo-quic 

### Summary
- extend op / feature test for htp arch (v69~v81)
- extend op / feature test for lpai / gpu backend

### Test plan
```bash
pytest backends/qualcomm/tests/rework/lpai/feature/v6/test.py --device f3c0531 --soc_model SM8850 --build_folder ./build-android/ --backend lpai
pytest backends/qualcomm/tests/rework/lpai/op/v6/test.py
```
```bash
pytest backends/qualcomm/tests/rework/gpu/feature/test.py --device f3c0531 --soc_model SM8650 --build_folder ./build-android/ --backend gpu
ytest backends/qualcomm/tests/rework/gpu/op/test.py --device f3c0531 --soc_model SM8650 --build_folder ./build-android/ --backend gpu
```

Author: haowhsu-quic

backends/qualcomm/_passes/backends/lpai/qnn_lpai_pass_manager.py

@@ -51,17 +51,18 @@ def get_passes_dependency_for_capture_program(cls):
             {
                 DecomposeHardsigmoid: [RemoveRedundancy],
                 DecomposeReciprocal: [RemoveRedundancy],
-                LpaiPartitionFallbackSupport: [TagQuantIO],
-                ResolveDebugHandle: [LpaiPartitionFallbackSupport],
+                LpaiPartitionFallbackSupport: [TagQuantIO, ResolveDebugHandle],
             }
         )
         return deps
 
     def _validate_edge_passes(self) -> None:
-        super()._validate_edge_passes()
         assert isinstance(
-            self.passes[-2], LpaiPartitionFallbackSupport
-        ), "Please ensure LpaiPartitionFallbackSupport is the last edge pass before ResolveDebugHandle."
+            self.passes[-2], ResolveDebugHandle
+        ), "Please ensure ResolveDebugHandle is the last edge pass before LpaiPartitionFallbackSupport."
+        assert isinstance(
+            self.passes[-1], LpaiPartitionFallbackSupport
+        ), "Please ensure LpaiPartitionFallbackSupport is the last pass."
 
     @classmethod
     def get_annotation_passes(cls):

backends/qualcomm/_passes/lpai_partition_fallback_support.py

@@ -254,7 +254,9 @@ def insert_partition_qdq(
                 output_dq_node.meta[QCOM_BYPASS_NODE] = True
         graph_module.graph.eliminate_dead_code()
 
-    def handle_back_to_back_nodes(self, graph_module: torch.fx.GraphModule):
+    def handle_back_to_back_nodes(
+        self, graph_module: torch.fx.GraphModule, unsupported_nodes: set[torch.fx.Node]
+    ):
         """
         This function takes care of following cases:
         1. When 2 contiguous fall back nodes ``a`` and ``b`` (both
@@ -279,6 +281,7 @@ def handle_back_to_back_nodes(self, graph_module: torch.fx.GraphModule):
                     input_node
                     for input_node in node.all_input_nodes
                     if input_node.op == "call_function"
+                    and input_node not in unsupported_nodes
                 ]
                 assert all(
                     input_node.target in dq_ops for input_node in input_call_func_nodes
@@ -327,7 +330,7 @@ def call(self, graph_module: torch.fx.GraphModule) -> PassResult:
         unsupported_nodes = self.get_unsupported_nodes(graph_module)
         for node in unsupported_nodes:
             self.insert_partition_qdq(graph_module, node)
-        self.handle_back_to_back_nodes(graph_module)
+        self.handle_back_to_back_nodes(graph_module, unsupported_nodes)
         graph_module.graph.eliminate_dead_code()
         graph_module.recompile()
-        return PassResult(graph_module, bool(unsupported_nodes))
+        return PassResult(graph_module, True)

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -323,9 +323,7 @@ def get_passes_dependency_for_capture_program(cls):
             RecomposePixelUnshuffle: [RemoveRedundancy],
             RecomposeRmsNorm: [RemoveRedundancy],
             TagQuantIO: [LayoutTransform],
-            ResolveDebugHandle: [
-                TagQuantIO
-            ],  # IMPORTANT: Please always ensure ResolveDebugHandle is the last executed pass.
+            ResolveDebugHandle: [TagQuantIO],
         }
 
     @classmethod

backends/qualcomm/builders/op_batch_norm.py

@@ -29,6 +29,7 @@ class BatchNorm(NodeVisitor):
     target = [
         "aten._native_batch_norm_legit_no_training.default",
         "aten._native_batch_norm_legit.no_stats",
+        "aten._native_batch_norm_legit_functional.default",
     ]
 
     def __init__(self, *args) -> None:

backends/qualcomm/qnn_preprocess.py

@@ -243,6 +243,7 @@ def preprocess_multimethod(  # noqa: C901
                             (handle_id := node.meta.get(DEBUG_HANDLE_KEY))
                             and QCOM_TENSOR_NAME in node.meta
                             and len(node.meta[QCOM_TENSOR_NAME]) == 1
+                            and node.op == "call_function"
                         ):
                             debug_handle_builder.insert_delegate_mapping_entry(
                                 handles=handle_id,

backends/qualcomm/quantizer/annotators/lpai_rules.py

@@ -129,7 +129,7 @@ class AvgPool2d(GeneralOpDef):
 
 # TODO: Batch_norm op cannot directly map to QNN OpBatchnorm due to the number of input doesn't match.
 @register_annotator(
-    [torch.ops.aten.batch_norm.default, torch.ops.aten.instance_norm.default],
+    [torch.ops.aten.batch_norm.default],
     qnn_op=None,
 )
 class BatchNorm(GeneralOpDef):
@@ -420,7 +420,8 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
             torch.ops.aten.topk.default,
             torch.ops.aten.sort.default,
         ):
-            out_act_quantization_spec = SharedQuantizationSpec(node.args[0])
+            # assign to None since they are not supported so far
+            out_act_quantization_spec = None
         node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
             output_qspec=out_act_quantization_spec,
             _annotated=True,
@@ -807,21 +808,6 @@ class ReluMinMax(GeneralOpDef):
     pass
 
 
-# TODO: Expand_as op cannot directly map to QNN OpTile due to the number of input doesn't match.
-@register_annotator(
-    [
-        torch.ops.aten.expand_as.default,
-    ],
-    qnn_op=None,
-)
-class ExpandAs(GeneralOpDef):
-    @staticmethod
-    def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
-        annotate_in_out_obs_sharing_op(node, quantization_config)
-        if not _is_annotated([node]):
-            annotate_single_in_share_out(node, quantization_config)
-
-
 @register_annotator(
     [
         torch.ops.aten.flatten.using_ints,
@@ -854,7 +840,6 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
             return
 
         act_node = node.args[0]
-        weight_node = node.args[2]
 
         # TODO current only support 16a16w
         annotate_input_qspec_map(
@@ -863,94 +848,23 @@ def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
             quantization_config.input_activation,
         )
 
-        annotate_input_qspec_map(
-            node,
-            weight_node,
-            quantization_config.input_activation,
-        )
+        if len(node.args) > 2 and node.args[2] is not None:
+            weight_node = node.args[2]
+            annotate_input_qspec_map(
+                node,
+                weight_node,
+                quantization_config.input_activation,
+            )
         nodes_to_mark_annotated = [node]
         annotate_output_qspec(node, quantization_config.output_activation)
         _mark_nodes_as_annotated(nodes_to_mark_annotated)
 
 
-# TODO: There is a bug in the BackendOpInfo library, so it is bypassed now.
-@register_annotator([torch.ops.aten.rsqrt.default], qnn_op=None)
-class Rsqrt(GeneralOpDef):
-    pass
-
-
 @register_annotator([torch.ops.aten.scaled_dot_product_attention.default], qnn_op=None)
 class ScaledDotProductAttention(GeneralOpDef):
     pass
 
 
-@register_annotator(
-    [
-        torch.ops.aten.scatter.src,
-        torch.ops.aten.scatter.value,
-        torch.ops.aten.scatter_add.default,
-        torch.ops.aten.scatter_reduce.two,
-    ],
-    qnn_op=None,
-)
-class ScatterElements(GeneralOpDef):
-    @staticmethod
-    def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
-        if _is_annotated([node]):
-            return
-
-        input_act = node.args[0]
-        if not isinstance(input_act, Node) or not _is_float_tensor(input_act):
-            return
-
-        input_qspec_map = {}
-        input_qspec_map[input_act] = quantization_config.input_activation
-
-        if (
-            len(node.args) > 3
-            and isinstance(node.args[3], Node)
-            and _is_float_tensor(node.args[3])
-        ):
-            input_qspec_map[node.args[3]] = SharedQuantizationSpec((input_act, node))
-
-        output_act_qspec = (
-            SharedQuantizationSpec((input_act, node))
-            if _is_float_tensor(node)
-            else None
-        )
-
-        if len(input_qspec_map) > 0 or output_act_qspec is not None:
-            node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
-                input_qspec_map=input_qspec_map,
-                output_qspec=output_act_qspec,
-                _annotated=True,
-            )
-
-
-@register_annotator([torch.ops.aten.sort.default], QnnConstants.OpTopK.op_name)
-class Sort(GeneralOpDef):
-    @staticmethod
-    def annotate(node: Node, quantization_config: QuantizationConfig) -> None:
-        if _is_annotated([node]):
-            return
-
-        input_qspec_map = {}
-        input_act_qspec = quantization_config.input_activation
-        out_act_quantization_spec = None
-        if input_act_qspec is not None:
-            if _is_float_tensor(node.args[0]):
-                input_act = node.args[0]
-                assert isinstance(input_act, Node)
-                input_qspec_map[input_act] = input_act_qspec
-                out_act_quantization_spec = SharedQuantizationSpec((input_act, node))
-
-        node.meta[Q_ANNOTATION_KEY] = QuantizationAnnotation(
-            input_qspec_map=input_qspec_map,
-            output_qspec=out_act_quantization_spec,
-            _annotated=True,
-        )
-
-
 @register_annotator(
     [torch.ops.aten.sigmoid, torch.ops.aten.sigmoid.default],
     QnnConstants.OpSigmoid.op_name,

backends/qualcomm/tests/models.py

@@ -2943,6 +2943,16 @@ def forward(self, x):
         )
 
 
+class ConvRelu(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(3, 8, kernel_size=3, padding=1)
+        self.relu = torch.nn.ReLU()
+
+    def forward(self, x):
+        return self.relu(self.conv(x))
+
+
 class TopKandIndex(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/rework/conftest.py

@@ -32,6 +32,7 @@
     get_qnn_context_binary_alignment,
     prepare_pt2e,
     QnnConfig,
+    QnnExecuTorchBackendType,
     QnnQuantizer,
     setup_common_args_and_variables,
     SimpleADB,
@@ -269,7 +270,7 @@ def qnn_config(global_setup, request):
             f'invalid configuration detected, fall back to emulator workload:\n"{e}"'
         )
         config = QnnConfig(
-            soc_model="unknown", build_folder="build-x86", compile_only=True
+            soc_model="unknown", build_folder="build-x86", enable_x86_64=True
         )
 
     return config
@@ -349,6 +350,7 @@ def invoke_remote(
     qnn_config: QnnConfig,
     executorch_prog: ExecutorchProgramManager,
     callback: callable,
+    inputs: Tuple[torch.Tensor] = None,
 ):
     with tempfile.TemporaryDirectory() as tmp_dir:
         pte_fname = f"{tmp_dir}/qnn_executorch_test.pte"
@@ -363,7 +365,7 @@ def invoke_remote(
             pte_path=[pte_fname],
             workspace=f"/data/local/tmp/{device_workspace}",
         )
-        adb.push()
+        adb.push(inputs=[inputs] if inputs is not None else None)
         callback(adb)
 
 
@@ -478,7 +480,23 @@ def export_and_verify(
     metrics: Metrics,
 ):
     with calibrate(module, [inputs], quantizer) as exported_module:
-        if quantizer is not None:
+        fake_tensors = (
+            [
+                node.meta["val"]
+                for node in exported_module.graph.nodes
+                if node.op == "call_function" and "val" in node.meta
+            ]
+            if quantizer
+            else []
+        )
+        dtypes = set()
+        for tensor in fake_tensors:
+            if isinstance(tensor, (tuple, list)):
+                dtypes.update([n.dtype for n in tensor])
+            else:
+                dtypes.add(tensor.dtype)
+
+        if quantizer and {torch.float, torch.float32} & dtypes:
             nodes = {node.target for node in exported_module.graph.nodes}
             q_and_dq = {
                 torch.ops.quantized_decomposed.quantize_per_tensor.default,
@@ -505,15 +523,34 @@ def export_and_verify(
         )
     )
     execution_plan = executorch_prog.executorch_program.execution_plan[0]
+
+    def validate():
+        match qnn_config.backend:
+            case QnnExecuTorchBackendType.kHtpBackend:
+                return len(execution_plan.operators) == 0
+            case QnnExecuTorchBackendType.kGpuBackend:
+                return len(execution_plan.operators) == 0
+            case QnnExecuTorchBackendType.kLpaiBackend:
+                aten_op_names = {
+                    op.name
+                    for op in execution_plan.operators
+                    if "quantize" not in op.name
+                }
+                return len(aten_op_names) == 0
+            case _:
+                return True
+
     assert all(
         [
-            len(execution_plan.delegates) == 1,
-            execution_plan.delegates[0].id == "QnnBackend",
-            len(execution_plan.operators) == 0,
+            (
+                len(execution_plan.delegates) == 1
+                and execution_plan.delegates[0].id == "QnnBackend"
+            ),
+            validate(),
         ]
     ), EXPECT_NOT_FULLY_DELEGATED
 
-    mode = "emulator" if qnn_config.build_folder == "build-x86" else "remote"
+    mode = "emulator" if qnn_config.enable_x86_64 else "remote"
     globals()[f"verify_output_{mode}"](
         module=module,
         inputs=inputs,

backends/qualcomm/tests/rework/gpu/conftest.py

@@ -3,3 +3,40 @@
 #
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
+
+from typing import Any
+
+import pytest
+
+from executorch.backends.qualcomm.export_utils import (
+    generate_gpu_compiler_spec,
+    generate_qnn_executorch_compiler_spec,
+    QcomChipset,
+)
+
+
+def with_gpu_context(func):
+    def wrapper(request, kwargs):
+        preserved = {k: kwargs.pop(k) for k in ["expected"]}
+        qnn_config = request.getfixturevalue("qnn_config")
+        fixtures = {
+            "quantizer": None,
+            "compile_spec": generate_qnn_executorch_compiler_spec(
+                soc_model=getattr(QcomChipset, qnn_config.soc_model),
+                backend_options=generate_gpu_compiler_spec(),
+                online_prepare=True,
+            ),
+        }
+        return func(request, fixtures | preserved)
+
+    return wrapper
+
+
+def enumerate_fp_dtype(metric: Any):
+    def wrapper(test_body):
+        return pytest.mark.parametrize(
+            "kwargs",
+            [pytest.param({"act": None, "expected": metric}, id="fp")],
+        )(test_body)
+
+    return wrapper