[Bugfix] Fix the uninitialized pin_memory flag for empty graph (#5609)

include/dgl/aten/coo.h

@@ -64,9 +64,11 @@ struct COOMatrix {
         data(darr),
         row_sorted(rsorted),
         col_sorted(csorted) {
+    if (!IsEmpty()) {
       is_pinned = (aten::IsNullArray(row) || row.IsPinned()) &&
                   (aten::IsNullArray(col) || col.IsPinned()) &&
                   (aten::IsNullArray(data) || data.IsPinned());
+    }
     CheckValidity();
   }
 
@@ -127,6 +129,11 @@ struct COOMatrix {
     CHECK_NO_OVERFLOW(row->dtype, num_cols);
   }
 
+  inline bool IsEmpty() const {
+    return aten::IsNullArray(row) && aten::IsNullArray(col) &&
+           aten::IsNullArray(data);
+  }
+
   /** @brief Return a copy of this matrix on the give device context. */
   inline COOMatrix CopyTo(const DGLContext& ctx) const {
     if (ctx == row->ctx) return *this;
@@ -138,6 +145,7 @@ struct COOMatrix {
 
   /** @brief Return a copy of this matrix in pinned (page-locked) memory. */
   inline COOMatrix PinMemory() {
+    if (!IsEmpty()) {
       if (is_pinned) return *this;
       auto new_coo = COOMatrix(
           num_rows, num_cols, row.PinMemory(), col.PinMemory(),
@@ -145,10 +153,14 @@ struct COOMatrix {
           col_sorted);
       CHECK(new_coo.is_pinned)
           << "An internal DGL error has occured while trying to pin a COO "
-           "matrix. Please file a bug at 'https://github.com/dmlc/dgl/issues' "
+             "matrix. Please file a bug at "
+             "'https://github.com/dmlc/dgl/issues' "
              "with the above stacktrace.";
       return new_coo;
     }
+    is_pinned = true;
+    return *this;
+  }
 
   /**
    * @brief Pin the row, col and data (if not Null) of the matrix.
@@ -159,6 +171,7 @@ struct COOMatrix {
    *       The context check is deferred to pinning the NDArray.
    */
   inline void PinMemory_() {
+    if (!IsEmpty()) {
       if (is_pinned) return;
       row.PinMemory_();
       col.PinMemory_();
@@ -167,6 +180,9 @@ struct COOMatrix {
       }
       is_pinned = true;
     }
+    is_pinned = true;
+    return;
+  }
 
   /**
    * @brief Unpin the row, col and data (if not Null) of the matrix.
@@ -176,6 +192,7 @@ struct COOMatrix {
    *       The context check is deferred to unpinning the NDArray.
    */
   inline void UnpinMemory_() {
+    if (!IsEmpty()) {
       if (!is_pinned) return;
       row.UnpinMemory_();
       col.UnpinMemory_();
@@ -184,6 +201,9 @@ struct COOMatrix {
       }
       is_pinned = false;
     }
+    is_pinned = false;
+    return;
+  }
 
   /**
    * @brief Record stream for the row, col and data (if not Null) of the matrix.

include/dgl/aten/csr.h

@@ -60,9 +60,11 @@ struct CSRMatrix {
         indices(iarr),
         data(darr),
         sorted(sorted_flag) {
+    if (!IsEmpty()) {
       is_pinned = (aten::IsNullArray(indptr) || indptr.IsPinned()) &&
                   (aten::IsNullArray(indices) || indices.IsPinned()) &&
                   (aten::IsNullArray(data) || data.IsPinned());
+    }
     CheckValidity();
   }
 
@@ -121,6 +123,11 @@ struct CSRMatrix {
     CHECK_EQ(indptr->shape[0], num_rows + 1);
   }
 
+  inline bool IsEmpty() const {
+    return aten::IsNullArray(indptr) && aten::IsNullArray(indices) &&
+           aten::IsNullArray(data);
+  }
+
   /** @brief Return a copy of this matrix on the give device context. */
   inline CSRMatrix CopyTo(const DGLContext& ctx) const {
     if (ctx == indptr->ctx) return *this;
@@ -131,16 +138,21 @@ struct CSRMatrix {
 
   /** @brief Return a copy of this matrix in pinned (page-locked) memory. */
   inline CSRMatrix PinMemory() {
+    if (!IsEmpty()) {
       if (is_pinned) return *this;
       auto new_csr = CSRMatrix(
           num_rows, num_cols, indptr.PinMemory(), indices.PinMemory(),
           aten::IsNullArray(data) ? data : data.PinMemory(), sorted);
       CHECK(new_csr.is_pinned)
           << "An internal DGL error has occured while trying to pin a CSR "
-           "matrix. Please file a bug at 'https://github.com/dmlc/dgl/issues' "
+             "matrix. Please file a bug at "
+             "'https://github.com/dmlc/dgl/issues' "
              "with the above stacktrace.";
       return new_csr;
     }
+    is_pinned = true;
+    return *this;
+  }
 
   /**
    * @brief Pin the indptr, indices and data (if not Null) of the matrix.
@@ -151,6 +163,7 @@ struct CSRMatrix {
    *       The context check is deferred to pinning the NDArray.
    */
   inline void PinMemory_() {
+    if (!IsEmpty()) {
       if (is_pinned) return;
       indptr.PinMemory_();
       indices.PinMemory_();
@@ -159,6 +172,9 @@ struct CSRMatrix {
       }
       is_pinned = true;
     }
+    is_pinned = true;
+    return;
+  }
 
   /**
    * @brief Unpin the indptr, indices and data (if not Null) of the matrix.
@@ -168,6 +184,7 @@ struct CSRMatrix {
    *       The context check is deferred to unpinning the NDArray.
    */
   inline void UnpinMemory_() {
+    if (!IsEmpty()) {
       if (!is_pinned) return;
       indptr.UnpinMemory_();
       indices.UnpinMemory_();
@@ -176,6 +193,9 @@ struct CSRMatrix {
       }
       is_pinned = false;
     }
+    is_pinned = false;
+    return;
+  }
 
   /**
    * @brief Record stream for the indptr, indices and data (if not Null) of the

tests/python/common/test_heterograph-index.py

@@ -68,6 +68,7 @@ def test_pin_memory(idtype):
 
     # Test pinning an empty homograph
     g2 = dgl.graph(([], []))
+    assert not g2.is_pinned()
     g2._graph = g2._graph.pin_memory()
     assert g2.is_pinned()
 

tests/python/common/test_heterograph.py

@@ -1259,6 +1259,7 @@ def test_pin_memory_(idtype):
 
     # test pin empty homograph
     g2 = dgl.graph(([], []))
+    assert not g2.is_pinned()
     g2.pin_memory_()
     assert g2.is_pinned()
     g2.unpin_memory_()