Repository: pytorch/pytorch
Pull request: Create captured inputs recursively for loop to resolve loop-carried dependencies across nested blocks
This PR fixes #7818 , the bug raises because we did not correctly create the captured inputs for nested loops. we changed the condition when we generate the captured input, recursively create a captured input if the use of ident crosses loops.


Problem statement:
[JIT][script] Bug in how loop-carried dependencies are captured across nested blocks
Code example:
```
import torch

@torch.jit.script
def capture_bug(x):
    output = 0
    for i in range(20):
        if i == 0:
            output = x.unsqueeze(0)
        else:
            output = torch.cat((output, x.unsqueeze(0)), dim=0)
    return output

print(capture_bug.graph)
```

Creates:
```
graph(%x : Dynamic) {
  %output.1 : Long() = prim::Constant[value={0}]()
  %2 : Long() = prim::Constant[value={20}]()
  %3 : Byte() = prim::Constant[value={1}]()
  %output : Dynamic = prim::Loop(%2, %3, %output.1)
    block0(%i : Dynamic, %15 : Dynamic) {
      %5 : Long() = prim::Constant[value={0}]()
      %6 : Dynamic = aten::eq(%i, %5)
      %output.4 : Dynamic = prim::If(%6)
        block0() {
          %output.2 : Dynamic = aten::unsqueeze[dim=0](%x)
          -> (%output.2)
        }
        block1() {
          %10 : Dynamic = aten::unsqueeze[dim=0](%x)
          %12 : Long() = prim::Constant[value={0}]()
          %output.3 : Dynamic = aten::cat(%output.1, %10, %12)
          -> (%output.3)
        }
      %16 : Byte() = prim::Constant[value={1}]()
      -> (%16, %output.4)
    }
  return (%output);
}
```

The first operand to `cat` in `block1()` should be a captured value, not `%output.1`

Hints from issue discussion:

Commit message: recursively generate captured values for nested loops

Patch to review:
diff --git a/test/expect/TestScript.test_if_for_in_range.expect b/test/expect/TestScript.test_if_for_in_range.expect
new file mode 100644
index 000000000000..a37c165318f7
--- /dev/null
+++ b/test/expect/TestScript.test_if_for_in_range.expect
@@ -0,0 +1,26 @@
+graph(%a.1 : Dynamic
+      %b.1 : Dynamic) {
+  %d.1 : Long() = prim::Constant[value={3}]()
+  %3 : Long() = prim::Constant[value={20}]()
+  %4 : Byte() = prim::Constant[value={1}]()
+  %a : Dynamic, %d : Long(), %b : Dynamic = prim::Loop(%3, %4, %a.1, %d.1, %b.1)
+    block0(%_ : Dynamic, %6 : Dynamic, %10 : Long(), %14 : Dynamic) {
+      %7 : Long() = prim::Constant[value={10}]()
+      %8 : Dynamic = aten::gt(%6, %7)
+      %a.3 : Dynamic, %b.3 : Dynamic, %d.3 : Long() = prim::If(%8)
+        block0() {
+          %9 : Long() = prim::Constant[value={3}]()
+          %a.2 : Dynamic = aten::add[alpha={1}](%9, %10)
+          -> (%a.2, %14, %10)
+        }
+        block1() {
+          %12 : Long() = prim::Constant[value={3}]()
+          %b.2 : Dynamic = aten::add[alpha={1}](%12, %10)
+          %d.2 : Long() = prim::Constant[value={4}]()
+          -> (%6, %b.2, %d.2)
+        }
+      %20 : Byte() = prim::Constant[value={1}]()
+      -> (%20, %a.3, %d.3, %b.3)
+    }
+  return (%d);
+}
diff --git a/test/test_jit.py b/test/test_jit.py
index e667520024be..bfefd42d404b 100644
--- a/test/test_jit.py
+++ b/test/test_jit.py
@@ -1317,6 +1317,20 @@ def func(a, b):
         inputs = self._make_scalar_vars([1, -1], torch.int64)
         self.checkScript(func, inputs, optimize=True)
 
+    def test_if_for_in_range(self):
+        def func(a, b):
+            d = 3
+            for _ in range(20):
+                if a > 10:
+                    a = 3 + d
+                else:
+                    b = 3 + d
+                    d = 4
+                c = a + b
+            return d
+        inputs = self._make_scalar_vars([1, -1], torch.int64)
+        self.checkScript(func, inputs, optimize=True)
+
     def test_if_noelse(self):
         def func(a, b):
             if a > 10:
@@ -1422,6 +1436,20 @@ def test_script_for_in_range_ast(zero):
 
         self.assertEqual(test_script_for_in_range_ast(*inputs), 161700)
 
+    def test_script_for_in_range_if_ast(self):
+        @torch.jit.script
+        def test_script_for_in_range_if_ast(x):
+            output = 0
+            for i in range(20):
+                if i == 0:
+                    output = x.unsqueeze(0)
+                else:
+                    output = torch.cat((output, x.unsqueeze(0)), dim=0)
+            return output
+        inputs = self._make_scalar_vars([0], torch.int64)
+
+        self.assertEqual(test_script_for_in_range_if_ast(*inputs).shape[0], 20)
+
     def test_script_bool_constant(self):
         script = '''
         def test_script_bool_constant():
diff --git a/torch/csrc/jit/script/compiler.cpp b/torch/csrc/jit/script/compiler.cpp
index 8f9f818e6239..837cd172f08c 100644
--- a/torch/csrc/jit/script/compiler.cpp
+++ b/torch/csrc/jit/script/compiler.cpp
@@ -111,6 +111,22 @@ struct Environment {
 
     return sv;
   }
+
+  SugaredValuePtr createCapturedInputIfNeeded(const SourceRange& loc, std::string ident) {
+    auto in_frame = findInThisFrame(ident);
+    if (in_frame)
+      return in_frame;
+
+    // recursively handles the case where parent blocks are also loops
+    auto from_parent = next ? next->createCapturedInputIfNeeded(loc, ident) : nullptr;
+
+    // recursively create the captured input if it is the loop block
+    if (from_parent && getBlockOwningKind() == prim::Loop) {
+      from_parent = createCapturedInput(from_parent->asValue(loc, method), ident);
+    }
+    return from_parent;
+  }
+
   Block* block() {
     return b;
   }
@@ -158,12 +174,8 @@ struct Environment {
         << " but is now being assigned to a value of type " << as_simple_value->type()->name();
       }
     }
-    if (as_simple_value &&
-        !findInThisFrame(name) &&
-        findInParentFrame(name) &&
-        getBlockOwningKind() == prim::Loop) {
-      createCapturedInput(as_simple_value, name);
-    }
+    if (as_simple_value)
+      createCapturedInputIfNeeded(loc, name);
     value_table[name] = std::move(value);
   }
 
@@ -175,14 +187,7 @@ struct Environment {
   }
 
   SugaredValuePtr getSugaredVar(const std::string& ident, SourceRange range, bool required=true) {
-    auto retval = findInThisFrame(ident);
-
-    if (!retval && (retval = findInParentFrame(ident)) &&
-        getBlockOwningKind() == prim::Loop) {
-      if(Value* simple_val = asSimple(retval)) {
-        retval = createCapturedInput(simple_val, ident);
-      }
-    }
+    auto retval = createCapturedInputIfNeeded(range, ident);
 
     if(!retval) {
       retval = resolver(ident);


Code review comments:
diff_hunk: @@ -111,6 +111,22 @@ struct Environment {
 
     return sv;
   }
+
+  SugaredValuePtr createCapturedInputIfNeeded(const SourceRange& loc, std::string ident) {
+    auto in_frame = findInThisFrame(ident);
+    if (in_frame)
+      return in_frame;
+
+    // recursively handles the case where parent blocks are also loops
+    auto from_parent = next ? next->createCapturedInputIfNeeded(loc, ident) : nullptr;
+
+    // recursively create the captured input if it is the loop block
+    if (from_parent && getBlockOwningKind() == prim::Loop) {
+      from_parent = createCapturedInput(from_parent->asValue(loc, method), ident);
path: torch/csrc/jit/script/compiler.cpp
text: @user1:
The asValue here should be replaced with:

```
if(Value* simple_val = asSimple(from_parent))
  from_parent = createCapturedInput(simple_val, ident);
```

Like it is in the original getSugaredVar code. Otherwise things like `self` won't be usable inside loops.