Repository: pytorch/pytorch

Issue:
DCE trick for inplace tracing doesn't work
Exposed by #4395, hotpatched in #4471.

The DCE trick says that if I have `y = f(x)`, and `f` is internally implemented as `g`, it's OK to trace both `g` and `f`. Recall the tracing algorithm is:

1. enter f(x)
2. compute its result y (recursively calling other possibly tracing functions)
3. trace y = f(x) (use the mapping to figure out the identity of x, and writing in a mapping that y binds to the result JIT IR value of this trace)
4. return from f

So when you run the example above, you'll do this:

```
# suppose x is mapped to %1
enter f(x)
enter g(x)
result of g is y
trace y = g(x a.k.a. %1) (mapping y to %2)
return from g
result of f is y
trace y = f(x a.k.a. %1) (remapping y to %3)
return from f
```

and end up with a trace like this:

```
%2 = g(%1)
%3 = f(%1)
... only %3 is live, because %2 was killed from the mapping...
```

Subsequent DCE will eliminate the invocation of `g` and you'll only see `f` in the final trace.

However, if f and g are inplace functions, the machinery breaks:

```
# suppose x is mapped to %1
enter f(x)
enter g(x)
result of g is x
trace x = g(x a.k.a. %1) (remapping x to %2)
return from g
result of f is x
trace x = f(x a.k.a. %2) (remapping x to %3)
return from f
```

resulting in:

```
%2 = g(%1)
%3 = f(%2) # OOPS
```

Proper TLS tracing will solve this problem. But maybe a more robust stopgap than #4471 is to add a little bit of extra logic for inplace functions. Essentially, what we need to do is save the original Value for a function before we make the inner recursive call, and use that as the Value when we determine the arguments of tracing, rather than doing it after we do the inner recursive call (in which case `x` could have been overwritten). In fact, I originally implemented the tracer this way but @colesbury convinced me that it wasn't necessary (and it was true, in the code at the time, it was not necessary.)

@houseroad, when you fix this, please work this discussion into a comment somewhere in the code.

CC @apaszke 
  
  


Discussion:

Resolving patch:
diff --git a/tools/autograd/gen_variable_type.py b/tools/autograd/gen_variable_type.py
--- a/tools/autograd/gen_variable_type.py
+++ b/tools/autograd/gen_variable_type.py
@@ -122,15 +122,22 @@
 RECORD_FUNCTION = CodeTemplate("""\
 profiler::RecordFunction profiler("${name}");""")
 
-RECORD_TRACE = CodeTemplate("""\
+PRE_RECORD_TRACE = CodeTemplate("""\
+jit::tracer::PreTraceInfo trace_info;
 if (jit::tracer::isTracing( ${tensor_args} )) {
-  jit::Node *n = jit::tracer::recordTrace( "${trace_name}", ${trace_inputs}, ${trace_outputs} );
+  trace_info = jit::tracer::preRecordTrace( "${trace_name}", ${trace_inputs} );
   ${record_attributes}
 }
 """)
 
+POST_RECORD_TRACE = CodeTemplate("""\
+if (trace_info.state != nullptr) {
+  jit::tracer::postRecordTrace( trace_info,  ${trace_outputs} );
+}
+""")
+
 RECORD_ATTRIBUTE = CodeTemplate("""\
-setattr(n, jit::Symbol("${name}"), ${name});""")
+setattr(trace_info.n, jit::Symbol("${name}"), ${name});""")
 
 
 def gen_variable_type(out, aten_declarations):
@@ -305,15 +312,15 @@ def emit_record_trace(env):
         # Operations involving Generator and Storage are not traceable
         # at the moment
         if any(arg['simple_type'] in {'Generator', 'Storage'} for arg in declaration['arguments']):
-            return ''
+            return ('', '')
         # We can't trace functions which don't have any Tensor or TensorList returns
         if 'Tensor' not in declaration['return_type']:
-            return ''
+            return ('', '')
         tensor_args = [arg for arg in arguments if arg['simple_type'] in {'Tensor', 'TensorList'}]
         if len(tensor_args) == 0:
-            return ''
+            return ('', '')
         if base_name in DONT_RECORD_TRACE:
-            return ''
+            return ('', '')
 
         # Note [clang-802.0.42 tuple overload bug]
         # ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -352,8 +359,6 @@ def emit_record_trace(env):
             if arg['simple_type'] in {'Tensor', 'TensorList'}:
                 continue
             local['record_attributes'].append(RECORD_ATTRIBUTE.substitute(name=arg['name']))
-        if not local['record_attributes']:
-            local['record_attributes'].append('(void)n;')
 
         local['trace_name'] = declaration['api_name']
         if local['trace_name'].endswith('_'):
@@ -361,7 +366,7 @@ def emit_record_trace(env):
         local['trace_outputs'] = get_trace_outputs(declaration)
 
         combined = nested_dict(local, nested_dict(env, declaration))
-        return RECORD_TRACE.substitute(combined)
+        return (PRE_RECORD_TRACE.substitute(combined), POST_RECORD_TRACE.substitute(combined))
 
     def declare_returned_variables():
         if modifies_arguments:
@@ -456,6 +461,10 @@ def emit_increment_version():
         body.extend(emit_check_inplace())
         body.extend(setup_derivative())
     body.append(declare_returned_variables())
+
+    pre_record_trace, post_record_trace = emit_record_trace(env)
+
+    body.append(pre_record_trace)
     body.append(emit_call(env))
     if requires_derivative:
         if inplace and is_view:
@@ -464,11 +473,9 @@ def emit_increment_version():
         # requires that the counter is incremented before it is called
         body.extend(emit_increment_version())
         body.append(emit_history())
-    # record_trace must appear before save_outputs so that saved outputs
-    # have their tracing_state saved
-    body.append(emit_record_trace(env))
     if requires_derivative:
         body.append(emit_save_outputs())
+    body.append(post_record_trace)
     body.append('return {};'.format(get_return_value()))
     return body
 
diff --git a/torch/onnx/__init__.py b/torch/onnx/__init__.py
--- a/torch/onnx/__init__.py
+++ b/torch/onnx/__init__.py
@@ -297,10 +297,6 @@ def _run_symbolic_function(g, n, inputs, aten=False):
         # See Note [Export inplace]
         if n.kind().endswith('_'):
             op_name = n.kind()[:-1]
-        elif n.kind().endswith('_forward'):
-            # NB: it seems the tracing generates duplicate ops,
-            # such as both threshold and threshold_forward, introduced in #4395.
-            return inputs
         else:
             op_name = n.kind()
         # Export ops in aten mode.


Test patch:
diff --git a/test/expect/TestJit.test_nested_inplace.expect b/test/expect/TestJit.test_nested_inplace.expect
new file mode 100644
--- /dev/null
+++ b/test/expect/TestJit.test_nested_inplace.expect
@@ -0,0 +1,4 @@
+graph(%0 : Double(2, 2)) {
+  %1 : Double(2, 2) = threshold[threshold={0}, value={0}](%0)
+  return (%1);
+}
diff --git a/test/test_jit.py b/test/test_jit.py
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -1165,6 +1165,11 @@ def forward(self, x):
         self.assertEqual(len(list(trace.graph().inputs())), 2)
         self.assertExpected(str(trace))
 
+    def test_nested_inplace(self):
+        x = Variable(torch.randn(2, 2))
+        trace, _ = torch.jit.trace(lambda x: F.threshold(x, 0, 0, inplace=True), (x,), nderivs=0)
+        self.assertExpectedTrace(trace)
+
 
 if __name__ == '__main__':
     run_tests()