Add PyObject preservation for UntypedStorage (pytorch#97470)

Part of pytorch#91395

Pull Request resolved: pytorch#97470
Approved by: https://github.com/ezyang

Author: kurtamohler

build_variables.bzl

@@ -888,6 +888,7 @@ libtorch_python_core_sources = [
     "torch/csrc/utils/python_dispatch.cpp",
     "torch/csrc/utils/python_symnode.cpp",
     "torch/csrc/utils/pybind.cpp",
+    "torch/csrc/utils/pyobject_preservation.cpp",
     "torch/csrc/utils/structseq.cpp",
     "torch/csrc/utils/tensor_apply.cpp",
     "torch/csrc/utils/tensor_dtypes.cpp",

c10/core/RefcountedDeleter.cpp

@@ -0,0 +1,79 @@
+#include <c10/core/RefcountedDeleter.h>
+
+#include <mutex>
+
+namespace c10 {
+
+void refcounted_deleter(void* ctx_) {
+  RefcountedDeleterContext& ctx =
+      *reinterpret_cast<RefcountedDeleterContext*>(ctx_);
+  ctx.refcount--;
+  if (ctx.refcount == 0) {
+    ctx.other_ctx = nullptr;
+    delete &ctx;
+  }
+}
+
+std::mutex replace_data_ptr_mutex;
+
+void maybeApplyRefcountedDeleter(c10::Storage storage) {
+  std::lock_guard<std::mutex> guard(replace_data_ptr_mutex);
+  c10::DataPtr& data_ptr = storage.mutable_data_ptr();
+
+  if ((void*)data_ptr.get_deleter() == (void*)&c10::refcounted_deleter) {
+    // Data pointer is already shared
+    return;
+  }
+
+  void* data = data_ptr.get();
+  void* other_ctx = data_ptr.get_context();
+  c10::DeleterFnPtr other_deleter = data_ptr.get_deleter();
+  c10::Device device = data_ptr.device();
+
+  // Release the context of the original DataPtr so that the data doesn't
+  // get deleted when the original DataPtr is replaced
+  data_ptr.release_context();
+
+  c10::RefcountedDeleterContext* refcount_ctx =
+      new c10::RefcountedDeleterContext(other_ctx, other_deleter);
+
+  c10::DataPtr new_data_ptr(
+      data,
+      reinterpret_cast<void*>(refcount_ctx),
+      &c10::refcounted_deleter,
+      device);
+  storage.set_data_ptr(std::move(new_data_ptr));
+}
+
+c10::Storage newStorageImplFromRefcountedDataPtr(c10::Storage storage) {
+  c10::maybeApplyRefcountedDeleter(storage);
+
+  c10::StorageImpl* storage_impl = storage.unsafeGetStorageImpl();
+
+  c10::DataPtr& data_ptr = storage.mutable_data_ptr();
+  c10::DataPtr new_data_ptr(
+      data_ptr.get(),
+      data_ptr.get_context(),
+      data_ptr.get_deleter(),
+      data_ptr.device());
+
+  // NOTE: This refcount increment should always happen immediately after
+  // `new_data_ptr` is created. No other lines of code should be added between
+  // them in the future, unless there's a very good reason for it, because if
+  // any errors are raised and `new_data_ptr` is deleted before the refcount is
+  // incremented, the refcount will get decremented and end up being one less
+  // than it should be.
+  reinterpret_cast<c10::RefcountedDeleterContext*>(data_ptr.get_context())
+      ->refcount++;
+
+  c10::Allocator* allocator = c10::GetAllocator(storage_impl->device_type());
+  c10::Storage new_storage = c10::make_intrusive<c10::StorageImpl>(
+      c10::StorageImpl::use_byte_size_t(),
+      storage_impl->nbytes(),
+      allocator,
+      /*resizable=*/storage_impl->resizable());
+  new_storage.set_data_ptr(std::move(new_data_ptr));
+  return new_storage;
+}
+
+} // namespace c10

c10/core/RefcountedDeleter.h

@@ -0,0 +1,50 @@
+#pragma once
+
+#include <c10/core/Storage.h>
+#include <c10/util/UniqueVoidPtr.h>
+
+#include <atomic>
+#include <memory>
+
+namespace c10 {
+
+// A RefcountedDeleterContext object is used as the `ctx` argument for DataPtr
+// to implement a shared DataPtr. Normally, a DataPtr is unique, but we use
+// this custom context and the `refcounted_deleter` function below to make the
+// DataPtr act like a non-unique DataPtr. This context object holds onto an
+// inner context and deleter function which handle the actual deletion of the
+// data when the refcount reaches 0.
+//
+// This shared DataPtr feature is only used when storages are shared between
+// multiple Python interpreters in MultiPy. Before storages had PyObject
+// preservation, interpreters could just share the same StorageImpl instance.
+// But now a StorageImpl can only be associated with one interpreter in order
+// to properly manage a zombie PyObject. So we share storages across Python
+// interpreters by creating a different StorageImpl instance for each one, but
+// they all point to the same data.
+struct C10_API RefcountedDeleterContext {
+  RefcountedDeleterContext(void* other_ctx, c10::DeleterFnPtr other_deleter)
+      : other_ctx(other_ctx, other_deleter), refcount(1) {}
+
+  std::unique_ptr<void, c10::DeleterFnPtr> other_ctx;
+  std::atomic_int refcount;
+};
+
+// `refcounted_deleter` is used as the `ctx_deleter` for DataPtr to implement
+// a shared DataPtr.
+//
+// Warning: This should only be called on a pointer to
+// a RefcountedDeleterContext that was allocated on the heap with `new`,
+// because when the refcount reaches 0, the context is deleted with `delete`
+C10_API void refcounted_deleter(void* ctx_);
+
+// If the storage's DataPtr does not use `refcounted_deleter`, replace it with
+// a DataPtr that does, so it can be shared between multiple StorageImpls
+C10_API void maybeApplyRefcountedDeleter(c10::Storage storage);
+
+// Create a new StorageImpl that points to the same data. If the original
+// StorageImpl's DataPtr does not use `refcounted_deleter`, it will be replaced
+// with one that does
+C10_API c10::Storage newStorageImplFromRefcountedDataPtr(c10::Storage storage);
+
+} // namespace c10

c10/core/SafePyObject.h

@@ -29,7 +29,7 @@ struct C10_API SafePyObject {
   SafePyObject& operator=(SafePyObject const&) = delete;
 
   ~SafePyObject() {
-    (*pyinterpreter_)->decref(data_, /*is_tensor*/ false);
+    (*pyinterpreter_)->decref(data_, /*has_pyobj_slot*/ false);
   }
 
   c10::impl::PyInterpreter& pyinterpreter() const {

c10/core/Storage.cpp

@@ -1,3 +1,18 @@
+#include <c10/core/RefcountedDeleter.h>
 #include <c10/core/Storage.h>
 
-namespace c10 {} // namespace c10
+namespace c10 {
+
+bool isSharedStorageAlias(const Storage& storage0, const Storage& storage1) {
+  c10::DeleterFnPtr deleter_expected = &c10::refcounted_deleter;
+  c10::DeleterFnPtr deleter0 = storage0.data_ptr().get_deleter();
+  c10::DeleterFnPtr deleter1 = storage1.data_ptr().get_deleter();
+
+  if ((deleter0 != deleter_expected) || (deleter1 != deleter_expected)) {
+    return false;
+  }
+
+  return storage0.data_ptr().get_context() == storage1.data_ptr().get_context();
+}
+
+} // namespace c10

c10/core/Storage.h

@@ -1,12 +1,22 @@
 #pragma once
 
 #include <c10/core/StorageImpl.h>
+#include <c10/util/ExclusivelyOwned.h>
 
 namespace c10 {
 
+struct Storage;
+
+C10_API bool isSharedStorageAlias(
+    const Storage& storage0,
+    const Storage& storage1);
+
 struct C10_API Storage {
  public:
   struct use_byte_size_t {};
+  struct unsafe_borrow_t {
+    explicit unsafe_borrow_t() = default;
+  };
 
   Storage() = default;
   Storage(c10::intrusive_ptr<StorageImpl> ptr)
@@ -40,6 +50,14 @@ struct C10_API Storage {
             allocator,
             resizable)) {}
 
+ protected:
+  explicit Storage(unsafe_borrow_t, const Storage& rhs)
+      : storage_impl_(c10::intrusive_ptr<c10::StorageImpl>::reclaim(
+            rhs.storage_impl_.get())) {}
+
+  friend MaybeOwnedTraits<Storage>;
+
+ public:
   // Legacy constructor for partially initialized (dtype or memory) storages
   // that can be temporarily created with Caffe2 APIs. See the note on top of
   // TensorImpl.h for details.
@@ -144,7 +162,9 @@ struct C10_API Storage {
   }
 
   bool is_alias_of(const Storage& other) const {
-    return storage_impl_ == other.storage_impl_;
+    return (
+        storage_impl_ == other.storage_impl_ ||
+        isSharedStorageAlias(*this, other));
   }
 
   void UniqueStorageShareExternalPointer(
@@ -175,4 +195,67 @@ struct C10_API Storage {
   c10::intrusive_ptr<StorageImpl> storage_impl_;
 };
 
+template <>
+struct MaybeOwnedTraits<c10::Storage> {
+  using owned_type = c10::Storage;
+  using borrow_type = c10::Storage;
+
+  static borrow_type createBorrow(const owned_type& from) {
+    return borrow_type(borrow_type::unsafe_borrow_t{}, from);
+  }
+
+  static void assignBorrow(borrow_type& lhs, const borrow_type& rhs) {
+    lhs.unsafeReleaseStorageImpl();
+    lhs = borrow_type(borrow_type::unsafe_borrow_t{}, rhs);
+  }
+
+  static void destroyBorrow(borrow_type& toDestroy) {
+    toDestroy.unsafeReleaseStorageImpl(); // "leak" it, but it was already +0.
+  }
+
+  static const owned_type& referenceFromBorrow(const borrow_type& borrow) {
+    return borrow;
+  }
+
+  static const owned_type* pointerFromBorrow(const borrow_type& borrow) {
+    return &borrow;
+  }
+
+  static bool debugBorrowIsValid(const borrow_type& /*borrow*/) {
+    return true;
+  }
+};
+
+template <>
+struct ExclusivelyOwnedTraits<c10::Storage> {
+  using repr_type = c10::Storage;
+  using pointer_type = c10::Storage*;
+  using const_pointer_type = const c10::Storage*;
+
+  static repr_type nullRepr() {
+    return c10::Storage();
+  }
+
+  template <class... Args>
+  static repr_type createInPlace(Args&&... args) {
+    return c10::Storage(std::forward<Args>(args)...);
+  }
+
+  static repr_type moveToRepr(c10::Storage&& x) {
+    return std::move(x);
+  }
+
+  static c10::Storage take(c10::Storage& x) {
+    return std::move(x);
+  }
+
+  static pointer_type getImpl(repr_type& x) {
+    return &x;
+  }
+
+  static const_pointer_type getImpl(const repr_type& x) {
+    return &x;
+  }
+};
+
 } // namespace c10

c10/core/StorageImpl.h

@@ -205,6 +205,14 @@ struct C10_API StorageImpl : public c10::intrusive_ptr_target {
     return received_cuda_;
   }
 
+  impl::PyObjectSlot* pyobj_slot() {
+    return &pyobj_slot_;
+  }
+
+  const impl::PyObjectSlot* pyobj_slot() const {
+    return &pyobj_slot_;
+  }
+
  private:
   DataPtr data_ptr_;
   SymInt size_bytes_;

c10/core/TensorImpl.cpp

@@ -73,9 +73,7 @@ void TensorImpl::_set_fw_grad(
   autograd_meta_->set_fw_grad(new_grad, self, level, is_inplace_op);
 }
 
-TensorImpl::~TensorImpl() {
-  pyobj_slot_.destroy_pyobj_if_needed();
-}
+TensorImpl::~TensorImpl() = default;
 
 TensorImpl::TensorImpl(
     Storage&& storage,
@@ -582,7 +580,7 @@ void TensorImpl::release_resources() {
   if (storage_) {
     storage_ = {};
   }
-  pyobj_slot_.destroy_pyobj_if_needed();
+  pyobj_slot_.maybe_destroy_pyobj();
 }
 
 #ifndef C10_DISABLE_TENSORIMPL_EXTENSIBILITY

c10/core/impl/PyInterpreter.cpp

@@ -10,7 +10,8 @@ struct NoopPyInterpreterVTable final : public PyInterpreterVTable {
     return "<unloaded interpreter>";
   }
 
-  void decref(PyObject* pyobj, bool is_tensor) const override {} // do nothing
+  void decref(PyObject* pyobj, bool has_pyobj_slot) const override {
+  } // do nothing
 
 #define PANIC(m)              \
   TORCH_INTERNAL_ASSERT(      \

c10/core/impl/PyInterpreter.h

@@ -127,8 +127,8 @@ struct C10_API PyInterpreterVTable {
   virtual std::string name() const = 0;
 
   // Run Py_DECREF on a PyObject.  We DO NOT assume the GIL is held on call
-  // See NOTE [PyInterpreter::decref takes an `is_tensor` arg]
-  virtual void decref(PyObject* pyobj, bool is_tensor) const = 0;
+  // See NOTE [PyInterpreter::decref takes a `has_pyobj_slot` arg]
+  virtual void decref(PyObject* pyobj, bool has_pyobj_slot) const = 0;
 
   // Perform a detach by deferring to the __torch_dispatch__ implementation of
   // detach, which will also arrange for the PyObject to get copied in this