sparse_matrix_coalesce.cc 1.02 KB
Newer Older
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
/**
 *  Copyright (c) 2022 by Contributors
 * @file sparse_matrix_coalesce.cc
 * @brief Operators related to sparse matrix coalescing.
 */
// clang-format off
#include <sparse/dgl_headers.h>
// clang-format on

#include <sparse/sparse_matrix.h>

#include "./utils.h"

namespace dgl {
namespace sparse {

c10::intrusive_ptr<SparseMatrix> SparseMatrix::Coalesce() {
  auto torch_coo = COOToTorchCOO(this->COOPtr(), this->value());
  auto coalesced_coo = torch_coo.coalesce();
20
21
  return SparseMatrix::FromCOO(
      coalesced_coo.indices(), coalesced_coo.values(), this->shape());
22
23
24
25
}

bool SparseMatrix::HasDuplicate() {
  aten::CSRMatrix dgl_csr;
czkkkkkk's avatar
czkkkkkk committed
26
27
28
  if (HasDiag()) {
    return false;
  }
29
30
31
32
33
34
35
36
37
38
39
40
  // The format for calculation will be chosen in the following order: CSR,
  // CSC. CSR is created if the sparse matrix only has CSC format.
  if (HasCSR() || !HasCSC()) {
    dgl_csr = CSRToOldDGLCSR(CSRPtr());
  } else {
    dgl_csr = CSRToOldDGLCSR(CSCPtr());
  }
  return aten::CSRHasDuplicate(dgl_csr);
}

}  // namespace sparse
}  // namespace dgl