Revert "Cleaner semantics for Reserve (pytorch#10261)"

This reverts commit e0d4357.

Author: ezyang

caffe2/core/tensor.h

@@ -231,17 +231,6 @@ class Tensor {
 
   virtual ~Tensor() noexcept {}
 
-  /**
-   * @brief Extend the outer-most dimension of this tensor
-   *        to dimension of `num`.
-   */
-  void ExtendTo(TIndex num, float growthPct, BaseContext* context) {
-    CAFFE_ENFORCE_GE_WITH_CALLER(dims_.size(), 1);
-    CAFFE_ENFORCE_GE_WITH_CALLER(growthPct, 0);
-    CAFFE_ENFORCE(context != nullptr, "Context must be provided.");
-    Extend(num - dims_[0], growthPct, context);
-  }
-
   /**
    * @brief Extends the outer-most dimension of this tensor by num elements,
    * preserving the existing data.
@@ -253,8 +242,6 @@ class Tensor {
    */
   void Extend(TIndex num, float growthPct, BaseContext* context) {
     CAFFE_ENFORCE_GE_WITH_CALLER(dims_.size(), 1);
-    CAFFE_ENFORCE_GE_WITH_CALLER(
-        num, 0, "`num` must be non-negative for Extend");
     auto newDims = dims_;
     newDims[0] += num;
     if (!data_) {
@@ -274,17 +261,30 @@ class Tensor {
     auto newCapacity = dims_;
     newCapacity[0] = std::max<size_t>(
         newDims[0], std::ceil(dims_[0] * (growthPct + 100) / 100));
+    Reserve(newCapacity, context);
+    dims_ = newDims;
+    size_ = newSize;
+  }
+
+  template <class T>
+  void Reserve(const std::vector<T>& newCapacity, BaseContext* context) {
+    auto newSize = std::accumulate(
+        newCapacity.begin(),
+        newCapacity.end(),
+        static_cast<TIndex>(1),
+        std::multiplies<TIndex>());
+    if (newSize * meta_.itemsize() <= capacity_) {
+      return;
+    }
     auto oldData = std::move(data_);
     auto oldSize = size_;
     auto oldDims = dims_;
     Resize(newCapacity);
     auto* newData = raw_mutable_data(meta_);
-    CAFFE_ENFORCE(
-        context != nullptr, "Context must be provided to Extend the tensor");
     context->CopyItemsSameDevice(meta_, oldSize, oldData.get(), newData);
+    dims_ = oldDims;
+    size_ = oldSize;
     reserved_ = true;
-    dims_ = newDims;
-    size_ = newSize;
   }
 
   /**
@@ -293,7 +293,7 @@ class Tensor {
    * This method guarantees that no re-allocations are carried out, which means
    * that the extra capacity after the end of the shurnk tensor is maintained.
    */
-  void ShrinkTo(TIndex outer_dim) {
+  void Shrink(TIndex outer_dim) {
     CAFFE_ENFORCE_WITH_CALLER(dims_.size() >= 1, "Tensor must be at least 1D");
     CAFFE_ENFORCE_WITH_CALLER(
         outer_dim <= dims_[0],
@@ -306,38 +306,6 @@ class Tensor {
         std::multiplies<TIndex>());
   }
 
-  /**
-   * @brief Reserve space for the underlying tensor.
-   *
-   * This must be called after Resize(), since we only specify the first
-   * dimension This does not copy over the old data to the newly allocated space
-   */
-  template <class T>
-  void ReserveSpace(const T& outer_dim) {
-    CAFFE_ENFORCE(
-        size_ != -1, "size should be initialized before calling ReserveSpace");
-    auto newCapacity = dims_;
-    newCapacity[0] = outer_dim;
-    auto newSize = std::accumulate(
-        newCapacity.begin(),
-        newCapacity.end(),
-        static_cast<TIndex>(1),
-        std::multiplies<TIndex>());
-    if (newSize * meta_.itemsize() <= capacity_) {
-      return;
-    }
-    // Old data is discarded
-    data_.reset();
-    auto oldSize = size_;
-    auto oldDims = dims_;
-    Resize(newCapacity);
-    // Allocate new memory and don't copy over the data
-    raw_mutable_data(meta_);
-    dims_ = oldDims;
-    size_ = oldSize;
-    reserved_ = true;
-  }
-
   /**
    * @brief Resizes a tensor.
    *
@@ -421,7 +389,7 @@ class Tensor {
     capacity_ = 0;
     // If reserved is true and we changed tensor memory then it is fine
     // to switch it to false, if Resize is called from Reserve and it triggers
-    // FreeMemory() then reserved_ will be set to true at end of ReserveSpace()
+    // FreeMemory() then reserved_ will be set to true at end of Reserve()
     reserved_ = false;
   }
 
@@ -772,10 +740,6 @@ class Tensor {
   TypeMeta meta_;
   std::shared_ptr<void> data_;
   size_t capacity_ = 0;
-  // we decide to keep reserved and it will
-  // live in Tensor after the split
-  // The logic is that if Extend() or ReserveSpace() were ever called,
-  // then subsequent Resize()s will not free up Storage.
   bool reserved_ = false;
   DeviceType device_type_ = CPU;
   // In case of chunk load we store how much data was already loaded

caffe2/experiments/operators/tt_pad_op.h

@@ -83,7 +83,7 @@ class TTPadGradientOp final : public Operator<Context> {
     auto dim1 = G.dim(1);
 
     if (old_dim0 < new_dim0) {
-      output->ShrinkTo(old_dim0);
+      output->Shrink(old_dim0);
     }
 
     return true;

caffe2/mobile/contrib/ulp2/ulp_neon.cc

@@ -537,7 +537,7 @@ void run2b1bConvIm2ColGEMM(QConvState* state,
   } else {
     CAFFE_ENFORCE_EQ(Y->dim32(0), divRoundUp(X.dim32(0) * OH * OW, kGEMMTileSize) * kGEMMTileSize);
     CAFFE_ENFORCE_EQ(Y->dim32(1), OC);
-    Y->ShrinkTo(X.dim32(0) * OH * OW);
+    Y->Shrink(X.dim32(0) * OH * OW);
     Y->Reshape(std::vector<TIndex>{{TIndex(X.dim(0)), TIndex(OH), TIndex(OW), TIndex(OC)}});
   }
 }

caffe2/operators/dataset_ops.cc

@@ -1004,7 +1004,7 @@ class TrimDatasetOp : public Operator<CPUContext> {
     // trim each column to the offset
     for (int col = 0; col < walker.fields().size(); ++col) {
       auto newOuterSize = walker.fields().at(col).offset();
-      Output(col)->ShrinkTo(newOuterSize);
+      Output(col)->Shrink(newOuterSize);
     }
     return true;
   }

caffe2/operators/last_n_window_collector.cc

@@ -46,7 +46,8 @@ class LastNWindowCollectorOp : public Operator<Context> {
       }
     }
 
-    auto num_entries = input.dims()[0];
+    auto dims = input.dims();
+    auto num_entries = dims[0];
 
     if (OutputSize() > NUM_VISITED) {
       auto* num_visited_tensor = Output(NUM_VISITED);
@@ -59,14 +60,8 @@ class LastNWindowCollectorOp : public Operator<Context> {
       *num_visited += num_entries;
     }
 
-    if (!output_initialized) {
-      auto dims = input.dims();
-      dims[0] = 0;
-      output->Resize(dims);
-      // pass meta to output
-      output->raw_mutable_data(input.meta());
-      output->ReserveSpace(numToCollect_);
-    }
+    dims[0] = numToCollect_;
+    output->Reserve(dims, &context_);
 
     if (num_entries == 0) {
       if (!output_initialized) {
@@ -78,14 +73,10 @@ class LastNWindowCollectorOp : public Operator<Context> {
 
     auto num_to_copy = std::min<int32_t>(num_entries, numToCollect_);
     auto output_batch_size = output_initialized ? output->dim(0) : 0;
-    auto output_num =
-        std::min<size_t>(numToCollect_, output_batch_size + num_to_copy);
-
-    // output_num is >= output_batch_size
-    if (output_num > output_batch_size) {
-      output->ExtendTo(output_num, 50, &context_);
+    dims[0] = std::min<size_t>(numToCollect_, output_batch_size + num_to_copy);
+    if (output_batch_size < numToCollect_) {
+      output->Resize(dims);
     }
-
     auto* output_data =
         static_cast<char*>(output->raw_mutable_data(input.meta()));
 

caffe2/operators/reservoir_sampling.cc

@@ -39,27 +39,23 @@ class ReservoirSamplingOp final : public Operator<Context> {
       }
     }
 
-    auto num_entries = input.dims()[0];
+    auto dims = input.dims();
+    auto num_entries = dims[0];
 
-    if (!output_initialized) {
-      // IMPORTANT: Force the output to have the right type before reserving,
-      // so that the output gets the right capacity
-      auto dims = input.dims();
-      dims[0] = 0;
-      output->Resize(dims);
-      output->raw_mutable_data(input.meta());
-      output->ReserveSpace(numToCollect_);
-    }
+    dims[0] = numToCollect_;
+    // IMPORTANT: Force the output to have the right type before reserving,
+    // so that the output gets the right capacity
+    output->raw_mutable_data(input.meta());
+    output->Reserve(dims, &context_);
 
     auto* pos_to_object =
         OutputSize() > POS_TO_OBJECT ? Output(POS_TO_OBJECT) : nullptr;
     if (pos_to_object) {
       if (!output_initialized) {
         // Cleaning up in case the reservoir got reset.
         pos_to_object->Resize(0);
-        pos_to_object->template mutable_data<int64_t>();
-        pos_to_object->ReserveSpace(numToCollect_);
       }
+      pos_to_object->Reserve(std::vector<TIndex>{numToCollect_}, &context_);
     }
 
     auto* object_to_pos_map = OutputSize() > OBJECT_TO_POS_MAP
@@ -100,14 +96,13 @@ class ReservoirSamplingOp final : public Operator<Context> {
     const auto num_new_entries = countNewEntries(unique_object_ids);
     auto num_to_copy = std::min<int32_t>(num_new_entries, numToCollect_);
     auto output_batch_size = output_initialized ? output->dim(0) : 0;
-    auto output_num =
-        std::min<size_t>(numToCollect_, output_batch_size + num_to_copy);
-    // output_num is >= output_batch_size
-    output->ExtendTo(output_num, 50, &context_);
-    if (pos_to_object) {
-      pos_to_object->ExtendTo(output_num, 50, &context_);
+    dims[0] = std::min<size_t>(numToCollect_, output_batch_size + num_to_copy);
+    if (output_batch_size < numToCollect_) {
+      output->Resize(dims);
+      if (pos_to_object) {
+        pos_to_object->Resize(dims[0]);
+      }
     }
-
     auto* output_data =
         static_cast<char*>(output->raw_mutable_data(input.meta()));
     auto* pos_to_object_data = pos_to_object

caffe2/operators/text_file_reader.cc

@@ -150,7 +150,7 @@ class TextFileReaderReadOp : public Operator<CPUContext> {
     }
 
     for (int i = 0; i < numFields; ++i) {
-      Output(i)->ShrinkTo(rowsRead);
+      Output(i)->Shrink(rowsRead);
     }
     return true;
   }