[PyTorch] Add IValue::IValue(std::vector<T>&&) ctors (#117769)

swolchok · pytorchmergebot · commit 414a1fd29f04 · 2024-01-19T21:11:11.000Z
There are two IValue constructors that take `const std::vector<T>&`. Add moving variants to allow callers to save on reference counting. Differential Revision: [D52879065](https://our.internmc.facebook.com/intern/diff/D52879065/) Pull Request resolved: #117769 Approved by: https://github.com/suo, https://github.com/Skylion007
diff --git a/aten/src/ATen/core/ivalue.h b/aten/src/ATen/core/ivalue.h
@@ -756,6 +756,8 @@ struct TORCH_API IValue final {
   IValue(at::ArrayRef<T> v);
   template <class T, enable_if_list_is_ivalue_constructible<T> = nullptr>
   IValue(const std::vector<T>& v);
+  template <class T, enable_if_list_is_ivalue_constructible<T> = nullptr>
+  IValue(std::vector<T>&& v);
   template <class T, size_t N>
   IValue(std::array<T, N> v);
 
@@ -772,6 +774,9 @@ struct TORCH_API IValue final {
   IValue(at::OptionalArrayRef<T> v);
   template <class T, enable_if_symint<T> = nullptr>
   IValue(const std::vector<T>& v);
+  template <class T, enable_if_symint<T> = nullptr>
+  IValue(std::vector<T>&& v);
+
 
   template <class T>
   using enable_if_ilist_is_ivalue_constructible = std::enable_if_t<
diff --git a/aten/src/ATen/core/ivalue_inl.h b/aten/src/ATen/core/ivalue_inl.h
@@ -2185,6 +2185,23 @@ template <class T, IValue::enable_if_symint<T>>
 inline IValue::IValue(const std::vector<T>& v) : IValue() {
   *this = IValue(at::ArrayRef<T>(v));
 }
+template <class T, IValue::enable_if_symint<T>>
+inline IValue::IValue(std::vector<T>&& v) : IValue() {
+  auto vi = c10::asIntArrayRefSlowOpt(v);
+  if (vi.has_value()) {
+    // This list is entirely integers; ensure it is typed as
+    // an IntList so toIntList works
+    *this = IValue(*vi);
+  } else {
+    // This list has SymInts; type it as a SymInt
+    *this = IValue(impl::toList<c10::SymInt>(c10::List<c10::SymInt>()));
+    auto list = to<c10::List<c10::SymInt>>();
+    list.reserve(v.size());
+    for (auto& e : v) {
+      list.push_back(std::move(e));
+    }
+  }
+}
 template <class T, IValue::enable_if_list_is_ivalue_constructible<T>>
 inline IValue::IValue(const std::vector<T>& v) : IValue(c10::List<T>()) {
   auto list = to<c10::List<T>>();
@@ -2193,6 +2210,22 @@ inline IValue::IValue(const std::vector<T>& v) : IValue(c10::List<T>()) {
     list.push_back(e);
   }
 }
+
+template <class T, IValue::enable_if_list_is_ivalue_constructible<T>>
+inline IValue::IValue(std::vector<T>&& v) : IValue(c10::List<T>()) {
+  auto list = to<c10::List<T>>();
+  list.reserve(v.size());
+  if constexpr (std::is_same_v<T, bool>) {
+    for (auto e : v) {
+      list.push_back(e);
+    }
+  } else {
+    for (auto& e : v) {
+      list.push_back(std::move(e));
+    }
+  }
+}
+
 template <class T, IValue::enable_if_list_is_ivalue_constructible<T>>
 inline IValue::IValue(c10::OptionalArrayRef<T> v) : IValue() {
   if (v.has_value()) {
diff --git a/aten/src/ATen/test/ivalue_test.cpp b/aten/src/ATen/test/ivalue_test.cpp
@@ -49,6 +49,14 @@ TEST(IValueTest, Basic) {
   ASSERT_TRUE(dlist.isNone());
   dlist = IValue(c10::List<double>({3.4}));
   ASSERT_TRUE(dlist.toDoubleVector() == std::vector<double>({3.4}));
+  dlist = IValue(std::vector<double>({3.3, 3.2}));
+  ASSERT_TRUE(dlist.toDoubleVector() == std::vector<double>({3.3, 3.2}));
+  IValue blist(std::vector<bool>{true, false});
+  ASSERT_TRUE(blist.isList());
+  const auto blistRef = blist.toListRef();
+  ASSERT_EQ(blistRef.size(), 2);
+  ASSERT_TRUE(blistRef[0].toBool());
+  ASSERT_FALSE(blistRef[1].toBool());
   IValue the_list(
       at::ivalue::Tuple::create({IValue(3.4), IValue(4), IValue(foo)}));
   ASSERT_EQ(foo.use_count(), 3);
