[bugfix] Explicitly unpin tensoradapter allocated arrays (#3997)

* Explicitly unpin tensoradapter allocated arrays

* Undo unrelated change

* Add unit test

* update unit test

include/dgl/runtime/ndarray.h

@@ -415,6 +415,8 @@ struct NDArray::Container {
   std::vector<int64_t> stride_;
   /*! \brief The internal array object */
   std::atomic<int> ref_counter_{0};
+
+  bool from_tensor_dispatcher_{false};
 };
 
 // implementations of inline functions

src/runtime/ndarray.cc

@@ -78,6 +78,12 @@ struct NDArray::Internal {
   // This enables us to create NDArray from memory allocated by other
   // frameworks that are DLPack compatible
   static void DLPackDeleter(NDArray::Container* ptr) {
+    if (ptr->from_tensor_dispatcher_) {
+       if (IsDataPinned(&(ptr->dl_tensor))) {
+        UnpinData(&(ptr->dl_tensor));
+      }
+    }
+
     DLManagedTensor* tensor = static_cast<DLManagedTensor*>(ptr->manager_ctx);
     if (tensor->deleter != nullptr) {
       (*tensor->deleter)(tensor);
@@ -210,8 +216,11 @@ NDArray NDArray::Empty(std::vector<int64_t> shape,
                        DLDataType dtype,
                        DLContext ctx) {
   TensorDispatcher* td = TensorDispatcher::Global();
-  if (td->IsAvailable())
-    return td->Empty(shape, dtype, ctx);
+  if (td->IsAvailable()) {
+    auto nd = td->Empty(shape, dtype, ctx);
+    nd.data_->from_tensor_dispatcher_ = true;
+    return nd;
+  }
 
   NDArray ret = Internal::Create(shape, dtype, ctx);
   // setup memory content

tests/pytorch/test_unpin_tensoradapter.py

+import backend as F
+import dgl
+import gc
+import unittest
+import torch
+
+
+@unittest.skipIf(F.ctx().type == 'cpu', reason='Pinning memory tests require GPU.')
+def test_unpin_tensoradapater():
+    # run a sufficient number of iterations such that the memory pool should be
+    # re-used
+    for j in range(3):
+        t = F.zerocopy_from_dlpack(dgl.ndarray.empty(
+            [10000, 10],
+            F.reverse_data_type_dict[F.float32],
+            ctx=dgl.utils.to_dgl_context(torch.device('cpu'))).to_dlpack()).zero_()
+        assert not F.is_pinned(t)
+        dgl.utils.pin_memory_inplace(t)
+        assert F.is_pinned(t)
+        del t