add argmin, add_ ops

Summary: add argmin, add_ ops  annontation

Differential Revision: D68811662

Author: Min Guo

backends/qualcomm/_passes/layout_transform.py

@@ -45,6 +45,7 @@ class LayoutTransform(ExportPass):
     layout_agnostic_ops = {
         exir_ops.edge.aten.abs.default,
         exir_ops.edge.aten.add.Tensor,
+        exir_ops.edge.aten.argmin.default,
         exir_ops.edge.aten.bmm.default,
         exir_ops.edge.aten.cat.default,
         exir_ops.edge.aten.ceil.default,

backends/qualcomm/builders/__init__.py

@@ -10,6 +10,7 @@
     op_adaptive_avg_pool2d,
     op_add,
     op_arange,
+    op_argmin,
     op_avg_pool2d,
     op_batch_norm,
     op_bmm,
@@ -82,6 +83,7 @@
     op_adaptive_avg_pool2d,
     op_add,
     op_arange,
+    op_argmin,
     op_avg_pool2d,
     op_batch_norm,
     op_bmm,

backends/qualcomm/builders/op_add.py

@@ -15,7 +15,7 @@
 
 @register_node_visitor
 class Add(NodeVisitor):
-    target = ["aten.add.Tensor"]
+    target = ["aten.add.Tensor", "aten.add_.Tensor"]
 
     def __init__(self, *args) -> None:
         super().__init__(*args)

backends/qualcomm/builders/op_argmin.py

@@ -0,0 +1,78 @@
+# 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.
+from typing import cast, Dict
+
+import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper
+import numpy as np
+import torch
+from executorch.backends.qualcomm.utils.constants import QCOM_AXIS_ORDER, QCOM_DATA
+
+from .node_visitor import NodeVisitor, register_node_visitor
+from .qnn_constants import OpArgmin, QNN_OP_PACKAGE_NAME_QTI_AISW
+
+
+@register_node_visitor
+class Argmin(NodeVisitor):
+    target = ["aten.argmin.default"]
+
+    def __init__(self, *args) -> None:
+        super().__init__(*args)
+
+    def define_node(
+        self,
+        node: torch.fx.Node,
+        nodes_to_wrappers: Dict[torch.fx.Node, PyQnnWrapper.TensorWrapper],
+    ) -> PyQnnWrapper.PyQnnOpWrapper:
+        input_node = node.args[0]
+        input_tensor = self.get_tensor(input_node, node)
+        argmin_inp_tensor_wrapper = self.define_tensor(
+            input_node,
+            node,
+            input_tensor,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+        argmin_input_tensors = [argmin_inp_tensor_wrapper]
+
+        output_tensor = self.get_tensor(node, node)
+        output_tensor_wrapper = self.define_tensor(
+            node,
+            node,
+            output_tensor,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+        argmin_output_tensors = [output_tensor_wrapper]
+
+        dim = cast(int, node.args[1])
+        if dim < 0:
+            dim = dim % len(input_tensor.shape)
+        if QCOM_AXIS_ORDER in node.meta:
+            dim = node.meta[QCOM_AXIS_ORDER].index(dim)
+
+        argmin_op = PyQnnWrapper.PyQnnOpWrapper(
+            node.name,
+            QNN_OP_PACKAGE_NAME_QTI_AISW,
+            OpArgmin.op_name,
+        )
+        argmin_op.AddInputTensors(argmin_input_tensors)
+        argmin_op.AddOutputTensors(argmin_output_tensors)
+
+        argmin_op.AddScalarParam(
+            OpArgmin.param_axis,
+            PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_UINT_32,
+            {QCOM_DATA: np.uint32(dim)},
+        )
+        
+        if len(node.args) > 2:
+            keep_dims = cast(bool, node.args[2])
+            argmin_op.AddScalarParam(
+                OpArgmin.param_keep_dims,
+                PyQnnWrapper.Qnn_DataType_t.QNN_DATATYPE_BOOL_8,
+                {QCOM_DATA: keep_dims},
+            )
+        
+        return argmin_op

backends/qualcomm/builders/qnn_constants.py

@@ -301,6 +301,13 @@ class OpReduceMean:
     param_keep_dims: str = "keep_dims"
 
 
+@dataclass(init=False, frozen=True)
+class OpArgmin:
+    op_name: str = "Argmin"
+    param_axis: str = "axis"
+    param_keep_dims: str = "keep_dims"
+
+
 @dataclass(init=False, frozen=True)
 class OpReduceSum:
     op_name: str = "ReduceSum"

backends/qualcomm/quantizer/annotators.py

@@ -162,7 +162,9 @@ def annotate_binary(node: Node, quantization_config: QuantizationConfig) -> None
     )
 
 
-@register_annotator([torch.ops.aten.add, torch.ops.aten.add.Tensor])
+@register_annotator(
+    [torch.ops.aten.add, torch.ops.aten.add_.Tensor, torch.ops.aten.add.Tensor]
+)
 def annotate_add(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_binary(node, quantization_config)
 
@@ -476,6 +478,10 @@ def annotate_softmax(node: Node, quantization_config: QuantizationConfig) -> Non
 def annotate_log_softmax(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_single_in_single_out(node, quantization_config)
 
+@register_annotator([torch.ops.aten.argmin.default])
+def annotate_argmin(node: Node, quantization_config: QuantizationConfig) -> None:
+    annotate_single_in_single_out(node, quantization_config)
+
 
 @register_annotator([torch.ops.aten.log.default])
 def annotate_log(node: Node, quantization_config: QuantizationConfig) -> None:

backends/qualcomm/tests/models.py

@@ -274,7 +274,6 @@ def forward(self, x):
         x = self.logsoftmax(x)
         return x
 
-
 class Conv2dAvgPool2d(torch.nn.Module):
     def __init__(self):
         super().__init__()
@@ -287,6 +286,18 @@ def forward(self, x):
         return self.pool(self.conv(x))
 
 
+class Conv2dArgmin(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv = torch.nn.Conv2d(
+            3, 16, 7, bias=True, stride=2, padding=3, dilation=1
+        )
+    
+    def forward(self, x):
+        x = self.conv(x)
+        return torch.argmin(x, dim=0, keepdim=True)
+        
+
 class Conv2dBnHardtanhMean(torch.nn.Module):
     def __init__(self):
         super(Conv2dBnHardtanhMean, self).__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -681,6 +681,10 @@ def test_qnn_backend_view(self):
         sample_input = (torch.randn([1, 8, 512]), torch.randn([1, 2, 8, 256]))
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_argmin(self):
+        module = Conv2dArgmin()  # noqa: F405   
+        sample_input = (torch.randn(16, 3, 16, 16),)
+        self.lower_module_and_test_output(module, sample_input)
 
 class TestQNNFloatingPointModel(TestQNN):
     # TODO: refactor to support different backends
@@ -704,7 +708,8 @@ def test_qnn_backend_chunk_add(self):
         torch.manual_seed(8)
         sample_input = (torch.randn(1, 2, 4, 2),)
         self.lower_module_and_test_output(module, sample_input)
-
+    
+   
     def test_qnn_backend_conv1d_relu_log_softmax(self):
         module = Conv1dReluLogSoftmax()  # noqa: F405
         sample_input = (torch.rand(1, 2, 28),)
@@ -1585,6 +1590,12 @@ def test_qnn_backend_view(self):
         sample_input = (torch.randn([1, 8, 512]), torch.randn([1, 2, 8, 256]))
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
+    
+    def test_qnn_backend_argmin(self):
+        module = Conv2dArgmin()  # noqa: F405   
+        sample_input = (torch.randn(16, 3, 16, 16),)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
 
 
 class TestQNNQuantizedModel(TestQNN):
@@ -1610,6 +1621,8 @@ def test_qnn_backend_chunk_add(self):
         sample_input = (torch.randn(1, 1, 4, 2),)
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
+    
+
 
     def test_qnn_backend_conv1d_relu_log_softmax(self):
         module = Conv1dReluLogSoftmax()  # noqa: F405