PyTorch/Caffe2 tensor interop in Python (#17190)

Summary:
Because of two separate python extensions with different pybind
instances I have to go through void* conversion. Since it's hidden from
user, it's fine.

New APIs added on C2 side:
- workspace.FetchTorch('blob')
- workspace.Workspace.current.blobs['blob'].to_torch()
- workspace.FeedBlob('blob', pytorch_tensor)

Works on CPU an GPU.

The only glitches are with resizing because of variable/tensor split.
But data sharing works properly.
Pull Request resolved: https://github.com/pytorch/pytorch/pull/17190

Reviewed By: ezyang

Differential Revision: D14163882

Pulled By: dzhulgakov

fbshipit-source-id: d18e5b8fcae026f393c842a1149e972515732de2

2 files changed, 60 insertions, 11 deletions

diff --git a/c10/test/util/intrusive_ptr_test.cpp b/c10/test/util/intrusive_ptr_test.cpp
index 3532cf81c5..c548cfdb76 100644
--- a/c10/test/util/intrusive_ptr_test.cpp
+++ b/c10/test/util/intrusive_ptr_test.cpp
@@ -1539,6 +1539,7 @@ TEST(IntrusivePtrTest, givenCopyAssignedPtr_whenReassigningCopy_thenIsUnique) {
 TEST(IntrusivePtrTest, givenPtr_whenReleasedAndReclaimed_thenDoesntCrash) {
   intrusive_ptr<SomeClass> obj = make_intrusive<SomeClass>();
   SomeClass* ptr = obj.release();
+  EXPECT_FALSE(obj.defined());
   intrusive_ptr<SomeClass> reclaimed = intrusive_ptr<SomeClass>::reclaim(ptr);
 }
 
@@ -1574,6 +1575,39 @@ TEST(IntrusivePtrTest, givenStackObject_whenReclaimed_thenCrashes) {
   EXPECT_ANY_THROW(ptr = intrusive_ptr<SomeClass>::reclaim(&obj));
 }
 
+TEST(IntrusivePtrTest, givenPtr_whenNonOwningReclaimed_thenDoesntCrash) {
+  intrusive_ptr<SomeClass> obj = make_intrusive<SomeClass>();
+  SomeClass* raw_ptr = obj.get();
+  EXPECT_TRUE(obj.defined());
+  intrusive_ptr<SomeClass> reclaimed =
+      intrusive_ptr<SomeClass>::unsafe_reclaim_from_nonowning(raw_ptr);
+  EXPECT_TRUE(reclaimed.defined());
+  EXPECT_EQ(reclaimed.get(), obj.get());
+}
+
+TEST(
+    IntrusivePtrTest,
+    givenPtr_whenNonOwningReclaimed_thenIsDestructedAtEnd) {
+  bool resourcesReleased = false;
+  bool wasDestructed = false;
+  {
+    intrusive_ptr<DestructableMock> outer;
+    {
+      intrusive_ptr<DestructableMock> inner =
+          make_intrusive<DestructableMock>(&resourcesReleased, &wasDestructed);
+      DestructableMock* raw_ptr = inner.get();
+      outer = intrusive_ptr<DestructableMock>::unsafe_reclaim_from_nonowning(
+          raw_ptr);
+    }
+    // inner is destructed
+    EXPECT_FALSE(resourcesReleased);
+    EXPECT_FALSE(wasDestructed);
+  }
+  // outer is destructed
+  EXPECT_TRUE(resourcesReleased);
+  EXPECT_TRUE(wasDestructed);
+}
+
 namespace {
 template <class T>
 struct IntrusiveAndWeak final {
diff --git a/c10/util/intrusive_ptr.h b/c10/util/intrusive_ptr.h
index 9f5e686927..b71cd7294a 100644
--- a/c10/util/intrusive_ptr.h
+++ b/c10/util/intrusive_ptr.h
@@ -73,14 +73,14 @@ class C10_API intrusive_ptr_target {
 // some other compilers don't know about -Wterminate or -Wexceptions and
 // will show a warning about unknown warning options otherwise.
 #ifdef _MSC_VER
-#  pragma warning(push)  
-#  pragma warning(disable: 4297) // function assumed not to throw an exception but does  
-#else  
-#  pragma GCC diagnostic push  
-#  pragma GCC diagnostic ignored "-Wpragmas"  
-#  pragma GCC diagnostic ignored "-Wunknown-warning-option"  
-#  pragma GCC diagnostic ignored "-Wterminate"  
-#  pragma GCC diagnostic ignored "-Wexceptions"  
+#  pragma warning(push)
+#  pragma warning(disable: 4297) // function assumed not to throw an exception but does
+#else
+#  pragma GCC diagnostic push
+#  pragma GCC diagnostic ignored "-Wpragmas"
+#  pragma GCC diagnostic ignored "-Wunknown-warning-option"
+#  pragma GCC diagnostic ignored "-Wterminate"
+#  pragma GCC diagnostic ignored "-Wexceptions"
 #endif
     AT_ASSERTM(
         refcount_.load() == 0,
@@ -89,9 +89,9 @@ class C10_API intrusive_ptr_target {
         weakcount_.load() == 0,
         "Tried to destruct an intrusive_ptr_target that still has weak_intrusive_ptr to it");
 #ifdef _MSC_VER
-#  pragma warning(pop)  
-#else  
-#  pragma GCC diagnostic pop  
+#  pragma warning(pop)
+#else
+#  pragma GCC diagnostic pop
 #endif
   }
 
@@ -362,6 +362,21 @@ class intrusive_ptr final {
 
     return result;
   }
+
+  /**
+   * Turn a **non-owning raw pointer** to an intrusive_ptr.
+   *
+   * This method is potentially dangerous (as it can mess up refcount).
+   */
+  static intrusive_ptr unsafe_reclaim_from_nonowning(TTarget* raw_ptr) {
+    // See Note [Stack allocated intrusive_ptr_target safety]
+    AT_ASSERTM(
+        raw_ptr == NullType::singleton() || raw_ptr->refcount_.load() > 0,
+        "intrusive_ptr: Can only reclaim pointers that are owned by someone");
+    auto ptr = reclaim(raw_ptr); // doesn't increase refcount
+    ptr.retain_();
+    return ptr;
+  }
 };
 
 template <