[fx2trt] exception for inplace op (#82)

Summary:
Pull Request resolved: https://github.com/pytorch/fx2trt/pull/82

`check_mutable_operations` does not work.
We override `create_node` function to throw exception for inplace

Reviewed By: jfix71

Differential Revision: D36556398

fbshipit-source-id: db35732007eae77fb40500753d7ffe26d43e659c

Author: Wei Wei

test/tracer/test_acc_tracer.py

@@ -2382,6 +2382,34 @@ def forward(self, a):
                 torch.equal(ref, res), f"Tensors at don't match {ref=} {res=}"
             )
 
+    def test_inplace_raise(self):
+        """
+        Test that encountering inplace is raised for exception
+        """
+
+        class TestModule(nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, a):
+                a = a + 2
+                a.sub_(3)
+                return a
+
+        m = TestModule()
+        in_a = torch.randn(5)
+        try:
+            acc_tracer.trace(
+                m,
+                [in_a],
+            )
+            self.fail("Shouldn't get here because exception should be thrown.")
+        except RuntimeError as e:
+            self.assertEqual(
+                "Tried to trace mutable operation sub_. FX only supports functional code",
+                str(e),
+            )
+
     def test_repeat_interleave(self):
         class TestModule(nn.Module):
             def __init__(self):

tracer/acc_tracer/acc_tracer.py

@@ -18,6 +18,7 @@
 from torch._sources import normalize_source_lines
 from torch.fx import Graph, Tracer
 from torch.fx.experimental.normalize import NormalizeArgs
+from torch.fx.node import Argument, Node, Target
 from torch.fx.passes import shape_prop
 
 
@@ -242,6 +243,42 @@ def trace(
         rewritten = _rewrite(root, ast_rewriter_allow_list, self.leaf_module_list)
         return super().trace(rewritten, concrete_args), rewritten
 
+    # override TraceBase's method
+    def create_node(
+        self,
+        kind: str,
+        target: Target,
+        args: Tuple[Argument, ...],
+        kwargs: Dict[str, Argument],
+        name: Optional[str] = None,
+        type_expr: Optional[Any] = None,
+    ) -> Node:
+        """
+        Inserts a graph node given target, args, kwargs, and name.
+
+        This method can be overridden to do extra checking, validation, or
+        modification of values used in node creation. For example, one might
+        want to disallow in-place operations from being recorded.
+        """
+
+        ## Hacky way to decide inplace ops
+        if type(target) != str:
+            name_target = target.__name__
+        else:
+            name_target = target
+
+        allow_list = ["and_", "or_"]  # python  operator.and_,  operator.or_
+        if (
+            name_target[-1] == "_"
+            and name_target[0] != "_"
+            and not (name_target in allow_list)
+        ):
+            raise RuntimeError(
+                f"Tried to trace mutable operation {name_target}. FX only supports functional code"
+            )
+
+        return self.graph.create_node(kind, target, args, kwargs, name, type_expr)
+
 
 # List of modules that need rewriting to be supported for tracing.
 DEFAULT_REWRITE_ALLOW_LIST = {