chore: import upstream snapshot with attribution

dgl_sparse/CMakeLists.txt

@@ -0,0 +1,93 @@
+cmake_minimum_required(VERSION 3.8)
+project(dgl_sparse C CXX)
+
+# Find PyTorch cmake files and PyTorch versions with the python interpreter $PYTHON_INTERP
+# ("python3" or "python" if empty)
+if(NOT PYTHON_INTERP)
+  find_program(PYTHON_INTERP NAMES python3 python)
+endif()
+message(STATUS "Using Python interpreter: ${PYTHON_INTERP}")
+file(TO_NATIVE_PATH ${CMAKE_CURRENT_SOURCE_DIR}/find_cmake.py FIND_CMAKE_PY)
+execute_process(
+  COMMAND ${PYTHON_INTERP} ${FIND_CMAKE_PY}
+  OUTPUT_VARIABLE TORCH_PREFIX_VER
+  OUTPUT_STRIP_TRAILING_WHITESPACE)
+message(STATUS "find_cmake.py output: ${TORCH_PREFIX_VER}")
+list(GET TORCH_PREFIX_VER 0 TORCH_PREFIX)
+list(GET TORCH_PREFIX_VER 1 TORCH_VER)
+message(STATUS "Configuring for PyTorch ${TORCH_VER}")
+string(REPLACE "." ";" TORCH_VERSION_LIST ${TORCH_VER})
+list(GET TORCH_VERSION_LIST 0 TORCH_VERSION_MAJOR)
+list(GET TORCH_VERSION_LIST 1 TORCH_VERSION_MINOR)
+
+set(SPARSE_LINKER_LIBS "")
+
+if(USE_CUDA)
+  add_definitions(-DDGL_USE_CUDA)
+  enable_language(CUDA)
+endif()
+
+# For windows, define NOMINMAX to avoid conflict with std::min/max
+if(MSVC)
+  add_definitions(-DNOMINMAX)
+endif()
+
+set(Torch_DIR "${TORCH_PREFIX}/Torch")
+message(STATUS "Setting directory to ${Torch_DIR}")
+find_package(Torch REQUIRED)
+set(CMAKE_C_FLAGS "${CMAKE_C_FLAGS} ${TORCH_C_FLAGS}")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} ${TORCH_CXX_FLAGS}")
+set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -O0 -g3 -ggdb")
+
+set(LIB_DGL_SPARSE_NAME "dgl_sparse_pytorch_${TORCH_VER}")
+list(APPEND SPARSE_LINKER_LIBS ${TORCH_LIBRARIES})
+
+set(SPARSE_DIR "${CMAKE_CURRENT_SOURCE_DIR}/src")
+set(SPARSE_INCLUDE "${CMAKE_CURRENT_SOURCE_DIR}/include")
+file(GLOB SPARSE_HEADERS ${SPARSE_INCLUDE})
+file(GLOB SPARSE_SRC
+  ${SPARSE_DIR}/*.cc
+  ${SPARSE_DIR}/cpu/*.cc
+)
+if(USE_CUDA)
+  file(GLOB SPARSE_CUDA_SRC
+    ${SPARSE_DIR}/cuda/*.cu
+  )
+  list(APPEND SPARSE_SRC ${SPARSE_CUDA_SRC})
+endif()
+
+add_library(${LIB_DGL_SPARSE_NAME} SHARED ${SPARSE_SRC} ${SPARSE_HEADERS})
+target_include_directories(
+  ${LIB_DGL_SPARSE_NAME} PRIVATE ${SPARSE_DIR} ${SPARSE_HEADERS})
+target_link_libraries(${LIB_DGL_SPARSE_NAME} ${SPARSE_LINKER_LIBS})
+target_compile_definitions(${LIB_DGL_SPARSE_NAME} PRIVATE TORCH_VERSION_MAJOR=${TORCH_VERSION_MAJOR})
+target_compile_definitions(${LIB_DGL_SPARSE_NAME} PRIVATE TORCH_VERSION_MINOR=${TORCH_VERSION_MINOR})
+
+target_include_directories(${LIB_DGL_SPARSE_NAME} PRIVATE "${CMAKE_SOURCE_DIR}/third_party/dmlc-core/include")
+message(STATUS "DGL include directories: ${DGL_INCLUDE_DIRS}")
+target_include_directories(${LIB_DGL_SPARSE_NAME} PRIVATE ${DGL_INCLUDE_DIRS})
+target_link_directories(${LIB_DGL_SPARSE_NAME} PRIVATE ${DGL_BUILD_DIR} "${DGL_BUILD_DIR}/third_party/dmlc-core")
+
+# The Torch CMake configuration only sets up the path for the MKL library when
+# using the conda distribution. The following is a workaround to address this
+# when using a standalone installation of MKL.
+if(DEFINED MKL_LIBRARIES)
+  target_link_directories(${LIB_DGL_SPARSE_NAME} PRIVATE ${MKL_ROOT}/lib/${MKL_ARCH})
+endif()
+if (EXTERNAL_DMLC_LIB_PATH)
+   # external dmlc requires OpenMP link
+   include(FindOpenMP)
+   if(OPENMP_FOUND)
+        set(CMAKE_C_FLAGS "${OpenMP_C_FLAGS} ${CMAKE_C_FLAGS}")
+        set(CMAKE_CXX_FLAGS "${OpenMP_CXX_FLAGS} ${CMAKE_CXX_FLAGS}")
+   endif(OPENMP_FOUND)	
+   message(STATUS "looking for dmlc library in ${EXTERNAL_DMLC_LIB_PATH}")
+   find_package(dmlc REQUIRED HINTS ${EXTERNAL_DMLC_LIB_PATH})
+   target_link_libraries(${LIB_DGL_SPARSE_NAME} dmlc::dmlc dgl)
+else (EXTERNAL_DMLC_LIB_PATH)
+   target_link_libraries(${LIB_DGL_SPARSE_NAME} dmlc dgl)
+endif()
+set(GOOGLE_TEST 0) # Turn off dmlc-core test
+
+# Configure dgl_sparse library to use C++17 standard for compatibility with PyTorch
+set_property(TARGET ${LIB_DGL_SPARSE_NAME} PROPERTY CXX_STANDARD 17)

dgl_sparse/build.bat

@@ -0,0 +1,33 @@
+REM Helper script to build DGL sparse libraries for PyTorch
+@ECHO OFF
+SETLOCAL EnableDelayedExpansion
+
+MD "%BINDIR%\dgl_sparse"
+DEL /S /Q build
+MD build
+PUSHD build
+
+IF x%1x == xx GOTO single
+COPY %BINDIR%\third_party\dmlc-core\Release\dmlc.lib %BINDIR%
+COPY %BINDIR%\Release\dgl.lib %BINDIR%
+
+FOR %%X IN (%*) DO (
+	DEL /S /Q *
+	"%CMAKE_COMMAND%" -DDGL_BUILD_DIR=%BINDIR% -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DDGL_INCLUDE_DIRS=%INCLUDEDIR: =;% -DUSE_CUDA=%USE_CUDA% -DPYTHON_INTERP=%%X .. -G "Visual Studio 16 2019" || EXIT /B 1
+	msbuild dgl_sparse.sln /m /nr:false || EXIT /B 1
+	COPY /Y Release\*.dll "%BINDIR%\dgl_sparse" || EXIT /B 1
+)
+
+GOTO end
+
+:single
+
+DEL /S /Q *
+"%CMAKE_COMMAND%" -DDGL_BUILD_DIR=%BINDIR% -DCMAKE_CONFIGURATION_TYPES=Release -DCUDA_TOOLKIT_ROOT_DIR="%CUDA_TOOLKIT_ROOT_DIR%" -DTORCH_CUDA_ARCH_LIST=%TORCH_CUDA_ARCH_LIST% -DUSE_CUDA=%USE_CUDA% -DDGL_INCLUDE_DIRS=%INCLUDEDIR: =;% .. -G "Visual Studio 16 2019" || EXIT /B 1
+msbuild dgl_sparse.sln /m /nr:false || EXIT /B 1
+COPY /Y Release\*.dll "%BINDIR%\dgl_sparse" || EXIT /B 1
+
+:end
+POPD
+
+ENDLOCAL

dgl_sparse/build.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+# Helper script to build dgl sparse libraries for PyTorch
+set -e
+
+mkdir -p build
+mkdir -p $BINDIR/dgl_sparse
+cd build
+
+if [ $(uname) = 'Darwin' ]; then
+	CPSOURCE=*.dylib
+else
+	CPSOURCE=*.so
+fi
+
+CMAKE_FLAGS="-DCUDA_TOOLKIT_ROOT_DIR=$CUDA_TOOLKIT_ROOT_DIR -DTORCH_CUDA_ARCH_LIST=$TORCH_CUDA_ARCH_LIST -DUSE_CUDA=$USE_CUDA -DEXTERNAL_DMLC_LIB_PATH=$EXTERNAL_DMLC_LIB_PATH"
+# CMake passes in the list of directories separated by spaces.  Here we replace them with semicolons.
+CMAKE_FLAGS="$CMAKE_FLAGS -DDGL_INCLUDE_DIRS=${INCLUDEDIR// /;} -DDGL_BUILD_DIR=$BINDIR"
+echo $CMAKE_FLAGS
+
+if [ $# -eq 0 ]; then
+	$CMAKE_COMMAND $CMAKE_FLAGS ..
+	make -j
+	cp -v $CPSOURCE $BINDIR/dgl_sparse
+else
+	for PYTHON_INTERP in $@; do
+		TORCH_VER=$($PYTHON_INTERP -c 'import torch; print(torch.__version__.split("+")[0])')
+		mkdir -p $TORCH_VER
+		cd $TORCH_VER
+		$CMAKE_COMMAND $CMAKE_FLAGS -DPYTHON_INTERP=$PYTHON_INTERP ../..
+		make -j
+		cp -v $CPSOURCE $BINDIR/dgl_sparse
+		cd ..
+	done
+fi

dgl_sparse/find_cmake.py

@@ -0,0 +1,11 @@
+import os
+
+import torch
+
+cmake_prefix_path = getattr(
+    torch.utils,
+    "cmake_prefix_path",
+    os.path.join(os.path.dirname(torch.__file__), "share", "cmake"),
+)
+version = torch.__version__.split("+")[0]
+print(";".join([cmake_prefix_path, version]))

dgl_sparse/include/sparse/dgl_headers.h

@@ -0,0 +1,43 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/dgl_headers.h
+ * @brief DGL headers used in the sparse library. This is a workaround to
+ * avoid the macro naming conflict between dmlc/logging.h and torch logger. This
+ * file includes all the DGL headers used in the sparse library and
+ * undefines logging macros defined in dmlc/logging.h. There are two rules to
+ * use this file. (1) All DGL headers used in the sparse library should be and
+ * only be registered in this file. (2) When including Pytorch headers, this
+ * file should be included in advance.
+ */
+#ifndef SPARSE_DGL_HEADERS_H_
+#define SPARSE_DGL_HEADERS_H_
+
+#include <dgl/aten/coo.h>
+#include <dgl/aten/csr.h>
+#include <dgl/kernel.h>
+#include <dgl/runtime/dlpack_convert.h>
+#include <dmlc/logging.h>
+
+#undef CHECK
+#undef CHECK_OP
+#undef CHECK_EQ
+#undef CHECK_NE
+#undef CHECK_LE
+#undef CHECK_LT
+#undef CHECK_GE
+#undef CHECK_GT
+#undef CHECK_NOTNULL
+#undef DCHECK
+#undef DCHECK_EQ
+#undef DCHECK_NE
+#undef DCHECK_LE
+#undef DCHECK_LT
+#undef DCHECK_GE
+#undef DCHECK_GT
+#undef DCHECK_NOTNULL
+#undef VLOG
+#undef LOG
+#undef DLOG
+#undef LOG_IF
+
+#endif  // SPARSE_DGL_HEADERS_H_

dgl_sparse/include/sparse/elementwise_op.h

@@ -0,0 +1,53 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/elementwise_op.h
+ * @brief DGL C++ sparse elementwise operators.
+ */
+#ifndef SPARSE_ELEMENTWISE_OP_H_
+#define SPARSE_ELEMENTWISE_OP_H_
+
+#include <sparse/sparse_matrix.h>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Adds two sparse matrices possibly with different sparsities.
+ *
+ * @param lhs_mat SparseMatrix
+ * @param rhs_mat SparseMatrix
+ *
+ * @return SparseMatrix
+ */
+c10::intrusive_ptr<SparseMatrix> SpSpAdd(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat);
+
+/**
+ * @brief Multiplies two sparse matrices possibly with different sparsities.
+ *
+ * @param lhs_mat SparseMatrix
+ * @param rhs_mat SparseMatrix
+ *
+ * @return SparseMatrix
+ */
+c10::intrusive_ptr<SparseMatrix> SpSpMul(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat);
+
+/**
+ * @brief Divides two sparse matrices with the same sparsity.
+ *
+ * @param lhs_mat SparseMatrix
+ * @param rhs_mat SparseMatrix
+ *
+ * @return SparseMatrix
+ */
+c10::intrusive_ptr<SparseMatrix> SpSpDiv(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_ELEMENTWISE_OP_H_

dgl_sparse/include/sparse/matrix_ops.h

@@ -0,0 +1,54 @@
+/**
+ *  Copyright (c) 2023 by Contributors
+ * @file sparse/matrix_ops.h
+ * @brief DGL C++ sparse matrix operators.
+ */
+#ifndef SPARSE_MATRIX_OPS_H_
+#define SPARSE_MATRIX_OPS_H_
+
+#include <sparse/sparse_matrix.h>
+
+#include <tuple>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Compute the intersection of two COO matrices. Return the intersection
+ * matrix, and the indices of the intersection in the left-hand-side and
+ * right-hand-side 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);
+
+/**
+ * @brief Compact sparse matrix by removing rows or columns without non-zero
+ * elements in the sparse matrix and relabeling indices of the dimension.
+ *
+ * This function serves a dual purpose: it allows you to reorganize the
+ * indices within a specific dimension (rows or columns) of the sparse matrix
+ * and, if needed, place certain 'leading_indices' at the beginning of the
+ * compact dimension.
+ *
+ * @param mat The sparse matrix to be compacted.
+ * @param dim The dimension to compact. Should be 0 or 1. Use 0 for row-wise
+ *        compaction and 1 for column-wise compaction.
+ * @param leading_indices An optional tensor containing row or column ids that
+ *        should be placed at the beginning of the compact dimension.
+ *
+ * @return A tuple containing the compacted sparse matrix and the index mapping
+ *         of the compact dimension from the new index to the original index.
+ */
+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);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_MATRIX_OPS_H_

dgl_sparse/include/sparse/reduction.h

@@ -0,0 +1,64 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/reduction.h
+ * @brief DGL C++ sparse matrix reduction operators.
+ */
+#ifndef SPARSE_REDUCTION_H_
+#define SPARSE_REDUCTION_H_
+
+#include <sparse/sparse_matrix.h>
+
+#include <string>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Reduces a sparse matrix along the specified sparse dimension.
+ *
+ * @param A The sparse matrix.
+ * @param dim The sparse dimension to reduce along.  Must be either 0 (rows) or
+ * 1 (columns).
+ * @param reduce The reduce operator.  Must be either "sum", "smin", "smax",
+ * "mean", or "sprod".
+ *
+ * @return Tensor
+ */
+torch::Tensor Reduce(
+    const c10::intrusive_ptr<SparseMatrix>& A, const std::string& reduce,
+    const torch::optional<int64_t>& dim = torch::nullopt);
+
+inline torch::Tensor ReduceSum(
+    const c10::intrusive_ptr<SparseMatrix>& A,
+    const torch::optional<int64_t>& dim = torch::nullopt) {
+  return Reduce(A, "sum", dim);
+}
+
+inline torch::Tensor ReduceMin(
+    const c10::intrusive_ptr<SparseMatrix>& A,
+    const torch::optional<int64_t>& dim = torch::nullopt) {
+  return Reduce(A, "smin", dim);
+}
+
+inline torch::Tensor ReduceMax(
+    const c10::intrusive_ptr<SparseMatrix>& A,
+    const torch::optional<int64_t>& dim = torch::nullopt) {
+  return Reduce(A, "smax", dim);
+}
+
+inline torch::Tensor ReduceMean(
+    const c10::intrusive_ptr<SparseMatrix>& A,
+    const torch::optional<int64_t>& dim = torch::nullopt) {
+  return Reduce(A, "smean", dim);
+}
+
+inline torch::Tensor ReduceProd(
+    const c10::intrusive_ptr<SparseMatrix>& A,
+    const torch::optional<int64_t>& dim = torch::nullopt) {
+  return Reduce(A, "sprod", dim);
+}
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_REDUCTION_H_

dgl_sparse/include/sparse/sddmm.h

@@ -0,0 +1,44 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/sddmm.h
+ * @brief DGL C++ SDDMM operator.
+ */
+#ifndef SPARSE_SDDMM_H_
+#define SPARSE_SDDMM_H_
+
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Perform a sampled matrix multiplication of a sparse matrix and two
+ * dense matrices. It calculates `sparse_mat * (mat1 @ mat2)`. The SDDMM can be
+ * batched, where the batch dimension is the last dimension for all input
+ * matrices.
+ *
+ * There are four cases for the input and output matrix shapes:
+ *   (1) (n, m), (n, k), (k, m), and (n, m);
+ *   (2) (n, m), (n,), and (m,), and (n, m);
+ *   (3) (n, m, b), (n, k, b), (k, m, b), and (n, m, b);
+ *   (4) (n, m), (n, k, b), (k, m, b), and (n, m, b);
+ *
+ * This function supports autograd for `mat1` and `mat2` but does not support
+ * high order gradient.
+ *
+ *
+ * @param sparse_mat The sparse matrix.
+ * @param mat1 The first dense matrix.
+ * @param mat2 The second dense matrix.
+ *
+ * @return SparseMatrix
+ */
+c10::intrusive_ptr<SparseMatrix> SDDMM(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, torch::Tensor mat1,
+    torch::Tensor mat2);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_SDDMM_H_

dgl_sparse/include/sparse/softmax.h

@@ -0,0 +1,33 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/softmax.h
+ * @brief DGL C++ Softmax operator
+ */
+#ifndef SPARSE_SOFTMAX_H_
+#define SPARSE_SOFTMAX_H_
+
+#include <sparse/sparse_matrix.h>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Apply softmax to the non-zero entries of the sparse matrix on the
+ * dimension dim. dim = 0 or 1 indicates column-wise or row-wise softmax
+ * respectively.
+ *
+ * This function supports autograd for the sparse matrix, but it does not
+ * support higher order gradient.
+ *
+ * @param sparse_mat The sparse matrix
+ * @param dim The dimension to apply softmax
+ *
+ * @return Sparse matrix
+ */
+c10::intrusive_ptr<SparseMatrix> Softmax(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat, int64_t dim);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_SOFTMAX_H_

dgl_sparse/include/sparse/sparse_format.h

@@ -0,0 +1,127 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/sparse_format.h
+ * @brief DGL C++ sparse format header.
+ */
+#ifndef SPARSE_SPARSE_FORMAT_H_
+#define SPARSE_SPARSE_FORMAT_H_
+
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <torch/custom_class.h>
+#include <torch/script.h>
+
+#include <memory>
+#include <utility>
+
+namespace dgl {
+namespace sparse {
+
+/** @brief SparseFormat enumeration. */
+enum SparseFormat { kCOO, kCSR, kCSC, kDiag };
+
+/** @brief COO sparse structure. */
+struct COO {
+  /** @brief The shape of the matrix. */
+  int64_t num_rows = 0, num_cols = 0;
+  /**
+   * @brief COO tensor of shape (2, nnz), stacking the row and column indices.
+   */
+  torch::Tensor indices;
+  /** @brief Whether the row indices are sorted. */
+  bool row_sorted = false;
+  /** @brief Whether the column indices per row are sorted. */
+  bool col_sorted = false;
+};
+
+/** @brief CSR sparse structure. */
+struct CSR {
+  /** @brief The dense shape of the matrix. */
+  int64_t num_rows = 0, num_cols = 0;
+  /** @brief CSR format index pointer array of the matrix. */
+  torch::Tensor indptr;
+  /** @brief CSR format index array of the matrix. */
+  torch::Tensor indices;
+  /** @brief Data index tensor. When it is null, assume it is from 0 to NNZ - 1.
+   */
+  torch::optional<torch::Tensor> value_indices;
+  /** @brief Whether the column indices per row are sorted. */
+  bool sorted = false;
+};
+
+struct Diag {
+  /** @brief The dense shape of the matrix. */
+  int64_t num_rows = 0, num_cols = 0;
+};
+
+/** @brief Convert an old DGL COO format to a COO in the sparse library. */
+std::shared_ptr<COO> COOFromOldDGLCOO(const aten::COOMatrix& dgl_coo);
+
+/** @brief Convert a COO in the sparse library to an old DGL COO matrix. */
+aten::COOMatrix COOToOldDGLCOO(const std::shared_ptr<COO>& coo);
+
+/** @brief Convert an old DGL CSR format to a CSR in the sparse library. */
+std::shared_ptr<CSR> CSRFromOldDGLCSR(const aten::CSRMatrix& dgl_csr);
+
+/** @brief Convert a CSR in the sparse library to an old DGL CSR matrix. */
+aten::CSRMatrix CSRToOldDGLCSR(const std::shared_ptr<CSR>& csr);
+
+/**
+ *  @brief Convert a COO and its nonzero values to a Torch COO matrix.
+ *  @param coo The COO format in the sparse library
+ *  @param value Values of the sparse matrix
+ *
+ *  @return Torch Sparse Tensor in COO format
+ */
+torch::Tensor COOToTorchCOO(
+    const std::shared_ptr<COO>& coo, torch::Tensor value);
+
+/** @brief Convert a CSR format to COO format. */
+std::shared_ptr<COO> CSRToCOO(const std::shared_ptr<CSR>& csr);
+
+/** @brief Convert a CSC format to COO format. */
+std::shared_ptr<COO> CSCToCOO(const std::shared_ptr<CSR>& csc);
+
+/** @brief Convert a COO format to CSR format. */
+std::shared_ptr<CSR> COOToCSR(const std::shared_ptr<COO>& coo);
+
+/** @brief Convert a CSC format to CSR format. */
+std::shared_ptr<CSR> CSCToCSR(const std::shared_ptr<CSR>& csc);
+
+/** @brief Convert a COO format to CSC format. */
+std::shared_ptr<CSR> COOToCSC(const std::shared_ptr<COO>& coo);
+
+/** @brief Convert a CSR format to CSC format. */
+std::shared_ptr<CSR> CSRToCSC(const std::shared_ptr<CSR>& csr);
+
+/** @brief Convert a Diag format to COO format. */
+std::shared_ptr<COO> DiagToCOO(
+    const std::shared_ptr<Diag>& diag,
+    const c10::TensorOptions& indices_options);
+
+/** @brief Convert a Diag format to CSR format. */
+std::shared_ptr<CSR> DiagToCSR(
+    const std::shared_ptr<Diag>& diag,
+    const c10::TensorOptions& indices_options);
+
+/** @brief Convert a Diag format to CSC format. */
+std::shared_ptr<CSR> DiagToCSC(
+    const std::shared_ptr<Diag>& diag,
+    const c10::TensorOptions& indices_options);
+
+/** @brief COO transposition. */
+std::shared_ptr<COO> COOTranspose(const std::shared_ptr<COO>& coo);
+
+/**
+ * @brief Sort the COO matrix by row and column indices.
+ * @return A pair of the sorted COO matrix and the permutation indices.
+ */
+std::pair<std::shared_ptr<COO>, torch::Tensor> COOSort(
+    const std::shared_ptr<COO>& coo);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_SPARSE_FORMAT_H_

dgl_sparse/include/sparse/sparse_matrix.h

@@ -0,0 +1,313 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/sparse_matrix.h
+ * @brief DGL C++ sparse matrix header.
+ */
+#ifndef SPARSE_SPARSE_MATRIX_H_
+#define SPARSE_SPARSE_MATRIX_H_
+
+// clang-format off
+#include <sparse/dgl_headers.h>
+// clang-format on
+
+#include <sparse/sparse_format.h>
+#include <torch/custom_class.h>
+#include <torch/script.h>
+
+#include <memory>
+#include <tuple>
+#include <utility>
+#include <vector>
+
+namespace dgl {
+namespace sparse {
+
+/** @brief SparseMatrix bound to Python.  */
+class SparseMatrix : public torch::CustomClassHolder {
+ public:
+  /**
+   * @brief General constructor to construct a sparse matrix for different
+   * sparse formats. At least one of the sparse formats should be provided,
+   * while others could be nullptrs.
+   *
+   * @param coo The COO format.
+   * @param csr The CSR format.
+   * @param csc The CSC format.
+   * @param value Value of the sparse matrix.
+   * @param shape Shape of the sparse matrix.
+   */
+  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);
+
+  /**
+   * @brief Construct a SparseMatrix from a COO format.
+   * @param coo The COO format
+   * @param value Values of the sparse matrix
+   * @param shape Shape of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromCOOPointer(
+      const std::shared_ptr<COO>& coo, torch::Tensor value,
+      const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Construct a SparseMatrix from a CSR format.
+   * @param csr The CSR format
+   * @param value Values of the sparse matrix
+   * @param shape Shape of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromCSRPointer(
+      const std::shared_ptr<CSR>& csr, torch::Tensor value,
+      const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Construct a SparseMatrix from a CSC format.
+   * @param csc The CSC format
+   * @param value Values of the sparse matrix
+   * @param shape Shape of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromCSCPointer(
+      const std::shared_ptr<CSR>& csc, torch::Tensor value,
+      const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Construct a SparseMatrix from a Diag format.
+   * @param diag The Diag format
+   * @param value Values of the sparse matrix
+   * @param shape Shape of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromDiagPointer(
+      const std::shared_ptr<Diag>& diag, torch::Tensor value,
+      const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Create a SparseMatrix from tensors in COO format.
+   * @param indices COO coordinates with shape (2, nnz).
+   * @param value Values of the sparse matrix.
+   * @param shape Shape of the sparse matrix.
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromCOO(
+      torch::Tensor indices, torch::Tensor value,
+      const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Create a SparseMatrix from tensors in CSR format.
+   * @param indptr Index pointer array of the CSR
+   * @param indices Indices array of the CSR
+   * @param value Values of the sparse matrix
+   * @param shape Shape of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromCSR(
+      torch::Tensor indptr, torch::Tensor indices, torch::Tensor value,
+      const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Create a SparseMatrix from tensors in CSC format.
+   * @param indptr Index pointer array of the CSC
+   * @param indices Indices array of the CSC
+   * @param value Values of the sparse matrix
+   * @param shape Shape of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromCSC(
+      torch::Tensor indptr, torch::Tensor indices, torch::Tensor value,
+      const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Create a SparseMatrix with Diag format.
+   * @param value Values of the sparse matrix
+   * @param shape Shape of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> FromDiag(
+      torch::Tensor value, const std::vector<int64_t>& shape);
+
+  /**
+   * @brief Create a SparseMatrix by selecting rows or columns based on provided
+   * indices.
+   *
+   * This function allows you to create a new SparseMatrix by selecting specific
+   * rows or columns from the original SparseMatrix based on the provided
+   * indices. The selection can be performed either row-wise or column-wise,
+   * determined by the 'dim' parameter.
+   *
+   * @param dim Select rows (dim=0) or columns (dim=1).
+   * @param ids A tensor containing the indices of the selected rows or columns.
+   *
+   * @return A new SparseMatrix containing the selected rows or columns.
+   *
+   * @note The 'dim' parameter should be either 0 (for row-wise selection) or 1
+   * (for column-wise selection).
+   * @note The 'ids' tensor should contain valid indices within the range of the
+   * original SparseMatrix's dimensions.
+   */
+  c10::intrusive_ptr<SparseMatrix> IndexSelect(int64_t dim, torch::Tensor ids);
+
+  /**
+   * @brief Create a SparseMatrix by selecting a range of rows or columns based
+   * on provided indices.
+   *
+   * This function allows you to create a new SparseMatrix by selecting a range
+   * of specific rows or columns from the original SparseMatrix based on the
+   * provided indices. The selection can be performed either row-wise or
+   * column-wise, determined by the 'dim' parameter.
+   *
+   * @param dim Select rows (dim=0) or columns (dim=1).
+   * @param start The starting index (inclusive) of the range.
+   * @param end The ending index (exclusive) of the range.
+   *
+   * @return A new SparseMatrix containing the selected range of rows or
+   * columns.
+   *
+   * @note The 'dim' parameter should be either 0 (for row-wise selection) or 1
+   * (for column-wise selection).
+   * @note The 'start' and 'end' indices should be valid indices within
+   * the valid range of the original SparseMatrix's dimensions.
+   */
+  c10::intrusive_ptr<SparseMatrix> RangeSelect(
+      int64_t dim, int64_t start, int64_t end);
+
+  /**
+   * @brief Create a SparseMatrix by sampling elements based on the specified
+   * dimension and sample count.
+   *
+   * If `ids` is provided, this function samples elements from the specified
+   * set of row or column IDs, resulting in a sparse matrix containing only
+   * the sampled rows or columns.
+   *
+   * @param dim Select rows (dim=0) or columns (dim=1) for sampling.
+   * @param fanout The number of elements to randomly sample from each row or
+   * column.
+   * @param ids An optional tensor containing row or column IDs from which to
+   * sample elements.
+   * @param replace Indicates whether repeated sampling of the same element
+   * is allowed. If True, repeated sampling is allowed; otherwise, it is not
+   * allowed.
+   * @param bias An optional boolean flag indicating whether to enable biasing
+   * during sampling. If True, the values of the sparse matrix will be used as
+   * bias weights, meaning that elements with higher values will be more likely
+   * to be sampled. Otherwise, all elements will be sampled uniformly,
+   * regardless of their value.
+   *
+   * @return A new SparseMatrix with the same shape as the original matrix
+   * containing the sampled elements.
+   *
+   * @note If 'replace = false' and there are fewer elements than 'fanout',
+   * all non-zero elements will be sampled.
+   * @note If 'ids' is not provided, the function will sample from
+   * all rows or columns.
+   */
+  c10::intrusive_ptr<SparseMatrix> Sample(
+      int64_t dim, int64_t fanout, torch::Tensor ids, bool replace, bool bias);
+
+  /**
+   * @brief Create a SparseMatrix from a SparseMatrix using new values.
+   * @param mat An existing sparse matrix
+   * @param value New values of the sparse matrix
+   *
+   * @return SparseMatrix
+   */
+  static c10::intrusive_ptr<SparseMatrix> ValLike(
+      const c10::intrusive_ptr<SparseMatrix>& mat, torch::Tensor value);
+
+  /** @return Value of the sparse matrix. */
+  inline torch::Tensor value() const { return value_; }
+  /** @return Shape of the sparse matrix. */
+  inline const std::vector<int64_t>& shape() const { return shape_; }
+  /** @return Number of non-zero values */
+  inline int64_t nnz() const { return value_.size(0); }
+  /** @return Non-zero value data type */
+  inline caffe2::TypeMeta dtype() const { return value_.dtype(); }
+  /** @return Device of the sparse matrix */
+  inline torch::Device device() const { return value_.device(); }
+
+  /** @return COO of the sparse matrix. The COO is created if not exists. */
+  std::shared_ptr<COO> COOPtr();
+  /** @return CSR of the sparse matrix. The CSR is created if not exists. */
+  std::shared_ptr<CSR> CSRPtr();
+  /** @return CSC of the sparse matrix. The CSC is created if not exists. */
+  std::shared_ptr<CSR> CSCPtr();
+  /**
+   * @return Diagonal format of the sparse matrix. An error will be raised if
+   * it does not have a diagonal format.
+   */
+  std::shared_ptr<Diag> DiagPtr();
+
+  /** @brief Check whether this sparse matrix has COO format. */
+  inline bool HasCOO() const { return coo_ != nullptr; }
+  /** @brief Check whether this sparse matrix has CSR format. */
+  inline bool HasCSR() const { return csr_ != nullptr; }
+  /** @brief Check whether this sparse matrix has CSC format. */
+  inline bool HasCSC() const { return csc_ != nullptr; }
+  /** @brief Check whether this sparse matrix has Diag format. */
+  inline bool HasDiag() const { return diag_ != nullptr; }
+
+  /** @return {row, col} tensors in the COO format. */
+  std::tuple<torch::Tensor, torch::Tensor> COOTensors();
+  /** @return Stacked row and col tensors in the COO format. */
+  torch::Tensor Indices();
+  /** @return {row, col, value_indices} tensors in the CSR format. */
+  std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+  CSRTensors();
+  /** @return {row, col, value_indices} tensors in the CSC format. */
+  std::tuple<torch::Tensor, torch::Tensor, torch::optional<torch::Tensor>>
+  CSCTensors();
+
+  /** @brief Return the transposition of the sparse matrix. It transposes the
+   * first existing sparse format by checking COO, CSR, and CSC.
+   */
+  c10::intrusive_ptr<SparseMatrix> Transpose() const;
+
+  /**
+   * @brief Return a new coalesced matrix.
+   *
+   * A coalesced sparse matrix satisfies the following properties:
+   *   - the indices of the non-zero elements are unique,
+   *   - the indices are sorted in lexicographical order.
+   *
+   * @return A coalesced sparse matrix.
+   */
+  c10::intrusive_ptr<SparseMatrix> Coalesce();
+
+  /**
+   * @brief Return true if this sparse matrix contains duplicate indices.
+   * @return A bool flag.
+   */
+  bool HasDuplicate();
+
+ private:
+  /** @brief Create the COO format for the sparse matrix internally */
+  void _CreateCOO();
+  /** @brief Create the CSR format for the sparse matrix internally */
+  void _CreateCSR();
+  /** @brief Create the CSC format for the sparse matrix internally */
+  void _CreateCSC();
+
+  // COO/CSC/CSR/Diag pointers. Nullptr indicates non-existence.
+  std::shared_ptr<COO> coo_;
+  std::shared_ptr<CSR> csr_, csc_;
+  std::shared_ptr<Diag> diag_;
+  // Value of the SparseMatrix
+  torch::Tensor value_;
+  // Shape of the SparseMatrix
+  const std::vector<int64_t> shape_;
+};
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  //  SPARSE_SPARSE_MATRIX_H_

dgl_sparse/include/sparse/spmm.h

@@ -0,0 +1,40 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/spmm.h
+ * @brief DGL C++ SpMM operator.
+ */
+#ifndef SPARSE_SPMM_H_
+#define SPARSE_SPMM_H_
+
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Perform a matrix multiplication of the sparse matrix and dense
+ * matrix. The SpMM can be batched, where the batch dimension is the last
+ * dimension for both sparse and dense matrices.
+ *
+ * There are three cases for sparse, dense, and output matrix shapes:
+ *   (1) (n, m), (m, k), and (n, k);
+ *   (2) (n, m), (m,), and (n,);
+ *   (3) (n, m, b), (m, k, b), and (n, k, b).
+ *
+ * This function supports autograd for both the sparse and dense matrix but does
+ * not support higher order gradient.
+ *
+ * @param sparse_mat The sparse matrix.
+ * @param dense_mat The dense matrix.
+ *
+ * @return Dense matrix.
+ */
+torch::Tensor SpMM(
+    const c10::intrusive_ptr<SparseMatrix>& sparse_mat,
+    torch::Tensor dense_mat);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_SPMM_H_

dgl_sparse/include/sparse/spspmm.h

@@ -0,0 +1,37 @@
+/**
+ *  Copyright (c) 2022 by Contributors
+ * @file sparse/spspmm.h
+ * @brief DGL C++ SpSpMM operator.
+ */
+#ifndef SPARSE_SPSPMM_H_
+#define SPARSE_SPSPMM_H_
+
+#include <sparse/sparse_matrix.h>
+#include <torch/script.h>
+
+namespace dgl {
+namespace sparse {
+
+/**
+ * @brief Perform a sparse-sparse matrix multiplication on matrices with
+ * possibly different sparsities. The two sparse matrices must have
+ * 1-D 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 supports autograd for both sparse matrices but does
+ * not support higher order gradient.
+ *
+ * @param lhs_mat The first sparse matrix of shape (n, m).
+ * @param rhs_mat The second sparse matrix of shape (m, k).
+ *
+ * @return Sparse matrix of shape (n, k).
+ */
+c10::intrusive_ptr<SparseMatrix> SpSpMM(
+    const c10::intrusive_ptr<SparseMatrix>& lhs_mat,
+    const c10::intrusive_ptr<SparseMatrix>& rhs_mat);
+
+}  // namespace sparse
+}  // namespace dgl
+
+#endif  // SPARSE_SPSPMM_H_

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_