[Opset13][PT FE] Update torch bitwise operators (#20339)

* Add opset-13 bitwise implementation

* Improvements in test

* Add transformation BitwiseOps->LogicalOps for bool

* Improve existing tests to better tests dtypes

* Disable transformatiions for supported bitwise ops

* Improvebitwise test inputs

* Update src/common/transformations/src/transformations/op_conversions/convert_bitwise_to_logical_bool.cpp

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>

* Update src/common/transformations/src/transformations/op_conversions/convert_bitwise_to_logical_bool.cpp

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>

* Update src/common/transformations/src/transformations/op_conversions/convert_bitwise_to_logical_bool.cpp

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>

* Update src/common/transformations/src/transformations/op_conversions/convert_bitwise_to_logical_bool.cpp

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>

* Update to REGISETR_PASS

---------

Co-authored-by: Katarzyna Mitrus <katarzyna.mitrus@intel.com>

tests/layer_tests/pytorch_tests/test_and.py

@@ -9,13 +9,11 @@ from pytorch_layer_test_class import PytorchLayerTest
 
 
 class TestAnd(PytorchLayerTest):
-
     def _prepare_input(self):
         return self.input_data
 
     def create_model_tensor_input(self):
         class aten_and_tensor(torch.nn.Module):
-
             def __init__(self) -> None:
                 super().__init__()
 
@@ -25,10 +23,9 @@ class TestAnd(PytorchLayerTest):
         ref_net = None
 
         return aten_and_tensor(), ref_net, "aten::__and__"
-    
+
     def create_model_bool_input(self):
         class aten_and_bool(torch.nn.Module):
-
             def __init__(self) -> None:
                 super().__init__()
 
@@ -39,18 +36,43 @@ class TestAnd(PytorchLayerTest):
 
         return aten_and_bool(), ref_net, "aten::__and__"
 
+    def create_model_int_input(self):
+        class aten_and_int(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def forward(self, int_a: int, int_b: int):
+                return int_a & int_b
+
+        ref_net = None
+
+        return aten_and_int(), ref_net, "aten::__and__"
+
     @pytest.mark.nightly
     @pytest.mark.precommit
     def test_and_tensor(self, ie_device, precision, ir_version):
-        self.input_data = (np.array([True, False, False], dtype=np.bool_), np.array(
-            [True, True, False], dtype=np.bool_))
-        self._test(*self.create_model_tensor_input(),
-                   ie_device, precision, ir_version)
+        self.input_data = (
+            np.array([True, False, False], dtype=np.bool_),
+            np.array([True, True, False], dtype=np.bool_),
+        )
+        self._test(*self.create_model_tensor_input(), ie_device, precision, ir_version)
 
     @pytest.mark.nightly
     @pytest.mark.precommit
     def test_and_bool(self, ie_device, precision, ir_version):
-        self.input_data = (np.array(True, dtype=np.bool_),
-                           np.array(True, dtype=np.bool_))
-        self._test(*self.create_model_bool_input(),
-                   ie_device, precision, ir_version)
+        self.input_data = (np.array(True, dtype=np.bool_), np.array(True, dtype=np.bool_))
+        self._test(*self.create_model_bool_input(), ie_device, precision, ir_version)
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_and_int(self, ie_device, precision, ir_version):
+        self.input_data = (np.array(3, dtype=np.int32), np.array(4, dtype=np.int32))
+        self._test(*self.create_model_int_input(), ie_device, precision, ir_version)
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_and_tensor(self, ie_device, precision, ir_version):
+        self.input_data = (np.array([3, 5, 8], dtype=np.int32), np.array([7, 11, 2], dtype=np.int32))
+        self._test(
+            *self.create_model_tensor_input(), ie_device, precision, ir_version, freeze_model=False, trace_model=True
+        )

tests/layer_tests/pytorch_tests/test_bitwise_not.py

@@ -1,29 +0,0 @@
-# Copyright (C) 2018-2023 Intel Corporation
-# SPDX-License-Identifier: Apache-2.0
-
-import pytest
-
-from pytorch_layer_test_class import PytorchLayerTest
-
-
-class TestBitwiseNot(PytorchLayerTest):
-    def _prepare_input(self):
-        import numpy as np
-        return ((np.random.randn(1, 5) > 0).astype(bool),)
-
-    def create_model(self):
-        import torch
-
-        class aten_bitwise_not(torch.nn.Module):
-
-            def forward(self, x):
-                return torch.bitwise_not(x)
-
-        ref_net = None
-
-        return aten_bitwise_not(), ref_net, "aten::bitwise_not"
-
-    @pytest.mark.nightly
-    @pytest.mark.precommit
-    def test_bitwise_not(self, ie_device, precision, ir_version):
-        self._test(*self.create_model(), ie_device, precision, ir_version)

tests/layer_tests/pytorch_tests/test_bitwise_ops.py

@@ -0,0 +1,132 @@
+# Copyright (C) 2018-2023 Intel Corporation
+# SPDX-License-Identifier: Apache-2.0
+
+import numpy as np
+import pytest
+import torch
+from pytorch_layer_test_class import PytorchLayerTest
+
+
+class TestBitwiseOp(PytorchLayerTest):
+    def _prepare_input(self, out, unary, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape):
+        choices = np.array([0, 1, 255, 7])
+        x = np.random.choice(choices, lhs_shape).astype(lhs_dtype)
+        if unary:
+            return (x,) if not out else (x, np.zeros_like(x).astype(lhs_dtype))
+        y = np.random.choice(choices, rhs_shape).astype(rhs_dtype)
+        if not out:
+            return x, y
+        return x, y, np.zeros_like(x).astype(lhs_dtype) + np.zeros_like(y).astype(rhs_dtype)
+
+    def create_model(self, op_name, out):
+        ops = {
+            "and": torch.bitwise_and,
+            "or": torch.bitwise_or,
+            "xor": torch.bitwise_xor,
+            "not": torch.bitwise_not,
+        }
+        op = ops[op_name]
+
+        class aten_bitwise(torch.nn.Module):
+            def __init__(self, op, out) -> None:
+                super().__init__()
+                self.op = op
+                if op == torch.bitwise_not:
+                    self.forward = self.forward_not
+                if out:
+                    self.forward = self.forward_out if not op == torch.bitwise_not else self.forward_not_out
+
+            def forward(self, tensor_a, tensor_b):
+                return self.op(tensor_a, tensor_b)
+
+            def forward_out(self, tensor_a, tensor_b, out):
+                return self.op(tensor_a, tensor_b, out=out), out
+
+            def forward_not(self, tensor_a):
+                return self.op(tensor_a)
+
+            def forward_not_out(self, tensor_a, out):
+                return self.op(tensor_a, out=out), out
+
+        ref_net = None
+
+        return aten_bitwise(op, out), ref_net, f"aten::bitwise_{op_name}"
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("op_type", ["and", "or", "not", "xor"])
+    @pytest.mark.parametrize("lhs_dtype", ["bool", "int32", "uint8", "int64"])
+    @pytest.mark.parametrize("rhs_dtype", ["bool", "int32", "uint8", "int64"])
+    @pytest.mark.parametrize(
+        ("lhs_shape", "rhs_shape"),
+        [
+            ([2, 3], [2, 3]),
+            ([2, 3], []),
+            ([], [2, 3]),
+        ],
+    )
+    @pytest.mark.parametrize("out", [False, True])
+    def test_bitwise_mixed_dtypes(
+        self, op_type, out, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape, ie_device, precision, ir_version
+    ):
+        self._test(
+            *self.create_model(op_type, out),
+            ie_device,
+            precision,
+            ir_version,
+            kwargs_to_prepare_input={
+                "out": out,
+                "unary": op_type == "not",
+                "lhs_dtype": lhs_dtype,
+                "rhs_dtype": rhs_dtype,
+                "lhs_shape": lhs_shape,
+                "rhs_shape": rhs_shape,
+            },
+            freeze_model=False,
+            trace_model=True,
+        )
+
+
+class TestBitwiseOperators(PytorchLayerTest):
+    def _prepare_input(self, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape):
+        choices = np.array([0, 1, 255, 7])
+        x = np.random.choice(choices, lhs_shape).astype(lhs_dtype)
+        y = np.random.choice(choices, rhs_shape).astype(rhs_dtype)
+        return x, y
+
+    def create_model(self):
+        class aten_bitwise(torch.nn.Module):
+            def forward(self, lhs, rhs):
+                return lhs & rhs, ~lhs, lhs | rhs, lhs ^ rhs
+
+        ref_net = None
+
+        return aten_bitwise(), ref_net, ("aten::__and__", "aten::bitwise_not", "aten::__or__", "aten::__xor__")
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    @pytest.mark.parametrize("lhs_dtype", ["bool", "int32"])
+    @pytest.mark.parametrize("rhs_dtype", ["bool", "int32"])
+    @pytest.mark.parametrize(
+        ("lhs_shape", "rhs_shape"),
+        [
+            ([2, 3], [2, 3]),
+            ([2, 3], []),
+            ([], [2, 3]),
+        ],
+    )
+    def test_bitwise_operators(self, lhs_dtype, rhs_dtype, lhs_shape, rhs_shape, ie_device, precision, ir_version):
+        self._test(
+            *self.create_model(),
+            ie_device,
+            precision,
+            ir_version,
+            kwargs_to_prepare_input={
+                "lhs_dtype": lhs_dtype,
+                "rhs_dtype": rhs_dtype,
+                "lhs_shape": lhs_shape,
+                "rhs_shape": rhs_shape,
+            },
+            trace_model=True,
+            freeze_model=False,
+        )

tests/layer_tests/pytorch_tests/test_or.py

@@ -1,29 +1,78 @@
 # Copyright (C) 2018-2023 Intel Corporation
 # SPDX-License-Identifier: Apache-2.0
 
+import numpy as np
 import pytest
+import torch
+
 from pytorch_layer_test_class import PytorchLayerTest
 
 
-class TestLog(PytorchLayerTest):
+class TestOr(PytorchLayerTest):
     def _prepare_input(self):
-        import numpy as np
-        return (np.random.randint(0, 255, (20, 30, 40, 50)),)
+        return self.input_data
 
-    def create_model(self):
-        import torch
+    def create_model_tensor_input(self):
+        class aten_or_tensor(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
 
-        class aten_or(torch.nn.Module):
-            def forward(self, x):
-                res = torch.ByteTensor(x.size()).zero_()
-                res[:, :, :, 1:] = res[:, :, :, 1:] | (x[:, :, :, 1:] != x[:, :, :, :-1])
-                res[:, :, :, :-1] = res[:, :, :, :-1] | (x[:, :, :, 1:] != x[:, :, :, :-1])
-                return res.float()
+            def forward(self, tensor_a, tensor_b):
+                return tensor_a | tensor_b
 
-        return aten_or(), None, "aten::__or__"
+        ref_net = None
+
+        return aten_or_tensor(), ref_net, "aten::__or__"
+
+    def create_model_bool_input(self):
+        class aten_or_bool(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def forward(self, bool_a: bool, bool_b: bool):
+                return bool_a | bool_b
+
+        ref_net = None
+
+        return aten_or_bool(), ref_net, "aten::__or__"
+
+    def create_model_int_input(self):
+        class aten_or_int(torch.nn.Module):
+            def __init__(self) -> None:
+                super().__init__()
+
+            def forward(self, int_a: int, int_b: int):
+                return int_a | int_b
+
+        ref_net = None
+
+        return aten_or_int(), ref_net, "aten::__or__"
 
     @pytest.mark.nightly
     @pytest.mark.precommit
-    def test_or(self, ie_device, precision, ir_version):
-        self._test(*self.create_model(), ie_device, precision, ir_version,
-                   dynamic_shapes=False, trace_model=True, use_convert_model=True)
+    def test_or_tensor(self, ie_device, precision, ir_version):
+        self.input_data = (
+            np.array([True, False, False], dtype=np.bool_),
+            np.array([True, True, False], dtype=np.bool_),
+        )
+        self._test(*self.create_model_tensor_input(), ie_device, precision, ir_version)
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_or_bool(self, ie_device, precision, ir_version):
+        self.input_data = (np.array(True, dtype=np.bool_), np.array(True, dtype=np.bool_))
+        self._test(*self.create_model_bool_input(), ie_device, precision, ir_version)
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_or_int(self, ie_device, precision, ir_version):
+        self.input_data = (np.array(3, dtype=np.int32), np.array(4, dtype=np.int32))
+        self._test(*self.create_model_int_input(), ie_device, precision, ir_version)
+
+    @pytest.mark.nightly
+    @pytest.mark.precommit
+    def test_or_tensor(self, ie_device, precision, ir_version):
+        self.input_data = (np.array([3, 5, 8], dtype=np.int32), np.array([7, 11, 2], dtype=np.int32))
+        self._test(
+            *self.create_model_tensor_input(), ie_device, precision, ir_version, freeze_model=False, trace_model=True
+        )

tests/layer_tests/pytorch_tests/test_xor.py

@@ -9,13 +9,11 @@ from pytorch_layer_test_class import PytorchLayerTest
 
 
 class TestXor(PytorchLayerTest):
-
     def _prepare_input(self):
         return self.input_data
 
     def create_model_tensor_input(self):
         class aten_xor_tensor(torch.nn.Module):
-
             def __init__(self) -> None:
                 super().__init__()
 
@@ -28,7 +26,6 @@ class TestXor(PytorchLayerTest):
 
     def create_model_bool_input(self):
         class aten_xor_bool(torch.nn.Module):
-
             def __init__(self) -> None:
                 super().__init__()
 
@@ -41,7 +38,6 @@ class TestXor(PytorchLayerTest):
 
     def create_model_int_input(self):
         class aten_xor_int(torch.nn.Module):
-
             def __init__(self) -> None:
                 super().__init__()
 
@@ -55,33 +51,28 @@ class TestXor(PytorchLayerTest):
     @pytest.mark.nightly
     @pytest.mark.precommit
     def test_xor_tensor(self, ie_device, precision, ir_version):
-        self.input_data = (np.array([True, False, False], dtype=np.bool_), np.array(
-            [True, True, False], dtype=np.bool_))
-        self._test(*self.create_model_tensor_input(),
-                   ie_device, precision, ir_version)
+        self.input_data = (
+            np.array([True, False, False], dtype=np.bool_),
+            np.array([True, True, False], dtype=np.bool_),
+        )
+        self._test(*self.create_model_tensor_input(), ie_device, precision, ir_version)
 
     @pytest.mark.nightly
     @pytest.mark.precommit
     def test_xor_bool(self, ie_device, precision, ir_version):
-        self.input_data = (np.array(True, dtype=np.bool_),
-                           np.array(True, dtype=np.bool_))
-        self._test(*self.create_model_bool_input(),
-                   ie_device, precision, ir_version)
+        self.input_data = (np.array(True, dtype=np.bool_), np.array(True, dtype=np.bool_))
+        self._test(*self.create_model_bool_input(), ie_device, precision, ir_version)
 
-    @pytest.mark.xfail(reason="bitwise_xor is not implemented")
     @pytest.mark.nightly
     @pytest.mark.precommit
     def test_xor_int(self, ie_device, precision, ir_version):
-        self.input_data = (np.array(3, dtype=np.int),
-                           np.array(4, dtype=np.int))
-        self._test(*self.create_model_int_input(),
-                   ie_device, precision, ir_version)
+        self.input_data = (np.array(3, dtype=np.int32), np.array(4, dtype=np.int32))
+        self._test(*self.create_model_int_input(), ie_device, precision, ir_version)
 
-    @pytest.mark.xfail(reason="bitwise_xor is not implemented")
     @pytest.mark.nightly
     @pytest.mark.precommit
     def test_xor_tensor(self, ie_device, precision, ir_version):
-        self.input_data = (np.array([3, 5, 8], dtype=np.int), np.array(
-            [7, 11, 2], dtype=np.int))
-        self._test(*self.create_model_tensor_input(),
-                   ie_device, precision, ir_version)
+        self.input_data = (np.array([3, 5, 8], dtype=np.int32), np.array([7, 11, 2], dtype=np.int32))
+        self._test(
+            *self.create_model_tensor_input(), ie_device, precision, ir_version, freeze_model=False, trace_model=True
+        )