chore: import upstream snapshot with attribution

dgl_sparse/src/cpu/matrix_ops_impl.cc

@@ -0,0 +1,10 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file cpu/matrix_ops_impl.cc
+ * @brief DGL C++ matrix operators.
+ */
+#include "./matrix_ops_impl.h"
+
+namespace dgl {
+namespace sparse {}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/elemenwise_op.cc

@@ -0,0 +1,151 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file elementwise_op.cc
+ * @brief DGL C++ sparse elementwise operator implementation.
+ */
+
+#include <sparse/elementwise_op.h>
+#include <sparse/matrix_ops.h>
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+#include <memory>
+
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+using namespace torch::autograd;
+
+c10::intrusive_ptr<SparseMatrix> SpSpAdd(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat) {
+  ElementwiseOpSanityCheck(lhs_mat, rhs_mat);
+  if (lhs_mat->HasDiag() && rhs_mat->HasDiag()) {
+    return SparseMatrix::FromDiagPointer(
+        lhs_mat->DiagPtr(), lhs_mat->value() + rhs_mat->value(),
+        lhs_mat->shape());
+  }
+  auto torch_lhs = COOToTorchCOO(lhs_mat->COOPtr(), lhs_mat->value());
+  auto torch_rhs = COOToTorchCOO(rhs_mat->COOPtr(), rhs_mat->value());
+  auto sum = (torch_lhs + torch_rhs).coalesce();
+  return SparseMatrix::FromCOO(sum.indices(), sum.values(), lhs_mat->shape());
+}
+
+class SpSpMulAutoGrad : public Function<SpSpMulAutoGrad> {
+ public:
+  static variable_list forward(
+      AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> lhs_mat,
+      torch::Tensor lhs_val, c10::intrusive_ptr<SparseMatrix> rhs_mat,
+      torch::Tensor rhs_val);
+
+  static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs);
+};
+
+variable_list SpSpMulAutoGrad::forward(
+    AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> lhs_mat,
+    torch::Tensor lhs_val, c10::intrusive_ptr<SparseMatrix> rhs_mat,
+    torch::Tensor rhs_val) {
+  std::shared_ptr<COO> intersection;
+  torch::Tensor lhs_indices, rhs_indices;
+  std::tie(intersection, lhs_indices, rhs_indices) =
+      COOIntersection(lhs_mat->COOPtr(), rhs_mat->COOPtr());
+  auto lhs_intersect_val = lhs_val.index_select(0, lhs_indices);
+  auto rhs_intersect_val = rhs_val.index_select(0, rhs_indices);
+  auto ret_val = lhs_intersect_val * rhs_intersect_val;
+  auto ret_mat =
+      SparseMatrix::FromCOOPointer(intersection, ret_val, lhs_mat->shape());
+
+  ctx->saved_data["lhs_require_grad"] = lhs_val.requires_grad();
+  ctx->saved_data["rhs_require_grad"] = rhs_val.requires_grad();
+  if (lhs_val.requires_grad()) {
+    ctx->saved_data["lhs_val_shape"] = lhs_val.sizes().vec();
+    ctx->saved_data["rhs_intersect_lhs"] =
+        SparseMatrix::ValLike(ret_mat, rhs_intersect_val);
+    ctx->saved_data["lhs_indices"] = lhs_indices;
+  }
+  if (rhs_val.requires_grad()) {
+    ctx->saved_data["rhs_val_shape"] = rhs_val.sizes().vec();
+    ctx->saved_data["lhs_intersect_rhs"] =
+        SparseMatrix::ValLike(ret_mat, lhs_intersect_val);
+    ctx->saved_data["rhs_indices"] = rhs_indices;
+  }
+  return {intersection->indices, ret_val};
+}
+
+tensor_list SpSpMulAutoGrad::backward(
+    AutogradContext* ctx, tensor_list grad_outputs) {
+  torch::Tensor lhs_val_grad, rhs_val_grad;
+  auto output_grad = grad_outputs[1];
+  if (ctx->saved_data["lhs_require_grad"].toBool()) {
+    auto rhs_intersect_lhs =
+        ctx->saved_data["rhs_intersect_lhs"].toCustomClass<SparseMatrix>();
+    const auto& lhs_val_shape = ctx->saved_data["lhs_val_shape"].toIntVector();
+    auto lhs_indices = ctx->saved_data["lhs_indices"].toTensor();
+    lhs_val_grad = torch::zeros(lhs_val_shape, output_grad.options());
+    auto intersect_grad = rhs_intersect_lhs->value() * output_grad;
+    lhs_val_grad.index_put_({lhs_indices}, intersect_grad);
+  }
+  if (ctx->saved_data["rhs_require_grad"].toBool()) {
+    auto lhs_intersect_rhs =
+        ctx->saved_data["lhs_intersect_rhs"].toCustomClass<SparseMatrix>();
+    const auto& rhs_val_shape = ctx->saved_data["rhs_val_shape"].toIntVector();
+    auto rhs_indices = ctx->saved_data["rhs_indices"].toTensor();
+    rhs_val_grad = torch::zeros(rhs_val_shape, output_grad.options());
+    auto intersect_grad = lhs_intersect_rhs->value() * output_grad;
+    rhs_val_grad.index_put_({rhs_indices}, intersect_grad);
+  }
+  return {torch::Tensor(), lhs_val_grad, torch::Tensor(), rhs_val_grad};
+}
+
+c10::intrusive_ptr<SparseMatrix> SpSpMul(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat) {
+  ElementwiseOpSanityCheck(lhs_mat, rhs_mat);
+  if (lhs_mat->HasDiag() && rhs_mat->HasDiag()) {
+    return SparseMatrix::FromDiagPointer(
+        lhs_mat->DiagPtr(), lhs_mat->value() * rhs_mat->value(),
+        lhs_mat->shape());
+  }
+  TORCH_CHECK(
+      !lhs_mat->HasDuplicate() && !rhs_mat->HasDuplicate(),
+      "Only support SpSpMul on sparse matrices without duplicate values")
+  auto results = SpSpMulAutoGrad::apply(
+      lhs_mat, lhs_mat->value(), rhs_mat, rhs_mat->value());
+  const auto& indices = results[0];
+  const auto& val = results[1];
+  return SparseMatrix::FromCOO(indices, val, lhs_mat->shape());
+}
+
+c10::intrusive_ptr<SparseMatrix> SpSpDiv(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat) {
+  ElementwiseOpSanityCheck(lhs_mat, rhs_mat);
+  if (lhs_mat->HasDiag() && rhs_mat->HasDiag()) {
+    return SparseMatrix::FromDiagPointer(
+        lhs_mat->DiagPtr(), lhs_mat->value() / rhs_mat->value(),
+        lhs_mat->shape());
+  }
+  std::shared_ptr<COO> sorted_lhs, sorted_rhs;
+  torch::Tensor lhs_sorted_perm, rhs_sorted_perm;
+  std::tie(sorted_lhs, lhs_sorted_perm) = COOSort(lhs_mat->COOPtr());
+  std::tie(sorted_rhs, rhs_sorted_perm) = COOSort(rhs_mat->COOPtr());
+  TORCH_CHECK(
+      !lhs_mat->HasDuplicate() && !rhs_mat->HasDuplicate(),
+      "Only support SpSpDiv on sparse matrices without duplicate values")
+  TORCH_CHECK(
+      torch::equal(sorted_lhs->indices, sorted_rhs->indices),
+      "Cannot divide two COO matrices with different sparsities.");
+  // This is to make sure the return matrix is in the same order as the lhs_mat
+  auto lhs_sorted_rperm = lhs_sorted_perm.argsort();
+  auto rhs_perm_on_lhs = rhs_sorted_perm.index_select(0, lhs_sorted_rperm);
+  auto lhs_value = lhs_mat->value();
+  auto rhs_value = rhs_mat->value().index_select(0, rhs_perm_on_lhs);
+  auto ret_val = lhs_value / rhs_value;
+  return SparseMatrix::FromCOOPointer(
+      lhs_mat->COOPtr(), ret_val, lhs_mat->shape());
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/matmul.cc

@@ -0,0 +1,195 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file matmul.cc
+ * @brief DGL sparse matrix multiplication functions.
+ */
+#include "./matmul.h"
+
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+torch::Tensor SpMMNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat,
+    torch::Tensor sparse_val, torch::Tensor dense_mat, bool transpose_sparse) {
+  const std::string op = "mul";
+  const std::string reduce = "sum";
+  const int64_t out_row =
+      transpose_sparse ? sparse_mat->shape()[1] : sparse_mat->shape()[0];
+  std::vector<int64_t> shape = {out_row, dense_mat.size(1)};
+  // Batched SpMM
+  if (sparse_val.dim() >= 2) {
+    shape = {out_row, dense_mat.size(1), sparse_val.size(1)};
+  }
+
+  auto ret = torch::zeros(shape, dense_mat.options());
+  auto dgl_sparse_val = TorchTensorToDGLArray(sparse_val);
+  auto dgl_dense_mat = TorchTensorToDGLArray(dense_mat);
+  auto dgl_ret = TorchTensorToDGLArray(ret);
+  if (!transpose_sparse) {
+    // The format for calculation will be chosen in the following order: CSR,
+    // COO. CSR is created if the sparse matrix only has CSC format.
+    if (sparse_mat->HasCSR() || !sparse_mat->HasCOO()) {
+      // sparse_mat->CSRPtr() will implicitly convert CSC to CSR format if CSR
+      // does not exist.
+      auto csr = CSRToOldDGLCSR(sparse_mat->CSRPtr());
+      aten::CSRSpMM(
+          op.c_str(), reduce.c_str(), csr, dgl_dense_mat, dgl_sparse_val,
+          dgl_ret, {});
+    } else {  // COO
+      // Use the reverse order of aten::COOSpMM because it calculates A^T @ X.
+      auto coo = COOToOldDGLCOO(sparse_mat->COOPtr());
+      coo = aten::COOTranspose(coo);
+      aten::COOSpMM(
+          op.c_str(), reduce.c_str(), coo, dgl_dense_mat, dgl_sparse_val,
+          dgl_ret, {});
+    }
+  } else {  // transpose_sparse
+    // The format for calculation will be chosen in the following order: CSC,
+    // COO. CSC is created if the sparse matrix only has CSR format.
+    if (sparse_mat->HasCSC() || !sparse_mat->HasCOO()) {
+      // sparse_mat->CSCPtr() will implicitly convert CSR to CSC format if CSR
+      // does not exist.
+      // Use CSC in DGL's CSRSpMM is equivalent as computing A^T @ X.
+      auto csc = CSRToOldDGLCSR(sparse_mat->CSCPtr());
+      aten::CSRSpMM(
+          op.c_str(), reduce.c_str(), csc, dgl_dense_mat, dgl_sparse_val,
+          dgl_ret, {});
+    } else {  // COO
+      // Use the reverse order of aten::COOSpMM because it calculates A^T @ X.
+      auto coo = COOToOldDGLCOO(sparse_mat->COOPtr());
+      aten::COOSpMM(
+          op.c_str(), reduce.c_str(), coo, dgl_dense_mat, dgl_sparse_val,
+          dgl_ret, {});
+    }
+  }
+  return ret;
+}
+
+torch::Tensor SDDMMNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor mat1,
+    torch::Tensor mat2_tr) {
+  const int64_t out_row = sparse_mat->nnz();
+  std::vector<int64_t> shape({out_row});
+  // Batched SDDMM
+  if (mat1.dim() >= 3) {
+    shape.push_back(mat1.size(2));
+    // (N, K, B) -> (N, B, K)
+    mat1 = mat1.transpose(1, 2);
+    // (M, K, B) -> (M, B, K)
+    mat2_tr = mat2_tr.transpose(1, 2);
+  }
+  auto ret = torch::zeros(shape, mat1.options());
+  const std::string op = "dot";
+  auto dgl_mat1 = TorchTensorToDGLArray(mat1);
+  auto dgl_mat2_tr = TorchTensorToDGLArray(mat2_tr);
+  auto dgl_ret = TorchTensorToDGLArray(ret);
+  // The format for calculation will be chosen in the following order: CSR,
+  // COO. CSR is created if the sparse matrix only has CSC format.
+  if (sparse_mat->HasCSR() || !sparse_mat->HasCOO()) {
+    // sparse_mat->CSRPtr() will implicitly convert CSC to CSR format if CSR
+    // does not exist.
+    auto csr = CSRToOldDGLCSR(sparse_mat->CSRPtr());
+    aten::CSRSDDMM(
+        op.c_str(), csr, dgl_mat1, dgl_mat2_tr, dgl_ret, 0 /* Lhs target: u */,
+        2 /* rhs target: v */);
+  } else {  // COO
+    auto coo = COOToOldDGLCOO(sparse_mat->COOPtr());
+    aten::COOSDDMM(
+        op.c_str(), coo, dgl_mat1, dgl_mat2_tr, dgl_ret, 0 /* Lhs target: u */,
+        2 /* rhs target: v */);
+  }
+  return ret;
+}
+
+torch::Tensor BroadcastOpNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    const std::string& op, int64_t dim) {
+  auto sparse_val = sparse_mat->value();
+  const int64_t out_row = sparse_mat->nnz();
+  const std::vector<int64_t> shape({out_row, sparse_val.size(1)});
+  auto ret = torch::zeros(shape, sparse_val.options());
+
+  auto dgl_sparse_val = TorchTensorToDGLArray(sparse_val);
+  auto dgl_dense_mat = TorchTensorToDGLArray(dense_mat);
+  auto dgl_ret = TorchTensorToDGLArray(ret);
+  // Setting dgl_rhs_target to 0 or 2 means using row or column coordinators
+  // to access dgl_dense_mat for each edge, respectively.
+  auto dgl_rhs_target = dim == 0 ? 2 : 0;
+
+  // The format for calculation will be chosen in the following order: COO, CSR
+  // . COO is created if the sparse matrix only has CSC format.
+  if (sparse_mat->HasCOO() || !sparse_mat->HasCSR()) {
+    // sparse_mat->COOPtr() will implicitly convert CSC to COO format if COO
+    // does not exist.
+    auto coo = COOToOldDGLCOO(sparse_mat->COOPtr());
+    aten::COOSDDMM(
+        op.c_str(), coo, dgl_sparse_val, dgl_dense_mat, dgl_ret,
+        1 /* Lhs target: e */, dgl_rhs_target);
+  } else {
+    auto csr = CSRToOldDGLCSR(sparse_mat->CSRPtr());
+    aten::CSRSDDMM(
+        op.c_str(), csr, dgl_sparse_val, dgl_dense_mat, dgl_ret,
+        1 /* Lhs target: e */, dgl_rhs_target);
+  }
+  return ret;
+}
+
+torch::Tensor BroadcastSubNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    int64_t dim) {
+  return BroadcastOpNoAutoGrad(sparse_mat, dense_mat, "sub", dim);
+}
+
+torch::Tensor BroadcastDivNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    int64_t dim) {
+  return BroadcastOpNoAutoGrad(sparse_mat, dense_mat, "div", dim);
+}
+
+torch::Tensor BroadcastMulNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    int64_t dim) {
+  return BroadcastOpNoAutoGrad(sparse_mat, dense_mat, "mul", dim);
+}
+
+c10::intrusive_ptr<SparseMatrix> SpSpMMNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat, torch::Tensor lhs_val,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat, torch::Tensor rhs_val,
+    bool lhs_transpose, bool rhs_transpose) {
+  aten::CSRMatrix lhs_dgl_csr, rhs_dgl_csr;
+  if (!lhs_transpose) {
+    lhs_dgl_csr = CSRToOldDGLCSR(lhs_mat->CSRPtr());
+  } else {
+    lhs_dgl_csr = CSRToOldDGLCSR(lhs_mat->CSCPtr());
+  }
+  if (!rhs_transpose) {
+    rhs_dgl_csr = CSRToOldDGLCSR(rhs_mat->CSRPtr());
+  } else {
+    rhs_dgl_csr = CSRToOldDGLCSR(rhs_mat->CSCPtr());
+  }
+  auto lhs_dgl_val = TorchTensorToDGLArray(lhs_val);
+  auto rhs_dgl_val = TorchTensorToDGLArray(rhs_val);
+  const int64_t ret_row =
+      lhs_transpose ? lhs_mat->shape()[1] : lhs_mat->shape()[0];
+  const int64_t ret_col =
+      rhs_transpose ? rhs_mat->shape()[0] : rhs_mat->shape()[1];
+  std::vector<int64_t> ret_shape({ret_row, ret_col});
+  aten::CSRMatrix ret_dgl_csr;
+  runtime::NDArray ret_val;
+  std::tie(ret_dgl_csr, ret_val) =
+      aten::CSRMM(lhs_dgl_csr, lhs_dgl_val, rhs_dgl_csr, rhs_dgl_val);
+  return SparseMatrix::FromCSRPointer(
+      CSRFromOldDGLCSR(ret_dgl_csr), DGLArrayToTorchTensor(ret_val), ret_shape);
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/matmul.h

@@ -0,0 +1,156 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file matmul.h
+ * @brief DGL sparse matrix multiplication functions.
+ */
+#ifndef DGL_SPARSE_MATMUL_H_
+#define DGL_SPARSE_MATMUL_H_
+
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+#include <string>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Perform a matrix multiplication of the sparse matrix and dense
+ * matrix. It uses the sparse formats of `sparse_mat` and non-zero values of
+ * `sparse_val` for SpMM. The `sparse_val` must be 1-dimensional. If the sparse
+ * matrix has shape (n, m), the dense matrix must have shape (m, k). And
+ * the returned dense matrix has shape (n, k).
+ *
+ * This function does not take care of autograd.
+ *
+ * @param sparse_mat The sparse matrix.
+ * @param sparse_val Non-zero values of the sparse matrix.
+ * @param dense_mat The dense matrix.
+ * @param transpose_sparse Whether the sparse_mat is transposed.
+ *
+ * @return Dense tensor.
+ */
+torch::Tensor SpMMNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat,
+    torch::Tensor sparse_val, torch::Tensor dense_mat, bool transpose_sparse);
+
+/**
+ * @brief Perform a sampled matrix multiplication of a sparse matrix and two
+ * dense matrices. It calculates `(mat1 @ mat2_tr^T) * spy(A)` and does consider
+ * the values of the sparse matrix. For efficiency, `mat2_tr` is the
+ * transposition of the matrix to be multiplied. If the sparse matrix has shape
+ * (n, m), `mat1` and `mat2_tr` must have shapes of `(n, k)` and `(m,
+ * k)`respectively. And the returned tensor has shape
+ * `(sparse_matrix->nnz(),)`.
+ *
+ * This function does not take care of autograd.
+ *
+ * @param sparse_mat The sparse matrix.
+ * @param mat1 The first dense matrix.
+ * @param mat2_tr Transposition of the second matrix.
+ *
+ * @return Dense tensor.
+ */
+torch::Tensor SDDMMNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor mat1,
+    torch::Tensor mat2_tr);
+
+/**
+ * @brief Broadcast the dense feature to the nonzero entries and then compute
+ * x_e = \phi(x_e, x_v) on the dimension dim, where x_e is the nonzero value,
+ * x_v is the dense feature, and \phi is add, sub, mul, or div. dim = 0 or 1
+ * means column-wise or row-wise broadcast respectively.
+ *
+ * This function does not take care of autograd.
+ *
+ * @param sparse_mat The sparse matrix with N rows and (nnz, D) nonzero values
+ * @param dense_mat Dense feature of shape (N, D)
+ * @param op Operator, can be add, sub, mul, or div
+ * @param dim The dimension to broadcast.
+ *
+ * @return Dense tensor of shape (nnz, D)
+ */
+torch::Tensor BroadcastOpNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    const std::string& op, int64_t dim);
+
+/**
+ * @brief Broadcast the dense feature to the nonzero entries and then compute
+ * x_e = x_e - x_v on the dimension dim, where x_e is the nonzero value, x_v is
+ * the dense feature. dim = 0 or 1 means column-wise or row-wise broadcast
+ * respectively.
+ *
+ * This function does not take care of autograd.
+ *
+ * @param sparse_mat The sparse matrix with N rows and (nnz, D) nonzero values
+ * @param dense_mat Dense feature of shape (N, D)
+ * @param dim The dimension to broadcast.
+ *
+ * @return Dense tensor of shape (nnz, D)
+ */
+torch::Tensor BroadcastSubNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    int64_t dim);
+
+/**
+ * @brief Broadcast the dense feature to the nonzero entries and then compute
+ * x_e = x_e / x_v on the dimension dim, where x_e is the nonzero value, x_v is
+ * the dense feature. dim = 0 or 1 means column-wise or row-wise broadcast
+ * respectively.
+ *
+ * This function does not take care of autograd.
+ *
+ * @param sparse_mat The sparse matrix with N rows and (nnz, D) nonzero values
+ * @param dense_mat Dense feature of shape (N, D)
+ * @param dim The dimension to broadcast.
+ *
+ * @return Dense tensor of shape (nnz, D)
+ */
+torch::Tensor BroadcastDivNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    int64_t dim);
+
+/**
+ * @brief Broadcast the dense feature to the nonzero entries and then compute
+ * x_e = x_e * x_v on the dimension dim, where x_e is the nonzero value, x_v is
+ * the dense feature. dim = 0 or 1 means column-wise or row-wise broadcast
+ * respectively.
+ *
+ * This function does not take care of autograd.
+ *
+ * @param sparse_mat The sparse matrix with N rows and (nnz, D) nonzero values
+ * @param dense_mat Dense feature of shape (N, D)
+ * @param dim The dimension to broadcast.
+ *
+ * @return Dense tensor of shape (nnz, D)
+ */
+torch::Tensor BroadcastMulNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor dense_mat,
+    int64_t dim);
+
+/**
+ * @brief Perform a sparse-sparse matrix multiplication with possibly different
+ * sparsities. The two sparse values must have 1-dimensional values. If the
+ * first sparse matrix has shape (n, m), the second sparse matrix must have
+ * shape (m, k), and the returned sparse matrix has shape (n, k).
+ *
+ * This function does not take care of autograd.
+ *
+ * @param lhs_mat The first sparse matrix of shape (n, m).
+ * @param lhs_val Sparse value for the first sparse matrix.
+ * @param rhs_mat The second sparse matrix of shape (m, k).
+ * @param rhs_val Sparse value for the second sparse matrix.
+ * @param lhs_transpose Whether the first matrix is transposed.
+ * @param rhs_transpose Whether the second matrix is transposed.
+ *
+ * @return Sparse matrix of shape (n, k).
+ */
+c10::intrusive_ptr<SparseMatrix> SpSpMMNoAutoGrad(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat, torch::Tensor lhs_val,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat, torch::Tensor rhs_val,
+    bool lhs_transpose, bool rhs_transpose);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // DGL_SPARSE_MATMUL_H_

dgl_sparse/src/matrix_ops.cc

@@ -0,0 +1,177 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file matrix_ops.cc
+ * @brief DGL C++ matrix operators.
+ */
+#include <sparse/matrix_ops.h>
+#include <torch/script.h>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Compute the intersection of two COO matrices. Return the intersection
+ * COO matrix, and the indices of the intersection in the left-hand-side and
+ * right-hand-side COO matrices.
+ *
+ * @param lhs The left-hand-side COO matrix.
+ * @param rhs The right-hand-side COO matrix.
+ *
+ * @return A tuple of COO matrix, lhs indices, and rhs indices.
+ */
+std::tuple<std::shared_ptr<COO>, torch::Tensor, torch::Tensor> COOIntersection(
+    const std::shared_ptr<COO>& lhs, const std::shared_ptr<COO>& rhs) {
+  // 1. Encode the two COO matrices into arrays of integers.
+  auto lhs_arr =
+      lhs->indices.index({0}) * lhs->num_cols + lhs->indices.index({1});
+  auto rhs_arr =
+      rhs->indices.index({0}) * rhs->num_cols + rhs->indices.index({1});
+  // 2. Concatenate the two arrays.
+  auto arr = torch::cat({lhs_arr, rhs_arr});
+  // 3. Unique the concatenated array.
+  torch::Tensor unique, inverse, counts;
+  std::tie(unique, inverse, counts) =
+      torch::unique_dim(arr, 0, false, true, true);
+  // 4. Find the indices of the counts greater than 1 in the unique array.
+  auto mask = counts > 1;
+  // 5. Map the inverse array to the original array to generate indices.
+  auto lhs_inverse = inverse.slice(0, 0, lhs_arr.numel());
+  auto rhs_inverse = inverse.slice(0, lhs_arr.numel(), arr.numel());
+  auto map_to_original = torch::empty_like(unique);
+  map_to_original.index_put_(
+      {lhs_inverse},
+      torch::arange(lhs_inverse.numel(), map_to_original.options()));
+  auto lhs_indices = map_to_original.index({mask});
+  map_to_original.index_put_(
+      {rhs_inverse},
+      torch::arange(rhs_inverse.numel(), map_to_original.options()));
+  auto rhs_indices = map_to_original.index({mask});
+  // 6. Decode the indices to get the intersection COO matrix.
+  auto ret_arr = unique.index({mask});
+  auto ret_indices = torch::stack(
+      {ret_arr.floor_divide(lhs->num_cols), ret_arr % lhs->num_cols}, 0);
+  auto ret_coo = std::make_shared<COO>(
+      COO{lhs->num_rows, lhs->num_cols, ret_indices, false, false});
+  return {ret_coo, lhs_indices, rhs_indices};
+}
+
+/** @brief Return the reverted mapping of a permutation. */
+static torch::Tensor RevertPermutation(const torch::Tensor& perm) {
+  auto rev_tensor = torch::empty_like(perm);
+  rev_tensor.index_put_(
+      {perm}, torch::arange(0, perm.numel(), rev_tensor.options()));
+  return rev_tensor;
+}
+
+/**
+ * @brief Compute the compact indices of row indices and leading indices. Return
+ * the compacted indices and the original row indices of compacted indices.
+ *
+ * @param row The row indices.
+ * @param leading_indices The leading indices.
+ *
+ * @return A tuple of compact indices, original indices.
+ */
+static std::tuple<torch::Tensor, torch::Tensor> CompactIndices(
+    const torch::Tensor& row,
+    const torch::optional<torch::Tensor>& leading_indices) {
+  torch::Tensor sorted, sort_indices, uniqued, unique_reverse_indices, counts;
+  // 1. Sort leading indices and row indices in ascending order.
+  int64_t n_leading_indices = 0;
+  if (leading_indices.has_value()) {
+    n_leading_indices = leading_indices.value().numel();
+    std::tie(sorted, sort_indices) =
+        torch::cat({leading_indices.value(), row}).sort();
+  } else {
+    std::tie(sorted, sort_indices) = row.sort();
+  }
+  // 2. Reverse sort indices.
+  auto sort_rev_indices = RevertPermutation(sort_indices);
+  // 3. Unique the sorted array.
+  std::tie(uniqued, unique_reverse_indices, counts) =
+      torch::unique_consecutive(sorted, true);
+  auto reverse_indices = unique_reverse_indices.index({sort_rev_indices});
+  auto n_uniqued = uniqued.numel();
+
+  // 4. Relabel the indices and map the inverse array to the original array.
+  auto split_indices = torch::full({n_uniqued}, -1, reverse_indices.options());
+
+  split_indices.index_put_(
+      {reverse_indices.slice(0, 0, n_leading_indices)},
+      torch::arange(0, n_leading_indices, split_indices.options()));
+
+  split_indices.index_put_(
+      {(split_indices == -1).nonzero().view(-1)},
+      torch::arange(n_leading_indices, n_uniqued, split_indices.options()));
+  // 5. Decode the indices to get the compact indices.
+  auto new_row = split_indices.index({reverse_indices.slice(
+      0, n_leading_indices, n_leading_indices + row.numel())});
+  return {new_row, uniqued.index({RevertPermutation(split_indices)})};
+}
+
+static std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> CompactCOO(
+    const c10::intrusive_ptr<SparseMatrix>& mat, int64_t dim,
+    const torch::optional<torch::Tensor>& leading_indices) {
+  torch::Tensor row, col;
+  auto coo = mat->COOTensors();
+  if (dim == 0)
+    std::tie(row, col) = coo;
+  else
+    std::tie(col, row) = coo;
+
+  torch::Tensor new_row, uniqued;
+  std::tie(new_row, uniqued) = CompactIndices(row, leading_indices);
+
+  if (dim == 0) {
+    auto ret = SparseMatrix::FromCOO(
+        torch::stack({new_row, col}, 0), mat->value(),
+        std::vector<int64_t>{uniqued.numel(), mat->shape()[1]});
+    return {ret, uniqued};
+  } else {
+    auto ret = SparseMatrix::FromCOO(
+        torch::stack({col, new_row}, 0), mat->value(),
+        std::vector<int64_t>{mat->shape()[0], uniqued.numel()});
+    return {ret, uniqued};
+  }
+}
+
+static std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> CompactCSR(
+    const c10::intrusive_ptr<SparseMatrix>& mat, int64_t dim,
+    const torch::optional<torch::Tensor>& leading_indices) {
+  std::shared_ptr<CSR> csr;
+  if (dim == 0)
+    csr = mat->CSCPtr();
+  else
+    csr = mat->CSRPtr();
+
+  torch::Tensor new_indices, uniqued;
+  std::tie(new_indices, uniqued) =
+      CompactIndices(csr->indices, leading_indices);
+
+  auto ret_value = mat->value();
+  if (csr->value_indices.has_value())
+    ret_value = mat->value().index_select(0, csr->value_indices.value());
+  if (dim == 0) {
+    auto ret = SparseMatrix::FromCSC(
+        csr->indptr, new_indices, ret_value,
+        std::vector<int64_t>{uniqued.numel(), mat->shape()[1]});
+    return {ret, uniqued};
+  } else {
+    auto ret = SparseMatrix::FromCSR(
+        csr->indptr, new_indices, ret_value,
+        std::vector<int64_t>{mat->shape()[0], uniqued.numel()});
+    return {ret, uniqued};
+  }
+}
+
+std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> Compact(
+    const c10::intrusive_ptr<SparseMatrix>& mat, int64_t dim,
+    const torch::optional<torch::Tensor>& leading_indices) {
+  if (mat->HasCOO()) {
+    return CompactCOO(mat, dim, leading_indices);
+  }
+  return CompactCSR(mat, dim, leading_indices);
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/matrix_ops_impl.h

@@ -0,0 +1,21 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file matrix_ops_impl.h
+ * @brief DGL C++ sparse matrix operator implementations.
+ */
+#ifndef DGL_SPARSE_MATRIX_OPS_IMPL_H_
+#define DGL_SPARSE_MATRIX_OPS_IMPL_H_
+
+#include <sparse/sparse_format.h>
+#include <sparse/sparse_matrix.h>
+
+#include <tuple>
+#include <vector>
+
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {}  // namespace sparse
+}  // namespace dgl
+
+#endif  // DGL_SPARSE_MATRIX_OPS_IMPL_H_

dgl_sparse/src/python_binding.cc

@@ -0,0 +1,63 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file python_binding.cc
+ * @brief DGL sparse library Python binding.
+ */
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <sparse/elementwise_op.h>
+#include <sparse/matrix_ops.h>
+#include <sparse/reduction.h>
+#include <sparse/sddmm.h>
+#include <sparse/softmax.h>
+#include <sparse/sparse_matrix.h>
+#include <sparse/spmm.h>
+#include <sparse/spspmm.h>
+#include <torch/custom_class.h>
+#include <torch/script.h>
+
+namespace dgl {
+namespace sparse {
+
+TORCH_LIBRARY(dgl_sparse, m) {
+  m.class_<SparseMatrix>("SparseMatrix")
+      .def("val", &SparseMatrix::value)
+      .def("nnz", &SparseMatrix::nnz)
+      .def("device", &SparseMatrix::device)
+      .def("shape", &SparseMatrix::shape)
+      .def("coo", &SparseMatrix::COOTensors)
+      .def("indices", &SparseMatrix::Indices)
+      .def("csr", &SparseMatrix::CSRTensors)
+      .def("csc", &SparseMatrix::CSCTensors)
+      .def("transpose", &SparseMatrix::Transpose)
+      .def("coalesce", &SparseMatrix::Coalesce)
+      .def("has_duplicate", &SparseMatrix::HasDuplicate)
+      .def("is_diag", &SparseMatrix::HasDiag)
+      .def("index_select", &SparseMatrix::IndexSelect)
+      .def("range_select", &SparseMatrix::RangeSelect)
+      .def("sample", &SparseMatrix::Sample);
+  m.def("from_coo", &SparseMatrix::FromCOO)
+      .def("from_csr", &SparseMatrix::FromCSR)
+      .def("from_csc", &SparseMatrix::FromCSC)
+      .def("from_diag", &SparseMatrix::FromDiag)
+      .def("spsp_add", &SpSpAdd)
+      .def("spsp_mul", &SpSpMul)
+      .def("spsp_div", &SpSpDiv)
+      .def("reduce", &Reduce)
+      .def("sum", &ReduceSum)
+      .def("smean", &ReduceMean)
+      .def("smin", &ReduceMin)
+      .def("smax", &ReduceMax)
+      .def("sprod", &ReduceProd)
+      .def("val_like", &SparseMatrix::ValLike)
+      .def("spmm", &SpMM)
+      .def("sddmm", &SDDMM)
+      .def("softmax", &Softmax)
+      .def("spspmm", &SpSpMM)
+      .def("compact", &Compact);
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/reduction.cc

@@ -0,0 +1,98 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file reduction.cc
+ * @brief DGL C++ sparse matrix reduction operator implementation.
+ */
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <sparse/elementwise_op.h>
+#include <sparse/reduction.h>
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+#include <string>
+#include <vector>
+
+namespace dgl {
+namespace sparse {
+
+namespace {
+
+torch::Tensor ReduceAlong(
+    const c10::intrusive_ptr<SparseMatrix>& A, const std::string& reduce,
+    int64_t dim) {
+  auto value = A->value();
+  auto coo = A->COOPtr();
+
+  std::string reduce_op;
+  if (reduce == "sum") {
+    reduce_op = "sum";
+  } else if (reduce == "smin") {
+    reduce_op = "amin";
+  } else if (reduce == "smax") {
+    reduce_op = "amax";
+  } else if (reduce == "smean") {
+    reduce_op = "mean";
+  } else if (reduce == "sprod") {
+    reduce_op = "prod";
+  } else {
+    TORCH_CHECK(false, "unknown reduce function ", reduce);
+    return torch::Tensor();
+  }
+
+  // Create the output tensor with shape
+  //
+  //   [A.num_rows if dim == 1 else A.num_cols] + A.val.shape[1:]
+  std::vector<int64_t> output_shape = value.sizes().vec();
+  std::vector<int64_t> view_dims(output_shape.size(), 1);
+  view_dims[0] = -1;
+  torch::Tensor idx;
+  if (dim == 0) {
+    output_shape[0] = coo->num_cols;
+    idx = coo->indices.index({1}).view(view_dims).expand_as(value);
+  } else if (dim == 1) {
+    output_shape[0] = coo->num_rows;
+    idx = coo->indices.index({0}).view(view_dims).expand_as(value);
+  }
+  torch::Tensor out = torch::zeros(output_shape, value.options());
+
+  if (dim == 0) {
+    out.scatter_reduce_(0, idx, value, reduce_op, false);
+  } else if (dim == 1) {
+    out.scatter_reduce_(0, idx, value, reduce_op, false);
+  }
+
+  return out;
+}
+
+torch::Tensor ReduceAll(
+    const c10::intrusive_ptr<SparseMatrix>& A, const std::string& reduce) {
+  if (reduce == "sum") {
+    return A->value().sum(0);
+  } else if (reduce == "smin") {
+    return A->value().amin(0);
+  } else if (reduce == "smax") {
+    return A->value().amax(0);
+  } else if (reduce == "smean") {
+    return A->value().mean(0);
+  } else if (reduce == "sprod") {
+    return A->value().prod(0);
+  }
+
+  TORCH_CHECK(false, "unknown reduce function ", reduce);
+  return torch::Tensor();
+}
+
+}  // namespace
+
+torch::Tensor Reduce(
+    const c10::intrusive_ptr<SparseMatrix>& A, const std::string& reduce,
+    const torch::optional<int64_t>& dim) {
+  return dim.has_value() ? ReduceAlong(A, reduce, dim.value())
+                         : ReduceAll(A, reduce);
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/sddmm.cc

@@ -0,0 +1,129 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sddmm.cc
+ * @brief DGL C++ sparse SDDMM operator implementation.
+ */
+#include <sparse/sparse_matrix.h>
+#include <sparse/spmm.h>
+#include <torch/script.h>
+
+#include <sstream>
+
+#include "./matmul.h"
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+using namespace torch::autograd;
+
+class SDDMMAutoGrad : public Function<SDDMMAutoGrad> {
+ public:
+  static torch::Tensor forward(
+      AutogradContext* ctx, const c10::intrusive_ptr<SparseMatrix>& sparse_mat,
+      torch::Tensor mat1, torch::Tensor mat2_tr);
+
+  static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs);
+};
+
+void _SDDMMSanityCheck(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor mat1,
+    torch::Tensor mat2) {
+  bool shape_check = true;
+  shape_check &= mat1.dim() == mat2.dim();
+  shape_check &= mat1.dim() <= 3;
+  shape_check &= sparse_mat->shape()[0] == mat1.size(0);
+  if (mat1.dim() == 3) {
+    shape_check &= sparse_mat->shape()[1] == mat2.size(1);
+    shape_check &= mat1.size(2) == mat2.size(2);
+    if (sparse_mat->value().dim() > 1) {
+      shape_check &= sparse_mat->value().size(1) == mat1.size(2);
+    }
+  } else {
+    shape_check &= sparse_mat->shape()[1] == mat2.size(mat2.dim() - 1);
+  }
+  if (mat1.dim() >= 2) {
+    shape_check &= mat1.size(1) == mat2.size(0);
+  }
+  if (!shape_check) {
+    std::stringstream error;
+    error << "SDDMM: Invalid input shapes. sparse_mat: "
+          << c10::IntArrayRef(sparse_mat->shape())
+          << ", sparse_val: " << sparse_mat->value().sizes()
+          << ", mat1: " << mat1.sizes() << ", mat2: " << mat2.sizes()
+          << ". Valid input shapes (sparse_mat, mat1, mat2) are: (1) (n, m), "
+             "(n, k), and (k, m); (2) (n, m), (n,), and (m,); (3) (n, m, b), "
+             "(n, k, b) and (k, m, b); (4) "
+             "(n, m), (n, k, b), and (k, m, b).";
+    TORCH_CHECK(false, error.str());
+  }
+  TORCH_CHECK(
+      mat1.dtype() == mat2.dtype(),
+      "SDDMM: the two dense matrices should have the same dtype.");
+  TORCH_CHECK(
+      mat1.device() == mat2.device() && sparse_mat->device() == mat2.device(),
+      "SDDMM: the two dense matrices and sparse matrix should on the same "
+      "device.");
+}
+
+torch::Tensor SDDMMAutoGrad::forward(
+    AutogradContext* ctx, const c10::intrusive_ptr<SparseMatrix>& sparse_mat,
+    torch::Tensor mat1, torch::Tensor mat2) {
+  auto mat2_tr = mat2.transpose(0, 1);
+  auto ret = SDDMMNoAutoGrad(sparse_mat, mat1, mat2_tr);
+  torch::Tensor cache_mat1, cache_mat2;
+  if (mat1.requires_grad()) {
+    cache_mat2 = mat2;
+  }
+  if (mat2.requires_grad()) {
+    cache_mat1 = mat1;
+  }
+  ctx->save_for_backward({cache_mat1, cache_mat2});
+  ctx->saved_data["mat1_requires_grad"] = mat1.requires_grad();
+  ctx->saved_data["mat2_requires_grad"] = mat2.requires_grad();
+  ctx->saved_data["sparse_mat"] = sparse_mat;
+  return ret;
+}
+
+tensor_list SDDMMAutoGrad::backward(
+    AutogradContext* ctx, tensor_list grad_outputs) {
+  auto saved = ctx->get_saved_variables();
+  auto mat1 = saved[0];
+  auto mat2 = saved[1];
+  auto sparse_mat = ctx->saved_data["sparse_mat"].toCustomClass<SparseMatrix>();
+  auto grad = grad_outputs[0];
+  torch::Tensor mat1_grad, mat2_grad;
+  if (ctx->saved_data["mat1_requires_grad"].toBool()) {
+    // SDDMM(M, A, B) = C. dA = SpMM(dC, B^T)
+    mat1_grad = SpMMNoAutoGrad(sparse_mat, grad, mat2.transpose(0, 1), false);
+  }
+  if (ctx->saved_data["mat2_requires_grad"].toBool()) {
+    // SDDMM(M, A, B) = C. dB = SpMM(dC^T, A)^T
+    auto mat2_tr_grad = SpMMNoAutoGrad(sparse_mat, grad, mat1, true);
+    mat2_grad = mat2_tr_grad.transpose(0, 1);
+  }
+  return {torch::Tensor(), mat1_grad, mat2_grad};
+}
+
+c10::intrusive_ptr<SparseMatrix> SDDMM(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor mat1,
+    torch::Tensor mat2) {
+  if (mat1.dim() == 1) {
+    mat1 = mat1.view({mat1.size(0), 1});
+  }
+  if (mat2.dim() == 1) {
+    mat2 = mat2.view({1, mat2.size(0)});
+  }
+  _SDDMMSanityCheck(sparse_mat, mat1, mat2);
+  auto val = SDDMMAutoGrad::apply(sparse_mat, mat1, mat2);
+  auto sparse_val = sparse_mat->value();
+  // Broadcast the sparse value in batched SDDMM.
+  if (sparse_val.dim() < val.dim()) {
+    sparse_val = sparse_val.unsqueeze(-1);
+  }
+  val = val * sparse_val;
+  return SparseMatrix::ValLike(sparse_mat, val);
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/softmax.cc

@@ -0,0 +1,96 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file softmax.cc
+ * @brief DGL C++ Softmax operator implementation
+ */
+
+#include <sparse/reduction.h>
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+#include "./matmul.h"
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+using namespace torch::autograd;
+
+class SoftmaxAutoGrad : public Function<SoftmaxAutoGrad> {
+ public:
+  static torch::Tensor forward(
+      AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> sparse_mat,
+      torch::Tensor sparse_val, int64_t dim);
+
+  static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs);
+};
+
+torch::Tensor SoftmaxAutoGrad::forward(
+    AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> sparse_mat,
+    torch::Tensor sparse_val, int64_t dim) {
+  // Reduce by columns with dim 1.
+  auto sparse_val_max = ReduceMax(sparse_mat, dim);
+  auto sparse_val_exp =
+      BroadcastSubNoAutoGrad(sparse_mat, sparse_val_max, dim).exp();
+  auto sparse_val_sum =
+      ReduceSum(SparseMatrix::ValLike(sparse_mat, sparse_val_exp), dim);
+  auto sparse_score = BroadcastDivNoAutoGrad(
+      SparseMatrix::ValLike(sparse_mat, sparse_val_exp), sparse_val_sum, dim);
+
+  const bool sparse_requires_grad = sparse_val.requires_grad();
+  torch::Tensor cache_sparse_score;
+  if (sparse_requires_grad) {
+    cache_sparse_score = sparse_score;
+  }
+  ctx->saved_data["sparse_matrix"] = sparse_mat;
+  ctx->saved_data["sparse_requires_grad"] = sparse_requires_grad;
+  ctx->saved_data["dim"] = dim;
+  ctx->save_for_backward({cache_sparse_score});
+  return sparse_score;
+}
+
+tensor_list SoftmaxAutoGrad::backward(
+    AutogradContext* ctx, tensor_list grad_outputs) {
+  auto saved = ctx->get_saved_variables();
+  auto sparse_score = saved[0];
+  auto output_grad = grad_outputs[0];
+
+  auto sparse_mat =
+      ctx->saved_data["sparse_matrix"].toCustomClass<SparseMatrix>();
+  const bool sparse_requires_grad =
+      ctx->saved_data["sparse_requires_grad"].toBool();
+  const int64_t dim = ctx->saved_data["dim"].toInt();
+
+  torch::Tensor sparse_val_grad;
+  if (sparse_requires_grad) {
+    auto sds = sparse_score * output_grad;
+    auto accum = ReduceSum(SparseMatrix::ValLike(sparse_mat, sds), dim);
+    sparse_val_grad =
+        sds - BroadcastMulNoAutoGrad(
+                  SparseMatrix::ValLike(sparse_mat, sparse_score), accum, dim);
+  }
+
+  return {torch::Tensor(), sparse_val_grad, torch::Tensor()};
+}
+
+c10::intrusive_ptr<SparseMatrix> Softmax(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, int64_t dim) {
+  auto sparse_val = sparse_mat->value();
+  bool expand_dim = false;
+  auto new_sparse_mat = sparse_mat;
+  if (sparse_val.dim() == 1) {
+    sparse_val = sparse_val.view({-1, 1});
+    expand_dim = true;
+    new_sparse_mat = SparseMatrix::ValLike(sparse_mat, sparse_val);
+  }
+
+  auto new_sparse_val = SoftmaxAutoGrad::apply(new_sparse_mat, sparse_val, dim);
+
+  if (expand_dim) {
+    new_sparse_val = new_sparse_val.view(-1);
+  }
+  return SparseMatrix::ValLike(sparse_mat, new_sparse_val);
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/sparse_format.cc

@@ -0,0 +1,156 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse_format.cc
+ * @brief DGL C++ sparse format implementations.
+ */
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <sparse/sparse_format.h>
+
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+std::shared_ptr<COO> COOFromOldDGLCOO(const aten::COOMatrix& dgl_coo) {
+  auto row = DGLArrayToTorchTensor(dgl_coo.row);
+  auto col = DGLArrayToTorchTensor(dgl_coo.col);
+  TORCH_CHECK(aten::IsNullArray(dgl_coo.data));
+  auto indices = torch::stack({row, col});
+  return std::make_shared<COO>(
+      COO{dgl_coo.num_rows, dgl_coo.num_cols, indices, dgl_coo.row_sorted,
+          dgl_coo.col_sorted});
+}
+
+aten::COOMatrix COOToOldDGLCOO(const std::shared_ptr<COO>& coo) {
+  auto row = TorchTensorToDGLArray(coo->indices.index({0}));
+  auto col = TorchTensorToDGLArray(coo->indices.index({1}));
+  return aten::COOMatrix(
+      coo->num_rows, coo->num_cols, row, col, aten::NullArray(),
+      coo->row_sorted, coo->col_sorted);
+}
+
+std::shared_ptr<CSR> CSRFromOldDGLCSR(const aten::CSRMatrix& dgl_csr) {
+  auto indptr = DGLArrayToTorchTensor(dgl_csr.indptr);
+  auto indices = DGLArrayToTorchTensor(dgl_csr.indices);
+  auto value_indices = DGLArrayToOptionalTorchTensor(dgl_csr.data);
+  return std::make_shared<CSR>(
+      CSR{dgl_csr.num_rows, dgl_csr.num_cols, indptr, indices, value_indices,
+          dgl_csr.sorted});
+}
+
+aten::CSRMatrix CSRToOldDGLCSR(const std::shared_ptr<CSR>& csr) {
+  auto indptr = TorchTensorToDGLArray(csr->indptr);
+  auto indices = TorchTensorToDGLArray(csr->indices);
+  auto data = OptionalTorchTensorToDGLArray(csr->value_indices);
+  return aten::CSRMatrix(
+      csr->num_rows, csr->num_cols, indptr, indices, data, csr->sorted);
+}
+
+torch::Tensor COOToTorchCOO(
+    const std::shared_ptr<COO>& coo, torch::Tensor value) {
+  torch::Tensor indices = coo->indices;
+  if (value.ndimension() == 2) {
+    return torch::sparse_coo_tensor(
+        indices, value, {coo->num_rows, coo->num_cols, value.size(1)});
+  } else {
+    return torch::sparse_coo_tensor(
+        indices, value, {coo->num_rows, coo->num_cols});
+  }
+}
+
+std::shared_ptr<COO> CSRToCOO(const std::shared_ptr<CSR>& csr) {
+  auto dgl_csr = CSRToOldDGLCSR(csr);
+  auto dgl_coo = aten::CSRToCOO(dgl_csr, csr->value_indices.has_value());
+  return COOFromOldDGLCOO(dgl_coo);
+}
+
+std::shared_ptr<COO> CSCToCOO(const std::shared_ptr<CSR>& csc) {
+  auto dgl_csc = CSRToOldDGLCSR(csc);
+  auto dgl_coo = aten::CSRToCOO(dgl_csc, csc->value_indices.has_value());
+  dgl_coo = aten::COOTranspose(dgl_coo);
+  return COOFromOldDGLCOO(dgl_coo);
+}
+
+std::shared_ptr<CSR> COOToCSR(const std::shared_ptr<COO>& coo) {
+  auto dgl_coo = COOToOldDGLCOO(coo);
+  auto dgl_csr = aten::COOToCSR(dgl_coo);
+  return CSRFromOldDGLCSR(dgl_csr);
+}
+
+std::shared_ptr<CSR> CSCToCSR(const std::shared_ptr<CSR>& csc) {
+  auto dgl_csc = CSRToOldDGLCSR(csc);
+  auto dgl_csr = aten::CSRTranspose(dgl_csc);
+  return CSRFromOldDGLCSR(dgl_csr);
+}
+
+std::shared_ptr<CSR> COOToCSC(const std::shared_ptr<COO>& coo) {
+  auto dgl_coo = COOToOldDGLCOO(coo);
+  auto dgl_coo_transpose = aten::COOTranspose(dgl_coo);
+  auto dgl_csc = aten::COOToCSR(dgl_coo_transpose);
+  return CSRFromOldDGLCSR(dgl_csc);
+}
+
+std::shared_ptr<CSR> CSRToCSC(const std::shared_ptr<CSR>& csr) {
+  auto dgl_csr = CSRToOldDGLCSR(csr);
+  auto dgl_csc = aten::CSRTranspose(dgl_csr);
+  return CSRFromOldDGLCSR(dgl_csc);
+}
+
+std::shared_ptr<COO> DiagToCOO(
+    const std::shared_ptr<Diag>& diag,
+    const c10::TensorOptions& indices_options) {
+  int64_t nnz = std::min(diag->num_rows, diag->num_cols);
+  auto indices = torch::arange(nnz, indices_options).repeat({2, 1});
+  return std::make_shared<COO>(
+      COO{diag->num_rows, diag->num_cols, indices, true, true});
+}
+
+std::shared_ptr<CSR> DiagToCSR(
+    const std::shared_ptr<Diag>& diag,
+    const c10::TensorOptions& indices_options) {
+  int64_t nnz = std::min(diag->num_rows, diag->num_cols);
+  auto indptr = torch::full(diag->num_rows + 1, nnz, indices_options);
+  auto nnz_range = torch::arange(nnz + 1, indices_options);
+  indptr.index_put_({nnz_range}, nnz_range);
+  auto indices = torch::arange(nnz, indices_options);
+  return std::make_shared<CSR>(
+      CSR{diag->num_rows, diag->num_cols, indptr, indices,
+          torch::optional<torch::Tensor>(), true});
+}
+
+std::shared_ptr<CSR> DiagToCSC(
+    const std::shared_ptr<Diag>& diag,
+    const c10::TensorOptions& indices_options) {
+  int64_t nnz = std::min(diag->num_rows, diag->num_cols);
+  auto indptr = torch::full(diag->num_cols + 1, nnz, indices_options);
+  auto nnz_range = torch::arange(nnz + 1, indices_options);
+  indptr.index_put_({nnz_range}, nnz_range);
+  auto indices = torch::arange(nnz, indices_options);
+  return std::make_shared<CSR>(
+      CSR{diag->num_cols, diag->num_rows, indptr, indices,
+          torch::optional<torch::Tensor>(), true});
+}
+
+std::shared_ptr<COO> COOTranspose(const std::shared_ptr<COO>& coo) {
+  auto dgl_coo = COOToOldDGLCOO(coo);
+  auto dgl_coo_tr = aten::COOTranspose(dgl_coo);
+  return COOFromOldDGLCOO(dgl_coo_tr);
+}
+
+std::pair<std::shared_ptr<COO>, torch::Tensor> COOSort(
+    const std::shared_ptr<COO>& coo) {
+  auto encoded_coo =
+      coo->indices.index({0}) * coo->num_cols + coo->indices.index({1});
+  torch::Tensor sorted, perm;
+  std::tie(sorted, perm) = encoded_coo.sort();
+  auto sorted_coo = std::make_shared<COO>(
+      COO{coo->num_rows, coo->num_cols, coo->indices.index_select(1, perm),
+          true, true});
+  return {sorted_coo, perm};
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/sparse_matrix.cc

@@ -0,0 +1,339 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse_matrix.cc
+ * @brief DGL C++ sparse matrix implementations.
+ */
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <c10/util/Logging.h>
+#include <sparse/elementwise_op.h>
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+SparseMatrix::SparseMatrix(
+    const std::shared_ptr<COO>& coo, const std::shared_ptr<CSR>& csr,
+    const std::shared_ptr<CSR>& csc, const std::shared_ptr<Diag>& diag,
+    torch::Tensor value, const std::vector<int64_t>& shape)
+    : coo_(coo),
+      csr_(csr),
+      csc_(csc),
+      diag_(diag),
+      value_(value),
+      shape_(shape) {
+  TORCH_CHECK(
+      coo != nullptr || csr != nullptr || csc != nullptr || diag != nullptr,
+      "At least one of CSR/COO/CSC/Diag is required to construct a "
+      "SparseMatrix.")
+  TORCH_CHECK(
+      shape.size() == 2, "The shape of a sparse matrix should be ",
+      "2-dimensional.");
+  // NOTE: Currently all the tensors of a SparseMatrix should on the same
+  // device. Do we allow the graph structure and values are on different
+  // devices?
+  if (coo != nullptr) {
+    TORCH_CHECK(coo->indices.dim() == 2);
+    TORCH_CHECK(coo->indices.size(0) == 2);
+    TORCH_CHECK(coo->indices.size(1) == value.size(0));
+    TORCH_CHECK(coo->indices.device() == value.device());
+  }
+  if (csr != nullptr) {
+    TORCH_CHECK(csr->indptr.dim() == 1);
+    TORCH_CHECK(csr->indices.dim() == 1);
+    TORCH_CHECK(csr->indptr.size(0) == shape[0] + 1);
+    TORCH_CHECK(csr->indices.size(0) == value.size(0));
+    TORCH_CHECK(csr->indptr.device() == value.device());
+    TORCH_CHECK(csr->indices.device() == value.device());
+  }
+  if (csc != nullptr) {
+    TORCH_CHECK(csc->indptr.dim() == 1);
+    TORCH_CHECK(csc->indices.dim() == 1);
+    TORCH_CHECK(csc->indptr.size(0) == shape[1] + 1);
+    TORCH_CHECK(csc->indices.size(0) == value.size(0));
+    TORCH_CHECK(csc->indptr.device() == value.device());
+    TORCH_CHECK(csc->indices.device() == value.device());
+  }
+  if (diag != nullptr) {
+    TORCH_CHECK(value.size(0) == std::min(diag->num_rows, diag->num_cols));
+  }
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromCOOPointer(
+    const std::shared_ptr<COO>& coo, torch::Tensor value,
+    const std::vector<int64_t>& shape) {
+  return c10::make_intrusive<SparseMatrix>(
+      coo, nullptr, nullptr, nullptr, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromCSRPointer(
+    const std::shared_ptr<CSR>& csr, torch::Tensor value,
+    const std::vector<int64_t>& shape) {
+  return c10::make_intrusive<SparseMatrix>(
+      nullptr, csr, nullptr, nullptr, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromCSCPointer(
+    const std::shared_ptr<CSR>& csc, torch::Tensor value,
+    const std::vector<int64_t>& shape) {
+  return c10::make_intrusive<SparseMatrix>(
+      nullptr, nullptr, csc, nullptr, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromDiagPointer(
+    const std::shared_ptr<Diag>& diag, torch::Tensor value,
+    const std::vector<int64_t>& shape) {
+  return c10::make_intrusive<SparseMatrix>(
+      nullptr, nullptr, nullptr, diag, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromCOO(
+    torch::Tensor indices, torch::Tensor value,
+    const std::vector<int64_t>& shape) {
+  auto coo =
+      std::make_shared<COO>(COO{shape[0], shape[1], indices, false, false});
+  return SparseMatrix::FromCOOPointer(coo, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromCSR(
+    torch::Tensor indptr, torch::Tensor indices, torch::Tensor value,
+    const std::vector<int64_t>& shape) {
+  auto csr = std::make_shared<CSR>(
+      CSR{shape[0], shape[1], indptr, indices, torch::optional<torch::Tensor>(),
+          false});
+  return SparseMatrix::FromCSRPointer(csr, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromCSC(
+    torch::Tensor indptr, torch::Tensor indices, torch::Tensor value,
+    const std::vector<int64_t>& shape) {
+  auto csc = std::make_shared<CSR>(
+      CSR{shape[1], shape[0], indptr, indices, torch::optional<torch::Tensor>(),
+          false});
+  return SparseMatrix::FromCSCPointer(csc, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::FromDiag(
+    torch::Tensor value, const std::vector<int64_t>& shape) {
+  auto diag = std::make_shared<Diag>(Diag{shape[0], shape[1]});
+  return SparseMatrix::FromDiagPointer(diag, value, shape);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::IndexSelect(
+    int64_t dim, torch::Tensor ids) {
+  auto id_array = TorchTensorToDGLArray(ids);
+  bool rowwise = dim == 0;
+  auto csr = rowwise ? this->CSRPtr() : this->CSCPtr();
+  auto slice_csr = dgl::aten::CSRSliceRows(CSRToOldDGLCSR(csr), id_array);
+  auto slice_value =
+      this->value().index_select(0, DGLArrayToTorchTensor(slice_csr.data));
+  // To prevent potential errors in future conversions to the COO format,
+  // where this array might be used as an initialization array for
+  // constructing COO representations, it is necessary to clear this array.
+  slice_csr.data = dgl::aten::NullArray();
+  auto ret = CSRFromOldDGLCSR(slice_csr);
+  if (rowwise) {
+    return SparseMatrix::FromCSRPointer(
+        ret, slice_value, {ret->num_rows, ret->num_cols});
+  } else {
+    return SparseMatrix::FromCSCPointer(
+        ret, slice_value, {ret->num_cols, ret->num_rows});
+  }
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::RangeSelect(
+    int64_t dim, int64_t start, int64_t end) {
+  bool rowwise = dim == 0;
+  auto csr = rowwise ? this->CSRPtr() : this->CSCPtr();
+  auto slice_csr = dgl::aten::CSRSliceRows(CSRToOldDGLCSR(csr), start, end);
+  auto slice_value =
+      this->value().index_select(0, DGLArrayToTorchTensor(slice_csr.data));
+  // To prevent potential errors in future conversions to the COO format,
+  // where this array might be used as an initialization array for
+  // constructing COO representations, it is necessary to clear this array.
+  slice_csr.data = dgl::aten::NullArray();
+  auto ret = CSRFromOldDGLCSR(slice_csr);
+  if (rowwise) {
+    return SparseMatrix::FromCSRPointer(
+        ret, slice_value, {ret->num_rows, ret->num_cols});
+  } else {
+    return SparseMatrix::FromCSCPointer(
+        ret, slice_value, {ret->num_cols, ret->num_rows});
+  }
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::Sample(
+    int64_t dim, int64_t fanout, torch::Tensor ids, bool replace, bool bias) {
+  bool rowwise = dim == 0;
+  auto id_array = TorchTensorToDGLArray(ids);
+  auto csr = rowwise ? this->CSRPtr() : this->CSCPtr();
+  // Slicing matrix.
+  auto slice_csr = dgl::aten::CSRSliceRows(CSRToOldDGLCSR(csr), id_array);
+  auto slice_value =
+      this->value().index_select(0, DGLArrayToTorchTensor(slice_csr.data));
+  // Reset value indices.
+  slice_csr.data = dgl::aten::NullArray();
+
+  auto prob =
+      bias ? TorchTensorToDGLArray(slice_value) : dgl::aten::NullArray();
+  auto slice_id =
+      dgl::aten::Range(0, id_array.NumElements(), 64, id_array->ctx);
+  // Sampling all rows on sliced matrix.
+  auto sample_coo =
+      dgl::aten::CSRRowWiseSampling(slice_csr, slice_id, fanout, prob, replace);
+  auto sample_value =
+      slice_value.index_select(0, DGLArrayToTorchTensor(sample_coo.data));
+  sample_coo.data = dgl::aten::NullArray();
+  auto ret = COOFromOldDGLCOO(sample_coo);
+  if (!rowwise) ret = COOTranspose(ret);
+  return SparseMatrix::FromCOOPointer(
+      ret, sample_value, {ret->num_rows, ret->num_cols});
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::ValLike(
+    const c10::intrusive_ptr<SparseMatrix>& mat, torch::Tensor value) {
+  TORCH_CHECK(
+      mat->value().size(0) == value.size(0), "The first dimension of ",
+      "the old values and the new values must be the same.");
+  TORCH_CHECK(
+      mat->value().device() == value.device(), "The device of the ",
+      "old values and the new values must be the same.");
+  const auto& shape = mat->shape();
+  if (mat->HasDiag()) {
+    return SparseMatrix::FromDiagPointer(mat->DiagPtr(), value, shape);
+  }
+  if (mat->HasCOO()) {
+    return SparseMatrix::FromCOOPointer(mat->COOPtr(), value, shape);
+  }
+  if (mat->HasCSR()) {
+    return SparseMatrix::FromCSRPointer(mat->CSRPtr(), value, shape);
+  }
+  TORCH_CHECK(mat->HasCSC(), "Invalid sparse format for ValLike.")
+  return SparseMatrix::FromCSCPointer(mat->CSCPtr(), value, shape);
+}
+
+std::shared_ptr<COO> SparseMatrix::COOPtr() {
+  if (coo_ == nullptr) {
+    _CreateCOO();
+  }
+  return coo_;
+}
+
+std::shared_ptr<CSR> SparseMatrix::CSRPtr() {
+  if (csr_ == nullptr) {
+    _CreateCSR();
+  }
+  return csr_;
+}
+
+std::shared_ptr<CSR> SparseMatrix::CSCPtr() {
+  if (csc_ == nullptr) {
+    _CreateCSC();
+  }
+  return csc_;
+}
+
+std::shared_ptr<Diag> SparseMatrix::DiagPtr() {
+  TORCH_CHECK(
+      diag_ != nullptr,
+      "Cannot get Diag sparse format from a non-diagonal sparse matrix");
+  return diag_;
+}
+
+std::tuple<torch::Tensor, torch::Tensor> SparseMatrix::COOTensors() {
+  auto coo = COOPtr();
+  return std::make_tuple(coo->indices.index({0}), coo->indices.index({1}));
+}
+
+torch::Tensor SparseMatrix::Indices() {
+  auto coo = COOPtr();
+  return coo->indices;
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+SparseMatrix::CSRTensors() {
+  auto csr = CSRPtr();
+  auto val = value();
+  return std::make_tuple(csr->indptr, csr->indices, csr->value_indices);
+}
+
+std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+SparseMatrix::CSCTensors() {
+  auto csc = CSCPtr();
+  return std::make_tuple(csc->indptr, csc->indices, csc->value_indices);
+}
+
+c10::intrusive_ptr<SparseMatrix> SparseMatrix::Transpose() const {
+  auto shape = shape_;
+  std::swap(shape[0], shape[1]);
+  auto value = value_;
+  if (HasDiag()) {
+    return SparseMatrix::FromDiag(value, shape);
+  } else if (HasCOO()) {
+    auto coo = COOTranspose(coo_);
+    return SparseMatrix::FromCOOPointer(coo, value, shape);
+  } else if (HasCSR()) {
+    return SparseMatrix::FromCSCPointer(csr_, value, shape);
+  } else {
+    return SparseMatrix::FromCSRPointer(csc_, value, shape);
+  }
+}
+
+void SparseMatrix::_CreateCOO() {
+  if (HasCOO()) return;
+  if (HasDiag()) {
+    auto indices_options = torch::TensorOptions()
+                               .dtype(torch::kInt64)
+                               .layout(torch::kStrided)
+                               .device(this->device());
+    coo_ = DiagToCOO(diag_, indices_options);
+  } else if (HasCSR()) {
+    coo_ = CSRToCOO(csr_);
+  } else if (HasCSC()) {
+    coo_ = CSCToCOO(csc_);
+  } else {
+    LOG(FATAL) << "SparseMatrix does not have any sparse format";
+  }
+}
+
+void SparseMatrix::_CreateCSR() {
+  if (HasCSR()) return;
+  if (HasDiag()) {
+    auto indices_options = torch::TensorOptions()
+                               .dtype(torch::kInt64)
+                               .layout(torch::kStrided)
+                               .device(this->device());
+    csr_ = DiagToCSR(diag_, indices_options);
+  } else if (HasCOO()) {
+    csr_ = COOToCSR(coo_);
+  } else if (HasCSC()) {
+    csr_ = CSCToCSR(csc_);
+  } else {
+    LOG(FATAL) << "SparseMatrix does not have any sparse format";
+  }
+}
+
+void SparseMatrix::_CreateCSC() {
+  if (HasCSC()) return;
+  if (HasDiag()) {
+    auto indices_options = torch::TensorOptions()
+                               .dtype(torch::kInt64)
+                               .layout(torch::kStrided)
+                               .device(this->device());
+    csc_ = DiagToCSC(diag_, indices_options);
+  } else if (HasCOO()) {
+    csc_ = COOToCSC(coo_);
+  } else if (HasCSR()) {
+    csc_ = CSRToCSC(csr_);
+  } else {
+    LOG(FATAL) << "SparseMatrix does not have any sparse format";
+  }
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/sparse_matrix_coalesce.cc

@@ -0,0 +1,40 @@
+/**
+ *  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();
+  return SparseMatrix::FromCOO(
+      coalesced_coo.indices(), coalesced_coo.values(), this->shape());
+}
+
+bool SparseMatrix::HasDuplicate() {
+  aten::CSRMatrix dgl_csr;
+  if (HasDiag()) {
+    return false;
+  }
+  // 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

dgl_sparse/src/spmm.cc

@@ -0,0 +1,131 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file spmm.cc
+ * @brief DGL C++ sparse SpMM operator implementation.
+ */
+
+#include <sparse/sddmm.h>
+#include <sparse/sparse_matrix.h>
+#include <sparse/spmm.h>
+#include <torch/script.h>
+
+#include <sstream>
+
+#include "./matmul.h"
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+using namespace torch::autograd;
+
+class SpMMAutoGrad : public Function<SpMMAutoGrad> {
+ public:
+  static torch::Tensor forward(
+      AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> sparse_mat,
+      torch::Tensor sparse_val, torch::Tensor dense_mat);
+
+  static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs);
+};
+
+void _SpMMSanityCheck(
+    c10::intrusive_ptr<SparseMatrix> sparse_mat, torch::Tensor sparse_val,
+    torch::Tensor dense_mat) {
+  const auto& sparse_mat_shape = sparse_mat->shape();
+  auto val_shape = sparse_val.sizes();
+  auto dense_shape = dense_mat.sizes();
+  bool shape_check = true;
+  shape_check &= sparse_mat_shape[1] == dense_shape[0];
+  shape_check &= val_shape.size() <= 2;
+  shape_check &= val_shape[0] == sparse_mat->nnz();
+  shape_check &= dense_shape.size() <= 3;
+  if (dense_shape.size() == 3 || val_shape.size() == 2) {
+    shape_check &= dense_shape.size() == val_shape.size() + 1;
+    shape_check &= dense_shape[2] == val_shape[1];
+  }
+  if (!shape_check) {
+    std::stringstream error;
+    error << "SpMM: Invalid input shapes. sparse_mat: "
+          << c10::IntArrayRef(sparse_mat->shape())
+          << ", sparse_val: " << sparse_mat->value().sizes()
+          << ", dense_mat: " << dense_mat.sizes()
+          << ". Valid input shapes (sparse_mat, dense_mat) are: (1) (n, m) and "
+             "(m, k); (2) (n, m) and (m,); (3) (n, m, b) and (m, k, b).";
+    TORCH_CHECK(false, error.str());
+  }
+  TORCH_CHECK(
+      sparse_val.dtype() == dense_mat.dtype(),
+      "SpMM: the non-zero values does not have the same dtype as the dense "
+      "matrix.");
+  TORCH_CHECK(
+      sparse_val.device() == sparse_mat->device() &&
+          sparse_val.device() == dense_mat.device(),
+      "SpMM: sparse matrix, non-zero values and the dense matrix should be "
+      "on the same device.");
+}
+
+torch::Tensor SpMMAutoGrad::forward(
+    AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> sparse_mat,
+    torch::Tensor sparse_val, torch::Tensor dense_mat) {
+  auto ret = SpMMNoAutoGrad(sparse_mat, sparse_val, dense_mat, false);
+
+  const bool sparse_requires_grad = sparse_val.requires_grad();
+  const bool dense_requires_grad = dense_mat.requires_grad();
+  torch::Tensor cache_sparse_val, cache_dense_mat;
+  if (dense_requires_grad) {
+    cache_sparse_val = sparse_val;
+  }
+  if (sparse_requires_grad) {
+    cache_dense_mat = dense_mat;
+  }
+  ctx->saved_data["sparse_matrix"] = sparse_mat;
+  ctx->saved_data["sparse_requires_grad"] = sparse_requires_grad;
+  ctx->saved_data["dense_requires_grad"] = dense_requires_grad;
+  ctx->save_for_backward({cache_sparse_val, cache_dense_mat});
+  return ret;
+}
+
+tensor_list SpMMAutoGrad::backward(
+    AutogradContext* ctx, tensor_list grad_outputs) {
+  auto saved = ctx->get_saved_variables();
+  auto sparse_val = saved[0];
+  auto dense_mat = saved[1];
+  auto output_grad = grad_outputs[0];
+
+  auto sparse_mat =
+      ctx->saved_data["sparse_matrix"].toCustomClass<SparseMatrix>();
+  const bool sparse_requires_grad =
+      ctx->saved_data["sparse_requires_grad"].toBool();
+  const bool dense_requires_grad =
+      ctx->saved_data["dense_requires_grad"].toBool();
+
+  torch::Tensor dense_mat_grad, sparse_val_grad;
+  if (sparse_requires_grad) {
+    // A @ B = C -> dA = dC @ (B^T)
+    sparse_val_grad = SDDMMNoAutoGrad(sparse_mat, output_grad, dense_mat);
+  }
+  if (dense_requires_grad) {
+    // A @ B = C -> dB = (A^T) @ dC
+    dense_mat_grad = SpMMNoAutoGrad(sparse_mat, sparse_val, output_grad, true);
+  }
+  return {torch::Tensor(), sparse_val_grad, dense_mat_grad};
+}
+
+torch::Tensor SpMM(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat,
+    torch::Tensor dense_mat) {
+  _SpMMSanityCheck(sparse_mat, sparse_mat->value(), dense_mat);
+  bool expand_dim = false;
+  if (dense_mat.dim() == 1) {
+    dense_mat = dense_mat.view({-1, 1});
+    expand_dim = true;
+  }
+  auto ret = SpMMAutoGrad::apply(sparse_mat, sparse_mat->value(), dense_mat);
+  if (expand_dim) {
+    ret = ret.view(-1);
+  }
+  return ret;
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/spspmm.cc

@@ -0,0 +1,172 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file spspmm.cc
+ * @brief DGL C++ sparse SpSpMM operator implementation.
+ */
+
+#include <sparse/sddmm.h>
+#include <sparse/sparse_matrix.h>
+#include <sparse/spspmm.h>
+#include <torch/script.h>
+
+#include "./matmul.h"
+#include "./utils.h"
+
+namespace dgl {
+namespace sparse {
+
+using namespace torch::autograd;
+
+class SpSpMMAutoGrad : public Function<SpSpMMAutoGrad> {
+ public:
+  static variable_list forward(
+      AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> lhs_mat,
+      torch::Tensor lhs_val, c10::intrusive_ptr<SparseMatrix> rhs_mat,
+      torch::Tensor rhs_val);
+
+  static tensor_list backward(AutogradContext* ctx, tensor_list grad_outputs);
+};
+
+void _SpSpMMSanityCheck(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat) {
+  const auto& lhs_shape = lhs_mat->shape();
+  const auto& rhs_shape = rhs_mat->shape();
+  TORCH_CHECK(
+      lhs_shape[1] == rhs_shape[0],
+      "SpSpMM: the second dim of lhs_mat should be equal to the first dim ",
+      "of the second matrix");
+  TORCH_CHECK(
+      lhs_mat->value().dim() == 1,
+      "SpSpMM: the value shape of lhs_mat should be 1-D");
+  TORCH_CHECK(
+      rhs_mat->value().dim() == 1,
+      "SpSpMM: the value shape of rhs_mat should be 1-D");
+  TORCH_CHECK(
+      lhs_mat->device() == rhs_mat->device(),
+      "SpSpMM: lhs_mat and rhs_mat should be on the same device");
+  TORCH_CHECK(
+      lhs_mat->dtype() == rhs_mat->dtype(),
+      "SpSpMM: lhs_mat and rhs_mat should have the same dtype");
+  TORCH_CHECK(
+      !lhs_mat->HasDuplicate(),
+      "SpSpMM does not support lhs_mat with duplicate indices. ",
+      "Call A = A.coalesce() to dedup first.");
+  TORCH_CHECK(
+      !rhs_mat->HasDuplicate(),
+      "SpSpMM does not support rhs_mat with duplicate indices. ",
+      "Call A = A.coalesce() to dedup first.");
+}
+
+// Mask select value of `mat` by `sub_mat`.
+torch::Tensor _CSRMask(
+    const c10::intrusive_ptr<SparseMatrix>& mat, torch::Tensor value,
+    const c10::intrusive_ptr<SparseMatrix>& sub_mat) {
+  auto csr = CSRToOldDGLCSR(mat->CSRPtr());
+  auto val = TorchTensorToDGLArray(value);
+  auto row = TorchTensorToDGLArray(sub_mat->COOPtr()->indices.index({0}));
+  auto col = TorchTensorToDGLArray(sub_mat->COOPtr()->indices.index({1}));
+  runtime::NDArray ret = aten::CSRGetFloatingData(csr, row, col, val, 0.);
+  return DGLArrayToTorchTensor(ret);
+}
+
+variable_list SpSpMMAutoGrad::forward(
+    AutogradContext* ctx, c10::intrusive_ptr<SparseMatrix> lhs_mat,
+    torch::Tensor lhs_val, c10::intrusive_ptr<SparseMatrix> rhs_mat,
+    torch::Tensor rhs_val) {
+  auto ret_mat =
+      SpSpMMNoAutoGrad(lhs_mat, lhs_val, rhs_mat, rhs_val, false, false);
+
+  ctx->saved_data["lhs_mat"] = lhs_mat;
+  ctx->saved_data["rhs_mat"] = rhs_mat;
+  ctx->saved_data["ret_mat"] = ret_mat;
+  ctx->saved_data["lhs_require_grad"] = lhs_val.requires_grad();
+  ctx->saved_data["rhs_require_grad"] = rhs_val.requires_grad();
+  ctx->save_for_backward({lhs_val, rhs_val});
+
+  auto csr = ret_mat->CSRPtr();
+  auto val = ret_mat->value();
+  TORCH_CHECK(!csr->value_indices.has_value());
+  return {csr->indptr, csr->indices, val};
+}
+
+tensor_list SpSpMMAutoGrad::backward(
+    AutogradContext* ctx, tensor_list grad_outputs) {
+  auto saved = ctx->get_saved_variables();
+  auto lhs_val = saved[0];
+  auto rhs_val = saved[1];
+  auto output_grad = grad_outputs[2];
+  auto lhs_mat = ctx->saved_data["lhs_mat"].toCustomClass<SparseMatrix>();
+  auto rhs_mat = ctx->saved_data["rhs_mat"].toCustomClass<SparseMatrix>();
+  auto ret_mat = ctx->saved_data["ret_mat"].toCustomClass<SparseMatrix>();
+  torch::Tensor lhs_val_grad, rhs_val_grad;
+
+  if (ctx->saved_data["lhs_require_grad"].toBool()) {
+    // A @ B = C -> dA = dC @ (B^T)
+    auto lhs_mat_grad =
+        SpSpMMNoAutoGrad(ret_mat, output_grad, rhs_mat, rhs_val, false, true);
+    lhs_val_grad = _CSRMask(lhs_mat_grad, lhs_mat_grad->value(), lhs_mat);
+  }
+  if (ctx->saved_data["rhs_require_grad"].toBool()) {
+    // A @ B = C -> dB = (A^T) @ dC
+    auto rhs_mat_grad =
+        SpSpMMNoAutoGrad(lhs_mat, lhs_val, ret_mat, output_grad, true, false);
+    rhs_val_grad = _CSRMask(rhs_mat_grad, rhs_mat_grad->value(), rhs_mat);
+  }
+  return {torch::Tensor(), lhs_val_grad, torch::Tensor(), rhs_val_grad};
+}
+
+c10::intrusive_ptr<SparseMatrix> DiagSpSpMM(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat) {
+  if (lhs_mat->HasDiag() && rhs_mat->HasDiag()) {
+    // Diag @ Diag
+    const int64_t m = lhs_mat->shape()[0];
+    const int64_t n = lhs_mat->shape()[1];
+    const int64_t p = rhs_mat->shape()[1];
+    const int64_t common_diag_len = std::min({m, n, p});
+    const int64_t new_diag_len = std::min(m, p);
+    auto slice = torch::indexing::Slice(0, common_diag_len);
+    auto new_val =
+        lhs_mat->value().index({slice}) * rhs_mat->value().index({slice});
+    new_val =
+        torch::constant_pad_nd(new_val, {0, new_diag_len - common_diag_len}, 0);
+    return SparseMatrix::FromDiag(new_val, {m, p});
+  }
+  if (lhs_mat->HasDiag() && !rhs_mat->HasDiag()) {
+    // Diag @ Sparse
+    auto row = rhs_mat->Indices().index({0});
+    auto val = lhs_mat->value().index_select(0, row) * rhs_mat->value();
+    return SparseMatrix::ValLike(rhs_mat, val);
+  }
+  if (!lhs_mat->HasDiag() && rhs_mat->HasDiag()) {
+    // Sparse @ Diag
+    auto col = lhs_mat->Indices().index({1});
+    auto val = rhs_mat->value().index_select(0, col) * lhs_mat->value();
+    return SparseMatrix::ValLike(lhs_mat, val);
+  }
+  TORCH_CHECK(
+      false,
+      "For DiagSpSpMM, at least one of the sparse matries need to have kDiag "
+      "format");
+  return c10::intrusive_ptr<SparseMatrix>();
+}
+
+c10::intrusive_ptr<SparseMatrix> SpSpMM(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat) {
+  _SpSpMMSanityCheck(lhs_mat, rhs_mat);
+  if (lhs_mat->HasDiag() || rhs_mat->HasDiag()) {
+    return DiagSpSpMM(lhs_mat, rhs_mat);
+  }
+  auto results = SpSpMMAutoGrad::apply(
+      lhs_mat, lhs_mat->value(), rhs_mat, rhs_mat->value());
+  std::vector<int64_t> ret_shape({lhs_mat->shape()[0], rhs_mat->shape()[1]});
+  auto indptr = results[0];
+  auto indices = results[1];
+  auto value = results[2];
+  return SparseMatrix::FromCSR(indptr, indices, value, ret_shape);
+}
+
+}  // namespace sparse
+}  // namespace dgl

dgl_sparse/src/utils.h

@@ -0,0 +1,84 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file utils.h
+ * @brief DGL C++ sparse API utilities
+ */
+#ifndef DGL_SPARSE_UTILS_H_
+#define DGL_SPARSE_UTILS_H_
+
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <ATen/DLConvertor.h>
+#include <sparse/sparse_matrix.h>
+#include <torch/custom_class.h>
+#include <torch/script.h>
+
+namespace dgl {
+namespace sparse {
+
+/** @brief Find a proper sparse format for two sparse matrices. It chooses
+ * COO if anyone of the sparse matrices has COO format. If none of them has
+ * COO, it tries CSR and CSC in the same manner. */
+inline static SparseFormat FindAnyExistingFormat(
+    const c10::intrusive_ptr<SparseMatrix>& A,
+    const c10::intrusive_ptr<SparseMatrix>& B) {
+  SparseFormat fmt;
+  if (A->HasCOO() || B->HasCOO()) {
+    fmt = SparseFormat::kCOO;
+  } else if (A->HasCSR() || B->HasCSR()) {
+    fmt = SparseFormat::kCSR;
+  } else {
+    fmt = SparseFormat::kCSC;
+  }
+  return fmt;
+}
+
+/** @brief Check whether two matrices has the same dtype and shape for
+ * elementwise operators. */
+inline static void ElementwiseOpSanityCheck(
+    const c10::intrusive_ptr<SparseMatrix>& A,
+    const c10::intrusive_ptr<SparseMatrix>& B) {
+  TORCH_CHECK(
+      A->value().dtype() == B->value().dtype(),
+      "Elementwise operators"
+      " do not support two sparse matrices with different dtypes.");
+  TORCH_CHECK(
+      A->shape()[0] == B->shape()[0] && A->shape()[1] == B->shape()[1],
+      "Elementwise operators do not support two sparse matrices with different"
+      " shapes.");
+}
+
+/** @brief Convert a Torch tensor to a DGL array. */
+inline static runtime::NDArray TorchTensorToDGLArray(torch::Tensor tensor) {
+  return runtime::DLPackConvert::FromDLPack(at::toDLPack(tensor.contiguous()));
+}
+
+/** @brief Convert a DGL array to a Torch tensor. */
+inline static torch::Tensor DGLArrayToTorchTensor(runtime::NDArray array) {
+  return at::fromDLPack(runtime::DLPackConvert::ToDLPack(array));
+}
+
+/** @brief Convert an optional Torch tensor to a DGL array. */
+inline static runtime::NDArray OptionalTorchTensorToDGLArray(
+    torch::optional<torch::Tensor> tensor) {
+  if (!tensor.has_value()) {
+    return aten::NullArray();
+  }
+  return TorchTensorToDGLArray(tensor.value());
+}
+
+/** @brief Convert a DGL array to an optional Torch tensor. */
+inline static torch::optional<torch::Tensor> DGLArrayToOptionalTorchTensor(
+    runtime::NDArray array) {
+  if (aten::IsNullArray(array)) {
+    return torch::optional<torch::Tensor>();
+  }
+  return torch::make_optional<torch::Tensor>(DGLArrayToTorchTensor(array));
+}
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // DGL_SPARSE_UTILS_H_