Qualcomm AI Engine Direct - Test Framework Refactor

Co-authored-by: @winskuo-quic, @chenweng-quic
- introduce pytest and reorganize the file architecture for
  finer-grained testing
- wider coverage of operator test with combinations of different
  precisions, codebase was changed accordingly
- add feature tests for HTP

Author: haowhsu-quic

.ci/scripts/test_model.sh

@@ -264,7 +264,7 @@ test_model_with_qnn() {
         ;;
   esac
 
-  "${PYTHON_EXECUTABLE}" -m examples.qualcomm.${SCRIPT_FOLDER}.${EXPORT_SCRIPT} -b ${CMAKE_OUTPUT_DIR} -m ${QNN_CHIPSET} --ci --compile_only $EXTRA_FLAGS
+  "${PYTHON_EXECUTABLE}" -m examples.qualcomm.${SCRIPT_FOLDER}.${EXPORT_SCRIPT} --build_folder ${CMAKE_OUTPUT_DIR} --soc_model ${QNN_CHIPSET} --ci --compile_only $EXTRA_FLAGS
   EXPORTED_MODEL=$(find "./${EXPORT_SCRIPT}" -type f -name "${MODEL_NAME}*.pte" -print -quit)
 }
 

backends/qualcomm/_passes/__init__.py

@@ -52,11 +52,9 @@
 from .recompose_pad_maxpool2d import RecomposePadMaxPool2d
 from .recompose_pixel_unshuffle import RecomposePixelUnshuffle
 from .recompose_rms_norm import RecomposeRmsNorm
-from .reduce_dynamic_range import ReduceDynamicRange
 from .remove_0d_tensor import Remove0DTensor
 from .remove_redundancy import RemoveRedundancy
 from .replace_arange_args import ReplaceArangeArgs
-from .replace_inf_values import ReplaceInfValues
 from .resolve_debug_handle import ResolveDebugHandle
 from .seq_mse import SeqMSE
 from .tag_quant_io import TagQuantIO
@@ -110,11 +108,9 @@
     RecomposePadMaxPool2d,
     RecomposePixelUnshuffle,
     RecomposeRmsNorm,
-    ReduceDynamicRange,
     Remove0DTensor,
     RemoveRedundancy,
     ReplaceArangeArgs,
-    ReplaceInfValues,
     ResolveDebugHandle,
     SeqMSE,
     TagQuantIO,

backends/qualcomm/_passes/decompose_col_im.py

@@ -30,6 +30,8 @@ def _decompose_im2col(self, graph_module: torch.fx.GraphModule):
             if node.target == self.im2col_op:
                 input_node = node.args[0]
                 kernel_size = node.args[1]
+                dilation = node.args[2]
+                padding = node.args[3]
                 stride = node.args[4]
                 batch_size = node.meta["val"].shape[0]
                 assert (
@@ -41,6 +43,12 @@ def _decompose_im2col(self, graph_module: torch.fx.GraphModule):
                 assert (
                     kernel_size[0] == kernel_size[1]
                 ), "im2col can only be converted when kernel height == width"
+                assert all(
+                    d == 1 for d in dilation
+                ), "col2im can only be converted when dilation equals to (1, 1)"
+                assert all(
+                    p == 0 for p in padding
+                ), "col2im can only be converted when padding equals to (0, 0)"
                 users = list(node.users.keys())
                 with graph_module.graph.inserting_after(input_node):
                     pixel_unshuffle_node = graph_module.graph.create_node(
@@ -77,6 +85,8 @@ def _decompose_col2im(self, graph_module: torch.fx.GraphModule):
                 input_node = node.args[0]
                 output_size = node.args[1]
                 kernel_size = node.args[2]
+                dilation = node.args[3]
+                padding = node.args[4]
                 stride = node.args[5]
                 batch_size = node.meta["val"].shape[0]
                 assert (
@@ -88,6 +98,13 @@ def _decompose_col2im(self, graph_module: torch.fx.GraphModule):
                 assert (
                     kernel_size[0] == kernel_size[1]
                 ), "col2im can only be converted when kernel height == width"
+                assert all(
+                    d == 1 for d in dilation
+                ), "col2im can only be converted when dilation equals to (1, 1)"
+                assert all(
+                    p == 0 for p in padding
+                ), "col2im can only be converted when padding equals to (0, 0)"
+
                 users = list(node.users.keys())
                 with graph_module.graph.inserting_after(input_node):
                     view_tensor = input_node.meta["val"].reshape(

backends/qualcomm/_passes/decompose_linalg_vector_norm.py

@@ -25,6 +25,12 @@ def forward(self, x):
             self.dim = 0
 
         x = torch.abs(x)
+
+        # QNN would not be able to compute pow where exponential is inf or -inf.
+        if self.exp == float("inf"):
+            return torch.amax(x, dim=self.dim, keepdim=self.keepdim)
+        if self.exp == float("-inf"):
+            return torch.amin(x, dim=self.dim, keepdim=self.keepdim)
         x = torch.pow(x, self.exp)
         x = torch.sum(x, dim=self.dim, keepdim=self.keepdim)
         return torch.pow(x, 1.0 / self.exp)

backends/qualcomm/_passes/decompose_remainder.py

@@ -23,8 +23,10 @@ def __init__(self):
         super(DecomposeRemainder, self).__init__()
         self.remainder_targets = {
             torch.ops.aten.remainder.Scalar,
+            torch.ops.aten.remainder.Scalar_Tensor,
             torch.ops.aten.remainder.Tensor,
             exir_ops.edge.aten.remainder.Scalar,
+            exir_ops.edge.aten.remainder.Scalar_Tensor,
             exir_ops.edge.aten.remainder.Tensor,
         }
 
@@ -35,7 +37,7 @@ def call(self, graph_module: torch.fx.GraphModule):
 
         for node in list(graph.nodes):
             if node.op == "call_function" and node.target in self.remainder_targets:
-                x_node = node.args[0]
+                x_arg = node.args[0]
                 y_arg = node.args[1]
                 is_edge = isinstance(node.target, EdgeOpOverload)
                 meta = node.meta
@@ -61,8 +63,21 @@ def call(self, graph_module: torch.fx.GraphModule):
                     else torch.ops.aten.sub.Tensor
                 )
 
-                is_scalar = not isinstance(y_arg, torch.fx.Node)
-                if is_scalar and is_edge:
+                is_x_scalar = not isinstance(x_arg, torch.fx.Node)
+                if is_x_scalar and is_edge:
+                    if x_arg not in const_cache:
+                        attr_name = get_new_attr_name_with_prefix("_remainder_const_")(
+                            graph_module
+                        )
+                        const_cache[x_arg] = get_const_node(
+                            graph, graph_module, attr_name, x_arg, node
+                        )
+                    x_node = const_cache[x_arg]
+                else:
+                    x_node = x_arg
+
+                is_y_scalar = not isinstance(y_arg, torch.fx.Node)
+                if is_y_scalar and is_edge:
                     if y_arg not in const_cache:
                         attr_name = get_new_attr_name_with_prefix("_remainder_const_")(
                             graph_module

backends/qualcomm/_passes/decompose_roll.py

@@ -16,7 +16,7 @@ def __init__(self, val_shape, shifts, dims):
         super().__init__()
         self.val_shape = val_shape
         if dims[0] is None:
-            self.shifts = [shifts[0] % torch.numel(torch.tensor(val_shape))]
+            self.shifts = [shifts[0] % torch.numel(torch.empty(val_shape))]
         else:
             self.shifts = [shift % val_shape[dim] for shift, dim in zip(shifts, dims)]
         self.dims = dims

backends/qualcomm/_passes/lift_constant_scalar_operands.py

@@ -33,11 +33,17 @@ class TensorOpInfo:
 
 SCALAR_OPS = {
     aten.eq.Scalar: TensorOpInfo(aten.eq.Tensor, False, False),
+    aten.eq.Tensor: TensorOpInfo(aten.eq.Tensor, False, False),
     aten.ge.Scalar: TensorOpInfo(aten.ge.Tensor, False, False),
+    aten.ge.Tensor: TensorOpInfo(aten.ge.Tensor, False, False),
     aten.gt.Scalar: TensorOpInfo(aten.gt.Tensor, False, False),
+    aten.gt.Tensor: TensorOpInfo(aten.gt.Tensor, False, False),
     aten.le.Scalar: TensorOpInfo(aten.le.Tensor, False, False),
+    aten.le.Tensor: TensorOpInfo(aten.le.Tensor, False, False),
     aten.lt.Scalar: TensorOpInfo(aten.lt.Tensor, False, False),
+    aten.lt.Tensor: TensorOpInfo(aten.lt.Tensor, False, False),
     aten.ne.Scalar: TensorOpInfo(aten.ne.Tensor, False, False),
+    aten.ne.Tensor: TensorOpInfo(aten.ne.Tensor, False, False),
     aten.add.Scalar: TensorOpInfo(aten.add.Tensor, False, False),
     aten.add_.Scalar: TensorOpInfo(aten.add_.Tensor, False, False),
     # For below cases, refer to LiftAddTensor Model in UT for sample
@@ -87,6 +93,7 @@ def _build_tensor_constant(
     ) -> TensorConstant:
         # For dtype, in some cases, we cannot use node.args[0] as scalar dtype.
         # Ex: Where op args[0] can be bool, however, we probably want args[1] and args[2] to be dtype same as node.meta["val"] instead of bool type
+
         first_arg = node.args[0]
         tensor = torch.tensor(
             const_val,

backends/qualcomm/_passes/qnn_pass_manager.py

@@ -57,11 +57,9 @@
     RecomposePadMaxPool2d,
     RecomposePixelUnshuffle,
     RecomposeRmsNorm,
-    ReduceDynamicRange,
     Remove0DTensor,
     RemoveRedundancy,
     ReplaceArangeArgs,
-    ReplaceInfValues,
     ResolveDebugHandle,
     TagQuantIO,
 )
@@ -226,7 +224,6 @@ def transform_for_to_edge_pipeline(
     # Before quantizer
     def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(RemoveRedundancy(quantization_capture=True))
-        self.add_pass(ReduceDynamicRange())
         self.add_pass(RecomposePixelUnshuffle(quantization_capture=True))
         self.add_pass(RecomposeRmsNorm(quantization_capture=True))
         self.add_pass(ReplaceArangeArgs())
@@ -255,7 +252,6 @@ def transform_for_annotation_pipeline(self, graph_module: GraphModule):
         self.add_pass(DecomposeSelectScatter())
         self.add_pass(DecomposeLinalgVectorNorm(quantization_capture=True))
         self.add_pass(DecomposeLogVariants())
-        self.add_pass(ReplaceInfValues())
         self.add_pass(LiftConstantScalarOperands())
         self.add_pass(InsertReshapeForReduceOps())
         return self._transform(graph_module)