feat: support more elementwise and unary converters

Author: zewenli98

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -1701,6 +1701,120 @@ def aten_ops_logical_xor(
     )
 
 
+def bitwise_type_validator(node: Node) -> bool:
+    targets = [
+        torch.ops.aten.bitwise_and.Tensor,
+        torch.ops.aten.bitwise_or.Tensor,
+        torch.ops.aten.bitwise_xor.Tensor,
+    ]
+    if node.target not in targets:
+        return False
+
+    lhs_val = node.args[0]
+    rhs_val = node.args[1]
+    lhs_meta = lhs_val.meta.get("tensor_meta")
+    rhs_meta = rhs_val.meta.get("tensor_meta")
+
+    if lhs_meta is None or rhs_meta is None:
+        return False
+
+    supported_type = [torch.bool, bool]
+    return lhs_meta.dtype in supported_type and rhs_meta.dtype in supported_type
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_and.Tensor, capability_validator=bitwise_type_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_and.Scalar)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_and.Scalar_Tensor)  # type: ignore[misc]
+def aten_ops_bitwise_and(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.elementwise.bitwise_and(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_or.Tensor, capability_validator=bitwise_type_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_or.Scalar)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_or.Scalar_Tensor)  # type: ignore[misc]
+def aten_ops_bitwise_or(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.elementwise.bitwise_or(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_xor.Tensor, capability_validator=bitwise_type_validator)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_xor.Scalar)  # type: ignore[misc]
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_xor.Scalar_Tensor)  # type: ignore[misc]
+def aten_ops_bitwise_xor(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.elementwise.bitwise_xor(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+        args[1],
+    )
+
+
+def bitwise_not_type_validator(node: Node) -> bool:
+    val = node.args[0]
+    val_meta = val.meta.get("tensor_meta")
+
+    if val_meta is None:
+        return False
+
+    supported_type = [torch.bool, bool]
+    return val_meta.dtype in supported_type
+
+
+@dynamo_tensorrt_converter(torch.ops.aten.bitwise_not.default, capability_validator=bitwise_not_type_validator)  # type: ignore[misc]
+@enforce_tensor_types(
+    {
+        0: (TRTTensor,),
+    }
+)  # type: ignore[misc]
+def aten_ops_bitwise_not(
+    ctx: ConversionContext,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    return impl.unary.bitwise_not(
+        ctx,
+        target,
+        SourceIR.ATEN,
+        name,
+        args[0],
+    )
+
+
 @dynamo_tensorrt_converter(torch.ops.aten.eq.Tensor)  # type: ignore[misc]
 @dynamo_tensorrt_converter(torch.ops.aten.eq.Scalar)  # type: ignore[misc]
 @enforce_tensor_types(

py/torch_tensorrt/dynamo/conversion/impl/elementwise/ops.py

@@ -423,6 +423,39 @@ def logical_xor(
     )
 
 
+def bitwise_and(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    lhs_val: Union[TRTTensor, int, bool, Sequence[Union[int, bool]]],
+    rhs_val: Union[TRTTensor, int, bool, Sequence[Union[int, bool]]],
+) -> TRTTensor:
+    return logical_and(ctx, target, source_ir, f"{name}_logical_and", lhs_val, rhs_val)
+
+
+def bitwise_or(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    lhs_val: Union[TRTTensor, int, bool, Sequence[Union[int, bool]]],
+    rhs_val: Union[TRTTensor, int, bool, Sequence[Union[int, bool]]],
+) -> TRTTensor:
+    return logical_or(ctx, target, source_ir, f"{name}_logical_or", lhs_val, rhs_val)
+
+
+def bitwise_xor(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    lhs_val: Union[TRTTensor, int, bool, Sequence[Union[int, bool]]],
+    rhs_val: Union[TRTTensor, int, bool, Sequence[Union[int, bool]]],
+) -> TRTTensor:
+    return logical_xor(ctx, target, source_ir, f"{name}_logical_xor", lhs_val, rhs_val)
+
+
 def eq(
     ctx: ConversionContext,
     target: Target,

py/torch_tensorrt/dynamo/conversion/impl/unary/ops.py

@@ -1,6 +1,7 @@
 from typing import Optional
 
 import tensorrt as trt
+import torch_tensorrt.dynamo.conversion.impl as impl
 from torch.fx.node import Target
 from torch_tensorrt.dynamo._SourceIR import SourceIR
 from torch_tensorrt.dynamo.conversion._ConversionContext import ConversionContext
@@ -336,6 +337,18 @@ def logical_not(
     )
 
 
+def bitwise_not(
+    ctx: ConversionContext,
+    target: Target,
+    source_ir: Optional[SourceIR],
+    name: str,
+    input_val: TRTTensor,
+) -> TRTTensor:
+    return impl.unary.logical_not(
+        ctx, target, source_ir, f"{name}_logical_not", input_val
+    )
+
+
 def sign(
     ctx: ConversionContext,
     target: Target,

tests/py/dynamo/conversion/test_bitwise_and_aten.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestBitwiseAndConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d", (5, 3)),
+            ("3d", (5, 3, 2)),
+        ]
+    )
+    def test_bitwise_and_tensor(self, _, shape):
+        class bitwise_and(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.bitwise_and.Tensor(lhs_val, rhs_val)
+
+        inputs = [
+            torch.randint(0, 2, shape, dtype=bool),
+            torch.randint(0, 2, shape, dtype=bool),
+        ]
+        self.run_test(
+            bitwise_and(),
+            inputs,
+            enable_passes=True,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

tests/py/dynamo/conversion/test_bitwise_not_aten.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestBitwiseNotConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d", (5, 3)),
+            ("3d", (5, 3, 2)),
+        ]
+    )
+    def test_bitwise_not_tensor(self, _, shape):
+        class bitwise_not(nn.Module):
+            def forward(self, val):
+                return torch.ops.aten.bitwise_not.default(val)
+
+        inputs = [
+            torch.randint(0, 2, shape, dtype=torch.bool),
+        ]
+        self.run_test(
+            bitwise_not(),
+            inputs,
+            enable_passes=True,
+            output_dtypes=[torch.bool],
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

tests/py/dynamo/conversion/test_bitwise_or_aten.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestBitwiseOrConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d", (5, 3)),
+            ("3d", (5, 3, 2)),
+        ]
+    )
+    def test_bitwise_or_tensor(self, _, shape):
+        class bitwise_or(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.bitwise_or.Tensor(lhs_val, rhs_val)
+
+        inputs = [
+            torch.randint(0, 2, shape, dtype=bool),
+            torch.randint(0, 2, shape, dtype=bool),
+        ]
+        self.run_test(
+            bitwise_or(),
+            inputs,
+            enable_passes=True,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()

tests/py/dynamo/conversion/test_bitwise_xor_aten.py

@@ -0,0 +1,33 @@
+import torch
+import torch.nn as nn
+from parameterized import parameterized
+from torch.testing._internal.common_utils import run_tests
+
+from .harness import DispatchTestCase
+
+
+class TestBitwiseXorConverter(DispatchTestCase):
+    @parameterized.expand(
+        [
+            ("2d", (5, 3)),
+            ("3d", (5, 3, 2)),
+        ]
+    )
+    def test_bitwise_xor_tensor(self, _, shape):
+        class bitwise_xor(nn.Module):
+            def forward(self, lhs_val, rhs_val):
+                return torch.ops.aten.bitwise_xor.Tensor(lhs_val, rhs_val)
+
+        inputs = [
+            torch.randint(0, 2, shape, dtype=bool),
+            torch.randint(0, 2, shape, dtype=bool),
+        ]
+        self.run_test(
+            bitwise_xor(),
+            inputs,
+            enable_passes=True,
+        )
+
+
+if __name__ == "__main__":
+    run_tests()