[JIT] fix alias assertion (#33952)

Summary:
This bug has been hit a couple times recently. We need to handle all bivariant types, not just optional, when asserting mutability/immutability of pointed-to elements in alias analysis.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/33952

Differential Revision: D20166025

Pulled By: eellison

fbshipit-source-id: cf3df9897a639641ef8303a08ba2b13523d01ef1

test/test_jit.py

@@ -7704,6 +7704,23 @@ a")
         g = torch.jit.last_executed_optimized_graph()
         self.assertEqual(next(g.outputs()).type().str(), "int[]")
 
+    def test_alias_covariant_type_containers(self):
+        @torch.jit.script
+        def foo(x):
+            # type: (bool)
+            if x: 
+                a = (None,)
+            else:
+                a = ([],)
+            return a
+
+        @torch.jit.script
+        def foo2(x, li):
+            # type: (bool, Tuple[Optional[List[Tensor]]])
+            if x:
+                li = (None,)
+            return li
+
     def test_while_write_outer_then_read(self):
         def func(a, b):
             while bool(a < 10):

torch/csrc/jit/ir/alias_analysis.cpp

@@ -769,8 +769,18 @@ void AliasDb::makePointerTo(const Value* from, const Value* to) {
     return;
   }
 
+  // covariant type containers can be point to types which are not
+  // also mutable/immutable because we unify the contained types
+  if (mutableType(from) != mutableType(to)) {
+    auto from_kind = from->type()->kind();
+    TORCH_INTERNAL_ASSERT(
+        from_kind == TypeKind::OptionalType ||
+        from_kind == TypeKind::FutureType || from_kind == TypeKind::TupleType);
+    return;
+  }
+
+  // both immutable
   if (!mutableType(from)) {
-    TORCH_INTERNAL_ASSERT(!mutableType(to));
     return;
   }
 
@@ -778,16 +788,7 @@ void AliasDb::makePointerTo(const Value* from, const Value* to) {
     return;
   }
 
-  // Special case: if `from` is an optional, `to` could be a None. Don't
-  // create a pointer in that case
-  if (from->type()->kind() == TypeKind::OptionalType &&
-      to->type()->kind() == TypeKind::NoneType) {
-    return;
-  }
-
-  // At this point, we should be dealing with two mutable types.
-  TORCH_INTERNAL_ASSERT(mutableType(from) && mutableType(to));
-
+  // At this point, we are dealing with two mutable types.
   auto fromEl = getOrCreateElement(from);
   auto toEl = getOrCreateElement(to);