Repository: pytorch/pytorch

Issue:
Freeing a deep computation graph causes a stack overflow
The code below will segfault due to a stack overflow. If the computation graph is very linear, like the below, such that freeing a Function leads to the immediate free of another Function, a stack overflow can occur. 

This stack overflow occurs some ~16% earlier on master than on v0.3.1.
```
import torch
from torch.autograd import Variable

def test(n):
    # In the second of two loops, the computation graph from the first is freed
    for j in range(0, 2):
        x = Variable(torch.FloatTensor(range(9)), requires_grad=True) 
        time_step = 0.002
        y = x.clone()
        
        # build deeply nested computation graph
        for i in range(n):
            y = y + y*time_step    
        print('Loop: {}'.format(j))

# Smallest n such that test(n) causes a stack overflow

# On 0.3.1 from conda install
test(65461)

# On master
test(52330)
```

I'm looking into fixing it.


Discussion:
To iterate is human, to recurse divine

:)
Meh, that's just how `shared_ptr`s work. We could do something like Python (basically limit the recursion depth, and accumulate remaining objects in a list once you hit the limit), but we'll need to be extra careful to construct all our pointers with the custom deleter. Shouldn't be too bad for perf if we use thread-local to-delete queues and recursion depth counters.

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
@@ -93,7 +93,7 @@
 """)
 
 ASSIGN_GRAD_FN = CodeTemplate("""\
-grad_fn = std::make_shared<${op}>(${op_ctor});
+grad_fn = std::shared_ptr<${op}>(new ${op}(${op_ctor}), deleteFunction);
 grad_fn->set_next_edges(collect_next_edges( ${args_with_derivatives} ));
 """)
 


Test patch:
diff --git a/test/test_autograd.py b/test/test_autograd.py
--- a/test/test_autograd.py
+++ b/test/test_autograd.py
@@ -611,6 +611,78 @@ def backward(ctx, grad_b):
 
         TestFn.apply(b).sum().backward()
 
+    def test_free_deep_graph(self):
+        def scope():
+            depth = 150000
+            x = torch.randn(1, requires_grad=True)
+            y = x.clone()
+
+            # build a "chain" computation graph
+            for i in range(depth):
+                y = y + y * 0.000001
+
+            # triggers graph deletion
+            del x
+
+        # Should not stack overflow
+        scope()
+
+    def test_free_deep_graph_complicated(self):
+        def scope():
+            depth = 100000
+            randchoice = torch.randint(2, [depth, 2])
+            x = torch.randn(1, requires_grad=True)
+            y = x.clone()
+
+            # Hold the two previous values
+            prev_values = [None, None]
+
+            # Build a "chain with skip connections" graph
+            for i in range(depth):
+                prev_tensors = [tensor for tensor in prev_values[:-1]
+                                if tensor is not None]
+                prev_values.append(y)
+                prev_values.pop(0)
+
+                # Definitely pick one tensor to add
+                y += y * 0.000001
+
+                # Possibly add other tensors
+                nprev = len(prev_tensors)
+                if nprev == 2:
+                    y += randchoice[depth].mul(torch.cat(prev_tensors)).sum()
+
+            # triggers graph deletion
+            del x
+
+        # Should not stack overflow
+        scope()
+
+    def test_free_deep_graph_pyfunction(self):
+        class MyOp(Function):
+            @staticmethod
+            def forward(ctx, tensor1, tensor2):
+                return tensor1 + tensor2
+
+            @staticmethod
+            def backward(ctx, grad_output):
+                return grad_output, grad_output
+
+        def scope():
+            depth = 150000
+            x = torch.randn(1, requires_grad=True)
+            y = x.clone()
+
+            # build deeply nested computation graph
+            for i in range(depth):
+                y = MyOp.apply(y, y)
+
+            # triggers graph deletion
+            del x
+
+        # Should not stack overflow
+        scope()
+
     def test_no_grad(self):
         x = torch.ones(5, 5, requires_grad=True)
         y = Variable(torch.ones(5, 5) * 4)