[Sparse] Reduce peak memory of SpSpAdd. (dmlc#5143)

* [Sparse] Reduce peak memory usage of SpSpAdd * Update * Update
XiaoYixin7132 · Jan 12, 2023 · 21c4c29 · 21c4c29
1 parent b743f76
commit 21c4c29
Showing 1 changed file with 9 additions and 3 deletions.
diff --git a/dgl_sparse/src/elemenwise_op.cc b/dgl_sparse/src/elemenwise_op.cc
@@ -22,9 +22,15 @@ c10::intrusive_ptr<SparseMatrix> SpSpAdd(
     const c10::intrusive_ptr<SparseMatrix>& A,
     const c10::intrusive_ptr<SparseMatrix>& B) {
   ElementwiseOpSanityCheck(A, B);
-  auto torch_A = COOToTorchCOO(A->COOPtr(), A->value());
-  auto torch_B = COOToTorchCOO(B->COOPtr(), B->value());
-  auto sum = (torch_A + torch_B).coalesce();
+  torch::Tensor sum;
+  {
+    // TODO(#5145) This is a workaround to reduce peak memory usage. It is no
+    // longer needed after we address #5145.
+    auto torch_A = COOToTorchCOO(A->COOPtr(), A->value());
+    auto torch_B = COOToTorchCOO(B->COOPtr(), B->value());
+    sum = torch_A + torch_B;
+  }
+  sum = sum.coalesce();
   auto indices = sum.indices();
   auto row = indices[0];
   auto col = indices[1];