From 9027b639bcadc96459a67552462a59663e2a1e96 Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Wed, 13 Sep 2023 07:56:02 +0000 Subject: [PATCH 1/8] add relabel python api --- dgl_sparse/include/sparse/macro.h | 64 +++++++++++++++++++++ dgl_sparse/include/sparse/matrix_ops.h | 5 ++ dgl_sparse/include/sparse/matrix_ops_impl.h | 22 +++++++ dgl_sparse/src/cpu/matrix_ops_impl.cc | 2 +- dgl_sparse/src/matrix_ops.cc | 11 ++++ dgl_sparse/src/matrix_ops_impl.h | 13 ----- python/dgl/sparse/sparse_matrix.py | 36 +++++------- 7 files changed, 117 insertions(+), 36 deletions(-) create mode 100644 dgl_sparse/include/sparse/macro.h create mode 100644 dgl_sparse/include/sparse/matrix_ops_impl.h delete mode 100644 dgl_sparse/src/matrix_ops_impl.h diff --git a/dgl_sparse/include/sparse/macro.h b/dgl_sparse/include/sparse/macro.h new file mode 100644 index 000000000000..31ca7a12ceb0 --- /dev/null +++ b/dgl_sparse/include/sparse/macro.h @@ -0,0 +1,64 @@ +/** + * Copyright (c) 2023 by Contributors + * @file macro.h + * @brief DGL C++ sparse API macros. + */ +#ifndef DGL_SPARSE_MACRO_H_ +#define DGL_SPARSE_MACRO_H_ + +namespace dgl { +namespace sparse { + +/** + * Dispatch according to device: + * + * DGL_SPARSE_XPU_SWITCH(tensor.device().type(), XPU, { + * // Now XPU is a placeholder for tensor.device().type() + * DeviceSpecificImplementation(...); + * }); + */ +#define DGL_SPARSE_XPU_SWITCH(val, XPU, op, ...) \ + do { \ + if ((val) == c10::DeviceType::CPU) { \ + constexpr auto XPU = c10::DeviceType::CPU; \ + { __VA_ARGS__ } \ + } else { \ + LOG(FATAL) << "Operator " << (op) << " does not support " \ + << c10::DeviceTypeName(val) << " device."; \ + } \ + } while (0) + +/** + * Dispatch according to ID type (either int32 or int64): + * + * DGL_SPARSE_ID_TYPE_SWITCH(tensor.dtype(), IdType, { + * // Now IdType is the type corresponding to data type of the tensor. + * // For instance, one can do this for a CPU array: + * IdType *data = static_cast(array.data_ptr()); + * }); + */ +#define DGL_SPARSE_ID_TYPE_SWITCH(val, IdType, op, ...) \ + do { \ + if ((val) == torch::kInt32) { \ + typedef int32_t IdType; \ + { __VA_ARGS__ } \ + } else if ((val) == torch::kInt64) { \ + typedef int64_t IdType; \ + { __VA_ARGS__ } \ + } else { \ + LOG(FATAL) << "Operator " << (op) << " does not support " \ + << (val).name() << " as ID dtype."; \ + } \ + } while (0) + +// Macro to dispatch according to device and index type. +#define DGL_SPARSE_COO_SWITCH(coo, XPU, IdType, op, ...) \ + DGL_SPARSE_XPU_SWITCH(coo->indices.device().type(), XPU, op, { \ + DGL_SPARSE_ID_TYPE_SWITCH( \ + (coo)->indices.dtype(), IdType, op, {{__VA_ARGS__}}); \ + }); + +} // namespace sparse +} // namespace dgl + +#endif // DGL_SPARSE_MACRO_H_ diff --git a/dgl_sparse/include/sparse/matrix_ops.h b/dgl_sparse/include/sparse/matrix_ops.h index 62dee82a79cd..32f8fe342300 100644 --- a/dgl_sparse/include/sparse/matrix_ops.h +++ b/dgl_sparse/include/sparse/matrix_ops.h @@ -7,6 +7,7 @@ #define SPARSE_MATRIX_OPS_H_ #include +#include #include @@ -26,6 +27,10 @@ namespace sparse { std::tuple, torch::Tensor, torch::Tensor> COOIntersection( const std::shared_ptr& lhs, const std::shared_ptr& rhs); +std::tuple, torch::Tensor> Relabel( + const c10::intrusive_ptr& mat, uint64_t dim, + torch::Tensor leading_indices); + } // namespace sparse } // namespace dgl diff --git a/dgl_sparse/include/sparse/matrix_ops_impl.h b/dgl_sparse/include/sparse/matrix_ops_impl.h new file mode 100644 index 000000000000..e62d64e4af87 --- /dev/null +++ b/dgl_sparse/include/sparse/matrix_ops_impl.h @@ -0,0 +1,22 @@ +/** + * 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 + +namespace dgl { +namespace sparse { + +template +std::tuple, torch::Tensor> RelabelImpl( + const c10::intrusive_ptr& mat, uint64_t dim, + torch::Tensor leading_indices) {} + +} // namespace sparse +} // namespace dgl + +#endif // DGL_SPARSE_MATRIX_OPS_IMPL_H_ diff --git a/dgl_sparse/src/cpu/matrix_ops_impl.cc b/dgl_sparse/src/cpu/matrix_ops_impl.cc index 9972a1d9d202..bfbf7360b433 100644 --- a/dgl_sparse/src/cpu/matrix_ops_impl.cc +++ b/dgl_sparse/src/cpu/matrix_ops_impl.cc @@ -3,7 +3,7 @@ * @file cpu/matrix_ops_impl.cc * @brief DGL C++ matrix operators. */ -#include "./matrix_ops_impl.h" +#include "sparse/matrix_ops_impl.h" namespace dgl { namespace sparse {} // namespace sparse diff --git a/dgl_sparse/src/matrix_ops.cc b/dgl_sparse/src/matrix_ops.cc index 6df5eb3af642..efd066b9f776 100644 --- a/dgl_sparse/src/matrix_ops.cc +++ b/dgl_sparse/src/matrix_ops.cc @@ -6,6 +6,9 @@ #include #include +#include "sparse/macro.h" +#include "sparse/matrix_ops_impl.h" + namespace dgl { namespace sparse { @@ -55,5 +58,13 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( return {ret_coo, lhs_indices, rhs_indices}; } +std::tuple, torch::Tensor> Relabel( + const c10::intrusive_ptr& mat, uint64_t dim, + torch::Tensor leading_indices) { + DGL_SPARSE_COO_SWITCH(mat->COOPtr(), XPU, IdType, "Relabel", { + return RelabelImpl(mat, dim, leading_indices); + }); +} + } // namespace sparse } // namespace dgl diff --git a/dgl_sparse/src/matrix_ops_impl.h b/dgl_sparse/src/matrix_ops_impl.h deleted file mode 100644 index a235b6cdaae8..000000000000 --- a/dgl_sparse/src/matrix_ops_impl.h +++ /dev/null @@ -1,13 +0,0 @@ -/** - * 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_ - -namespace dgl { -namespace sparse {} -} // namespace dgl - -#endif // DGL_SPARSE_MATRIX_OPS_IMPL_H_ diff --git a/python/dgl/sparse/sparse_matrix.py b/python/dgl/sparse/sparse_matrix.py index cb85e1435ffa..a31cdaefcb53 100644 --- a/python/dgl/sparse/sparse_matrix.py +++ b/python/dgl/sparse/sparse_matrix.py @@ -680,13 +680,13 @@ def sample( self.c_sparse_matrix.sample(dim, fanout, ids, replace, bias) ) - def compact( + def relabel( self, dim: int, leading_indices: Optional[torch.Tensor] = None, ): - """Compact sparse matrix by removing rows or columns without non-zero - elements in the sparse matrix and relabeling indices of the dimension. + """Relabels indices of a dimension and remove rows or columns without + non-zero elements in the sparse matrix. This function serves a dual purpose: it allows you to reorganize the indices within a specific dimension (rows or columns) of the sparse @@ -697,10 +697,7 @@ def compact( of relabeled indices remains the same as the original order, except that rows or columns without non-zero elements are removed. When 'leading_indices' are provided, they are positioned at the start of the - relabeled dimension. To be precise, all rows selected by the specified - indices will be remapped from 0 to length(indices) - 1. Rows that are not - selected and contain any non-zero elements will be positioned after those - remapped rows while maintaining their original order. + relabeled dimension. This function mimics 'dgl.to_block', a method used to compress a sampled subgraph by eliminating redundant nodes. The 'leading_indices' parameter @@ -730,27 +727,22 @@ def compact( >>> indices = torch.tensor([[0, 2], [1, 2]]) >>> A = dglsp.spmatrix(indices) - >>> print(A.to_dense()) - tensor([[0., 1., 0.], - [0., 0., 0.], - [0., 0., 1.]]) - Case 1: Compact rows without indices. + Case 1: Relabel rows without indices. - >>> B, original_rows = A.compact(dim=0, leading_indices=None) - >>> print(B.to_dense()) - tensor([[0., 1., 0.], - [0., 0., 1.]]) + >>> B, original_rows = A.relabel(dim=0, leading_indices=None) + >>> print(B) + SparseMatrix(indices=tensor([[0, 1], [1, 2]]), + shape=(2, 3), nnz=2) >>> print(original_rows) torch.Tensor([0, 2]) - Case 2: Compact rows with indices. + Case 2: Relabel rows with indices. - >>> B, original_rows = A.compact(dim=0, leading_indices=[1, 2]) - >>> print(B.to_dense()) - tensor([[0., 0., 0.], - [0., 0., 1.], - [0., 1., 0.],]) + >>> B, original_rows = A.relabel(dim=0, leading_indices=[1, 2]) + >>> print(B) + SparseMatrix(indices=tensor([[1, 2], [2, 1]]), + shape=(3, 3), nnz=2) >>> print(original_rows) torch.Tensor([1, 2, 0]) """ From e7aeeccbc92edb417236ac280d20a47ebea21d15 Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Thu, 14 Sep 2023 09:34:16 +0000 Subject: [PATCH 2/8] add description --- dgl_sparse/include/sparse/matrix_ops.h | 18 ++++++++++++++++++ dgl_sparse/src/matrix_ops.cc | 18 ++++++++++++++++++ 2 files changed, 36 insertions(+) diff --git a/dgl_sparse/include/sparse/matrix_ops.h b/dgl_sparse/include/sparse/matrix_ops.h index 32f8fe342300..cc14a1c33abb 100644 --- a/dgl_sparse/include/sparse/matrix_ops.h +++ b/dgl_sparse/include/sparse/matrix_ops.h @@ -27,6 +27,24 @@ namespace sparse { std::tuple, torch::Tensor, torch::Tensor> COOIntersection( const std::shared_ptr& lhs, const std::shared_ptr& rhs); +/** + * @brief Relabels indices of a dimension and removes rows or columns without + * non-zero elements in the sparse matrix. + * + * 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 + * relabeled dimension. + * + * @param mat The sparse matrix to be relabeled. + * @param dim The dimension to relabel. Should be 0 or 1. Use 0 for row-wise + * relabeling and 1 for column-wise relabeling. + * @param leading_indices An optional tensor containing row or column ids that + * should be placed at the beginning of the relabeled dimension. + * + * @return A tuple containing the relabeled sparse matrix and the index mapping + * of the relabeled dimension from the new index to the original index. + */ std::tuple, torch::Tensor> Relabel( const c10::intrusive_ptr& mat, uint64_t dim, torch::Tensor leading_indices); diff --git a/dgl_sparse/src/matrix_ops.cc b/dgl_sparse/src/matrix_ops.cc index efd066b9f776..ddb8afde559f 100644 --- a/dgl_sparse/src/matrix_ops.cc +++ b/dgl_sparse/src/matrix_ops.cc @@ -58,6 +58,24 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( return {ret_coo, lhs_indices, rhs_indices}; } +/** + * @brief Relabels indices of a dimension and removes rows or columns without + * non-zero elements in the sparse matrix. + * + * 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 + * relabeled dimension. + * + * @param mat The sparse matrix to be relabeled. + * @param dim The dimension to relabel. Should be 0 or 1. Use 0 for row-wise + * relabeling and 1 for column-wise relabeling. + * @param leading_indices An optional tensor containing row or column ids that + * should be placed at the beginning of the relabeled dimension. + * + * @return A tuple containing the relabeled sparse matrix and the index mapping + * of the relabeled dimension from the new index to the original index. + */ std::tuple, torch::Tensor> Relabel( const c10::intrusive_ptr& mat, uint64_t dim, torch::Tensor leading_indices) { From d9c4173c5fca4935de93bdfadb429b4de1cd4f4d Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Fri, 15 Sep 2023 02:30:50 +0000 Subject: [PATCH 3/8] fix name --- dgl_sparse/include/sparse/matrix_ops.h | 6 ++-- dgl_sparse/include/sparse/matrix_ops_impl.h | 20 +++++++++++- dgl_sparse/src/matrix_ops.cc | 10 +++--- python/dgl/sparse/sparse_matrix.py | 36 +++++++++++++-------- 4 files changed, 49 insertions(+), 23 deletions(-) diff --git a/dgl_sparse/include/sparse/matrix_ops.h b/dgl_sparse/include/sparse/matrix_ops.h index cc14a1c33abb..ef8144fe77fd 100644 --- a/dgl_sparse/include/sparse/matrix_ops.h +++ b/dgl_sparse/include/sparse/matrix_ops.h @@ -28,8 +28,8 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( const std::shared_ptr& lhs, const std::shared_ptr& rhs); /** - * @brief Relabels indices of a dimension and removes rows or columns without - * non-zero elements in the sparse matrix. + * @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 @@ -45,7 +45,7 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( * @return A tuple containing the relabeled sparse matrix and the index mapping * of the relabeled dimension from the new index to the original index. */ -std::tuple, torch::Tensor> Relabel( +std::tuple, torch::Tensor> Compact( const c10::intrusive_ptr& mat, uint64_t dim, torch::Tensor leading_indices); diff --git a/dgl_sparse/include/sparse/matrix_ops_impl.h b/dgl_sparse/include/sparse/matrix_ops_impl.h index e62d64e4af87..168e685bd94b 100644 --- a/dgl_sparse/include/sparse/matrix_ops_impl.h +++ b/dgl_sparse/include/sparse/matrix_ops_impl.h @@ -11,8 +11,26 @@ namespace dgl { namespace sparse { +/** + * @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 + * relabeled dimension. + * + * @param mat The sparse matrix to be relabeled. + * @param dim The dimension to relabel. Should be 0 or 1. Use 0 for row-wise + * relabeling and 1 for column-wise relabeling. + * @param leading_indices An optional tensor containing row or column ids that + * should be placed at the beginning of the relabeled dimension. + * + * @return A tuple containing the relabeled sparse matrix and the index mapping + * of the relabeled dimension from the new index to the original index. + */ template -std::tuple, torch::Tensor> RelabelImpl( +std::tuple, torch::Tensor> CompactImpl( const c10::intrusive_ptr& mat, uint64_t dim, torch::Tensor leading_indices) {} diff --git a/dgl_sparse/src/matrix_ops.cc b/dgl_sparse/src/matrix_ops.cc index ddb8afde559f..ee318210aa6b 100644 --- a/dgl_sparse/src/matrix_ops.cc +++ b/dgl_sparse/src/matrix_ops.cc @@ -59,8 +59,8 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( } /** - * @brief Relabels indices of a dimension and removes rows or columns without - * non-zero elements in the sparse matrix. + * @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 @@ -76,11 +76,11 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( * @return A tuple containing the relabeled sparse matrix and the index mapping * of the relabeled dimension from the new index to the original index. */ -std::tuple, torch::Tensor> Relabel( +std::tuple, torch::Tensor> Compact( const c10::intrusive_ptr& mat, uint64_t dim, torch::Tensor leading_indices) { - DGL_SPARSE_COO_SWITCH(mat->COOPtr(), XPU, IdType, "Relabel", { - return RelabelImpl(mat, dim, leading_indices); + DGL_SPARSE_COO_SWITCH(mat->COOPtr(), XPU, IdType, "Compact", { + return CompactImpl(mat, dim, leading_indices); }); } diff --git a/python/dgl/sparse/sparse_matrix.py b/python/dgl/sparse/sparse_matrix.py index a31cdaefcb53..cb85e1435ffa 100644 --- a/python/dgl/sparse/sparse_matrix.py +++ b/python/dgl/sparse/sparse_matrix.py @@ -680,13 +680,13 @@ def sample( self.c_sparse_matrix.sample(dim, fanout, ids, replace, bias) ) - def relabel( + def compact( self, dim: int, leading_indices: Optional[torch.Tensor] = None, ): - """Relabels indices of a dimension and remove rows or columns without - non-zero elements in the sparse matrix. + """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 @@ -697,7 +697,10 @@ def relabel( of relabeled indices remains the same as the original order, except that rows or columns without non-zero elements are removed. When 'leading_indices' are provided, they are positioned at the start of the - relabeled dimension. + relabeled dimension. To be precise, all rows selected by the specified + indices will be remapped from 0 to length(indices) - 1. Rows that are not + selected and contain any non-zero elements will be positioned after those + remapped rows while maintaining their original order. This function mimics 'dgl.to_block', a method used to compress a sampled subgraph by eliminating redundant nodes. The 'leading_indices' parameter @@ -727,22 +730,27 @@ def relabel( >>> indices = torch.tensor([[0, 2], [1, 2]]) >>> A = dglsp.spmatrix(indices) + >>> print(A.to_dense()) + tensor([[0., 1., 0.], + [0., 0., 0.], + [0., 0., 1.]]) - Case 1: Relabel rows without indices. + Case 1: Compact rows without indices. - >>> B, original_rows = A.relabel(dim=0, leading_indices=None) - >>> print(B) - SparseMatrix(indices=tensor([[0, 1], [1, 2]]), - shape=(2, 3), nnz=2) + >>> B, original_rows = A.compact(dim=0, leading_indices=None) + >>> print(B.to_dense()) + tensor([[0., 1., 0.], + [0., 0., 1.]]) >>> print(original_rows) torch.Tensor([0, 2]) - Case 2: Relabel rows with indices. + Case 2: Compact rows with indices. - >>> B, original_rows = A.relabel(dim=0, leading_indices=[1, 2]) - >>> print(B) - SparseMatrix(indices=tensor([[1, 2], [2, 1]]), - shape=(3, 3), nnz=2) + >>> B, original_rows = A.compact(dim=0, leading_indices=[1, 2]) + >>> print(B.to_dense()) + tensor([[0., 0., 0.], + [0., 0., 1.], + [0., 1., 0.],]) >>> print(original_rows) torch.Tensor([1, 2, 0]) """ From aea728168bb382d0deec43c76b3b03ca40808c9a Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Fri, 15 Sep 2023 02:38:28 +0000 Subject: [PATCH 4/8] fix lint --- dgl_sparse/include/sparse/macro.h | 6 +++--- dgl_sparse/include/sparse/matrix_ops_impl.h | 8 +++++--- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/dgl_sparse/include/sparse/macro.h b/dgl_sparse/include/sparse/macro.h index 31ca7a12ceb0..7b55cc828c75 100644 --- a/dgl_sparse/include/sparse/macro.h +++ b/dgl_sparse/include/sparse/macro.h @@ -3,8 +3,8 @@ * @file macro.h * @brief DGL C++ sparse API macros. */ -#ifndef DGL_SPARSE_MACRO_H_ -#define DGL_SPARSE_MACRO_H_ +#ifndef SPARSE_MACRO_H_ +#define SPARSE_MACRO_H_ namespace dgl { namespace sparse { @@ -61,4 +61,4 @@ namespace sparse { } // namespace sparse } // namespace dgl -#endif // DGL_SPARSE_MACRO_H_ +#endif // SPARSE_MACRO_H_ diff --git a/dgl_sparse/include/sparse/matrix_ops_impl.h b/dgl_sparse/include/sparse/matrix_ops_impl.h index 168e685bd94b..05a6173436ce 100644 --- a/dgl_sparse/include/sparse/matrix_ops_impl.h +++ b/dgl_sparse/include/sparse/matrix_ops_impl.h @@ -3,11 +3,13 @@ * @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_ +#ifndef SPARSE_MATRIX_OPS_IMPL_H_ +#define SPARSE_MATRIX_OPS_IMPL_H_ #include +#include + namespace dgl { namespace sparse { @@ -37,4 +39,4 @@ std::tuple, torch::Tensor> CompactImpl( } // namespace sparse } // namespace dgl -#endif // DGL_SPARSE_MATRIX_OPS_IMPL_H_ +#endif // SPARSE_MATRIX_OPS_IMPL_H_ From 3681fe7a9aa326258e07271979d246e2077bc843 Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Fri, 15 Sep 2023 03:00:19 +0000 Subject: [PATCH 5/8] fix return value --- dgl_sparse/include/sparse/matrix_ops_impl.h | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/dgl_sparse/include/sparse/matrix_ops_impl.h b/dgl_sparse/include/sparse/matrix_ops_impl.h index 05a6173436ce..cf746d55d2b3 100644 --- a/dgl_sparse/include/sparse/matrix_ops_impl.h +++ b/dgl_sparse/include/sparse/matrix_ops_impl.h @@ -34,7 +34,10 @@ namespace sparse { template std::tuple, torch::Tensor> CompactImpl( const c10::intrusive_ptr& mat, uint64_t dim, - torch::Tensor leading_indices) {} + torch::Tensor leading_indices) { + // Place holder only. + return {mat, leading_indices}; +} } // namespace sparse } // namespace dgl From 5fb362368dd8cd63db143b7c692774f327baf715 Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Fri, 15 Sep 2023 07:26:34 +0000 Subject: [PATCH 6/8] fix desxription --- dgl_sparse/include/sparse/matrix_ops.h | 16 +++++++-------- dgl_sparse/{include/sparse => src}/macro.h | 3 ++- dgl_sparse/src/matrix_ops.cc | 20 +------------------ .../{include/sparse => src}/matrix_ops_impl.h | 16 +++++++-------- 4 files changed, 19 insertions(+), 36 deletions(-) rename dgl_sparse/{include/sparse => src}/macro.h (96%) rename dgl_sparse/{include/sparse => src}/matrix_ops_impl.h (69%) diff --git a/dgl_sparse/include/sparse/matrix_ops.h b/dgl_sparse/include/sparse/matrix_ops.h index ef8144fe77fd..6dc41ac548df 100644 --- a/dgl_sparse/include/sparse/matrix_ops.h +++ b/dgl_sparse/include/sparse/matrix_ops.h @@ -34,19 +34,19 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( * 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 - * relabeled dimension. + * compact dimension. * - * @param mat The sparse matrix to be relabeled. - * @param dim The dimension to relabel. Should be 0 or 1. Use 0 for row-wise - * relabeling and 1 for column-wise relabeling. + * @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 relabeled dimension. + * should be placed at the beginning of the compact dimension. * - * @return A tuple containing the relabeled sparse matrix and the index mapping - * of the relabeled dimension from the new index to the original index. + * @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, torch::Tensor> Compact( - const c10::intrusive_ptr& mat, uint64_t dim, + const c10::intrusive_ptr& mat, int64_t dim, torch::Tensor leading_indices); } // namespace sparse diff --git a/dgl_sparse/include/sparse/macro.h b/dgl_sparse/src/macro.h similarity index 96% rename from dgl_sparse/include/sparse/macro.h rename to dgl_sparse/src/macro.h index 7b55cc828c75..e6bf43391e24 100644 --- a/dgl_sparse/include/sparse/macro.h +++ b/dgl_sparse/src/macro.h @@ -10,7 +10,8 @@ namespace dgl { namespace sparse { /** - * Dispatch according to device: + * Dispatch an operator to a templated implementation function + * according to its device: * * DGL_SPARSE_XPU_SWITCH(tensor.device().type(), XPU, { * // Now XPU is a placeholder for tensor.device().type() diff --git a/dgl_sparse/src/matrix_ops.cc b/dgl_sparse/src/matrix_ops.cc index ee318210aa6b..cf6b71f01f58 100644 --- a/dgl_sparse/src/matrix_ops.cc +++ b/dgl_sparse/src/matrix_ops.cc @@ -6,7 +6,7 @@ #include #include -#include "sparse/macro.h" +#include "./macro.h" #include "sparse/matrix_ops_impl.h" namespace dgl { @@ -58,24 +58,6 @@ std::tuple, torch::Tensor, torch::Tensor> COOIntersection( return {ret_coo, lhs_indices, rhs_indices}; } -/** - * @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 - * relabeled dimension. - * - * @param mat The sparse matrix to be relabeled. - * @param dim The dimension to relabel. Should be 0 or 1. Use 0 for row-wise - * relabeling and 1 for column-wise relabeling. - * @param leading_indices An optional tensor containing row or column ids that - * should be placed at the beginning of the relabeled dimension. - * - * @return A tuple containing the relabeled sparse matrix and the index mapping - * of the relabeled dimension from the new index to the original index. - */ std::tuple, torch::Tensor> Compact( const c10::intrusive_ptr& mat, uint64_t dim, torch::Tensor leading_indices) { diff --git a/dgl_sparse/include/sparse/matrix_ops_impl.h b/dgl_sparse/src/matrix_ops_impl.h similarity index 69% rename from dgl_sparse/include/sparse/matrix_ops_impl.h rename to dgl_sparse/src/matrix_ops_impl.h index cf746d55d2b3..271efdbe3d17 100644 --- a/dgl_sparse/include/sparse/matrix_ops_impl.h +++ b/dgl_sparse/src/matrix_ops_impl.h @@ -20,20 +20,20 @@ namespace sparse { * 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 - * relabeled dimension. + * compact dimension. * - * @param mat The sparse matrix to be relabeled. - * @param dim The dimension to relabel. Should be 0 or 1. Use 0 for row-wise - * relabeling and 1 for column-wise relabeling. + * @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 relabeled dimension. + * should be placed at the beginning of the compact dimension. * - * @return A tuple containing the relabeled sparse matrix and the index mapping - * of the relabeled dimension from the new index to the original index. + * @return A tuple containing the compacted sparse matrix and the index mapping + * of the compact dimension from the new index to the original index. */ template std::tuple, torch::Tensor> CompactImpl( - const c10::intrusive_ptr& mat, uint64_t dim, + const c10::intrusive_ptr& mat, int64_t dim, torch::Tensor leading_indices) { // Place holder only. return {mat, leading_indices}; From ea6b14ba3024e936247513f7c7271771e1af3ee2 Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Fri, 15 Sep 2023 07:33:02 +0000 Subject: [PATCH 7/8] fix headand lint --- dgl_sparse/src/cpu/matrix_ops_impl.cc | 2 +- dgl_sparse/src/macro.h | 6 +++--- dgl_sparse/src/matrix_ops.cc | 2 +- dgl_sparse/src/matrix_ops_impl.h | 6 +++--- 4 files changed, 8 insertions(+), 8 deletions(-) diff --git a/dgl_sparse/src/cpu/matrix_ops_impl.cc b/dgl_sparse/src/cpu/matrix_ops_impl.cc index bfbf7360b433..9972a1d9d202 100644 --- a/dgl_sparse/src/cpu/matrix_ops_impl.cc +++ b/dgl_sparse/src/cpu/matrix_ops_impl.cc @@ -3,7 +3,7 @@ * @file cpu/matrix_ops_impl.cc * @brief DGL C++ matrix operators. */ -#include "sparse/matrix_ops_impl.h" +#include "./matrix_ops_impl.h" namespace dgl { namespace sparse {} // namespace sparse diff --git a/dgl_sparse/src/macro.h b/dgl_sparse/src/macro.h index e6bf43391e24..dde3c5e95752 100644 --- a/dgl_sparse/src/macro.h +++ b/dgl_sparse/src/macro.h @@ -3,8 +3,8 @@ * @file macro.h * @brief DGL C++ sparse API macros. */ -#ifndef SPARSE_MACRO_H_ -#define SPARSE_MACRO_H_ +#ifndef DGL_SPARSE_MACRO_H_ +#define DGL_SPARSE_MACRO_H_ namespace dgl { namespace sparse { @@ -62,4 +62,4 @@ namespace sparse { } // namespace sparse } // namespace dgl -#endif // SPARSE_MACRO_H_ +#endif // DGL_SPARSE_MACRO_H_ diff --git a/dgl_sparse/src/matrix_ops.cc b/dgl_sparse/src/matrix_ops.cc index cf6b71f01f58..4316fda84b41 100644 --- a/dgl_sparse/src/matrix_ops.cc +++ b/dgl_sparse/src/matrix_ops.cc @@ -7,7 +7,7 @@ #include #include "./macro.h" -#include "sparse/matrix_ops_impl.h" +#include "./matrix_ops_impl.h" namespace dgl { namespace sparse { diff --git a/dgl_sparse/src/matrix_ops_impl.h b/dgl_sparse/src/matrix_ops_impl.h index 271efdbe3d17..7ec10a421ba6 100644 --- a/dgl_sparse/src/matrix_ops_impl.h +++ b/dgl_sparse/src/matrix_ops_impl.h @@ -3,8 +3,8 @@ * @file matrix_ops_impl.h * @brief DGL C++ sparse matrix operator implementations. */ -#ifndef SPARSE_MATRIX_OPS_IMPL_H_ -#define SPARSE_MATRIX_OPS_IMPL_H_ +#ifndef DGL_SPARSE_MATRIX_OPS_IMPL_H_ +#define DGL_SPARSE_MATRIX_OPS_IMPL_H_ #include @@ -42,4 +42,4 @@ std::tuple, torch::Tensor> CompactImpl( } // namespace sparse } // namespace dgl -#endif // SPARSE_MATRIX_OPS_IMPL_H_ +#endif // DGL_SPARSE_MATRIX_OPS_IMPL_H_ From c5afdea772d74ca9ba5ce87dcce540e67ce8144a Mon Sep 17 00:00:00 2001 From: xiangyuzhi Date: Fri, 15 Sep 2023 09:22:48 +0000 Subject: [PATCH 8/8] fix comments --- dgl_sparse/src/macro.h | 14 +++++++------- dgl_sparse/src/matrix_ops_impl.h | 18 ------------------ 2 files changed, 7 insertions(+), 25 deletions(-) diff --git a/dgl_sparse/src/macro.h b/dgl_sparse/src/macro.h index dde3c5e95752..5eccbddc83a8 100644 --- a/dgl_sparse/src/macro.h +++ b/dgl_sparse/src/macro.h @@ -18,14 +18,14 @@ namespace sparse { * DeviceSpecificImplementation(...); * }); */ -#define DGL_SPARSE_XPU_SWITCH(val, XPU, op, ...) \ +#define DGL_SPARSE_XPU_SWITCH(device, XPU, op, ...) \ do { \ - if ((val) == c10::DeviceType::CPU) { \ + if ((device) == c10::DeviceType::CPU) { \ constexpr auto XPU = c10::DeviceType::CPU; \ { __VA_ARGS__ } \ } else { \ LOG(FATAL) << "Operator " << (op) << " does not support " \ - << c10::DeviceTypeName(val) << " device."; \ + << c10::DeviceTypeName(device) << " device."; \ } \ } while (0) @@ -38,17 +38,17 @@ namespace sparse { * IdType *data = static_cast(array.data_ptr()); * }); */ -#define DGL_SPARSE_ID_TYPE_SWITCH(val, IdType, op, ...) \ +#define DGL_SPARSE_ID_TYPE_SWITCH(dtype, IdType, op, ...) \ do { \ - if ((val) == torch::kInt32) { \ + if ((dtype) == torch::kInt32) { \ typedef int32_t IdType; \ { __VA_ARGS__ } \ - } else if ((val) == torch::kInt64) { \ + } else if ((dtype) == torch::kInt64) { \ typedef int64_t IdType; \ { __VA_ARGS__ } \ } else { \ LOG(FATAL) << "Operator " << (op) << " does not support " \ - << (val).name() << " as ID dtype."; \ + << (dtype).name() << " as ID dtype."; \ } \ } while (0) diff --git a/dgl_sparse/src/matrix_ops_impl.h b/dgl_sparse/src/matrix_ops_impl.h index 7ec10a421ba6..af04b3496a6c 100644 --- a/dgl_sparse/src/matrix_ops_impl.h +++ b/dgl_sparse/src/matrix_ops_impl.h @@ -13,24 +13,6 @@ namespace dgl { namespace sparse { -/** - * @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. - */ template std::tuple, torch::Tensor> CompactImpl( const c10::intrusive_ptr& mat, int64_t dim,