reciprocal lowering pass

Author: apbose

py/torch_tensorrt/dynamo/conversion/aten_ops_converters.py

@@ -8,6 +8,8 @@
 
 from torch_tensorrt.fx.types import TRTNetwork, TRTTensor
 from torch_tensorrt.dynamo.conversion import SourceIR, impl
+from torch_tensorrt.dynamo.conversion.converter_utils import cast_trt_tensor
+from torch_tensorrt.dynamo.conversion.converter_utils import cast_int_int_div_trt_tensor
 
 _LOGGER: logging.Logger = logging.getLogger(__name__)
 
@@ -54,6 +56,20 @@ def aten_ops_div(
         "input": args[0],
         "other": args[1],
     }
+    # If both are TRTTensor, both are cast to float32
+    if isinstance(args[0], TRTTensor) and isinstance(args[1], TRTTensor):
+        kwargs_new["input"], kwargs_new["other"] = cast_int_int_div_trt_tensor(
+            network, kwargs_new["input"], kwargs_new["other"]
+        )
+    # If one is TRTTensor, it is cast to float32
+    elif isinstance(args[0], TRTTensor) and (
+        kwargs_new["input"].dtype == trt.int8 or kwargs_new["input"].dtype == trt.int32
+    ):
+        kwargs_new["input"] = cast_trt_tensor(network, kwargs_new["input"], trt.float32)
+    elif isinstance(args[1], TRTTensor) and (
+        kwargs_new["other"].dtype == trt.int8 or kwargs_new["other"].dtype == trt.int32
+    ):
+        kwargs_new["other"] = cast_trt_tensor(network, kwargs_new["other"], trt.float32)
     rounding_mode = kwargs.get("rounding_mode")
     if rounding_mode is None:
         return acc_ops_converters.acc_ops_div(network, target, None, kwargs_new, name)

py/torch_tensorrt/dynamo/conversion/converter_utils.py

@@ -5,6 +5,8 @@
     TRTNetwork,
     TRTTensor,
 )
+import torch_tensorrt as trt
+from typing import List
 
 
 def dynamic_unsupported(node: torch.fx.Node) -> bool:
@@ -65,6 +67,28 @@ def cast_trt_tensor(
         return input_val
 
 
+def cast_int_int_div_trt_tensor(
+    network: TRTNetwork,
+    lhs_val: TRTTensor,
+    rhs_val: TRTTensor,
+) -> List[TRTTensor]:
+    """
+    Given two `int` data type TRT Tensor to div operation, cast the TRT Tensor to float type
+    Args:
+        network (TRTNetwork): A TensorRT network
+        lhs_val (TRTTensor): A TRT Tensor numerator
+        rhs_val (TRTTensor): A TRT Tensor numerator
+    Returns:
+        A list of lhs_val and rhs_val casted to the approriate datatype
+    """
+    if (lhs_val.dtype == trt.int8 or lhs_val.dtype == trt.int32) and (
+        rhs_val.dtype == trt.int8 or rhs_val.dtype == trt.int32
+    ):
+        lhs_val = cast_trt_tensor(network, lhs_val, trt.float32)
+        rhs_val = cast_trt_tensor(network, rhs_val, trt.float32)
+    return list((lhs_val, rhs_val))
+
+
 def broadcastable(
     a: TRTTensor,
     b: TRTTensor,

py/torch_tensorrt/dynamo/lowering/_decompositions.py

@@ -70,5 +70,12 @@ def addmm_replacement(
     )
 
 
+@register_decomposition(torch.ops.aten.reciprocal.default, registry=DECOMPOSITIONS)
+def reciprocal_replacement(
+    input_: torch.Tensor,
+) -> torch.Tensor:
+    return torch.div(1, input_)
+
+
 def get_decompositions():
     return DECOMPOSITIONS

tests/py/dynamo/backend/test_decompositions.py

@@ -78,15 +78,14 @@ def forward(self, x):
                 return y
 
         # Operations expected to be removed in the traced graph after decompositions
-        expected_ops = {torch.ops.aten.sqrt.default, torch.ops.aten.reciprocal.default}
-        unexpected_ops = {torch.ops.aten.rsqrt.default}
+        expected_ops = {torch.ops.aten.sqrt.default, torch.ops.aten.div.Tensor}
+        unexpected_ops = {
+            torch.ops.aten.rsqrt.default,
+            torch.ops.aten.reciprocal.default,
+        }
 
         inputs = [
-            torch.randint(
-                1,
-                10,
-                (5,),
-            ),
+            torch.randint(1, 10, (5,), dtype=torch.int32),
         ]
 
         fx_graph = torch.fx.symbolic_trace(Rsqrt())
@@ -182,6 +181,69 @@ def forward(self, x, y, z):
             f"AddMM TRT outputs don't match with the original model.",
         )
 
+    def test_lowering_reciprocal(self):
+        class Reciprocal(torch.nn.Module):
+            def __init__(self, *args, **kwargs) -> None:
+                super().__init__(*args, **kwargs)
+
+            def forward(self, x):
+                y = torch.ops.aten.reciprocal.default(x)
+                return y
+
+        # Operations expected to be removed in the traced graph after decompositions
+        expected_ops = {torch.ops.aten.div.Tensor}
+        unexpected_ops = {torch.ops.aten.reciprocal.default}
+
+        inputs = [
+            torch.randn(
+                5,
+            ).cuda()
+        ]
+
+        fx_graph = torch.fx.symbolic_trace(Reciprocal())
+        unexpected_ops_seen, expected_ops_unseen = lower_graph_testing(
+            fx_graph,
+            inputs,
+            expected_ops=expected_ops,
+            unexpected_ops=unexpected_ops,
+            min_block_size=1,
+        )
+
+        self.assertEquals(
+            len(unexpected_ops_seen),
+            0,
+            f"The following unexpected ops were encountered: {unexpected_ops_seen}",
+        )
+
+        self.assertEquals(
+            len(expected_ops_unseen),
+            0,
+            f"The following expected ops were not encountered: {expected_ops_unseen}",
+        )
+
+        torch._dynamo.reset()
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_tensorrt",
+            inputs,
+            min_block_size=1,
+            pass_through_build_failures=True,
+        )
+        optimized_model_results = optimized_model(*inputs).detach().cpu()
+        torch_model_results = fx_graph(*inputs).detach().cpu()
+
+        max_diff = float(
+            torch.max(torch.abs(optimized_model_results - torch_model_results))
+        )
+        self.assertAlmostEqual(
+            max_diff,
+            0,
+            DECIMALS_OF_AGREEMENT,
+            f"Reciprocal TRT outputs don't match with the original model.",
+        )
+
 
 if __name__ == "__main__":
     run_tests()

tests/py/dynamo/converters/test_binary_ops_aten.py

@@ -75,6 +75,23 @@ def forward(self, x):
         inputs = [torch.rand(1, 1) + 1]
         self.run_test(m, inputs, expected_ops={expected_op})
 
+    @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in elementwise_ops])
+    def test_elementwise_ops_mismatched_dtypes(
+        self, name, orig_op: Callable, expected_op
+    ):
+        class TestModule(nn.Module):
+            def __init__(self, orig_op):
+                super().__init__()
+                self.orig_op = orig_op
+
+            def forward(self, x):
+                return self.orig_op(x.int(), x)
+
+        m = TestModule(orig_op)
+        # Avoid dividing by 0.
+        inputs = [2 * torch.rand(1, 1, dtype=torch.float) + 1]
+        self.run_test(m, inputs, expected_ops={expected_op})
+
     @parameterized.expand([(op[1].__name__, op[0], op[1]) for op in elementwise_ops])
     def test_elementwise_ops_with_one_constant(
         self, name, orig_op: Callable, expected_op
@@ -114,6 +131,41 @@ def forward(self, x):
         inputs = [torch.randn(2, 2)]
         self.run_test(m, inputs, expected_ops={expected_op})
 
+    @parameterized.expand([((lambda x, y: x / y), torch.ops.aten.div.Tensor)])
+    def test_elementwise_op_div_with_two_ints(self, orig_op: Callable, expected_op):
+        class TestModule(nn.Module):
+            def __init__(self, orig_op):
+                super().__init__()
+                self.orig_op = orig_op
+
+            def forward(self, x):
+                return self.orig_op(x, x + 1)
+
+        m = TestModule(orig_op)
+        inputs = [torch.randint(1, 10, (5,), dtype=torch.int32)]
+        self.run_test(m, inputs, expected_ops={expected_op})
+
+    @parameterized.expand([((lambda x, y: x / y), torch.ops.aten.div.Tensor)])
+    def test_elementwise_op_div_with_one_int_one_constant(
+        self, orig_op: Callable, expected_op
+    ):
+        class TestModule(nn.Module):
+            def __init__(self, orig_op):
+                super().__init__()
+                self.constant1 = torch.nn.Parameter(
+                    torch.randn(
+                        5,
+                    )
+                )
+                self.orig_op = orig_op
+
+            def forward(self, x):
+                return self.orig_op(x, self.constant1)
+
+        m = TestModule(orig_op)
+        inputs = [torch.randint(1, 10, (5,), dtype=torch.int32)]
+        self.run_test(m, inputs, expected_ops={expected_op})
+
     # Dynamic shape test
     @parameterized.expand(
         [