Add to.dtype and neg ops (#8041)

Summary:
Pull Request resolved: #8041

#8041
add neg ops and annonation of to.dtype

Reviewed By: Andriyluck

Differential Revision: D68815927

Author: billmguo

backends/qualcomm/_passes/layout_transform.py

@@ -75,6 +75,7 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.mean.dim,
         exir_ops.edge.aten.minimum.default,
         exir_ops.edge.aten.mul.Tensor,
+        exir_ops.edge.aten.neg.default,
         exir_ops.edge.aten.pow.Tensor_Scalar,
         exir_ops.edge.aten.prelu.default,
         exir_ops.edge.aten.repeat.default,

backends/qualcomm/builders/__init__.py

@@ -47,6 +47,7 @@
     op_mean_dim,
     op_min,
     op_mul,
+    op_neg,
     op_pad,
     op_pow,
     op_prelu,
@@ -120,6 +121,7 @@
     op_mean_dim,
     op_min,
     op_mul,
+    op_neg,
     op_pad,
     op_pow,
     op_prelu,

backends/qualcomm/builders/op_neg.py

@@ -0,0 +1,53 @@
+# 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 Dict
+
+import executorch.backends.qualcomm.python.PyQnnWrapperAdaptor as PyQnnWrapper
+import torch
+
+from .node_visitor import NodeVisitor, register_node_visitor
+from .qnn_constants import OpElementWiseNeg, QNN_OP_PACKAGE_NAME_QTI_AISW
+
+
+@register_node_visitor
+class Neg(NodeVisitor):
+    target = ["aten.neg.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)
+        neg_inp_tensor_wrapper = self.define_tensor(
+            input_node,
+            node,
+            input_tensor,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+        neg_input_tensors = [neg_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,
+        )
+        neg_output_tensors = [output_tensor_wrapper]
+        neg_op = PyQnnWrapper.PyQnnOpWrapper(
+            node.name,
+            QNN_OP_PACKAGE_NAME_QTI_AISW,
+            OpElementWiseNeg.op_name,
+        )
+        neg_op.AddInputTensors(neg_input_tensors)
+        neg_op.AddOutputTensors(neg_output_tensors)
+        return neg_op

backends/qualcomm/builders/qnn_constants.py

@@ -145,6 +145,11 @@ class OpElementWiseMultiply:
     op_name: str = "ElementWiseMultiply"
 
 
+@dataclass(init=False, frozen=True)
+class OpElementWiseNeg:
+    op_name: str = "ElementWiseNeg"
+
+
 @dataclass(init=False, frozen=True)
 class OpElementWiseNeuron:
     op_name: str = "ElementWiseNeuron"

backends/qualcomm/quantizer/annotators.py

@@ -162,9 +162,7 @@ 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])
 def annotate_add(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_binary(node, quantization_config)
 
@@ -398,6 +396,11 @@ def annotate_max_pool2d_with_indices(
     annotate_single_in_single_out(node, quantization_config)
 
 
+@register_annotator([torch.ops.aten.neg.default])
+def annotate_neg(node: Node, quantization_config: QuantizationConfig) -> None:
+    annotate_single_in_single_out(node, quantization_config)
+
+
 @register_annotator([torch.ops.aten.adaptive_avg_pool2d.default])
 def annotate_adaptive_avgpool2d(
     node: Node, quantization_config: QuantizationConfig
@@ -479,8 +482,6 @@ def annotate_log_softmax(node: Node, quantization_config: QuantizationConfig) ->
     annotate_single_in_single_out(node, quantization_config)
 
 
-
-
 @register_annotator([torch.ops.aten.log.default])
 def annotate_log(node: Node, quantization_config: QuantizationConfig) -> None:
     annotate_single_in_single_out(node, quantization_config)

backends/qualcomm/tests/models.py

@@ -894,6 +894,14 @@ def forward(self, x):
         return attn_output
 
 
+class Neg(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x):
+        return torch.neg(x)
+
+
 class Pad(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -111,7 +111,7 @@ def test_qnn_backend_arange(self):
         for i, module in enumerate(modules):
             with self.subTest(i=i):
                 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),)
@@ -546,6 +546,11 @@ def test_qnn_backend_minimum(self):
         sample_input = (torch.randn(1, 2, 3, 4), torch.randn(2, 3, 4))
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_neg(self):
+        module = Neg()  # noqa: F405
+        sample_input = (torch.randn(1, 4, 16, 16),)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_pad(self):
         module = Pad()  # noqa: F405
         sample_input = (torch.randn([1, 8, 128]),)
@@ -1434,6 +1439,12 @@ def test_qnn_backend_minimum(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_neg(self):
+        module = Neg()  # noqa: F405
+        sample_input = (torch.randn(1, 4, 16, 16),)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_pad(self):
         module = Pad()  # noqa: F405
         sample_input = (torch.randn([1, 8, 128]),)
@@ -1603,8 +1614,6 @@ def test_qnn_backend_view(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
-   
-
 
 class TestQNNQuantizedModel(TestQNN):
     # TODO: refactor to support different backends