[fx2trt] support for ne, logical_not, logical_and (#75444)

Summary:
X-link: pytorch/pytorch#75444

as titled
1. support logical_and, logical_not
2. replace eq,gt,lt with python operator in acc_ops due to the fact that torch op needs input to be torch.Tensor but python op does not
3. add more test cases
4. add individual ne op without using combination of existing ops since there are limitations. For ex, in lowering it to equal+logical_not. It will fail the last test case in test_ne.py. The failure reason is that logical_not needs the input to be a tensor.
Also we can not use equal+operator.not since not is not tracable in FX with the error "symbolically traced variables cannot be used as inputs to control flow"
We also can not use equal+operator.invert since operator.invert(True)=-2

(Note: this ignores all push blocking failures!)

Reviewed By: 842974287

Differential Revision: D35232917

fbshipit-source-id: d4601a6883c977caa263f67b9db86cbc862d4780

Author: Wei Wei

fx/converters/acc_ops_converters.py

@@ -1237,6 +1237,87 @@ def acc_ops_minimum(
     )
 
 
+
+@tensorrt_converter(acc_ops.logical_not)
+def acc_ops_logical_not(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    input_val = kwargs["input"]
+    operation_type = trt.UnaryOperation.NOT
+    # cast to bool type
+    if input_val.dtype in (trt.float32, trt.float16, trt.int32):
+        input_val = type_cast(network, target, f"{name}_input", input_val, trt.bool)
+    return add_unary_layer(network, input_val, operation_type, target, name)
+
+
+@tensorrt_converter(acc_ops.logical_and, no_implicit_batch_dim=True)
+@tensorrt_converter(acc_ops.bitwise_and, no_implicit_batch_dim=True)
+def acc_ops_logical_and(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    if network.has_implicit_batch_dimension:
+        raise RuntimeError("The `logical_and` function should be called with explicit batch dimension.")
+
+    input_t = kwargs["input"]
+    other_t =  kwargs["other"]
+    # we only support both inputs are bool type
+    if target == acc_ops.bitwise_and:
+        def check_is_bool(input_t):
+            if isinstance(input_t, TRTTensor):
+                assert input_t.dtype == trt.bool, "We currently do not support input is non-bool"
+            elif isinstance(input_t, torch.Tensor):
+                assert input_t.dtype == torch.bool, "We currently do not support input is non-bool"
+            else:
+                assert isinstance(input_t. bool), "We currently do not support input is non-bool"
+
+        check_is_bool(input_t)
+        check_is_bool(other_t)
+
+    input_t = get_trt_tensor(network, input_t, f"{name}_input_t")
+    other_t = get_trt_tensor(network, other_t, f"{name}_other_t")
+
+    if input_t.dtype != trt.bool:
+        input_t = type_cast(network, target, f"{name}_input", input_t, trt.bool)
+    if other_t.dtype != trt.bool:
+        other_t = type_cast(network, target, f"{name}_other", other_t, trt.bool)
+    return add_binary_elementwise_layer(
+        network, input_t, other_t, trt.ElementWiseOperation.AND, target, name
+    )
+
+
+@tensorrt_converter(acc_ops.ne, no_implicit_batch_dim=True)
+def acc_ops_ne(
+    network: TRTNetwork,
+    target: Target,
+    args: Tuple[Argument, ...],
+    kwargs: Dict[str, Argument],
+    name: str,
+) -> Union[TRTTensor, Sequence[TRTTensor]]:
+    if network.has_implicit_batch_dimension:
+        raise RuntimeError("The `ne` function should be called with explicit batch dimension.")
+
+    input_t = kwargs["input"]
+    other_t =  kwargs["other"]
+
+    input_t = get_trt_tensor(network, input_t, f"{name}_input_t")
+    other_t = get_trt_tensor(network, other_t, f"{name}_other_t")
+
+    input_t, other_t = dtype_uniform(network, target, name, input_t, other_t)
+    eq_t = add_binary_elementwise_layer(
+        network, input_t, other_t, trt.ElementWiseOperation.EQUAL, target, name
+    )
+
+    return add_unary_layer(network, eq_t, trt.UnaryOperation.NOT, target, name)
+
+
 @tensorrt_converter(acc_ops.eq, no_implicit_batch_dim=True)
 def acc_ops_eq(
     network: TRTNetwork,
@@ -1250,14 +1331,11 @@ def acc_ops_eq(
 
     input_t = kwargs["input"]
     other_t =  kwargs["other"]
-    if isinstance(other_t, (torch.Tensor, bool)):
-        if isinstance(other_t, bool):
-            other_t = int(other_t)
-        elif other_t.dtype == torch.bool:
-            other_t = other_t.to(torch.int32)
+
+    input_t = get_trt_tensor(network, input_t, f"{name}_input_t")
     other_t = get_trt_tensor(network, other_t, f"{name}_other_t")
-    input_t, other_t = dtype_uniform(network, target, name, input_t, other_t)
 
+    input_t, other_t = dtype_uniform(network, target, name, input_t, other_t)
     return add_binary_elementwise_layer(
         network, input_t, other_t, trt.ElementWiseOperation.EQUAL, target, name
     )
@@ -1276,12 +1354,10 @@ def acc_ops_gt(
 
     input_t = kwargs["input"]
     other_t =  kwargs["other"]
-    if isinstance(other_t, (torch.Tensor, bool)):
-        if isinstance(other_t, bool):
-            other_t = int(other_t)
-        elif other_t.dtype == torch.bool:
-            other_t = other_t.to(torch.int32)
+
+    input_t = get_trt_tensor(network, input_t, f"{name}_input_t")
     other_t = get_trt_tensor(network, other_t, f"{name}_other_t")
+
     input_t, other_t = dtype_uniform(network, target, name, input_t, other_t)
     return add_binary_elementwise_layer(
         network, input_t, other_t, trt.ElementWiseOperation.GREATER, target, name
@@ -1301,14 +1377,11 @@ def acc_ops_lt(
 
     input_t = kwargs["input"]
     other_t =  kwargs["other"]
-    if isinstance(other_t, (torch.Tensor, bool)):
-        if isinstance(other_t, bool):
-            other_t = int(other_t)
-        elif other_t.dtype == torch.bool:
-            other_t = other_t.to(torch.int32)
+
+    input_t = get_trt_tensor(network, input_t, f"{name}_input_t")
     other_t = get_trt_tensor(network, other_t, f"{name}_other_t")
-    input_t, other_t = dtype_uniform(network, target, name, input_t, other_t)
 
+    input_t, other_t = dtype_uniform(network, target, name, input_t, other_t)
     return add_binary_elementwise_layer(
         network, input_t, other_t, trt.ElementWiseOperation.LESS, target, name
     )

fx/converters/converter_utils.py

@@ -8,7 +8,16 @@
 import tensorrt as trt
 import torch
 from torch.fx.node import Target, Argument
-from fx2trt_oss.fx.types import *  # noqa: F403
+from fx2trt_oss.fx.types import (
+    TRTNetwork,
+    TRTTensor,
+    TRTLayer,
+    TRTPluginFieldCollection,
+    TRTPlugin,
+    TRTDataType,
+    TRTElementWiseOp,
+    Shape
+)
 from fx2trt_oss.fx.utils import torch_dtype_from_trt
 
 
@@ -223,6 +232,15 @@ def get_trt_tensor(
     Returns:
         A TensorRT ITensor that represents the given value.
     """
+    # TRT can not add constant for bool type. We do a work around to 1) cast it to int and 2)cast to bool later
+    # This is useful for logical operations which require input to be bool type
+    if isinstance(input_val, bool):
+        input_val = int(input_val)
+    if isinstance(input_val, torch.Tensor) and input_val.dtype == torch.bool:
+        input_val = input_val.to(torch.int32)
+    if isinstance(input_val, torch.Tensor) and input_val.dtype == torch.int64:
+        input_val = input_val.to(torch.int32)
+
     if isinstance(input_val, (torch.Tensor, int, float)):
         return create_constant(network, input_val, name, dtype)
     elif not isinstance(input_val, TRTTensor):
@@ -439,6 +457,8 @@ def add_unary_layer(
         )
     layer = network.add_unary(input_val, operation_type)
     set_layer_name(layer, target, name)
+    output = layer.get_output(0)
+    output.name = output.name + "_" + target.__name__
     return layer.get_output(0)
 
 
@@ -672,7 +692,7 @@ def get_python_op_from_trt_elementwise_op(trt_op: TRTElementWiseOp) -> Callable[
     else:
         raise RuntimeError(f"{trt_op} is not supported yet!")
 
-def dtype_uniform(network, target, name, input, other):
+def dtype_uniform(network: TRTNetwork, target: Target, name: str, input: TRTTensor, other: TRTTensor):
     table = {trt.bool:0, trt.int32:1, trt.float16:2, trt.float32:3}
     input_dtype = input.dtype
     other_dtype = other.dtype
@@ -697,3 +717,12 @@ def dtype_uniform(network, target, name, input, other):
         set_layer_name(layer_o, target, f"{name}_other_dtype_change")
         other = layer_o.get_output(0)
     return input, other
+
+def type_cast(network: TRTNetwork, target: Target, name: str, input: TRTTensor, cast_type: TRTDataType):
+    """
+    This function helps to cast the input type to cast_type
+    """
+    layer_i = network.add_identity(input)
+    layer_i.set_output_type(0, cast_type)
+    set_layer_name(layer_i, target, f"{name}_dtype_change")
+    return layer_i.get_output(0)

fx/fx2trt.py

@@ -297,7 +297,7 @@ def output(self, target, args, kwargs):
             raise RuntimeError("TensorRT requires all outputs to be Tensor!")
 
         for i, output in enumerate(outputs):
-            if any(op_name in output.name.split("_") for op_name in ("eq", "gt", "lt", "or", "xor")):
+            if any(op_name in output.name.split("_") for op_name in ("eq", "gt", "lt", "or", "xor", "and", "not", "ne")):
                 output_bool = True
             else:
                 output_bool = False

test/converters/acc_op/test_eq.py

@@ -127,6 +127,21 @@ def forward(self, x):
         ]
         self.run_test(Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim = False)
 
+class TestConstInputConverter(AccTestCase):
+    def test_eq(self):
+        class Eq(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return  x.shape[0] == 4
+
+        input = torch.randn(3,4)
+        inputs = [
+            input,
+        ]
+        self.run_test(Eq(), inputs, expected_ops={acc_ops.eq}, test_implicit_batch_dim = False)
+
 
 if __name__ == '__main__':
     run_tests()

test/converters/acc_op/test_gt.py

@@ -124,5 +124,22 @@ def forward(self, x):
         ]
         self.run_test(Eq(), inputs, expected_ops={acc_ops.gt}, test_implicit_batch_dim = False)
 
+
+class TestConstInputConverter(AccTestCase):
+    def test_gt(self):
+        class Gt(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return  x.shape[0] > 4
+
+        input = torch.randn(3,4)
+        inputs = [
+            input,
+        ]
+        self.run_test(Gt(), inputs, expected_ops={acc_ops.gt}, test_implicit_batch_dim = False)
+
+
 if __name__ == '__main__':
     run_tests()

test/converters/acc_op/test_logical_and.py

@@ -0,0 +1,95 @@
+import torch
+import fx2trt_oss.tracer.acc_tracer.acc_ops as acc_ops
+from torch.testing._internal.common_fx2trt import AccTestCase
+from torch.testing._internal.common_utils import run_tests
+from parameterized import parameterized
+
+class TestAndMethodSimpleConverter(AccTestCase):
+    @parameterized.expand(
+        [
+            ("rand_2d_float_bool", torch.randn(3,4), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_int_bool", torch.randn(3,4).to(torch.int), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_bool_bool", torch.randn(3,4).to(torch.bool), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_float_int", torch.randn(3,4).to(torch.float), torch.randn(3,4).to(torch.int)),
+            ("rand_2d_float_single_bool", torch.randn(3,4), torch.tensor(0).to(torch.bool)),
+            ("rand_2d_int_single_bool", torch.randn(3,4).to(torch.int), torch.tensor(0).to(torch.bool)),
+            ("rand_2d_bool_single_bool", torch.randn(3,4).to(torch.bool), torch.tensor(0).to(torch.bool)),
+        ]
+    )
+    def test_and(self, _, input, other):
+        class And(torch.nn.Module):
+            def forward(self, x, y):
+                return x.logical_and(y)
+
+        inputs = [
+            input,
+            other,
+        ]
+        self.run_test(And(), inputs, expected_ops={acc_ops.logical_and}, test_implicit_batch_dim = False)
+
+class TestAndFunctionSimpleConverter(AccTestCase):
+    @parameterized.expand(
+        [
+            ("rand_2d_float_bool", torch.randn(3,4), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_int_bool", torch.randn(3,4).to(torch.int), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_bool_bool", torch.randn(3,4).to(torch.bool), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_float_int", torch.randn(3,4).to(torch.float), torch.randn(3,4).to(torch.int)),
+            ("rand_2d_float_single_bool", torch.randn(3,4), torch.tensor(0).to(torch.bool)),
+            ("rand_2d_int_single_bool", torch.randn(3,4).to(torch.int), torch.tensor(0).to(torch.bool)),
+            ("rand_2d_bool_single_bool", torch.randn(3,4).to(torch.bool), torch.tensor(0).to(torch.bool)),
+        ]
+    )
+    def test_and(self, _, input, other):
+        class And(torch.nn.Module):
+            def forward(self, x, y):
+                return torch.logical_and(x, y)
+
+        inputs = [
+            input,
+            other,
+        ]
+        self.run_test(And(), inputs, expected_ops={acc_ops.logical_and}, test_implicit_batch_dim = False)
+
+class TestAndOperatorSimpleConverter(AccTestCase):
+    @parameterized.expand(
+        [
+            ("rand_2d_bool_bool", torch.randn(3,4).to(torch.bool), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_bool_single_bool", torch.randn(3,4).to(torch.bool), torch.tensor(0).to(torch.bool)),
+        ]
+    )
+    def test_and(self, _, input, other):
+        class And(torch.nn.Module):
+            def forward(self, x, y):
+                return x & y
+
+        inputs = [
+            input,
+            other,
+        ]
+        self.run_test(And(), inputs, expected_ops={acc_ops.bitwise_and}, test_implicit_batch_dim = False)
+
+
+class TestAndOperatorConstantConverter(AccTestCase):
+    @parameterized.expand(
+        [
+            ("rand_2d_bool_bool", torch.randn(3,4).to(torch.bool), torch.randn(3,4).to(torch.bool)),
+            ("rand_2d_bool_single_bool", torch.randn(3,4).to(torch.bool), torch.tensor(0).to(torch.bool)),
+        ]
+    )
+    def test_and(self, _, input, other):
+        class And(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.other = other
+
+            def forward(self, x):
+                return  x & self.other
+
+        inputs = [
+            input,
+        ]
+        self.run_test(And(), inputs, expected_ops={acc_ops.bitwise_and}, test_implicit_batch_dim = False)
+
+
+if __name__ == '__main__':
+    run_tests()

test/converters/acc_op/test_lt.py

@@ -124,5 +124,22 @@ def forward(self, x):
         ]
         self.run_test(Eq(), inputs, expected_ops={acc_ops.lt}, test_implicit_batch_dim = False)
 
+
+class TestConstInputConverter(AccTestCase):
+    def test_lt(self):
+        class Lt(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, x):
+                return  x.shape[0] < 4
+
+        input = torch.randn(3,4)
+        inputs = [
+            input,
+        ]
+        self.run_test(Lt(), inputs, expected_ops={acc_ops.lt}, test_implicit_batch_dim = False)
+
+
 if __name__ == '__main__':
     run_tests()