chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
@@ -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_
|
||||
Reference in New Issue
Block a user