give up on making TupleElements noncopyable on "[PyTorch] RFC: Add tu…

…ple inline storage" I noticed a bunch of time being spent heap-allocating Tuples in the unpickler. 1-, 2-, and 3-element Tuples are apparently common enough that they get their own bytecode instructions, so I decided to try also giving them their own representation. We store up to 3 IValues inline in `Tuple` rather than doing a second heap allocation for a `std::vector<IValue>`. Differential Revision: [D30592622](https://our.internmc.facebook.com/intern/diff/D30592622/) **NOTE FOR REVIEWERS**: This PR has internal Facebook specific changes or comments, please review them on [Phabricator](https://our.internmc.facebook.com/intern/diff/D30592622/)! cc pietern mrshenli pritamdamania87 zhaojuanmao satgera rohan-varma gqchen aazzolini osalpekar jiayisuse @SciPioneer @H-Huang gcramer23 [ghstack-poisoned]
pytorch · swolchok · Aug 26, 2021 · Aug 27, 2021 · Aug 27, 2021 · Aug 27, 2021
commit 236cf76efc7060a2028cca10bfa06f44800f42c7
diff --git a/aten/src/ATen/core/ivalue_inl.h b/aten/src/ATen/core/ivalue_inl.h
@@ -312,10 +312,60 @@ struct TORCH_API TupleElements {
     }
   }
 
-  // Simply not implemented; no particular reason not to implement in
-  // the future except that it seems unnecessary.
-  TupleElements(const TupleElements&) = delete;
-  TupleElements& operator=(const TupleElements&) = delete;
+  // It would be nice to make this noncopyable to prevent people from
+  // writing code like `auto output =
+  // forward(...).toTuple()->elements()` (which does refcount bumps on
+  // each element, unlike the more efficient but verbose
+  // ```
+  // auto outputIntrusivePtr = forward(...).toTuple();
+  // const auto& output = outputIntrusivePtr->elements();
+  // ```
+  // ), but there is simply an overwhelming amount of code that does
+  // it the inefficient way.
+
+  TupleElements(const TupleElements& rhs)
+  : inlineSize_(rhs.inlineSize_) {
+    if (rhs.inlineSize_) {
+      for (const auto  ii : c10::irange(inlineSize_)) {
+        new (&elementsInline_[ii]) IValue(rhs.elementsInline_[ii]);
+      }
+    } else {
+      new (&elementsVector_) std::vector<IValue>(rhs.elementsVector_);
+    }
+  }
+
+  TupleElements& operator=(const TupleElements& rhs) {
+    if (inlineSize_) {
+      if (rhs.inlineSize_) {
+        for (const auto ii : c10::irange(std::min(inlineSize_, rhs.inlineSize_))) {
+          elementsInline_[ii] = rhs.elementsInline_[ii];
+        }
+        if (rhs.inlineSize_ > inlineSize_) {
+          for (const auto ii : c10::irange(inlineSize_, rhs.inlineSize_)) {
+            new (&elementsInline_[ii]) IValue(rhs.elementsInline_[ii]);
+          }
+        } else {
+          for (const auto ii : c10::irange(rhs.inlineSize_, inlineSize_)) {
+            elementsInline_[ii].~IValue();
+          }
+        }
+      } else {
+        destroyInline();
+        new (&elementsVector_) std::vector<IValue>(rhs.elementsVector_);
+      }
+    } else {
+      if (rhs.inlineSize_) {
+        elementsVector_.~vector();
+        for (const auto ii : c10::irange(rhs.inlineSize_)) {
+          new (&elementsInline_[ii]) IValue(rhs.elementsInline_[ii]);
+        }
+      } else {
+        elementsVector_ = rhs.elementsVector_;
+      }
+    }
+    inlineSize_ = rhs.inlineSize_;
+    return *this;
+  }
 
   TupleElements(TupleElements&& rhs) noexcept
   : inlineSize_(rhs.inlineSize_) {

diff --git a/aten/src/ATen/test/ivalue_test.cpp b/aten/src/ATen/test/ivalue_test.cpp
@@ -811,23 +811,31 @@ TEST(TupleElementsTest, Resize) {
   }
 }
 
-TEST(TupleElementsTest, MoveConstruct) {
+TEST(TupleElementsTest, CopyAndMoveConstruct) {
   int idx = 0;
   for (auto fromFactory : factories) {
     auto toMoveFrom = fromFactory();
     TupleElements movedInto(std::move(toMoveFrom));
-    validateTupleElements(movedInto, expectedContents[idx++]);
+    validateTupleElements(movedInto, expectedContents[idx]);
+    auto toCopyFrom = fromFactory();
+    TupleElements copiedInto(toCopyFrom);
+    validateTupleElements(copiedInto, expectedContents[idx]);
+    idx++;
   }
 }
 
-TEST(TupleElementsTest, MoveAssign) {
+TEST(TupleElementsTest, CopyAndMoveAssign) {
   int fromIdx = 0;
   for (auto fromFactory : factories) {
     for (auto toFactory : factories) {
       auto from = fromFactory();
       auto to = toFactory();
+      auto copyFrom = fromFactory();
+      auto toCopy = toFactory();
       to = std::move(from);
       validateTupleElements(to, expectedContents[fromIdx]);
+      toCopy = copyFrom;
+      validateTupleElements(toCopy, expectedContents[fromIdx]);
     }
     fromIdx++;
   }