fix: Remove input aliasing with builtin ops

- Add replacements for inplace builtin operators with their out-of-place
equivalents
- Add utility to automatically perform replacement prior to AOT tracing
- Add test cases to verify inplace operators are replaced accurately

Author: gs-olive

py/torch_tensorrt/dynamo/backend/backends.py

@@ -13,7 +13,10 @@
 from torch._ops import OpOverload
 from torch_tensorrt.dynamo import CompilationSettings
 from torch_tensorrt.dynamo.compile import compile_module
-from torch_tensorrt.dynamo.lowering._decompositions import get_decompositions
+from torch_tensorrt.dynamo.lowering import (
+    get_decompositions,
+    replace_builtin_inplace_ops,
+)
 from torch_tensorrt.dynamo.lowering._pre_aot_lowering import pre_aot_substitutions
 from torch_tensorrt.dynamo.utils import parse_dynamo_kwargs
 
@@ -74,6 +77,7 @@ def _pretraced_backend(
         with unittest.mock.patch.object(
             fake_mode, "allow_non_fake_inputs", True
         ), fake_mode:
+            replace_builtin_inplace_ops(gm)
             # Invoke AOTAutograd to translate operators to aten
             graph_module = aot_export_for_compile(
                 gm,

py/torch_tensorrt/dynamo/lowering/__init__.py

@@ -2,4 +2,5 @@
 from ._fusers import *  # noqa: F401
 from ._pre_aot_lowering import SUBSTITUTION_REGISTRY  # noqa: F401
 from ._pre_aot_lowering import register_substitution  # noqa: F401
+from ._replace_inplace_ops import replace_builtin_inplace_ops
 from .substitutions import *  # noqa: F401

py/torch_tensorrt/dynamo/lowering/_replace_inplace_ops.py

@@ -0,0 +1,50 @@
+import logging
+import operator
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+BUILTIN_TRANSLATION = {
+    operator.ipow: operator.pow,
+    operator.imul: operator.mul,
+    operator.imatmul: operator.matmul,
+    operator.ifloordiv: operator.floordiv,
+    operator.itruediv: operator.truediv,
+    operator.imod: operator.mod,
+    operator.iadd: operator.add,
+    operator.isub: operator.sub,
+    operator.ilshift: operator.lshift,
+    operator.irshift: operator.rshift,
+    operator.iand: operator.and_,
+    operator.ixor: operator.xor,
+    operator.ior: operator.or_,
+}
+
+
+def replace_builtin_inplace_ops(gm: torch.fx.GraphModule) -> None:
+    """Replaces inplace builtins from Python's operator class
+
+    Replaces inplace builtins with out-of-place equivalent ops
+    """
+    for node in gm.graph.nodes:
+        # If a node uses one of the inplace builtins
+        # Replace it with its out-of-place equivalent
+        if node.target in BUILTIN_TRANSLATION:
+            out_of_place_op = BUILTIN_TRANSLATION[node.target]
+
+            # Replace inplace operator node and delete
+            with gm.graph.inserting_before(node):
+                out_of_place = gm.graph.call_function(
+                    out_of_place_op,
+                    args=node.args,
+                    kwargs=node.kwargs,
+                )
+
+            logger.debug(f"Replacing {node.target} with {out_of_place.target}")
+
+            node.replace_all_uses_with(out_of_place)
+            gm.graph.erase_node(node)
+
+    gm.graph.lint()
+    gm.recompile()

tests/py/dynamo/backend/test_specialized_models.py

@@ -157,5 +157,86 @@ def forward(self, x):
         torch._dynamo.reset()
 
 
+class TestInputModifications(TestCase):
+    def test_input_modifications_add(self):
+        class InplaceAdd(torch.nn.Module):
+            def forward(self, x):
+                x += 3
+                y = x + 1
+                return y
+
+        inputs = [
+            torch.rand(
+                3,
+                5,
+                7,
+            ).cuda(),
+        ]
+
+        fx_graph = torch.fx.symbolic_trace(InplaceAdd())
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            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,
+            msg=f"InplaceAdd TRT outputs don't match with the original model.",
+        )
+        torch._dynamo.reset()
+
+    def test_input_modifications_mul(self):
+        class InplaceMul(torch.nn.Module):
+            def forward(self, x):
+                x *= 5.0
+                x *= 1.9
+                y = x + 1
+                y /= 1.3
+                return y
+
+        inputs = [
+            torch.rand(
+                1,
+                3,
+                5,
+                7,
+            ).cuda(),
+        ]
+
+        fx_graph = torch.fx.symbolic_trace(InplaceMul())
+
+        # Validate that the results between Torch and Torch-TRT are similar
+        optimized_model = torch_tensorrt.compile(
+            fx_graph,
+            "torch_compile",
+            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,
+            msg=f"InplaceMul TRT outputs don't match with the original model.",
+        )
+        torch._dynamo.reset()
+
+
 if __name__ == "__main__":
     run_tests()