Add weak pointer and finalizer support directly to THStorage. (#9148)

Summary:
The underlying use-case is the file descriptor to storage cache in
torch.multiprocessing.reductions.  Previously, this was implemented by wrapping
an existing allocator with a "weak ref" allocator which also knew to null out
the weak reference when the storage died.  This is terribly oblique, and
prevents us from refactoring the allocators to get rid of per-storage allocator
state.

So instead of going through this fiasco, we instead directly implement weak
pointers and finalizers in THStorage.  Weak pointers to THStorage retain the
THStorage struct, but not the data_ptr.  When all strong references die,
data_ptr dies and the finalizers get invoked.

There is one major hazard in this patch, which is what happens if you
repeatedly call _weak_ref on a storage.  For cleanliness, we no longer
shove our grubby fingers into the finalizer struct to see if there is already
a Python object for the weak reference and return it; we just create a new one
(no one is checking these Python objects for identity).  This means if you
keep calling it, we'll keep piling on finalizers.  That's bad! But I am
not going to fix it until it is actually a problem for someone, because
then we need to add another caching layer.

Signed-off-by: Edward Z. Yang <[email protected]>
Pull Request resolved: pytorch/pytorch#9148

Differential Revision: D8729106

Pulled By: ezyang

fbshipit-source-id: 69710ca3b7c7e05069090e1b263f8b6b9f1cf72f

Author: ezyang

aten/src/TH/THStorage.cpp

@@ -14,27 +14,56 @@
 #include "generic/THStorageCopy.cpp"
 #include "THGenerateHalfType.h"
 
+// Free a non-weak pointer to THStorage
 void THStorage_free(THStorage *storage) {
   AT_ASSERT(storage->backend == at::kCPU);
 
   if (!storage) {
     return;
   }
 
-  if ((storage->flag & TH_STORAGE_REFCOUNTED) && (storage->refcount.load() > 0)) {
+  if (storage->flag & TH_STORAGE_REFCOUNTED) {
     if (--storage->refcount == 0) {
+      if (storage->finalizer) {
+        (*storage->finalizer)();
+      }
+      storage->finalizer.~unique_ptr<THFinalizer>();
       if (storage->flag & TH_STORAGE_FREEMEM) {
         static_cast<THAllocator*>(storage->allocatorVoidPtr)->free(storage->allocatorContext, storage->data_ptr);
       }
       if (storage->flag & TH_STORAGE_VIEW) {
         THStorage_free(storage->view);
       }
-      storage->refcount.~atomic<int>();
-      THFree(storage);
+      THStorage_weakFree(storage);
     }
   }
 }
 
+// Manually retains a weak reference
+void THStorage_weakRetain(THStorage *weak_storage) {
+  weak_storage->weakcount++;
+}
+
+// Releases a weak reference
+void THStorage_weakFree(THStorage *weak_storage) {
+  if (--weak_storage->weakcount == 0) {
+    weak_storage->refcount.~atomic<int>();
+    weak_storage->weakcount.~atomic<int>();
+    THFree(weak_storage);
+  }
+}
+
+// Given a weak reference, returns a strong reference to a storage (which must
+// be freed when done) or null if the storage is already dead.
+THStorage* THStorage_weakLock(THStorage *weak_storage) {
+  for (;;) {
+    int refcount = weak_storage->refcount.load();
+    if (refcount == 0) return nullptr;
+    if (weak_storage->refcount.compare_exchange_strong(refcount, refcount + 1)) break;
+  }
+  return weak_storage;
+}
+
 THDescBuff THLongStorage_sizeDesc(const THLongStorage *size) {
   return _THSizeDesc(THLongStorage_data(size), size->size);
 }
@@ -89,6 +118,8 @@ THStorage* THStorage_newWithAllocator(at::ScalarType scalar_type, ptrdiff_t size
   storage->data_ptr = allocator->malloc(allocatorContext, at::elementSize(scalar_type)*size);
   storage->size = size;
   new (&storage->refcount) std::atomic<int>(1);
+  new (&storage->weakcount) std::atomic<int>(1); // from the strong reference
+  new (&storage->finalizer) std::unique_ptr<THFinalizer>(nullptr);
   storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE | TH_STORAGE_FREEMEM;
   storage->allocatorVoidPtr = allocator;
   storage->allocatorContext = allocatorContext;
@@ -140,19 +171,6 @@ void THStorage_retain(THStorage *storage)
   }
 }
 
-int THStorage_retainIfLive(THStorage *storage)
-{
-  // TODO: Check if TH_STORAGE_REFCOUNTED?
-  int refcount = storage->refcount.load();
-  while (refcount > 0) {
-    if (storage->refcount.compare_exchange_strong(refcount, refcount + 1)) {
-      return 1;
-    }
-    refcount = storage->refcount.load();
-  }
-  return 0;
-}
-
 THStorage* THStorage_newWithData(at::ScalarType scalar_type, void *data, ptrdiff_t size)
 {
   return THStorage_newWithDataAndAllocator(scalar_type, data, size,
@@ -168,7 +186,9 @@ THStorage* THStorage_newWithDataAndAllocator(at::ScalarType scalar_type,
   storage->scalar_type = scalar_type;
   storage->data_ptr = data;
   storage->size = size;
-  storage->refcount = 1;
+  new (&storage->refcount) std::atomic<int>(1);
+  new (&storage->weakcount) std::atomic<int>(1); // from the strong reference
+  new (&storage->finalizer) std::unique_ptr<THFinalizer>(nullptr);
   storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE | TH_STORAGE_FREEMEM;
   storage->allocatorVoidPtr = allocator;
   storage->allocatorContext = allocatorContext;
@@ -227,6 +247,7 @@ void THStorage_swap(THStorage *storage1, THStorage *storage2)
     // don't swap refcount!
     SWAP(allocatorVoidPtr);
     SWAP(allocatorContext);
+    SWAP(finalizer);
     SWAP(view);
     SWAP(device);
 #undef SWAP

aten/src/TH/THStorage.h

@@ -19,6 +19,7 @@
 
 // This exists to have a data-type independent way of freeing (necessary for THPPointer).
 TH_API void THStorage_free(THStorage *storage);
+TH_API void THStorage_weakFree(THStorage *storage);
 
 TH_API THDescBuff THLongStorage_sizeDesc(const THLongStorage *size);
 TH_API THLongStorage *THLongStorage_newInferSize(THLongStorage *size, ptrdiff_t nElement);

aten/src/TH/THStorage.hpp

@@ -10,16 +10,45 @@
 #include "THTypeConversion.hpp"
 #include <atomic>
 
+// Note [Weak references for intrusive refcounting]
+// ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+// Here's the scheme:
+//
+//  - refcount == number of strong references to the object
+//    weakcount == number of weak references to the object,
+//      plus one more if refcount > 0
+//
+//  - THStorage stays live as long as there are any strong
+//    or weak pointers to it (weakcount > 0, since strong
+//    references count as a +1 to weakcount)
+//
+//  - finalizers are called and data_ptr is deallocated when refcount == 0
+//
+//  - Once refcount == 0, it can never again be > 0 (the transition
+//    from > 0 to == 0 is monotonic)
+//
+//  - When you access THStorage via a weak pointer, you must
+//    atomically increment the use count, if it is greater than 0.
+//    If it is not, you must report that the storage is dead.
+//
+
+struct THFinalizer {
+  virtual void operator()() = 0;
+  virtual ~THFinalizer() {};
+};
+
 typedef struct THStorage
 {
     at::Backend backend; // kCPU or kCUDA only
     at::ScalarType scalar_type;
     void *data_ptr;
     ptrdiff_t size;
     std::atomic<int> refcount;
+    std::atomic<int> weakcount;
     char flag;
     void *allocatorVoidPtr; // Either THDeviceAllocator or THCDeviceAllocator
     void *allocatorContext;
+    std::unique_ptr<THFinalizer> finalizer;
     struct THStorage *view;
     int device;
 
@@ -51,11 +80,14 @@ THStorage* THStorage_newWithMapping(at::ScalarType scalar_type, const char *file
 void THStorage_setFlag(THStorage *storage, const char flag);
 void THStorage_clearFlag(THStorage *storage, const char flag);
 void THStorage_retain(THStorage *storage);
-int THStorage_retainIfLive(THStorage *storage);
 THStorage* THStorage_newWithData(at::ScalarType scalar_type, void *data, ptrdiff_t size);
 THStorage* THStorage_newWithDataAndAllocator(at::ScalarType scalar_type,
                                              void* data, ptrdiff_t size,
                                              THAllocator* allocator,
                                              void* allocatorContext);
 void THStorage_resize(THStorage *storage, ptrdiff_t size);
 void THStorage_swap(THStorage *storage1, THStorage *storage2);
+
+void THStorage_weakRetain(THStorage *weak_storage);
+void THStorage_weakFree(THStorage *weak_storage);
+THStorage* THStorage_weakLock(THStorage *weak_storage);

aten/src/TH/generic/THStorage.cpp

@@ -95,11 +95,6 @@ void THStorage_(retain)(THStorage *storage)
   THStorage_retain(storage);
 }
 
-int THStorage_(retainIfLive)(THStorage *storage)
-{
-  return THStorage_retainIfLive(storage);
-}
-
 void THStorage_(free)(THStorage *storage)
 {
   THStorage_free(storage);

aten/src/TH/generic/THStorage.h

@@ -66,9 +66,6 @@ TH_API void THStorage_(clearFlag)(THStorage *storage, const char flag);
 TH_API void THStorage_(retain)(THStorage *storage);
 TH_API void THStorage_(swap)(THStorage *storage1, THStorage *storage2);
 
-/* used by StorageSharing */
-TH_API int THStorage_(retainIfLive)(THStorage *storage);
-
 /* might differ with other API (like CUDA) */
 TH_API void THStorage_(free)(THStorage *storage);
 TH_API void THStorage_(resize)(THStorage *storage, ptrdiff_t size);

aten/src/THC/THCStorage.cpp

@@ -34,6 +34,8 @@ THCStorage* THCStorage_newWithAllocator(THCState *state,
   THCStorage *storage = (THCStorage*)THAlloc(sizeof(THCStorage));
   memset(storage, 0, sizeof(THCStorage));
   new (&storage->refcount) std::atomic<int>(1);
+  new (&storage->weakcount) std::atomic<int>(1);
+  new (&storage->finalizer) std::unique_ptr<THFinalizer>(nullptr);
   storage->backend = at::kCUDA;
   storage->scalar_type = scalar_type;
   storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE | TH_STORAGE_FREEMEM;
@@ -60,25 +62,25 @@ THCStorage* THCStorage_newWithAllocator(THCState *state,
   return storage;
 }
 
-void THCStorage_free(THCState *state, THCStorage *self)
+void THCStorage_free(THCState *state, THCStorage *storage)
 {
-  AT_ASSERT(self->backend == at::kCUDA);
-
-  if(!(self->flag & TH_STORAGE_REFCOUNTED))
-    return;
+  AT_ASSERT(storage->backend == at::kCUDA);
 
-  if (--self->refcount == 0)
-  {
-    if(self->flag & TH_STORAGE_FREEMEM) {
-      auto* thc_device_allocator = static_cast<THCDeviceAllocator*>(self->allocatorVoidPtr);
-      THCudaCheck(
-        (*thc_device_allocator->free)(self->allocatorContext, self->data_ptr));
-    }
-    if(self->flag & TH_STORAGE_VIEW) {
-      THCStorage_free(state, self->view);
+  if (storage->flag & TH_STORAGE_REFCOUNTED) {
+    if (--storage->refcount == 0) {
+      if (storage->finalizer) {
+        (*storage->finalizer)();
+      }
+      storage->finalizer.~unique_ptr<THFinalizer>();
+      if (storage->flag & TH_STORAGE_FREEMEM) {
+        auto* thc_device_allocator = static_cast<THCDeviceAllocator*>(storage->allocatorVoidPtr);
+        THCudaCheck((*thc_device_allocator->free)(storage->allocatorContext, storage->data_ptr));
+      }
+      if (storage->flag & TH_STORAGE_VIEW) {
+        THCStorage_free(state, storage->view);
+      }
+      THStorage_weakFree(storage);
     }
-    self->refcount.~atomic<int>();
-    THFree(self);
   }
 }
 
@@ -174,7 +176,9 @@ THCStorage* THCStorage_newWithDataAndAllocator(
   storage->scalar_type = scalar_type;
   storage->data_ptr = data;
   storage->size = size;
-  storage->refcount = 1;
+  new (&storage->refcount) std::atomic<int>(1);
+  new (&storage->weakcount) std::atomic<int>(1);
+  new (&storage->finalizer) std::unique_ptr<THFinalizer>(nullptr);
   storage->flag = TH_STORAGE_REFCOUNTED | TH_STORAGE_RESIZABLE | TH_STORAGE_FREEMEM;
   storage->allocatorVoidPtr = allocator;
   storage->allocatorContext = allocatorContext;

aten/src/THC/generic/THCStorage.cpp

@@ -122,11 +122,6 @@ void THCStorage_(retain)(THCState *state, THCStorage *self)
   THStorage_retain(self);
 }
 
-int THCStorage_(retainIfLive)(THCState *state, THCStorage *storage)
-{
-  return THStorage_retainIfLive(storage);
-}
-
 void THCStorage_(free)(THCState *state, THCStorage *self)
 {
   THCStorage_free(state, self);

aten/src/THC/generic/THCStorage.h

@@ -53,9 +53,6 @@ THC_API void THCStorage_(setFlag)(THCState *state, THCStorage *storage, const ch
 THC_API void THCStorage_(clearFlag)(THCState *state, THCStorage *storage, const char flag);
 THC_API void THCStorage_(retain)(THCState *state, THCStorage *storage);
 
-/* used by StorageSharing */
-THC_API int THCStorage_(retainIfLive)(THCState *state, THCStorage *storage);
-
 THC_API void THCStorage_(free)(THCState *state, THCStorage *storage);
 THC_API void THCStorage_(resize)(THCState *state, THCStorage *storage, ptrdiff_t size);
 THC_API void THCStorage_(fill)(THCState *state, THCStorage *storage, real value);