[JIT] Cast return values of functions returning Any

Meghan Lele · Meghan Lele · commit 64d973e4efde · 2020-08-20T17:51:50.000-07:00
**Summary** This commit modifies IR generation to insert explicit cast that cast each return value to `Any` when a function is annotated as returning `Any`. This precludes the failure in type unification (see below) that caused this issue. Issue #41962 reported that the use of an `Any` return type in combination with different code paths returning values of different types causes a segmentation fault. This is because the exit transform pass tries to unify the different return types, fails, but silently sets the type of the if node to c10::nullopt. This causes problems later in shape analysis when that type object is dereferenced. **Test Plan** This commit adds a unit test that checks that a function similar to the one in #41962 can be scripted and executed. **Fixes** This commit fixes #41962. ghstack-source-id: 77d7dbf Pull Request resolved: #42259
diff --git a/test/test_jit_py3.py b/test/test_jit_py3.py
@@ -542,6 +542,20 @@ def forward(self, x):
         # Check that ignored method is still intact.
         self.assertEqual(inp, n.ignored_method(inp))
 
+    def test_if_returning_any(self):
+        """
+        Check that an if statement can return different
+        types early from each branch when the return
+        type of the function is Any.
+        """
+        def if_function(inp: torch.Tensor) -> Any:
+            if inp.shape[0] == 1:
+                return inp * inp
+            else:
+                return "str"
+
+        self.checkScript(if_function, (torch.randn(5),))
+
     def test_export_opnames_interface(self):
         global OneTwoModule
 
diff --git a/torch/csrc/jit/frontend/ir_emitter.cpp b/torch/csrc/jit/frontend/ir_emitter.cpp
@@ -996,7 +996,13 @@ struct to_ir {
       result_type = merged_result_type.value();
     }
     AT_ASSERT(result_type);
+
     def_stack_.back().merged_return_type_ = result_type;
+
+    if (result_type == AnyType::get() && result->type() != AnyType::get()) {
+      result = graph->insertUncheckedCast(result, result_type);
+    }
+
     graph->insertNode(graph->create(prim::ReturnStmt, {result}, 0));
     exit_blocks.insert(environment_stack->block());
   }

Original file line number	Diff line number	Diff line change
`@@ -996,7 +996,13 @@ struct to_ir {`
`996`	`996`	`result_type = merged_result_type.value();`
`997`	`997`	`}`
`998`	`998`	`AT_ASSERT(result_type);`
	`999`	`+`
`999`	`1000`	`def_stack_.back().merged_return_type_ = result_type;`
	`1001`	`+`
	`1002`	`+ if (result_type == AnyType::get() && result->type() != AnyType::get()) {`
	`1003`	`+ result = graph->insertUncheckedCast(result, result_type);`
	`1004`	`+ }`
	`1005`	`+`
`1000`	`1006`	`graph->insertNode(graph->create(prim::ReturnStmt, {result}, 0));`
`1001`	`1007`	`exit_blocks.insert(environment_stack->block());`
`1002`	`1008`	`}`