deepCopy also copies type information of lists

smessmer · smessmer · commit 82e8cae632da · 2019-07-23T14:57:51.000-07:00
Differential Revision: [D16449220](https://our.internmc.facebook.com/intern/diff/D16449220/) ghstack-source-id: 87041700 Pull Request resolved: #23271
diff --git a/aten/src/ATen/core/List.h b/aten/src/ATen/core/List.h
@@ -435,6 +435,10 @@ class List final {
   // TODO Test use_count
   size_t use_count() const;
 
+  // private API for now because the return type will change to TypePtr
+  // instead of optional<TypePtr> once types are mandatory.
+  optional<TypePtr> _elementType() const;
+
 private:
   explicit List(c10::intrusive_ptr<detail::ListImpl<StorageT>>&& elements);
   friend struct IValue;
diff --git a/aten/src/ATen/core/List_inl.h b/aten/src/ATen/core/List_inl.h
@@ -289,4 +289,9 @@ size_t List<T>::use_count() const {
   return impl_.use_count();
 }
 
+template<class T>
+optional<TypePtr> List<T>::_elementType() const {
+  return impl_->elementType;
+}
+
 }
diff --git a/torch/csrc/jit/passes/utils/check_alias_annotation.cpp b/torch/csrc/jit/passes/utils/check_alias_annotation.cpp
@@ -25,12 +25,22 @@ IValue deepCopy(const IValue& self) {
 
   // Lists of ivalues should recursively deep copy their contents
   if (self.isGenericList()) {
-    auto newList = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
-    newList.reserve(self.toGenericListRef().size());
-    for (const IValue& value : self.toGenericListRef()) {
-      newList.push_back(deepCopy(value));
+    auto source = std::move(self).toGenericList();
+    auto deepCopyGenericList = [&] (c10::impl::GenericList& dest) {
+      dest.reserve(source.size());
+      for (const IValue& value : source) {
+        dest.push_back(deepCopy(value));
+      }
+    };
+    if (source._elementType().has_value()) {
+      auto newList = c10::impl::GenericList(*source._elementType());
+      deepCopyGenericList(newList);
+      return newList;
+    } else {
+      auto newList = c10::impl::GenericList(c10::impl::deprecatedUntypedList());
+      deepCopyGenericList(newList);
+      return newList;
     }
-    return newList;
   }
 
   // Regular lists can copy assign

Original file line number	Diff line number	Diff line change
`@@ -289,4 +289,9 @@ size_t List<T>::use_count() const {`
`289`	`289`	`return impl_.use_count();`
`290`	`290`	`}`
`291`	`291`
	`292`	`+template<class T>`
	`293`	`+optional<TypePtr> List<T>::_elementType() const {`
	`294`	`+ return impl_->elementType;`
	`295`	`+}`
	`296`	`+`
`292`	`297`	`}`