Add logical and op (pytorch#13342)

Summary:

As title, add the logical and op for an internal model

Rollback Plan:

Differential Revision: D80122607

Author: cccclai

backends/qualcomm/_passes/layout_transform.py

@@ -89,6 +89,7 @@ class LayoutTransform(ExportPass):
         exir_ops.edge.aten.le.Tensor,
         exir_ops.edge.aten.linear.default,
         exir_ops.edge.aten.log.default,
+        exir_ops.edge.aten.logical_and.default,
         exir_ops.edge.aten.logical_not.default,
         exir_ops.edge.aten.lt.Scalar,
         exir_ops.edge.aten.lt.Tensor,

backends/qualcomm/builders/__init__.py

@@ -52,6 +52,7 @@
     op_linear,
     op_log,
     op_log_softmax,
+    op_logical_and,
     op_logical_not,
     op_lt,
     op_matmul,
@@ -148,6 +149,7 @@
     op_le,
     op_linear,
     op_log,
+    op_logical_and,
     op_logical_not,
     op_log_softmax,
     op_lt,

backends/qualcomm/builders/op_logical_and.py

@@ -0,0 +1,66 @@
+# 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
+from .node_visitor_manager import register_node_visitor
+from .qnn_constants import OpElementWiseAnd, QNN_OP_PACKAGE_NAME_QTI_AISW
+
+
+@register_node_visitor
+class And(NodeVisitor):
+    target = ["aten.logical_and.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_1 = self.get_node(node.args[0])
+        input_tensor_1 = self.get_tensor(input_node_1, node)
+        input_tensor_wrapper_1 = self.define_tensor(
+            input_node_1,
+            node,
+            input_tensor_1,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+
+        input_node_2 = self.get_node(node.args[1])
+        input_tensor_2 = self.get_tensor(input_node_2, node)
+        input_tensor_wrapper_2 = self.define_tensor(
+            input_node_2,
+            node,
+            input_tensor_2,
+            PyQnnWrapper.Qnn_TensorType_t.QNN_TENSOR_TYPE_NATIVE,
+            nodes_to_wrappers,
+        )
+
+        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,
+        )
+
+        logical_and_op = PyQnnWrapper.PyQnnOpWrapper(
+            node.name,
+            QNN_OP_PACKAGE_NAME_QTI_AISW,
+            OpElementWiseAnd.op_name,
+        )
+        logical_and_op.AddInputTensors([input_tensor_wrapper_1, input_tensor_wrapper_2])
+        logical_and_op.AddOutputTensors([output_tensor_wrapper])
+
+        return logical_and_op

backends/qualcomm/builders/qnn_constants.py

@@ -107,11 +107,6 @@ class OpElementWiseAdd:
     op_name: str = "ElementWiseAdd"
 
 
-@dataclass(init=False, frozen=True)
-class OpElementWiseAnd:
-    op_name: str = "ElementWiseAnd"
-
-
 @dataclass(init=False, frozen=True)
 class OpElementWiseAtan:
     op_name: str = "ElementWiseAtan"
@@ -200,6 +195,11 @@ class OpElementWiseNeuron:
     param_beta: str = "beta"
 
 
+@dataclass(init=False, frozen=True)
+class OpElementWiseAnd:
+    op_name: str = "ElementWiseAnd"
+
+
 @dataclass(init=False, frozen=True)
 class OpElementWiseNot:
     op_name: str = "ElementWiseNot"

backends/qualcomm/tests/models.py

@@ -1108,6 +1108,14 @@ def forward(self, x):
         return torch.log(x)
 
 
+class LogicalAnd(torch.nn.Module):
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x, y):
+        return torch.logical_and(x, y)
+
+
 class LogicalNot(torch.nn.Module):
     def __init__(self):
         super().__init__()

backends/qualcomm/tests/test_qnn_delegate.py

@@ -856,6 +856,14 @@ def test_qnn_backend_log(self):
         sample_input = (torch.rand([1, 2, 3, 4]),)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_logical_and(self):
+        module = LogicalAnd()  # noqa: F405
+        input1 = torch.tensor([True, False, True, False])
+        input2 = torch.tensor([True, True, False, False])
+        sample_input = (input1, input2)
+        self.lower_module_and_test_output(module, sample_input)
+
+
     def test_qnn_backend_logical_not(self):
         module = LogicalNot()  # noqa: F405
         sample_input = (torch.rand([1, 2, 3, 4]),)
@@ -2244,6 +2252,14 @@ def test_qnn_backend_log(self):
         module = self.get_qdq_module(module, sample_input)
         self.lower_module_and_test_output(module, sample_input)
 
+    def test_qnn_backend_logical_and(self):
+        module = LogicalAnd()  # noqa: F405
+        input1 = torch.tensor([True, False, True, False])
+        input2 = torch.tensor([True, True, False, False])
+        sample_input = (input1, input2)
+        module = self.get_qdq_module(module, sample_input)
+        self.lower_module_and_test_output(module, sample_input)
+
     def test_qnn_backend_logical_not(self):
         module = LogicalNot()  # noqa: F405
         sample_input = (torch.rand([1, 2, 3, 4]),)