Fix and add testing for nullptr allocator in c2->pt conversion (#16857)

Summary:
Fixes the bug for when tensor is created on Caffe2 side, then passed to PT and resized. Now we just initialize allocator correctly.

Note that the code in raw_mutable_data() is still necessary because of non-resizable tensors.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/16857

Reviewed By: houseroad

Differential Revision: D14019469

Pulled By: dzhulgakov

fbshipit-source-id: 14d3a3b946d718bbab747ea376903646b885706a

diff --git a/aten/src/ATen/test/tensor_interop_test.cpp b/aten/src/ATen/test/tensor_interop_test.cpp
index 6016142..495fe0f 100644 (file)
--- a/aten/src/ATen/test/tensor_interop_test.cpp
+++ b/aten/src/ATen/test/tensor_interop_test.cpp
@@ -35,6 +35,25 @@ TEST(Caffe2ToPytorch, Simple) {
   }
 }
 
+TEST(Caffe2ToPytorch, ExternalData) {
+  caffe2::Tensor c2_tensor = caffe2::empty({4, 4}, at::kLong);
+  int64_t buf[16];
+  for (int64_t i = 0; i < 16; i++) {
+    buf[i] = i;
+  }
+  c2_tensor.ShareExternalPointer(buf, 16);
+
+  // If the buffer is allocated externally, we can still pass tensor around,
+  // but we can't resize its storage using PT APIs
+  at::Tensor at_tensor(c2_tensor);
+  auto it = at_tensor.data<int64_t>();
+  for (int64_t i = 0; i < 16; i++) {
+    ASSERT_EQ(it[i], i);
+  }
+  ASSERT_FALSE(at_tensor.storage().resizable());
+  ASSERT_ANY_THROW(at_tensor.resize_({7,7}));
+}
+
 TEST(Caffe2ToPytorch, Op) {
   caffe2::Tensor c2_tensor(caffe2::CPU);
   c2_tensor.Resize(3, 3);
@@ -80,6 +99,40 @@ TEST(Caffe2ToPytorch, PartiallyInitialized) {
   }
 }
 
+TEST(Caffe2ToPytorch, MutualResizes) {
+  caffe2::Tensor c2_tensor = caffe2::empty({5, 5}, at::kFloat);
+  auto data = c2_tensor.mutable_data<float>();
+  for (int64_t i = 0; i < 25; i++) {
+    data[i] = 0;
+  }
+
+  at::Tensor at_tensor(c2_tensor);
+
+  // change is visible
+  at_tensor[0][0] = 123;
+  ASSERT_EQ(c2_tensor.mutable_data<float>()[0], 123);
+
+  // resize PT tensor in smaller direction - storage is preserved
+  at_tensor.resize_({4, 4});
+  c2_tensor.mutable_data<float>()[1] = 234;
+  ASSERT_EQ(at_tensor[0][1].item().to<float>(), 234);
+
+  // resize PT tensor in larger direction - storage is preserved
+  at_tensor.resize_({6, 6});
+  c2_tensor.mutable_data<float>()[2] = 345;
+  ASSERT_EQ(at_tensor[0][2].item().to<float>(), 345);
+  ASSERT_EQ(c2_tensor.sizes()[0], 6);
+  ASSERT_EQ(c2_tensor.sizes()[1], 6);
+
+  // resize Caffe2 tensor - semantics are to NOT preserve the data, but the
+  // TensorImpl is still shared
+  c2_tensor.Resize(7, 7);
+  c2_tensor.mutable_data<float>()[3] = 456;
+  ASSERT_EQ(at_tensor[0][3].item().to<float>(), 456);
+  ASSERT_EQ(at_tensor.sizes()[0], 7);
+  ASSERT_EQ(at_tensor.sizes()[1], 7);
+}
+
 TEST(PytorchToCaffe2, Op) {
   caffe2::Workspace workspace;
   caffe2::NetDef net;

diff --git a/aten/src/TH/generic/THStorage.cpp b/aten/src/TH/generic/THStorage.cpp
index 1ec5062..db45be9 100644 (file)
--- a/aten/src/TH/generic/THStorage.cpp
+++ b/aten/src/TH/generic/THStorage.cpp
@@ -120,7 +120,7 @@ THStorage* THStorage_(newWithDataAndAllocator)(at::DataPtr&& data, ptrdiff_t siz
       size,
       std::move(data),
       allocator,
-      true).release();
+      allocator != nullptr).release();
   return storage;
 }
 

diff --git a/aten/src/THC/generic/THCStorage.cpp b/aten/src/THC/generic/THCStorage.cpp
index 99edeaa..b5495e1 100644 (file)
--- a/aten/src/THC/generic/THCStorage.cpp
+++ b/aten/src/THC/generic/THCStorage.cpp
@@ -122,7 +122,7 @@ THCStorage* THCStorage_(newWithDataAndAllocator)(
       size,
       std::move(data),
       allocator,
-      true).release();
+      allocator != nullptr).release();
   return storage;
 }
 

diff --git a/c10/core/Allocator.h b/c10/core/Allocator.h
index af4dd92..398f47e 100644 (file)
--- a/c10/core/Allocator.h
+++ b/c10/core/Allocator.h
@@ -183,6 +183,7 @@ struct C10_API InefficientStdFunctionContext {
 
 } // namespace c10
 
+// TODO: move to c10
 namespace caffe2 {
 
 /** Set the allocator for DeviceType `t`. The passed in allocator pointer is

diff --git a/c10/core/StorageImpl.h b/c10/core/StorageImpl.h
index 81523b7..67bac37 100644 (file)
--- a/c10/core/StorageImpl.h
+++ b/c10/core/StorageImpl.h
@@ -21,6 +21,10 @@ struct C10_API StorageImpl final : public c10::intrusive_ptr_target {
         numel_(numel),
         resizable_(resizable),
         allocator_(allocator) {
+    if (resizable) {
+      AT_ASSERTM(
+          allocator_, "For resizable storage, allocator must be provided");
+    }
     if (numel > 0) {
       if (data_type_.id() == caffe2::TypeIdentifier::uninitialized()) {
         AT_ERROR(
@@ -45,8 +49,12 @@ struct C10_API StorageImpl final : public c10::intrusive_ptr_target {
       : StorageImpl(device, caffe2::TypeMeta()) {}
 
   StorageImpl(at::Device device, caffe2::TypeMeta data_type)
-      : StorageImpl(data_type, 0, at::DataPtr(nullptr, device), nullptr, true) {
-  }
+      : StorageImpl(
+            data_type,
+            0,
+            at::DataPtr(nullptr, device),
+            caffe2::GetAllocator(device.type()),
+            true) {}
 
   StorageImpl& operator=(StorageImpl&& other) = default;
   StorageImpl& operator=(const StorageImpl&) = delete;
@@ -171,6 +179,10 @@ struct C10_API StorageImpl final : public c10::intrusive_ptr_target {
   }
 
   void set_resizable(bool resizable) {
+    if (resizable) {
+      // We need an allocator to be resizable
+      AT_ASSERT(allocator_);
+    }
     resizable_ = resizable;
   }
 
@@ -208,6 +220,8 @@ struct C10_API StorageImpl final : public c10::intrusive_ptr_target {
     // capacity() might not return the value that was set here, if itemsize does
     // not evenly divide it.
     numel_ = capacity / data_type_.itemsize();
+    allocator_ = nullptr;
+    resizable_ = false;
   }
 
  private:

diff --git a/c10/core/TensorImpl.h b/c10/core/TensorImpl.h
index f97798d..e0d88df 100644 (file)
--- a/c10/core/TensorImpl.h
+++ b/c10/core/TensorImpl.h
@@ -1135,7 +1135,12 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
     } else {
       int64_t numel = capacity / data_type.itemsize();
       // Create a new Storage
-      storage_ = Storage(data_type, numel, std::move(data_ptr), nullptr, true);
+      storage_ = Storage(
+          data_type,
+          numel,
+          std::move(data_ptr),
+          /*allocator=*/nullptr,
+          /*resizable=*/false);
       data_type_ = data_type;
       storage_offset_ = 0;
     }
@@ -1178,7 +1183,10 @@ struct C10_API TensorImpl : public c10::intrusive_ptr_target {
         return storage_.data();
       }
       const Allocator* allocator = storage_.allocator();
-      // TODO: Get rid of StaticContext
+      // Storage might have nullptr allocator in rare cases, for example, if
+      // an external memory segment has been wrapped with Tensor and we don't
+      // know how to reallocate it. However, in order to preserve legacy C2
+      // behavior, we allow reallocating the memory using default allocator.
       if (allocator == nullptr) {
         allocator = caffe2::GetAllocator(storage_.device_type());
       }