restore caffe2 strides (pytorch#12883)

Summary:
Pull Request resolved: pytorch#12883

Attempting to do this again. last try broke oss ci: D10421896

Reallocation of strides_ if there's no change in dim seems to cause the error that broke internal flow last time. This fixes that. Found a potential race condition in caffe2 counter ops that might be the cause, we will investigate that.

Reviewed By: ezyang

Differential Revision: D10469960

fbshipit-source-id: 478186ff0d2f3dba1fbff6231db715322418d79c

Author: Roy Li

aten/src/ATen/core/TensorImpl.cpp

@@ -46,20 +46,13 @@ IntList TensorImpl::sizes() const {
 }
 
 IntList TensorImpl::strides() const {
-  AT_ASSERTM(strides_,
-             "Caffe2 tensors don't (yet) have meaningful strides and cannot "
-             "be used in PyTorch.");
   return IntList{strides_.get(), sizes_.size()};
 }
 
 bool TensorImpl::compute_contiguous() const {
   bool is_contiguous = true;
   if (is_empty())
     return is_contiguous;
-  if (!strides_) {
-    // Special case for Caffe2 tensors which don't have strides set.
-    return true;
-  }
   int64_t z = 1;
   for (int64_t d = dim() - 1; d >= 0; d--) {
     if (size(d) != 1) {
@@ -90,9 +83,6 @@ int64_t TensorImpl::size(int64_t d) const {
 }
 
 int64_t TensorImpl::stride(int64_t d) const {
-  AT_ASSERTM(strides_,
-             "Caffe2 tensors don't (yet) have meaningful strides and cannot "
-             "be used in PyTorch.");
   d = at::maybe_wrap_dim(d, dim(), false);
   return strides_[d];
 }

aten/src/ATen/core/TensorImpl.h

@@ -166,6 +166,8 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
   explicit TensorImpl(at::Storage storage) : storage_(std::move(storage)), storage_offset_(0) {
     AT_ASSERT(storage_);
     data_type_ = storage_.dtype();
+    strides_.reset(new int64_t[1]);
+    strides_[0] = 1;
   }
 
   TensorImpl(const TensorImpl&) = default;
@@ -317,15 +319,17 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
     // assumes that old values are preserved
     auto old_dim = sizes_.size();
     sizes_.resize(ndim);
-    auto new_strides = c10::guts::make_unique<int64_t[]>(ndim);
-    for (size_t i = 0; i < std::min(old_dim, static_cast<size_t>(ndim)); i++) {
-      new_strides[i] = strides_[i];
-    }
-    for (size_t i = old_dim; i < static_cast<size_t>(ndim); i++) {
-      // If ndim < old_dim, this loop never executes
-      new_strides[i] = 0;
+    if (old_dim != sizes_.size()) {
+      auto new_strides = c10::guts::make_unique<int64_t[]>(ndim);
+      for (size_t i = 0; i < std::min(old_dim, static_cast<size_t>(ndim)); i++) {
+        new_strides[i] = strides_[i];
+      }
+      for (size_t i = old_dim; i < static_cast<size_t>(ndim); i++) {
+        // If ndim < old_dim, this loop never executes
+        new_strides[i] = 0;
+      }
+      strides_ = std::move(new_strides);
     }
-    strides_ = std::move(new_strides);
     refresh_numel();
     refresh_contiguous();
   }
@@ -337,8 +341,6 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
   }
 
   virtual void set_stride(int64_t dim, int64_t new_stride) {
-    AT_ASSERTM(strides_, "Caffe2 tensors don't have meaningful strides and "
-                         "cannot be used in PyTorch");
     strides_[dim] = new_stride;
     refresh_numel();
     refresh_contiguous();
@@ -630,8 +632,9 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
         " The old caffe2 mixes Reshape and Resize but this behavior has "
         "been changed. If you find this error, most likely you will need "
         "to change corresponding code from Reshape to Resize.");
+    auto old_dim = sizes_.size();
     sizes_ = dims;
-    update_to_contiguous_strides();
+    update_to_contiguous_strides(old_dim);
   }
 
   /**
@@ -801,13 +804,14 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
       typename = typename std::enable_if<std::is_integral<T>::value>::type>
   bool SetDimsTemplate(at::ArrayRef<T> src) {
     auto old_numel = numel_;
+    auto old_dim = sizes_.size();
     sizes_.resize(src.size());
     int64_t new_numel = 1;
     for (size_t i = 0; i < src.size(); ++i) {
       new_numel *= src[i];
       sizes_[i] = src[i];
     }
-    update_to_contiguous_strides();
+    update_to_contiguous_strides(old_dim);
     numel_ = new_numel;
     return numel_ != old_numel;
   }
@@ -844,8 +848,17 @@ struct CAFFE2_API TensorImpl : public c10::intrusive_ptr_target {
     return SetDims(at::IntList{d0, d1, d2, d3});
   }
 
-  inline void update_to_contiguous_strides() {
-    strides_.reset();
+  inline void update_to_contiguous_strides(size_t old_dim) {
+    if (old_dim != sizes_.size()) {
+      strides_ = c10::guts::make_unique<int64_t[]>(sizes_.size());
+    }
+    if (dim() > 0) {
+      int last_idx = dim() - 1;
+      strides_[last_idx] = 1;
+      for (auto i = last_idx - 1; i >= 0; --i) {
+        strides_[i] = strides_[i + 1] * std::max<int64_t>(sizes_[i + 1], 1);
+      }
+    }
     is_contiguous_ = true;
   }
 

aten/src/ATen/core/TensorImpl_test.cpp

@@ -0,0 +1,7 @@
+#include "gtest/gtest.h"
+#include "caffe2/core/tensor.h"
+
+TEST(TensorImplTest, Caffe2Constructor) {
+  caffe2::Tensor tensor(caffe2::CPU);
+  ASSERT_EQ(tensor.strides()[0], 1);
+}