Skip to content

GitLab

Unverified Commit 8ac27dad authored by Hongzhi (Steve), Chen's avatar Hongzhi (Steve), Chen Committed by GitHub
Browse files

[Misc] clang-format auto fix. (#4824) 



* [Misc] clang-format auto fix.

* blabla

* ablabla

* blabla
Co-authored-by: default avatarSteve <ubuntu@ip-172-31-34-29.ap-northeast-1.compute.internal>
parent bcd37684
Expand all Show whitespace changes
Inline Side-by-side
Showing with 1707 additions and 1700 deletions
include/dgl/aten/array_ops.h
View file @ 8ac27dad
......@@ -10,10 +10,11 @@
#define DGL_ATEN_ARRAY_OPS_H_
#include <algorithm>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include <tuple>
#include <string>
#include "./types.h"
namespace dgl {
......@@ -24,7 +25,8 @@ namespace aten {
//////////////////////////////////////////////////////////////////////
/** @return A special array to represent null. */
inline NDArray NullArray(const DGLDataType& dtype = DGLDataType{kDGLInt, 64, 1},
inline NDArray NullArray(
const DGLDataType& dtype = DGLDataType{kDGLInt, 64, 1},
const DGLContext& ctx = DGLContext{kDGLCPU, 0}) {
return NDArray::Empty({0}, dtype, ctx);
}
......@@ -32,9 +34,7 @@ inline NDArray NullArray(const DGLDataType& dtype = DGLDataType{kDGLInt, 64, 1},
/**
* @return Whether the input array is a null array.
*/
inline bool IsNullArray(NDArray array) {
return array->shape[0] == 0;
}
inline bool IsNullArray(NDArray array) { return array->shape[0] == 0; }
/**
* @brief Create a new id array with given length
......@@ -43,8 +43,8 @@ inline bool IsNullArray(NDArray array) {
* @param nbits The number of integer bits
* @return id array
*/
IdArray NewIdArray(int64_t length,
DGLContext ctx = DGLContext{kDGLCPU, 0},
IdArray NewIdArray(
int64_t length, DGLContext ctx = DGLContext{kDGLCPU, 0},
uint8_t nbits = 64);
/**
......@@ -55,8 +55,8 @@ IdArray NewIdArray(int64_t length,
* @return the id array
*/
template <typename T>
IdArray VecToIdArray(const std::vector<T>& vec,
uint8_t nbits = 64,
IdArray VecToIdArray(
const std::vector<T>& vec, uint8_t nbits = 64,
DGLContext ctx = DGLContext{kDGLCPU, 0});
/**
......@@ -148,7 +148,7 @@ IdArray NonZero(BoolArray bool_arr);
* @brief Return the data under the index. In numpy notation, A[I]
* @tparam ValueType The type of return value.
*/
template<typename ValueType>
template <typename ValueType>
ValueType IndexSelect(NDArray array, int64_t index);
/**
......@@ -187,10 +187,11 @@ NDArray Scatter(NDArray array, IdArray indices);
void Scatter_(IdArray index, NDArray value, NDArray out);
/**
* @brief Repeat each element a number of times. Equivalent to np.repeat(array, repeats)
* @brief Repeat each element a number of times. Equivalent to np.repeat(array,
* repeats)
* @param array A 1D vector
* @param repeats A 1D integer vector for number of times to repeat for each element in
* \c array. Must have the same shape as \c array.
* @param repeats A 1D integer vector for number of times to repeat for each
* element in \c array. Must have the same shape as \c array.
*/
NDArray Repeat(NDArray array, IdArray repeats);
......@@ -253,12 +254,13 @@ inline bool IsValidIdArray(const dgl::runtime::NDArray& arr) {
*
* The packed tensor would be [1, 2, 3, 4, 5].
*
* The length tensor would be [2, 3], i.e. the length of each sequence before padding.
* The length tensor would be [2, 3], i.e. the length of each sequence before
* padding.
*
* The offset tensor would be [0, 2], i.e. the offset to the packed tensor for each
* sequence (before padding)
* The offset tensor would be [0, 2], i.e. the offset to the packed tensor for
* each sequence (before padding)
*/
template<typename ValueType>
template <typename ValueType>
std::tuple<NDArray, IdArray, IdArray> Pack(NDArray array, ValueType pad_value);
/**
......@@ -295,8 +297,9 @@ std::pair<NDArray, IdArray> ConcatSlices(NDArray array, IdArray lengths);
* @brief Return the cumulative summation (or inclusive sum) of the input array.
*
* The first element out[0] is equal to the first element of the input array
* array[0]. The rest elements are defined recursively, out[i] = out[i-1] + array[i].
* Hence, the result array length is the same as the input array length.
* array[0]. The rest elements are defined recursively, out[i] = out[i-1] +
* array[i]. Hence, the result array length is the same as the input array
* length.
*
* If prepend_zero is true, then the first element is zero and the result array
* length is the input array length plus one. This is useful for creating
......@@ -320,17 +323,18 @@ IdArray NonZero(NDArray array);
/**
* @brief Sort the ID vector in ascending order.
*
* It performs both sort and arg_sort (returning the sorted index). The sorted index
* is always in int64.
* It performs both sort and arg_sort (returning the sorted index). The sorted
* index is always in int64.
*
* @param array Input array.
* @param num_bits The number of bits used in key comparison. For example, if the data type
* of the input array is int32_t and `num_bits = 8`, it only uses bits in index
* range [0, 8) for sorting. Setting it to a small value could
* speed up the sorting if the underlying sorting algorithm is radix sort (e.g., on GPU).
* Setting it to zero (default value) means using all the bits for comparison.
* On CPU, it currently has no effect.
* @return A pair of arrays: sorted values and sorted index to the original position.
* @param num_bits The number of bits used in key comparison. For example, if
* the data type of the input array is int32_t and `num_bits = 8`, it only uses
* bits in index range [0, 8) for sorting. Setting it to a small value could
* speed up the sorting if the underlying sorting algorithm is
* radix sort (e.g., on GPU). Setting it to zero (default value) means using all
* the bits for comparison. On CPU, it currently has no effect.
* @return A pair of arrays: sorted values and sorted index to the original
* position.
*/
std::pair<IdArray, IdArray> Sort(IdArray array, int num_bits = 0);
......@@ -341,9 +345,7 @@ std::string ToDebugString(NDArray array);
// inline implementations
template <typename T>
IdArray VecToIdArray(const std::vector<T>& vec,
uint8_t nbits,
DGLContext ctx) {
IdArray VecToIdArray(const std::vector<T>& vec, uint8_t nbits, DGLContext ctx) {
IdArray ret = NewIdArray(vec.size(), DGLContext{kDGLCPU, 0}, nbits);
if (nbits == 32) {
std::copy(vec.begin(), vec.end(), static_cast<int32_t*>(ret->data));
......@@ -367,7 +369,8 @@ inline DGLContext GetContextOf(const std::vector<IdArray>& arrays) {
first = false;
result = array->ctx;
} else {
CHECK_EQ(array->ctx, result) << "Context of the input arrays are different";
CHECK_EQ(array->ctx, result)
<< "Context of the input arrays are different";
}
}
return result;
......
include/dgl/aten/coo.h
View file @ 8ac27dad
......@@ -9,14 +9,16 @@
#include <dmlc/io.h>
#include <dmlc/serializer.h>
#include <vector>
#include <utility>
#include <tuple>
#include <string>
#include "./types.h"
#include <tuple>
#include <utility>
#include <vector>
#include "./array_ops.h"
#include "./spmat.h"
#include "./macro.h"
#include "./spmat.h"
#include "./types.h"
namespace dgl {
namespace aten {
......@@ -52,9 +54,9 @@ struct COOMatrix {
/** @brief default constructor */
COOMatrix() = default;
/** @brief constructor */
COOMatrix(int64_t nrows, int64_t ncols, IdArray rarr, IdArray carr,
IdArray darr = NullArray(), bool rsorted = false,
bool csorted = false)
COOMatrix(
int64_t nrows, int64_t ncols, IdArray rarr, IdArray carr,
IdArray darr = NullArray(), bool rsorted = false, bool csorted = false)
: num_rows(nrows),
num_cols(ncols),
row(rarr),
......@@ -79,8 +81,9 @@ struct COOMatrix {
// Convert to a SparseMatrix object that can return to python.
SparseMatrix ToSparseMatrix() const {
return SparseMatrix(static_cast<int32_t>(SparseFormat::kCOO), num_rows,
num_cols, {row, col, data}, {row_sorted, col_sorted});
return SparseMatrix(
static_cast<int32_t>(SparseFormat::kCOO), num_rows, num_cols,
{row, col, data}, {row_sorted, col_sorted});
}
bool Load(dmlc::Stream* fs) {
......@@ -122,12 +125,12 @@ struct COOMatrix {
}
/** @brief Return a copy of this matrix on the give device context. */
inline COOMatrix CopyTo(const DGLContext &ctx) const {
if (ctx == row->ctx)
return *this;
return COOMatrix(num_rows, num_cols, row.CopyTo(ctx), col.CopyTo(ctx),
aten::IsNullArray(data) ? data : data.CopyTo(ctx),
row_sorted, col_sorted);
inline COOMatrix CopyTo(const DGLContext& ctx) const {
if (ctx == row->ctx) return *this;
return COOMatrix(
num_rows, num_cols, row.CopyTo(ctx), col.CopyTo(ctx),
aten::IsNullArray(data) ? data : data.CopyTo(ctx), row_sorted,
col_sorted);
}
/**
......@@ -139,8 +142,7 @@ struct COOMatrix {
* The context check is deferred to pinning the NDArray.
*/
inline void PinMemory_() {
if (is_pinned)
return;
if (is_pinned) return;
row.PinMemory_();
col.PinMemory_();
if (!aten::IsNullArray(data)) {
......@@ -157,8 +159,7 @@ struct COOMatrix {
* The context check is deferred to unpinning the NDArray.
*/
inline void UnpinMemory_() {
if (!is_pinned)
return;
if (!is_pinned) return;
row.UnpinMemory_();
col.UnpinMemory_();
if (!aten::IsNullArray(data)) {
......@@ -183,25 +184,25 @@ struct COOMatrix {
///////////////////////// COO routines //////////////////////////
/** @brief Return true if the value (row, col) is non-zero */
bool COOIsNonZero(COOMatrix , int64_t row, int64_t col);
bool COOIsNonZero(COOMatrix, int64_t row, int64_t col);
/**
* @brief Batched implementation of COOIsNonZero.
* @note This operator allows broadcasting (i.e, either row or col can be of length 1).
* @note This operator allows broadcasting (i.e, either row or col can be of
* length 1).
*/
runtime::NDArray COOIsNonZero(COOMatrix , runtime::NDArray row, runtime::NDArray col);
runtime::NDArray COOIsNonZero(
COOMatrix, runtime::NDArray row, runtime::NDArray col);
/** @brief Return the nnz of the given row */
int64_t COOGetRowNNZ(COOMatrix , int64_t row);
runtime::NDArray COOGetRowNNZ(COOMatrix , runtime::NDArray row);
int64_t COOGetRowNNZ(COOMatrix, int64_t row);
runtime::NDArray COOGetRowNNZ(COOMatrix, runtime::NDArray row);
/** @brief Return the data array of the given row */
std::pair<runtime::NDArray, runtime::NDArray>
COOGetRowDataAndIndices(COOMatrix , int64_t row);
std::pair<runtime::NDArray, runtime::NDArray> COOGetRowDataAndIndices(
COOMatrix, int64_t row);
/** @brief Whether the COO matrix contains data */
inline bool COOHasData(COOMatrix csr) {
return !IsNullArray(csr.data);
}
inline bool COOHasData(COOMatrix csr) { return !IsNullArray(csr.data); }
/**
* @brief Check whether the COO is sorted.
......@@ -217,11 +218,12 @@ std::pair<bool, bool> COOIsSorted(COOMatrix coo);
/**
* @brief Get the data and the row,col indices for each returned entries.
*
* The operator supports matrix with duplicate entries and all the matched entries
* will be returned. The operator assumes there is NO duplicate (row, col) pair
* in the given input. Otherwise, the returned result is undefined.
* The operator supports matrix with duplicate entries and all the matched
* entries will be returned. The operator assumes there is NO duplicate (row,
* col) pair in the given input. Otherwise, the returned result is undefined.
*
* @note This operator allows broadcasting (i.e, either row or col can be of length 1).
* @note This operator allows broadcasting (i.e, either row or col can be of
* length 1).
* @param mat Sparse matrix
* @param rows Row index
* @param cols Column index
......@@ -230,10 +232,13 @@ std::pair<bool, bool> COOIsSorted(COOMatrix coo);
std::vector<runtime::NDArray> COOGetDataAndIndices(
COOMatrix mat, runtime::NDArray rows, runtime::NDArray cols);
/** @brief Get data. The return type is an ndarray due to possible duplicate entries. */
/** @brief Get data. The return type is an ndarray due to possible duplicate
* entries. */
inline runtime::NDArray COOGetAllData(COOMatrix mat, int64_t row, int64_t col) {
IdArray rows = VecToIdArray<int64_t>({row}, mat.row->dtype.bits, mat.row->ctx);
IdArray cols = VecToIdArray<int64_t>({col}, mat.row->dtype.bits, mat.row->ctx);
IdArray rows =
VecToIdArray<int64_t>({row}, mat.row->dtype.bits, mat.row->ctx);
IdArray cols =
VecToIdArray<int64_t>({col}, mat.row->dtype.bits, mat.row->ctx);
const auto& rst = COOGetDataAndIndices(mat, rows, cols);
return rst[2];
}
......@@ -241,18 +246,20 @@ inline runtime::NDArray COOGetAllData(COOMatrix mat, int64_t row, int64_t col) {
/**
* @brief Get the data for each (row, col) pair.
*
* The operator supports matrix with duplicate entries but only one matched entry
* will be returned for each (row, col) pair. Support duplicate input (row, col)
* pairs.
* The operator supports matrix with duplicate entries but only one matched
* entry will be returned for each (row, col) pair. Support duplicate input
* (row, col) pairs.
*
* @note This operator allows broadcasting (i.e, either row or col can be of length 1).
* @note This operator allows broadcasting (i.e, either row or col can be of
* length 1).
*
* @param mat Sparse matrix.
* @param rows Row index.
* @param cols Column index.
* @return Data array. The i^th element is the data of (rows[i], cols[i])
*/
runtime::NDArray COOGetData(COOMatrix mat, runtime::NDArray rows, runtime::NDArray cols);
runtime::NDArray COOGetData(
COOMatrix mat, runtime::NDArray rows, runtime::NDArray cols);
/** @brief Return a transposed COO matrix */
COOMatrix COOTranspose(COOMatrix coo);
......@@ -301,14 +308,15 @@ COOMatrix COOSliceRows(COOMatrix coo, runtime::NDArray rows);
* @param cols The col index to select
* @return submatrix
*/
COOMatrix COOSliceMatrix(COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
COOMatrix COOSliceMatrix(
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
/** @return True if the matrix has duplicate entries */
bool COOHasDuplicate(COOMatrix coo);
/**
* @brief Deduplicate the entries of a sorted COO matrix, replacing the data with the
* number of occurrences of the row-col coordinates.
* @brief Deduplicate the entries of a sorted COO matrix, replacing the data
* with the number of occurrences of the row-col coordinates.
*/
std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo);
......@@ -316,11 +324,13 @@ std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo);
* @brief Sort the indices of a COO matrix in-place.
*
* The function sorts row indices in ascending order. If sort_column is true,
* col indices are sorted in ascending order too. The data array of the returned COOMatrix
* stores the shuffled index which could be used to fetch edge data.
* col indices are sorted in ascending order too. The data array of the returned
* COOMatrix stores the shuffled index which could be used to fetch edge data.
*
* Complexity: O(N*log(N)) time and O(1) space, where N is the number of nonzeros.
* TODO(minjie): The time complexity could be improved to O(N) by using a O(N) space.
* Complexity: O(N*log(N)) time and O(1) space, where N is the number of
* nonzeros.
* TODO(minjie): The time complexity could be improved to O(N) by using a O(N)
* space.
*
* @param mat The coo matrix to sort.
* @param sort_column True if column index should be sorted too.
......@@ -331,31 +341,32 @@ void COOSort_(COOMatrix* mat, bool sort_column = false);
* @brief Sort the indices of a COO matrix.
*
* The function sorts row indices in ascending order. If sort_column is true,
* col indices are sorted in ascending order too. The data array of the returned COOMatrix
* stores the shuffled index which could be used to fetch edge data.
* col indices are sorted in ascending order too. The data array of the returned
* COOMatrix stores the shuffled index which could be used to fetch edge data.
*
* Complexity: O(N*log(N)) time and O(1) space, where N is the number of nonzeros.
* TODO(minjie): The time complexity could be improved to O(N) by using a O(N) space.
* Complexity: O(N*log(N)) time and O(1) space, where N is the number of
* nonzeros.
* TODO(minjie): The time complexity could be improved to O(N) by using a O(N)
* space.
*
* @param mat The input coo matrix
* @param sort_column True if column index should be sorted too.
* @return COO matrix with index sorted.
*/
inline COOMatrix COOSort(COOMatrix mat, bool sort_column = false) {
if ((mat.row_sorted && !sort_column) || mat.col_sorted)
return mat;
COOMatrix ret(mat.num_rows, mat.num_cols,
mat.row.Clone(), mat.col.Clone(),
COOHasData(mat)? mat.data.Clone() : mat.data,
mat.row_sorted, mat.col_sorted);
if ((mat.row_sorted && !sort_column) || mat.col_sorted) return mat;
COOMatrix ret(
mat.num_rows, mat.num_cols, mat.row.Clone(), mat.col.Clone(),
COOHasData(mat) ? mat.data.Clone() : mat.data, mat.row_sorted,
mat.col_sorted);
COOSort_(&ret, sort_column);
return ret;
}
/**
* @brief Remove entries from COO matrix by entry indices (data indices)
* @return A new COO matrix as well as a mapping from the new COO entries to the old COO
* entries.
* @return A new COO matrix as well as a mapping from the new COO entries to the
* old COO entries.
*/
COOMatrix COORemove(COOMatrix coo, IdArray entries);
......@@ -365,10 +376,12 @@ COOMatrix COORemove(COOMatrix coo, IdArray entries);
* @param new_row_ids the new row Ids (the index is the old row Id)
* @param new_col_ids the new column Ids (the index is the old col Id).
*/
COOMatrix COOReorder(COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
COOMatrix COOReorder(
COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
/**
* @brief Randomly select a fixed number of non-zero entries along each given row independently.
* @brief Randomly select a fixed number of non-zero entries along each given
* row independently.
*
* The function performs random choices along each row independently.
* The picked indices are returned in the form of a COO matrix.
......@@ -400,15 +413,12 @@ COOMatrix COOReorder(COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArr
* Should be of the same length as the data array.
* If an empty array is provided, assume uniform.
* @param replace True if sample with replacement
* @return A COOMatrix storing the picked row and col indices. Its data field stores the
* the index of the picked elements in the value array.
* @return A COOMatrix storing the picked row and col indices. Its data field
* stores the the index of the picked elements in the value array.
*/
COOMatrix COORowWiseSampling(
COOMatrix mat,
IdArray rows,
int64_t num_samples,
NDArray prob_or_mask = NDArray(),
bool replace = true);
COOMatrix mat, IdArray rows, int64_t num_samples,
NDArray prob_or_mask = NDArray(), bool replace = true);
/**
* @brief Randomly select a fixed number of non-zero entries for each edge type
......@@ -433,8 +443,8 @@ COOMatrix COORowWiseSampling(
* COOMatrix coo = ...;
* IdArray rows = ... ; // [0, 3]
* std::vector<int64_t> num_samples = {2, 2, 2};
* COOMatrix sampled = COORowWisePerEtypeSampling(coo, rows, eid2etype_offset, num_samples,
* FloatArray(), false);
* COOMatrix sampled = COORowWisePerEtypeSampling(coo, rows, eid2etype_offset,
* num_samples, FloatArray(), false);
* // possible sampled coo matrix:
* // sampled.num_rows = 4
* // sampled.num_cols = 4
......@@ -450,20 +460,18 @@ COOMatrix COORowWiseSampling(
* Should be of the same length as the data array.
* If an empty array is provided, assume uniform.
* @param replace True if sample with replacement
* @return A COOMatrix storing the picked row and col indices. Its data field stores the
* the index of the picked elements in the value array.
* @return A COOMatrix storing the picked row and col indices. Its data field
* stores the the index of the picked elements in the value array.
* @note The edges of the entire graph must be ordered by their edge types.
*/
COOMatrix COORowWisePerEtypeSampling(
COOMatrix mat,
IdArray rows,
const std::vector<int64_t>& eid2etype_offset,
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
const std::vector<int64_t>& num_samples,
const std::vector<NDArray>& prob_or_mask,
bool replace = true);
const std::vector<NDArray>& prob_or_mask, bool replace = true);
/**
* @brief Select K non-zero entries with the largest weights along each given row.
* @brief Select K non-zero entries with the largest weights along each given
* row.
*
* The function performs top-k selection along each row independently.
* The picked indices are returned in the form of a COO matrix.
......@@ -492,18 +500,15 @@ COOMatrix COORowWisePerEtypeSampling(
* @param mat Input COO matrix.
* @param rows Rows to sample from.
* @param k The K value.
* @param weight Weight associated with each entry. Should be of the same length as the
* data array. If an empty array is provided, assume uniform.
* @param ascending If true, elements are sorted by ascending order, equivalent to find
* the K smallest values. Otherwise, find K largest values.
* @return A COOMatrix storing the picked row and col indices. Its data field stores the
* the index of the picked elements in the value array.
* @param weight Weight associated with each entry. Should be of the same length
* as the data array. If an empty array is provided, assume uniform.
* @param ascending If true, elements are sorted by ascending order, equivalent
* to find the K smallest values. Otherwise, find K largest values.
* @return A COOMatrix storing the picked row and col indices. Its data field
* stores the the index of the picked elements in the value array.
*/
COOMatrix COORowWiseTopk(
COOMatrix mat,
IdArray rows,
int64_t k,
NDArray weight,
COOMatrix mat, IdArray rows, int64_t k, NDArray weight,
bool ascending = false);
/**
......@@ -535,8 +540,7 @@ COOMatrix COORowWiseTopk(
* COOMatrix_C.num_rows : 3
* COOMatrix_C.num_cols : 4
*/
COOMatrix UnionCoo(
const std::vector<COOMatrix>& coos);
COOMatrix UnionCoo(const std::vector<COOMatrix>& coos);
/**
* @brief DisjointUnion a list COOMatrix into one COOMatrix.
......@@ -566,12 +570,13 @@ COOMatrix UnionCoo(
* COOMatrix_C.num_cols : 5
*
* @param coos The input list of coo matrix.
* @param src_offset A list of integers recording src vertix id offset of each Matrix in coos
* @param src_offset A list of integers recording dst vertix id offset of each Matrix in coos
* @param src_offset A list of integers recording src vertix id offset of each
* Matrix in coos
* @param src_offset A list of integers recording dst vertix id offset of each
* Matrix in coos
* @return The combined COOMatrix.
*/
COOMatrix DisjointUnionCoo(
const std::vector<COOMatrix>& coos);
COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos);
/**
* @brief COOMatrix toSimple.
......@@ -591,8 +596,8 @@ COOMatrix DisjointUnionCoo(
* edge_map = [0, 0, 0, 1, 1, 2, 3, 4, 4, 4, 4]
*
* @return The simplified COOMatrix
* The count recording the number of duplicated edges from the original graph.
* The edge mapping from the edge IDs of original graph to those of the
* The count recording the number of duplicated edges from the original
* graph. The edge mapping from the edge IDs of original graph to those of the
* returned graph.
*/
std::tuple<COOMatrix, IdArray, IdArray> COOToSimple(const COOMatrix& coo);
......@@ -642,11 +647,10 @@ std::tuple<COOMatrix, IdArray, IdArray> COOToSimple(const COOMatrix& coo);
* @return A list of COOMatrixes representing each disjoint components.
*/
std::vector<COOMatrix> DisjointPartitionCooBySizes(
const COOMatrix &coo,
const uint64_t batch_size,
const std::vector<uint64_t> &edge_cumsum,
const std::vector<uint64_t> &src_vertex_cumsum,
const std::vector<uint64_t> &dst_vertex_cumsum);
const COOMatrix& coo, const uint64_t batch_size,
const std::vector<uint64_t>& edge_cumsum,
const std::vector<uint64_t>& src_vertex_cumsum,
const std::vector<uint64_t>& dst_vertex_cumsum);
/**
* @brief Slice a contiguous chunk from a COOMatrix
......@@ -684,10 +688,9 @@ std::vector<COOMatrix> DisjointPartitionCooBySizes(
* @return COOMatrix representing the chunk.
*/
COOMatrix COOSliceContiguousChunk(
const COOMatrix &coo,
const std::vector<uint64_t> &edge_range,
const std::vector<uint64_t> &src_vertex_range,
const std::vector<uint64_t> &dst_vertex_range);
const COOMatrix& coo, const std::vector<uint64_t>& edge_range,
const std::vector<uint64_t>& src_vertex_range,
const std::vector<uint64_t>& dst_vertex_range);
/**
* @brief Create a LineGraph of input coo
......@@ -716,10 +719,11 @@ COOMatrix COOSliceContiguousChunk(
* [0, 1, 1, 0, 0]]
*
* @param coo COOMatrix to create the LineGraph
* @param backtracking whether the pair of (v, u) (u, v) edges are treated as linked
* @param backtracking whether the pair of (v, u) (u, v) edges are treated as
* linked
* @return LineGraph in COO format
*/
COOMatrix COOLineGraph(const COOMatrix &coo, bool backtracking);
COOMatrix COOLineGraph(const COOMatrix& coo, bool backtracking);
} // namespace aten
} // namespace dgl
......
include/dgl/aten/csr.h
View file @ 8ac27dad
......@@ -223,8 +223,10 @@ bool CSRIsSorted(CSRMatrix csr);
std::vector<runtime::NDArray> CSRGetDataAndIndices(
CSRMatrix, runtime::NDArray rows, runtime::NDArray cols);
/* @brief Get data. The return type is an ndarray due to possible duplicate
* entries. */
/**
* @brief Get data. The return type is an ndarray due to possible duplicate
* entries.
*/
inline runtime::NDArray CSRGetAllData(CSRMatrix mat, int64_t row, int64_t col) {
const auto& nbits = mat.indptr->dtype.bits;
const auto& ctx = mat.indptr->ctx;
......
include/dgl/aten/macro.h
View file @ 8ac27dad
This diff is collapsed.
include/dgl/graph_serializer.h
View file @ 8ac27dad
......@@ -10,7 +10,8 @@
#include <memory>
namespace dgl {
// Util class to call the private/public empty constructor, which is needed for serialization
// Util class to call the private/public empty constructor, which is needed for
// serialization
class Serializer {
public:
template <typename T>
......
include/dgl/runtime/ndarray.h
View file @ 8ac27dad
This diff is collapsed.
include/dgl/runtime/parallel_for.h
View file @ 8ac27dad
......@@ -7,19 +7,18 @@
#define DGL_RUNTIME_PARALLEL_FOR_H_
#include <dmlc/omp.h>
#include <algorithm>
#include <string>
#include <atomic>
#include <cstdlib>
#include <exception>
#include <vector>
#include <atomic>
#include <string>
#include <utility>
#include <vector>
namespace {
int64_t divup(int64_t x, int64_t y) {
return (x + y - 1) / y;
}
}
int64_t divup(int64_t x, int64_t y) { return (x + y - 1) / y; }
} // namespace
namespace dgl {
namespace runtime {
......@@ -37,9 +36,7 @@ struct DefaultGrainSizeT {
}
}
size_t operator()() {
return grain_size;
}
size_t operator()() { return grain_size; }
};
} // namespace
......@@ -48,7 +45,9 @@ inline size_t compute_num_threads(size_t begin, size_t end, size_t grain_size) {
if (omp_in_parallel() || end - begin <= grain_size || end - begin == 1)
return 1;
return std::min(static_cast<int64_t>(omp_get_max_threads()), divup(end - begin, grain_size));
return std::min(
static_cast<int64_t>(omp_get_max_threads()),
divup(end - begin, grain_size));
#else
return 1;
#endif
......@@ -60,15 +59,13 @@ static DefaultGrainSizeT default_grain_size;
* @brief OpenMP-based parallel for loop.
*
* It requires each thread's workload to have at least \a grain_size elements.
* The loop body will be a function that takes in two arguments \a begin and \a end, which
* stands for the starting (inclusive) and ending index (exclusive) of the workload.
* The loop body will be a function that takes in two arguments \a begin and \a
* end, which stands for the starting (inclusive) and ending index (exclusive)
* of the workload.
*/
template <typename F>
void parallel_for(
const size_t begin,
const size_t end,
const size_t grain_size,
F&& f) {
const size_t begin, const size_t end, const size_t grain_size, F&& f) {
if (begin >= end) {
return;
}
......@@ -89,13 +86,11 @@ void parallel_for(
try {
f(begin_tid, end_tid);
} catch (...) {
if (!err_flag.test_and_set())
eptr = std::current_exception();
if (!err_flag.test_and_set()) eptr = std::current_exception();
}
}
}
if (eptr)
std::rethrow_exception(eptr);
if (eptr) std::rethrow_exception(eptr);
#else
f(begin, end);
#endif
......@@ -110,22 +105,21 @@ void parallel_for(
* parallel for pragma with static scheduling.
*/
template <typename F>
void parallel_for(
const size_t begin,
const size_t end,
F&& f) {
void parallel_for(const size_t begin, const size_t end, F&& f) {
parallel_for(begin, end, default_grain_size(), std::forward<F>(f));
}
/**
* @brief OpenMP-based two-stage parallel reduction.
*
* The first-stage reduction function \a f works in parallel. Each thread's workload has
* at least \a grain_size elements. The loop body will be a function that takes in
* the starting index (inclusive), the ending index (exclusive), and the reduction identity.
* The first-stage reduction function \a f works in parallel. Each thread's
* workload has at least \a grain_size elements. The loop body will be a
* function that takes in the starting index (inclusive), the ending index
* (exclusive), and the reduction identity.
*
* The second-stage reduction function \a sf is a binary function working in the main
* thread. It aggregates the partially reduced result computed from each thread.
* The second-stage reduction function \a sf is a binary function working in the
* main thread. It aggregates the partially reduced result computed from each
* thread.
*
* Example to compute a parallelized max reduction of an array \c a:
*
......@@ -134,11 +128,9 @@ void parallel_for(
* 100, // ending index
* 1, // grain size
* -std::numeric_limits<float>::infinity, // identity
* [&a] (int begin, int end, float ident) { // first-stage partial reducer
* float result = ident;
* for (int i = begin; i < end; ++i)
* result = std::max(result, a[i]);
* return result;
* [&a] (int begin, int end, float ident) { // first-stage partial
* reducer float result = ident; for (int i = begin; i < end; ++i) result =
* std::max(result, a[i]); return result;
* },
* [] (float result, float partial_result) {
* return std::max(result, partial_result);
......@@ -146,12 +138,8 @@ void parallel_for(
*/
template <typename DType, typename F, typename SF>
DType parallel_reduce(
const size_t begin,
const size_t end,
const size_t grain_size,
const DType ident,
const F& f,
const SF& sf) {
const size_t begin, const size_t end, const size_t grain_size,
const DType ident, const F& f, const SF& sf) {
if (begin >= end) {
return ident;
}
......@@ -174,17 +162,14 @@ DType parallel_reduce(
try {
results[tid] = f(begin_tid, end_tid, ident);
} catch (...) {
if (!err_flag.test_and_set())
eptr = std::current_exception();
if (!err_flag.test_and_set()) eptr = std::current_exception();
}
}
}
if (eptr)
std::rethrow_exception(eptr);
if (eptr) std::rethrow_exception(eptr);
DType out = ident;
for (int64_t i = 0; i < num_threads; ++i)
out = sf(out, results[i]);
for (int64_t i = 0; i < num_threads; ++i) out = sf(out, results[i]);
return out;
}
......
src/array/array.cc
View file @ 8ac27dad
This diff is collapsed.
src/array/array_op.h
View file @ 8ac27dad
......@@ -8,9 +8,10 @@
#include <dgl/array.h>
#include <dgl/graph_traversal.h>
#include <vector>
#include <tuple>
#include <utility>
#include <vector>
namespace dgl {
namespace aten {
......@@ -82,7 +83,8 @@ template <DGLDeviceType XPU, typename IdType>
bool CSRIsNonZero(CSRMatrix csr, int64_t row, int64_t col);
template <DGLDeviceType XPU, typename IdType>
runtime::NDArray CSRIsNonZero(CSRMatrix csr, runtime::NDArray row, runtime::NDArray col);
runtime::NDArray CSRIsNonZero(
CSRMatrix csr, runtime::NDArray row, runtime::NDArray col);
template <DGLDeviceType XPU, typename IdType>
bool CSRHasDuplicate(CSRMatrix csr);
......@@ -104,19 +106,21 @@ bool CSRIsSorted(CSRMatrix csr);
template <DGLDeviceType XPU, typename IdType, typename DType>
runtime::NDArray CSRGetData(
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols, bool return_eids,
runtime::NDArray weights, DType filler);
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols,
bool return_eids, runtime::NDArray weights, DType filler);
template <DGLDeviceType XPU, typename IdType, typename DType>
runtime::NDArray CSRGetData(
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols,
runtime::NDArray weights, DType filler) {
return CSRGetData<XPU, IdType, DType>(csr, rows, cols, false, weights, filler);
return CSRGetData<XPU, IdType, DType>(
csr, rows, cols, false, weights, filler);
}
template <DGLDeviceType XPU, typename IdType>
NDArray CSRGetData(CSRMatrix csr, NDArray rows, NDArray cols) {
return CSRGetData<XPU, IdType, IdType>(csr, rows, cols, true, NullArray(rows->dtype), -1);
return CSRGetData<XPU, IdType, IdType>(
csr, rows, cols, true, NullArray(rows->dtype), -1);
}
template <DGLDeviceType XPU, typename IdType>
......@@ -141,20 +145,23 @@ template <DGLDeviceType XPU, typename IdType>
CSRMatrix CSRSliceRows(CSRMatrix csr, runtime::NDArray rows);
template <DGLDeviceType XPU, typename IdType>
CSRMatrix CSRSliceMatrix(CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
CSRMatrix CSRSliceMatrix(
CSRMatrix csr, runtime::NDArray rows, runtime::NDArray cols);
template <DGLDeviceType XPU, typename IdType>
void CSRSort_(CSRMatrix* csr);
template <DGLDeviceType XPU, typename IdType, typename TagType>
std::pair<CSRMatrix, NDArray> CSRSortByTag(
const CSRMatrix &csr, IdArray tag_array, int64_t num_tags);
const CSRMatrix& csr, IdArray tag_array, int64_t num_tags);
template <DGLDeviceType XPU, typename IdType>
CSRMatrix CSRReorder(CSRMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
CSRMatrix CSRReorder(
CSRMatrix csr, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
template <DGLDeviceType XPU, typename IdType>
COOMatrix COOReorder(COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
COOMatrix COOReorder(
COOMatrix coo, runtime::NDArray new_row_ids, runtime::NDArray new_col_ids);
template <DGLDeviceType XPU, typename IdType>
CSRMatrix CSRRemove(CSRMatrix csr, IdArray entries);
......@@ -162,15 +169,16 @@ CSRMatrix CSRRemove(CSRMatrix csr, IdArray entries);
// FloatType is the type of probability data.
template <DGLDeviceType XPU, typename IdType, typename DType>
COOMatrix CSRRowWiseSampling(
CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask, bool replace);
CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
bool replace);
// FloatType is the type of probability data.
template <DGLDeviceType XPU, typename IdType, typename DType>
COOMatrix CSRRowWisePerEtypeSampling(
CSRMatrix mat, IdArray rows,
const std::vector<int64_t>& eid2etype_offset,
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
const std::vector<int64_t>& num_samples,
const std::vector<NDArray>& prob_or_mask, bool replace, bool rowwise_etype_sorted);
const std::vector<NDArray>& prob_or_mask, bool replace,
bool rowwise_etype_sorted);
template <DGLDeviceType XPU, typename IdType>
COOMatrix CSRRowWiseSamplingUniform(
......@@ -178,9 +186,9 @@ COOMatrix CSRRowWiseSamplingUniform(
template <DGLDeviceType XPU, typename IdType>
COOMatrix CSRRowWisePerEtypeSamplingUniform(
CSRMatrix mat, IdArray rows,
const std::vector<int64_t>& eid2etype_offset,
const std::vector<int64_t>& num_samples, bool replace, bool rowwise_etype_sorted);
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
const std::vector<int64_t>& num_samples, bool replace,
bool rowwise_etype_sorted);
// FloatType is the type of weight data.
template <DGLDeviceType XPU, typename IdType, typename DType>
......@@ -189,21 +197,13 @@ COOMatrix CSRRowWiseTopk(
template <DGLDeviceType XPU, typename IdType, typename FloatType>
COOMatrix CSRRowWiseSamplingBiased(
CSRMatrix mat,
IdArray rows,
int64_t num_samples,
NDArray tag_offset,
FloatArray bias,
bool replace);
CSRMatrix mat, IdArray rows, int64_t num_samples, NDArray tag_offset,
FloatArray bias, bool replace);
template <DGLDeviceType XPU, typename IdType>
std::pair<IdArray, IdArray> CSRGlobalUniformNegativeSampling(
const CSRMatrix& csr,
int64_t num_samples,
int num_trials,
bool exclude_self_loops,
bool replace,
double redundancy);
const CSRMatrix& csr, int64_t num_samples, int num_trials,
bool exclude_self_loops, bool replace, double redundancy);
// Union CSRMatrixes
template <DGLDeviceType XPU, typename IdType>
......@@ -218,7 +218,8 @@ template <DGLDeviceType XPU, typename IdType>
bool COOIsNonZero(COOMatrix coo, int64_t row, int64_t col);
template <DGLDeviceType XPU, typename IdType>
runtime::NDArray COOIsNonZero(COOMatrix coo, runtime::NDArray row, runtime::NDArray col);
runtime::NDArray COOIsNonZero(
COOMatrix coo, runtime::NDArray row, runtime::NDArray col);
template <DGLDeviceType XPU, typename IdType>
bool COOHasDuplicate(COOMatrix coo);
......@@ -230,15 +231,16 @@ template <DGLDeviceType XPU, typename IdType>
runtime::NDArray COOGetRowNNZ(COOMatrix coo, runtime::NDArray row);
template <DGLDeviceType XPU, typename IdType>
std::pair<runtime::NDArray, runtime::NDArray>
COOGetRowDataAndIndices(COOMatrix coo, int64_t row);
std::pair<runtime::NDArray, runtime::NDArray> COOGetRowDataAndIndices(
COOMatrix coo, int64_t row);
template <DGLDeviceType XPU, typename IdType>
std::vector<runtime::NDArray> COOGetDataAndIndices(
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
template <DGLDeviceType XPU, typename IdType>
runtime::NDArray COOGetData(COOMatrix mat, runtime::NDArray rows, runtime::NDArray cols);
runtime::NDArray COOGetData(
COOMatrix mat, runtime::NDArray rows, runtime::NDArray cols);
template <DGLDeviceType XPU, typename IdType>
COOMatrix COOTranspose(COOMatrix coo);
......@@ -253,7 +255,8 @@ template <DGLDeviceType XPU, typename IdType>
COOMatrix COOSliceRows(COOMatrix coo, runtime::NDArray rows);
template <DGLDeviceType XPU, typename IdType>
COOMatrix COOSliceMatrix(COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
COOMatrix COOSliceMatrix(
COOMatrix coo, runtime::NDArray rows, runtime::NDArray cols);
template <DGLDeviceType XPU, typename IdType>
std::pair<COOMatrix, IdArray> COOCoalesce(COOMatrix coo);
......@@ -273,13 +276,13 @@ COOMatrix COORemove(COOMatrix coo, IdArray entries);
// FloatType is the type of probability data.
template <DGLDeviceType XPU, typename IdType, typename DType>
COOMatrix COORowWiseSampling(
COOMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask, bool replace);
COOMatrix mat, IdArray rows, int64_t num_samples, NDArray prob_or_mask,
bool replace);
// FloatType is the type of probability data.
template <DGLDeviceType XPU, typename IdType, typename DType>
COOMatrix COORowWisePerEtypeSampling(
COOMatrix mat, IdArray rows,
const std::vector<int64_t>& eid2etype_offset,
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
const std::vector<int64_t>& num_samples,
const std::vector<NDArray>& prob_or_mask, bool replace);
......@@ -289,8 +292,7 @@ COOMatrix COORowWiseSamplingUniform(
template <DGLDeviceType XPU, typename IdType>
COOMatrix COORowWisePerEtypeSamplingUniform(
COOMatrix mat, IdArray rows,
const std::vector<int64_t>& eid2etype_offset,
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
const std::vector<int64_t>& num_samples, bool replace);
// FloatType is the type of weight data.
......@@ -313,14 +315,12 @@ template <DGLDeviceType XPU, typename IdType>
Frontiers DGLDFSEdges(const CSRMatrix& csr, IdArray source);
template <DGLDeviceType XPU, typename IdType>
Frontiers DGLDFSLabeledEdges(const CSRMatrix& csr,
IdArray source,
const bool has_reverse_edge,
const bool has_nontree_edge,
const bool return_labels);
Frontiers DGLDFSLabeledEdges(
const CSRMatrix& csr, IdArray source, const bool has_reverse_edge,
const bool has_nontree_edge, const bool return_labels);
template <DGLDeviceType XPU, typename IdType>
COOMatrix COOLineGraph(const COOMatrix &coo, bool backtracking);
COOMatrix COOLineGraph(const COOMatrix& coo, bool backtracking);
} // namespace impl
} // namespace aten
......
src/array/cpu/gather_mm.cc
View file @ 8ac27dad
/**
/**
* Copyright (c) 2020 by Contributors
* @file kernel/cpu/gaher_mm.cc
* @brief GatherMM C APIs and definitions.
*/
#include "./gather_mm.h"
#include <dgl/array.h>
namespace dgl {
......@@ -11,81 +12,72 @@ namespace aten {
/** @brief Generalized SegmentMM. */
template <int XPU, typename IdType, typename DType>
void SegmentMM(const NDArray A,
const NDArray B,
NDArray C,
const NDArray seglen_A,
void SegmentMM(
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
bool a_trans, bool b_trans) {
LOG(FATAL) << "Unsupported CPU kernel for SegmentMM.";
}
template <int XPU, typename IdType, typename DType>
void SegmentMMBackwardB(const NDArray A,
const NDArray dC,
NDArray dB,
const NDArray seglen) {
void SegmentMMBackwardB(
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen) {
LOG(FATAL) << "Unsupported CPU kernel for SegmentMMBackwardB.";
}
/** @brief Generalized GatherMM. */
template <int XPU, typename IdType, typename DType>
void GatherMM(const NDArray A,
const NDArray B,
NDArray C,
const NDArray idx_a,
void GatherMM(
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b) {
LOG(FATAL) << "Unsupported CPU kernel for GatherMM.";
}
/** @brief Generalized GatherMM_scatter. */
template <int XPU, typename IdType, typename DType>
void GatherMMScatter(const NDArray A,
const NDArray B,
NDArray C,
const NDArray idx_a,
const NDArray idx_b,
const NDArray idx_c) {
void GatherMMScatter(
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b, const NDArray idx_c) {
LOG(FATAL) << "Unsupported CPU kernel for GatherMM.";
}
template void GatherMM<kDGLCPU, int32_t, float>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b);
template void GatherMM<kDGLCPU, int64_t, float>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b);
template void GatherMM<kDGLCPU, int32_t, double>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b);
template void GatherMM<kDGLCPU, int64_t, double>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b);
template void GatherMMScatter<kDGLCPU, int32_t, float>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b, const NDArray idx_c);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b, const NDArray idx_c);
template void GatherMMScatter<kDGLCPU, int64_t, float>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b, const NDArray idx_c);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b, const NDArray idx_c);
template void GatherMMScatter<kDGLCPU, int32_t, double>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b, const NDArray idx_c);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b, const NDArray idx_c);
template void GatherMMScatter<kDGLCPU, int64_t, double>(
const NDArray A, const NDArray B, NDArray C,
const NDArray idx_a, const NDArray idx_b, const NDArray idx_c);
const NDArray A, const NDArray B, NDArray C, const NDArray idx_a,
const NDArray idx_b, const NDArray idx_c);
template void SegmentMM<kDGLCPU, int32_t, float>(
const NDArray A, const NDArray B, NDArray C,
const NDArray seglen_A, bool a_trans, bool b_trans);
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
bool a_trans, bool b_trans);
template void SegmentMM<kDGLCPU, int64_t, float>(
const NDArray A, const NDArray B, NDArray C,
const NDArray seglen_A, bool a_trans, bool b_trans);
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
bool a_trans, bool b_trans);
template void SegmentMM<kDGLCPU, int32_t, double>(
const NDArray A, const NDArray B, NDArray C,
const NDArray seglen_A, bool a_trans, bool b_trans);
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
bool a_trans, bool b_trans);
template void SegmentMM<kDGLCPU, int64_t, double>(
const NDArray A, const NDArray B, NDArray C,
const NDArray seglen_A, bool a_trans, bool b_trans);
const NDArray A, const NDArray B, NDArray C, const NDArray seglen_A,
bool a_trans, bool b_trans);
template void SegmentMMBackwardB<kDGLCPU, int32_t, float>(
const NDArray A, const NDArray dC, NDArray dB, const NDArray seglen);
......
src/array/cpu/rowwise_pick.h
View file @ 8ac27dad
......@@ -7,13 +7,14 @@
#define DGL_ARRAY_CPU_ROWWISE_PICK_H_
#include <dgl/array.h>
#include <dmlc/omp.h>
#include <dgl/runtime/parallel_for.h>
#include <functional>
#include <dmlc/omp.h>
#include <algorithm>
#include <functional>
#include <memory>
#include <string>
#include <vector>
#include <memory>
namespace dgl {
namespace aten {
......@@ -41,10 +42,10 @@ namespace impl {
template <typename IdxType>
using PickFn = std::function<void(
IdxType rowid, IdxType off, IdxType len, IdxType num_picks,
const IdxType* col, const IdxType* data,
IdxType* out_idx)>;
const IdxType* col, const IdxType* data, IdxType* out_idx)>;
// User-defined function for determining the number of elements to pick from one row.
// User-defined function for determining the number of elements to pick from one
// row.
//
// The column indices of the given row are stored in
// [col + off, col + off + len)
......@@ -63,8 +64,8 @@ using PickFn = std::function<void(
// \param data Pointer of the data indices.
template <typename IdxType>
using NumPicksFn = std::function<IdxType(
IdxType rowid, IdxType off, IdxType len,
const IdxType* col, const IdxType* data)>;
IdxType rowid, IdxType off, IdxType len, const IdxType* col,
const IdxType* data)>;
// User-defined function for picking elements from a range within a row.
//
......@@ -73,7 +74,8 @@ using NumPicksFn = std::function<IdxType(
//
// Similarly, the data indices are stored in
// data[off+et_idx[et_offset+i])]
// Data index pointer could be NULL, which means data[i] == off+et_idx[et_offset+i])
// Data index pointer could be NULL, which means data[i] ==
// off+et_idx[et_offset+i])
//
// *ATTENTION*: This function will be invoked concurrently. Please make sure
// it is thread-safe.
......@@ -93,15 +95,17 @@ using EtypeRangePickFn = std::function<void(
const IdxType* eid, IdxType* out_idx)>;
// Template for picking non-zero values row-wise. The implementation utilizes
// OpenMP parallelization on rows because each row performs computation independently.
// OpenMP parallelization on rows because each row performs computation
// independently.
template <typename IdxType>
COOMatrix CSRRowWisePick(CSRMatrix mat, IdArray rows,
int64_t num_picks, bool replace, PickFn<IdxType> pick_fn,
NumPicksFn<IdxType> num_picks_fn) {
COOMatrix CSRRowWisePick(
CSRMatrix mat, IdArray rows, int64_t num_picks, bool replace,
PickFn<IdxType> pick_fn, NumPicksFn<IdxType> num_picks_fn) {
using namespace aten;
const IdxType* indptr = static_cast<IdxType*>(mat.indptr->data);
const IdxType* indices = static_cast<IdxType*>(mat.indices->data);
const IdxType* data = CSRHasData(mat)? static_cast<IdxType*>(mat.data->data) : nullptr;
const IdxType* data =
CSRHasData(mat) ? static_cast<IdxType*>(mat.data->data) : nullptr;
const IdxType* rows_data = static_cast<IdxType*>(rows->data);
const int64_t num_rows = rows->shape[0];
const auto& ctx = mat.indptr->ctx;
......@@ -115,30 +119,36 @@ COOMatrix CSRRowWisePick(CSRMatrix mat, IdArray rows,
// have at least num_picks number of nnz when replace is false.
//
// If the check holds, remove -1 elements by remove_if operation, which simply
// moves valid elements to the head of arrays and create a view of the original
// array. The implementation consumes a little extra memory than the actual requirement.
// moves valid elements to the head of arrays and create a view of the
// original array. The implementation consumes a little extra memory than the
// actual requirement.
//
// Otherwise, directly use the row and col arrays to construct the result COO matrix.
// Otherwise, directly use the row and col arrays to construct the result COO
// matrix.
//
// [02/29/2020 update]: OMP is disabled for now since batch-wise parallelism is more
// [02/29/2020 update]: OMP is disabled for now since batch-wise parallelism
// is more
// significant. (minjie)
// Do not use omp_get_max_threads() since that doesn't work for compiling without OpenMP.
// Do not use omp_get_max_threads() since that doesn't work for compiling
// without OpenMP.
const int num_threads = runtime::compute_num_threads(0, num_rows, 1);
std::vector<int64_t> global_prefix(num_threads + 1, 0);
// TODO(BarclayII) Using OMP parallel directly instead of using runtime::parallel_for
// does not handle exceptions well (directly aborts when an exception pops up).
// It runs faster though because there is less scheduling. Need to handle
// exceptions better.
// TODO(BarclayII) Using OMP parallel directly instead of using
// runtime::parallel_for does not handle exceptions well (directly aborts when
// an exception pops up). It runs faster though because there is less
// scheduling. Need to handle exceptions better.
IdArray picked_row, picked_col, picked_idx;
#pragma omp parallel num_threads(num_threads)
{
const int thread_id = omp_get_thread_num();
const int64_t start_i = thread_id * (num_rows/num_threads) +
const int64_t start_i =
thread_id * (num_rows / num_threads) +
std::min(static_cast<int64_t>(thread_id), num_rows % num_threads);
const int64_t end_i = (thread_id + 1) * (num_rows/num_threads) +
const int64_t end_i =
(thread_id + 1) * (num_rows / num_threads) +
std::min(static_cast<int64_t>(thread_id + 1), num_rows % num_threads);
assert(thread_id + 1 < num_threads || end_i == num_rows);
......@@ -149,7 +159,7 @@ COOMatrix CSRRowWisePick(CSRMatrix mat, IdArray rows,
local_prefix[0] = 0;
for (int64_t i = start_i; i < end_i; ++i) {
// build prefix-sum
const int64_t local_i = i-start_i;
const int64_t local_i = i - start_i;
const IdxType rid = rows_data[i];
IdxType len = num_picks_fn(
rid, indptr[rid], indptr[rid + 1] - indptr[rid], indices, data);
......@@ -157,8 +167,8 @@ COOMatrix CSRRowWisePick(CSRMatrix mat, IdArray rows,
}
global_prefix[thread_id + 1] = local_prefix[num_local];
#pragma omp barrier
#pragma omp master
#pragma omp barrier
#pragma omp master
{
for (int t = 0; t < num_threads; ++t) {
global_prefix[t + 1] += global_prefix[t];
......@@ -168,7 +178,7 @@ COOMatrix CSRRowWisePick(CSRMatrix mat, IdArray rows,
picked_idx = IdArray::Empty({global_prefix[num_threads]}, idtype, ctx);
}
#pragma omp barrier
#pragma omp barrier
IdxType* picked_rdata = picked_row.Ptr<IdxType>();
IdxType* picked_cdata = picked_col.Ptr<IdxType>();
IdxType* picked_idata = picked_idx.Ptr<IdxType>();
......@@ -180,14 +190,15 @@ COOMatrix CSRRowWisePick(CSRMatrix mat, IdArray rows,
const IdxType off = indptr[rid];
const IdxType len = indptr[rid + 1] - off;
if (len == 0)
continue;
if (len == 0) continue;
const int64_t local_i = i - start_i;
const int64_t row_offset = thread_offset + local_prefix[local_i];
const int64_t num_picks = thread_offset + local_prefix[local_i + 1] - row_offset;
const int64_t num_picks =
thread_offset + local_prefix[local_i + 1] - row_offset;
pick_fn(rid, off, len, num_picks, indices, data, picked_idata + row_offset);
pick_fn(
rid, off, len, num_picks, indices, data, picked_idata + row_offset);
for (int64_t j = 0; j < num_picks; ++j) {
const IdxType picked = picked_idata[row_offset + j];
picked_rdata[row_offset + j] = rid;
......@@ -200,25 +211,25 @@ COOMatrix CSRRowWisePick(CSRMatrix mat, IdArray rows,
const int64_t new_len = global_prefix.back();
return COOMatrix(
mat.num_rows,
mat.num_cols,
mat.num_rows, mat.num_cols,
picked_row.CreateView({new_len}, picked_row->dtype),
picked_col.CreateView({new_len}, picked_row->dtype),
picked_idx.CreateView({new_len}, picked_row->dtype));
}
// Template for picking non-zero values row-wise. The implementation utilizes
// OpenMP parallelization on rows because each row performs computation independently.
// OpenMP parallelization on rows because each row performs computation
// independently.
template <typename IdxType, typename DType>
COOMatrix CSRRowWisePerEtypePick(CSRMatrix mat, IdArray rows,
const std::vector<int64_t>& eid2etype_offset,
COOMatrix CSRRowWisePerEtypePick(
CSRMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
const std::vector<int64_t>& num_picks, bool replace,
bool rowwise_etype_sorted, EtypeRangePickFn<IdxType> pick_fn,
const std::vector<NDArray>& prob_or_mask) {
using namespace aten;
const IdxType* indptr = mat.indptr.Ptr<IdxType>();
const IdxType* indices = mat.indices.Ptr<IdxType>();
const IdxType* eid = CSRHasData(mat)? mat.data.Ptr<IdxType>() : nullptr;
const IdxType* eid = CSRHasData(mat) ? mat.data.Ptr<IdxType>() : nullptr;
const IdxType* rows_data = rows.Ptr<IdxType>();
const int64_t num_rows = rows->shape[0];
const auto& ctx = mat.indptr->ctx;
......@@ -229,8 +240,8 @@ COOMatrix CSRRowWisePerEtypePick(CSRMatrix mat, IdArray rows,
std::vector<IdArray> picked_idxs(rows->shape[0]);
// Check if the number of picks have the same value.
// If so, we can potentially speed up if we have a node with total number of neighbors
// less than the given number of picks with replace=False.
// If so, we can potentially speed up if we have a node with total number of
// neighbors less than the given number of picks with replace=False.
bool same_num_pick = true;
int64_t num_pick_value = num_picks[0];
for (int64_t num_pick : num_picks) {
......@@ -267,9 +278,10 @@ COOMatrix CSRRowWisePerEtypePick(CSRMatrix mat, IdArray rows,
for (int64_t j = 0; j < len; ++j) {
const IdxType homogenized_eid = eid ? eid[off + j] : off + j;
auto it = std::upper_bound(
eid2etype_offset.begin(), eid2etype_offset.end(), homogenized_eid);
eid2etype_offset.begin(), eid2etype_offset.end(),
homogenized_eid);
const IdxType heterogenized_etype = it - eid2etype_offset.begin() - 1;
const IdxType heterogenized_eid = \
const IdxType heterogenized_eid =
homogenized_eid - eid2etype_offset[heterogenized_etype];
if (!has_probs || IsNullArray(prob_or_mask[heterogenized_etype])) {
......@@ -305,17 +317,21 @@ COOMatrix CSRRowWisePerEtypePick(CSRMatrix mat, IdArray rows,
for (int64_t j = 0; j < len; ++j) {
const IdxType homogenized_eid = eid ? eid[off + j] : off + j;
auto it = std::upper_bound(
eid2etype_offset.begin(), eid2etype_offset.end(), homogenized_eid);
eid2etype_offset.begin(), eid2etype_offset.end(),
homogenized_eid);
const IdxType heterogenized_etype = it - eid2etype_offset.begin() - 1;
const IdxType heterogenized_eid = \
const IdxType heterogenized_eid =
homogenized_eid - eid2etype_offset[heterogenized_etype];
et[j] = heterogenized_etype;
et_eid[j] = heterogenized_eid;
}
if (!rowwise_etype_sorted) // the edge type is sorted, not need to sort it
std::sort(et_idx.begin(), et_idx.end(),
[&et](IdxType i1, IdxType i2) {return et[i1] < et[i2];});
CHECK_LT(et[et_idx[len - 1]], num_etypes) << "etype values exceed the number of fanouts";
if (!rowwise_etype_sorted) // the edge type is sorted, not need to sort
// it
std::sort(
et_idx.begin(), et_idx.end(),
[&et](IdxType i1, IdxType i2) { return et[i1] < et[i2]; });
CHECK_LT(et[et_idx[len - 1]], num_etypes)
<< "etype values exceed the number of fanouts";
IdxType cur_et = et[et_idx[0]];
int64_t et_offset = 0;
......@@ -333,14 +349,18 @@ COOMatrix CSRRowWisePerEtypePick(CSRMatrix mat, IdArray rows,
// fast path, select all
for (int64_t k = 0; k < et_len; ++k) {
const IdxType eid_offset = off + et_idx[et_offset + k];
const IdxType homogenized_eid = eid ? eid[eid_offset] : eid_offset;
const IdxType homogenized_eid =
eid ? eid[eid_offset] : eid_offset;
auto it = std::upper_bound(
eid2etype_offset.begin(), eid2etype_offset.end(), homogenized_eid);
const IdxType heterogenized_etype = it - eid2etype_offset.begin() - 1;
const IdxType heterogenized_eid = \
eid2etype_offset.begin(), eid2etype_offset.end(),
homogenized_eid);
const IdxType heterogenized_etype =
it - eid2etype_offset.begin() - 1;
const IdxType heterogenized_eid =
homogenized_eid - eid2etype_offset[heterogenized_etype];
if (!has_probs || IsNullArray(prob_or_mask[heterogenized_etype])) {
if (!has_probs ||
IsNullArray(prob_or_mask[heterogenized_etype])) {
// No probability, select all
rows.push_back(rid);
cols.push_back(indices[eid_offset]);
......@@ -357,32 +377,31 @@ COOMatrix CSRRowWisePerEtypePick(CSRMatrix mat, IdArray rows,
}
}
} else {
IdArray picked_idx = Full(-1, num_picks[cur_et], sizeof(IdxType) * 8, ctx);
IdArray picked_idx =
Full(-1, num_picks[cur_et], sizeof(IdxType) * 8, ctx);
IdxType* picked_idata = picked_idx.Ptr<IdxType>();
// need call random pick
pick_fn(off, et_offset, cur_et,
et_len, et_idx, et_eid,
eid, picked_idata);
pick_fn(
off, et_offset, cur_et, et_len, et_idx, et_eid, eid,
picked_idata);
for (int64_t k = 0; k < num_picks[cur_et]; ++k) {
const IdxType picked = picked_idata[k];
if (picked == -1)
continue;
if (picked == -1) continue;
rows.push_back(rid);
cols.push_back(indices[off+et_idx[et_offset+picked]]);
cols.push_back(indices[off + et_idx[et_offset + picked]]);
if (eid) {
idx.push_back(eid[off+et_idx[et_offset+picked]]);
idx.push_back(eid[off + et_idx[et_offset + picked]]);
} else {
idx.push_back(off+et_idx[et_offset+picked]);
idx.push_back(off + et_idx[et_offset + picked]);
}
}
}
if (j+1 == len)
break;
if (j + 1 == len) break;
// next etype
cur_et = et[et_idx[j+1]];
et_offset = j+1;
cur_et = et[et_idx[j + 1]];
et_offset = j + 1;
et_len = 1;
} else {
et_len++;
......@@ -402,31 +421,32 @@ COOMatrix CSRRowWisePerEtypePick(CSRMatrix mat, IdArray rows,
IdArray picked_row = Concat(picked_rows);
IdArray picked_col = Concat(picked_cols);
IdArray picked_idx = Concat(picked_idxs);
return COOMatrix(mat.num_rows, mat.num_cols,
picked_row, picked_col, picked_idx);
return COOMatrix(
mat.num_rows, mat.num_cols, picked_row, picked_col, picked_idx);
}
// Template for picking non-zero values row-wise. The implementation first slices
// out the corresponding rows and then converts it to CSR format. It then performs
// row-wise pick on the CSR matrix and rectifies the returned results.
// Template for picking non-zero values row-wise. The implementation first
// slices out the corresponding rows and then converts it to CSR format. It then
// performs row-wise pick on the CSR matrix and rectifies the returned results.
template <typename IdxType>
COOMatrix COORowWisePick(COOMatrix mat, IdArray rows,
int64_t num_picks, bool replace, PickFn<IdxType> pick_fn,
NumPicksFn<IdxType> num_picks_fn) {
COOMatrix COORowWisePick(
COOMatrix mat, IdArray rows, int64_t num_picks, bool replace,
PickFn<IdxType> pick_fn, NumPicksFn<IdxType> num_picks_fn) {
using namespace aten;
const auto& csr = COOToCSR(COOSliceRows(mat, rows));
const IdArray new_rows = Range(0, rows->shape[0], rows->dtype.bits, rows->ctx);
const IdArray new_rows =
Range(0, rows->shape[0], rows->dtype.bits, rows->ctx);
const auto& picked = CSRRowWisePick<IdxType>(
csr, new_rows, num_picks, replace, pick_fn, num_picks_fn);
return COOMatrix(mat.num_rows, mat.num_cols,
return COOMatrix(
mat.num_rows, mat.num_cols,
IndexSelect(rows, picked.row), // map the row index to the correct one
picked.col,
picked.data);
picked.col, picked.data);
}
// Template for picking non-zero values row-wise. The implementation first slices
// out the corresponding rows and then converts it to CSR format. It then performs
// row-wise pick on the CSR matrix and rectifies the returned results.
// Template for picking non-zero values row-wise. The implementation first
// slices out the corresponding rows and then converts it to CSR format. It then
// performs row-wise pick on the CSR matrix and rectifies the returned results.
template <typename IdxType, typename DType>
COOMatrix COORowWisePerEtypePick(
COOMatrix mat, IdArray rows, const std::vector<int64_t>& eid2etype_offset,
......@@ -435,13 +455,15 @@ COOMatrix COORowWisePerEtypePick(
const std::vector<NDArray>& prob_or_mask) {
using namespace aten;
const auto& csr = COOToCSR(COOSliceRows(mat, rows));
const IdArray new_rows = Range(0, rows->shape[0], rows->dtype.bits, rows->ctx);
const IdArray new_rows =
Range(0, rows->shape[0], rows->dtype.bits, rows->ctx);
const auto& picked = CSRRowWisePerEtypePick<IdxType, DType>(
csr, new_rows, eid2etype_offset, num_picks, replace, false, pick_fn, prob_or_mask);
return COOMatrix(mat.num_rows, mat.num_cols,
csr, new_rows, eid2etype_offset, num_picks, replace, false, pick_fn,
prob_or_mask);
return COOMatrix(
mat.num_rows, mat.num_cols,
IndexSelect(rows, picked.row), // map the row index to the correct one
picked.col,
picked.data);
picked.col, picked.data);
}
} // namespace impl
......
src/array/cpu/rowwise_sampling.cc
View file @ 8ac27dad
This diff is collapsed.
src/array/cpu/rowwise_topk.cc
View file @ 8ac27dad
......@@ -3,8 +3,9 @@
* @file array/cpu/rowwise_topk.cc
* @brief rowwise topk
*/
#include <numeric>
#include <algorithm>
#include <numeric>
#include "./rowwise_pick.h"
namespace dgl {
......@@ -14,9 +15,9 @@ namespace {
template <typename IdxType>
inline NumPicksFn<IdxType> GetTopkNumPicksFn(int64_t k) {
NumPicksFn<IdxType> num_picks_fn = [k]
(IdxType rowid, IdxType off, IdxType len,
const IdxType* col, const IdxType* data) {
NumPicksFn<IdxType> num_picks_fn = [k](IdxType rowid, IdxType off,
IdxType len, const IdxType* col,
const IdxType* data) {
const int64_t max_num_picks = (k == -1) ? len : k;
return std::min(static_cast<IdxType>(max_num_picks), len);
};
......@@ -26,28 +27,28 @@ inline NumPicksFn<IdxType> GetTopkNumPicksFn(int64_t k) {
template <typename IdxType, typename DType>
inline PickFn<IdxType> GetTopkPickFn(NDArray weight, bool ascending) {
const DType* wdata = static_cast<DType*>(weight->data);
PickFn<IdxType> pick_fn = [ascending, wdata]
(IdxType rowid, IdxType off, IdxType len, IdxType num_picks,
const IdxType* col, const IdxType* data,
IdxType* out_idx) {
PickFn<IdxType> pick_fn = [ascending, wdata](
IdxType rowid, IdxType off, IdxType len,
IdxType num_picks, const IdxType* col,
const IdxType* data, IdxType* out_idx) {
std::function<bool(IdxType, IdxType)> compare_fn;
if (ascending) {
if (data) {
compare_fn = [wdata, data] (IdxType i, IdxType j) {
compare_fn = [wdata, data](IdxType i, IdxType j) {
return wdata[data[i]] < wdata[data[j]];
};
} else {
compare_fn = [wdata] (IdxType i, IdxType j) {
compare_fn = [wdata](IdxType i, IdxType j) {
return wdata[i] < wdata[j];
};
}
} else {
if (data) {
compare_fn = [wdata, data] (IdxType i, IdxType j) {
compare_fn = [wdata, data](IdxType i, IdxType j) {
return wdata[data[i]] > wdata[data[j]];
};
} else {
compare_fn = [wdata] (IdxType i, IdxType j) {
compare_fn = [wdata](IdxType i, IdxType j) {
return wdata[i] > wdata[j];
};
}
......
src/array/cpu/sddmm.cc
View file @ 8ac27dad
......@@ -4,6 +4,7 @@
* @brief SDDMM C APIs and definitions.
*/
#include "./sddmm.h"
#include <dgl/array.h>
namespace dgl {
......@@ -25,7 +26,7 @@ namespace aten {
} \
} while (0)
#define SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, ...)\
#define SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, ...) \
do { \
if ((lhs_target) == 0) { \
constexpr int LhsTarget = 0; \
......@@ -41,35 +42,26 @@ namespace aten {
} \
} while (0)
/** @brief Generalized SDDMM on Csr format. */
template <int XPU, typename IdType, typename DType>
void SDDMMCsr(const std::string& op,
const BcastOff& bcast,
const CSRMatrix& csr,
NDArray lhs,
NDArray rhs,
NDArray out,
int lhs_target,
int rhs_target) {
void SDDMMCsr(
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target) {
SWITCH_OP(op, Op, {
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
cpu::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(bcast, csr, lhs, rhs, out);
cpu::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(
bcast, csr, lhs, rhs, out);
});
});
}
/** @brief Generalized SDDMM on Csr format with Heterograph support. */
template <int XPU, typename IdType, typename DType>
void SDDMMCsrHetero(const std::string& op,
const BcastOff& bcast,
const std::vector<CSRMatrix>& vec_csr,
const std::vector<NDArray>& vec_lhs,
const std::vector<NDArray>& vec_rhs,
std::vector<NDArray> vec_out,
int lhs_target,
int rhs_target,
const std::vector<dgl_type_t>& lhs_nid,
void SDDMMCsrHetero(
const std::string& op, const BcastOff& bcast,
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& vec_lhs,
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& lhs_nid,
const std::vector<dgl_type_t>& rhs_nid) {
SWITCH_OP(op, Op, {
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
......@@ -79,7 +71,8 @@ void SDDMMCsrHetero(const std::string& op,
NDArray lhs = vec_lhs[lhs_nid[etype]];
NDArray rhs = vec_rhs[rhs_nid[etype]];
NDArray out = vec_out[etype];
cpu::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(bcast, csr, lhs, rhs, out);
cpu::SDDMMCsr<IdType, DType, Op, LhsTarget, RhsTarget>(
bcast, csr, lhs, rhs, out);
}
});
});
......@@ -87,78 +80,62 @@ void SDDMMCsrHetero(const std::string& op,
template void SDDMMCsr<kDGLCPU, int32_t, float>(
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCsr<kDGLCPU, int64_t, float>(
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCsr<kDGLCPU, int32_t, double>(
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCsr<kDGLCPU, int64_t, double>(
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCsrHetero<kDGLCPU, int32_t, float>(
const std::string& op, const BcastOff& bcast,
const std::vector<CSRMatrix>& vec_csr,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
template void SDDMMCsrHetero<kDGLCPU, int64_t, float>(
const std::string& op, const BcastOff& bcast,
const std::vector<CSRMatrix>& vec_csr,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
template void SDDMMCsrHetero<kDGLCPU, int32_t, double>(
const std::string& op, const BcastOff& bcast,
const std::vector<CSRMatrix>& vec_csr,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
template void SDDMMCsrHetero<kDGLCPU, int64_t, double>(
const std::string& op, const BcastOff& bcast,
const std::vector<CSRMatrix>& vec_csr,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<CSRMatrix>& vec_csr, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
/** @brief Generalized SDDMM on Coo format. */
template <int XPU, typename IdType, typename DType>
void SDDMMCoo(const std::string& op,
const BcastOff& bcast,
const COOMatrix& coo,
NDArray lhs,
NDArray rhs,
NDArray out,
int lhs_target,
int rhs_target) {
void SDDMMCoo(
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target) {
SWITCH_OP(op, Op, {
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
cpu::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(bcast, coo, lhs, rhs, out);
cpu::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(
bcast, coo, lhs, rhs, out);
});
});
}
/** @brief Generalized SDDMM on Coo format with Heterograph support. */
template <int XPU, typename IdType, typename DType>
void SDDMMCooHetero(const std::string& op,
const BcastOff& bcast,
const std::vector<COOMatrix>& vec_coo,
const std::vector<NDArray>& vec_lhs,
const std::vector<NDArray>& vec_rhs,
std::vector<NDArray> vec_out,
int lhs_target,
int rhs_target,
const std::vector<dgl_type_t>& lhs_nid,
void SDDMMCooHetero(
const std::string& op, const BcastOff& bcast,
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& vec_lhs,
const std::vector<NDArray>& vec_rhs, std::vector<NDArray> vec_out,
int lhs_target, int rhs_target, const std::vector<dgl_type_t>& lhs_nid,
const std::vector<dgl_type_t>& rhs_nid) {
SWITCH_OP(op, Op, {
SWITCH_TARGET(lhs_target, rhs_target, LhsTarget, RhsTarget, {
......@@ -168,7 +145,8 @@ void SDDMMCooHetero(const std::string& op,
NDArray lhs = vec_lhs[lhs_nid[etype]];
NDArray rhs = vec_rhs[rhs_nid[etype]];
NDArray out = vec_out[etype];
cpu::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(bcast, coo, lhs, rhs, out);
cpu::SDDMMCoo<IdType, DType, Op, LhsTarget, RhsTarget>(
bcast, coo, lhs, rhs, out);
}
});
});
......@@ -176,48 +154,40 @@ void SDDMMCooHetero(const std::string& op,
template void SDDMMCoo<kDGLCPU, int32_t, float>(
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCoo<kDGLCPU, int64_t, float>(
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCoo<kDGLCPU, int32_t, double>(
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCoo<kDGLCPU, int64_t, double>(
const std::string& op, const BcastOff& bcast, const COOMatrix& coo,
NDArray lhs, NDArray rhs, NDArray out,
int lhs_target, int rhs_target);
NDArray lhs, NDArray rhs, NDArray out, int lhs_target, int rhs_target);
template void SDDMMCooHetero<kDGLCPU, int32_t, float>(
const std::string& op, const BcastOff& bcast,
const std::vector<COOMatrix>& vec_coo,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
template void SDDMMCooHetero<kDGLCPU, int64_t, float>(
const std::string& op, const BcastOff& bcast,
const std::vector<COOMatrix>& vec_coo,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
template void SDDMMCooHetero<kDGLCPU, int32_t, double>(
const std::string& op, const BcastOff& bcast,
const std::vector<COOMatrix>& vec_coo,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
template void SDDMMCooHetero<kDGLCPU, int64_t, double>(
const std::string& op, const BcastOff& bcast,
const std::vector<COOMatrix>& vec_coo,
const std::vector<NDArray>& lhs, const std::vector<NDArray>& rhs,
std::vector<NDArray> out, int lhs_target, int rhs_target,
const std::vector<dgl_type_t>& in_eid,
const std::vector<COOMatrix>& vec_coo, const std::vector<NDArray>& lhs,
const std::vector<NDArray>& rhs, std::vector<NDArray> out, int lhs_target,
int rhs_target, const std::vector<dgl_type_t>& in_eid,
const std::vector<dgl_type_t>& out_eid);
} // namespace aten
......
src/array/cpu/segment_reduce.cc
View file @ 8ac27dad
......@@ -4,8 +4,11 @@
* @brief Segment reduce C APIs and definitions.
*/
#include "./segment_reduce.h"
#include <dgl/array.h>
#include <string>
#include "./spmm_binary_ops.h"
namespace dgl {
......@@ -14,10 +17,7 @@ namespace aten {
/** @brief Segment Reduce operator. */
template <int XPU, typename IdType, typename DType>
void SegmentReduce(
const std::string& op,
NDArray feat,
NDArray offsets,
NDArray out,
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
NDArray arg) {
if (op == "sum") {
cpu::SegmentSum<IdType, DType>(feat, offsets, out);
......@@ -36,73 +36,47 @@ void SegmentReduce(
/** @brief Scatter Add.*/
template <int XPU, typename IdType, typename DType>
void ScatterAdd(NDArray feat,
NDArray idx,
NDArray out) {
void ScatterAdd(NDArray feat, NDArray idx, NDArray out) {
cpu::ScatterAdd<IdType, DType>(feat, idx, out);
}
/** @brief Update gradients for reduce operator max/min on heterogeneous graph.*/
/** @brief Update gradients for reduce operator max/min on heterogeneous
* graph.*/
template <int XPU, typename IdType, typename DType>
void UpdateGradMinMax_hetero(const HeteroGraphPtr& g,
const std::string& op,
const std::vector<NDArray>& feat,
const std::vector<NDArray>& idx,
const std::vector<NDArray>& idx_etype,
std::vector<NDArray>* out) {
void UpdateGradMinMax_hetero(
const HeteroGraphPtr& g, const std::string& op,
const std::vector<NDArray>& feat, const std::vector<NDArray>& idx,
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out) {
cpu::UpdateGradMinMax_hetero<IdType, DType>(g, op, feat, idx, idx_etype, out);
}
/** @brief Backward function of segment cmp.*/
template <int XPU, typename IdType, typename DType>
void BackwardSegmentCmp(
NDArray feat,
NDArray arg,
NDArray out) {
void BackwardSegmentCmp(NDArray feat, NDArray arg, NDArray out) {
cpu::BackwardSegmentCmp<IdType, DType>(feat, arg, out);
}
template void SegmentReduce<kDGLCPU, int32_t, float>(
const std::string &op,
NDArray feat,
NDArray offsets,
NDArray out,
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
NDArray arg);
template void SegmentReduce<kDGLCPU, int64_t, float>(
const std::string &op,
NDArray feat,
NDArray offsets,
NDArray out,
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
NDArray arg);
template void SegmentReduce<kDGLCPU, int32_t, double>(
const std::string &op,
NDArray feat,
NDArray offsets,
NDArray out,
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
NDArray arg);
template void SegmentReduce<kDGLCPU, int64_t, double>(
const std::string &op,
NDArray feat,
NDArray offsets,
NDArray out,
const std::string& op, NDArray feat, NDArray offsets, NDArray out,
NDArray arg);
template void ScatterAdd<kDGLCPU, int32_t, float>(
NDArray feat,
NDArray idx,
NDArray out);
NDArray feat, NDArray idx, NDArray out);
template void ScatterAdd<kDGLCPU, int64_t, float>(
NDArray feat,
NDArray idx,
NDArray out);
NDArray feat, NDArray idx, NDArray out);
template void ScatterAdd<kDGLCPU, int32_t, double>(
NDArray feat,
NDArray idx,
NDArray out);
NDArray feat, NDArray idx, NDArray out);
template void ScatterAdd<kDGLCPU, int64_t, double>(
NDArray feat,
NDArray arg,
NDArray out);
NDArray feat, NDArray arg, NDArray out);
template void UpdateGradMinMax_hetero<kDGLCPU, int32_t, float>(
const HeteroGraphPtr& g, const std::string& op,
......@@ -122,21 +96,13 @@ template void UpdateGradMinMax_hetero<kDGLCPU, int64_t, double>(
const std::vector<NDArray>& idx_etype, std::vector<NDArray>* out);
template void BackwardSegmentCmp<kDGLCPU, int32_t, float>(
NDArray feat,
NDArray arg,
NDArray out);
NDArray feat, NDArray arg, NDArray out);
template void BackwardSegmentCmp<kDGLCPU, int64_t, float>(
NDArray feat,
NDArray arg,
NDArray out);
NDArray feat, NDArray arg, NDArray out);
template void BackwardSegmentCmp<kDGLCPU, int32_t, double>(
NDArray feat,
NDArray arg,
NDArray out);
NDArray feat, NDArray arg, NDArray out);
template void BackwardSegmentCmp<kDGLCPU, int64_t, double>(
NDArray feat,
NDArray arg,
NDArray out);
NDArray feat, NDArray arg, NDArray out);
} // namespace aten
} // namespace dgl
src/array/cpu/spmm.cc
View file @ 8ac27dad
......@@ -4,6 +4,7 @@
* @brief SPMM C APIs and definitions.
*/
#include "./spmm.h"
#include <dgl/array.h>
namespace dgl {
......@@ -11,12 +12,9 @@ namespace aten {
/** @brief Generalized SpMM on Csr format. */
template <int XPU, typename IdType, typename DType>
void SpMMCsr(const std::string& op, const std::string& reduce,
const BcastOff& bcast,
const CSRMatrix& csr,
NDArray ufeat,
NDArray efeat,
NDArray out,
void SpMMCsr(
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux) {
const int64_t dim = bcast.out_len;
if (reduce == "sum") {
......@@ -25,13 +23,15 @@ void SpMMCsr(const std::string& op, const std::string& reduce,
});
} else if (reduce == "max" || reduce == "min") {
SWITCH_OP(op, Op, {
DType *out_off = out.Ptr<DType>();
DType* out_off = out.Ptr<DType>();
if (reduce == "max") {
std::fill(out_off, out_off + csr.num_rows * dim, cpu::op::Max<DType>::zero);
std::fill(
out_off, out_off + csr.num_rows * dim, cpu::op::Max<DType>::zero);
cpu::SpMMCmpCsr<IdType, DType, Op, cpu::op::Max<DType>>(
bcast, csr, ufeat, efeat, out, out_aux[0], out_aux[1]);
} else {
std::fill(out_off, out_off + csr.num_rows * dim, cpu::op::Min<DType>::zero);
std::fill(
out_off, out_off + csr.num_rows * dim, cpu::op::Min<DType>::zero);
cpu::SpMMCmpCsr<IdType, DType, Op, cpu::op::Min<DType>>(
bcast, csr, ufeat, efeat, out, out_aux[0], out_aux[1]);
}
......@@ -43,12 +43,11 @@ void SpMMCsr(const std::string& op, const std::string& reduce,
/** @brief Generalized SpMM on Csr format. */
template <int XPU, typename IdType, typename DType>
void SpMMCsrHetero(const std::string& op, const std::string& reduce,
const BcastOff& bcast,
void SpMMCsrHetero(
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const std::vector<CSRMatrix>& vec_csr,
const std::vector<NDArray>& vec_ufeat,
const std::vector<NDArray>& vec_efeat,
std::vector<NDArray>* vec_out,
const std::vector<NDArray>& vec_efeat, std::vector<NDArray>* vec_out,
std::vector<std::vector<NDArray>>* out_aux,
const std::vector<dgl_type_t>& ufeat_node_tids,
const std::vector<dgl_type_t>& out_node_tids) {
......@@ -60,8 +59,10 @@ void SpMMCsrHetero(const std::string& op, const std::string& reduce,
const dgl_type_t src_id = ufeat_node_tids[etype];
const dgl_type_t dst_id = out_node_tids[etype];
CSRMatrix csr = vec_csr[etype];
NDArray ufeat = (vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
NDArray efeat = (vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
NDArray ufeat =
(vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
NDArray efeat =
(vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
NDArray out = (*vec_out)[dst_id];
cpu::SpMMSumCsr<IdType, DType, Op>(bcast, csr, ufeat, efeat, out);
}
......@@ -71,21 +72,27 @@ void SpMMCsrHetero(const std::string& op, const std::string& reduce,
std::vector<bool> updated((*vec_out).size(), false);
// TODO(Israt): use vector updated to fill(out...) too
for (dgl_type_t etype = 0; etype < ufeat_node_tids.size(); ++etype) {
DType *out_off = (*vec_out)[out_node_tids[etype]].Ptr<DType>();
DType* out_off = (*vec_out)[out_node_tids[etype]].Ptr<DType>();
if (reduce == "max")
std::fill(out_off, out_off + vec_csr[etype].num_rows * dim, cpu::op::Max<DType>::zero);
std::fill(
out_off, out_off + vec_csr[etype].num_rows * dim,
cpu::op::Max<DType>::zero);
else
std::fill(out_off, out_off + vec_csr[etype].num_rows * dim, cpu::op::Min<DType>::zero);
std::fill(
out_off, out_off + vec_csr[etype].num_rows * dim,
cpu::op::Min<DType>::zero);
const dgl_type_t dst_id = out_node_tids[etype];
if (!updated[dst_id]) {
updated[dst_id] = true;
if (Op::use_lhs) {
IdType *argu_ntype = (*out_aux)[2][dst_id].Ptr<IdType>();
std::fill(argu_ntype, argu_ntype + vec_csr[etype].num_rows * dim, -1);
IdType* argu_ntype = (*out_aux)[2][dst_id].Ptr<IdType>();
std::fill(
argu_ntype, argu_ntype + vec_csr[etype].num_rows * dim, -1);
}
if (Op::use_rhs) {
IdType *arge_etype = (*out_aux)[3][dst_id].Ptr<IdType>();
std::fill(arge_etype, arge_etype + vec_csr[etype].num_rows * dim, -1);
IdType* arge_etype = (*out_aux)[3][dst_id].Ptr<IdType>();
std::fill(
arge_etype, arge_etype + vec_csr[etype].num_rows * dim, -1);
}
}
}
......@@ -94,17 +101,21 @@ void SpMMCsrHetero(const std::string& op, const std::string& reduce,
const dgl_type_t src_id = ufeat_node_tids[etype];
const dgl_type_t dst_id = out_node_tids[etype];
CSRMatrix csr = vec_csr[etype];
NDArray ufeat = (vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
NDArray efeat = (vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
NDArray ufeat =
(vec_ufeat.size() == 0) ? NullArray() : vec_ufeat[src_id];
NDArray efeat =
(vec_efeat.size() == 0) ? NullArray() : vec_efeat[etype];
NDArray out = (*vec_out)[dst_id];
if (reduce == "max") {
cpu::SpMMCmpCsrHetero<IdType, DType, Op, cpu::op::Max<DType>>(
bcast, csr, ufeat, efeat, out, (*out_aux)[0][dst_id], (*out_aux)[1][dst_id],
(*out_aux)[2][dst_id], (*out_aux)[3][dst_id], src_id, etype);
bcast, csr, ufeat, efeat, out, (*out_aux)[0][dst_id],
(*out_aux)[1][dst_id], (*out_aux)[2][dst_id],
(*out_aux)[3][dst_id], src_id, etype);
} else {
cpu::SpMMCmpCsrHetero<IdType, DType, Op, cpu::op::Min<DType>>(
bcast, csr, ufeat, efeat, out, (*out_aux)[0][dst_id], (*out_aux)[1][dst_id],
(*out_aux)[2][dst_id], (*out_aux)[3][dst_id], src_id, etype);
bcast, csr, ufeat, efeat, out, (*out_aux)[0][dst_id],
(*out_aux)[1][dst_id], (*out_aux)[2][dst_id],
(*out_aux)[3][dst_id], src_id, etype);
}
}
});
......@@ -114,119 +125,104 @@ void SpMMCsrHetero(const std::string& op, const std::string& reduce,
}
template void SpMMCsr<kDGLCPU, int32_t, float>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
template void SpMMCsr<kDGLCPU, int64_t, float>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
template void SpMMCsr<kDGLCPU, int32_t, double>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
template void SpMMCsr<kDGLCPU, int64_t, double>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const CSRMatrix& csr, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
template void SpMMCsrHetero<kDGLCPU, int32_t, float>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const std::vector<CSRMatrix>& csr,
const std::vector<NDArray>& ufeat, const std::vector<NDArray>& efeat,
std::vector<NDArray>* out, std::vector<std::vector<NDArray>>* out_aux,
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
std::vector<std::vector<NDArray>>* out_aux,
const std::vector<dgl_type_t>& ufeat_node_tids,
const std::vector<dgl_type_t>& out_node_tids);
template void SpMMCsrHetero<kDGLCPU, int64_t, float>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const std::vector<CSRMatrix>& csr,
const std::vector<NDArray>& ufeat, const std::vector<NDArray>& efeat,
std::vector<NDArray>* out, std::vector<std::vector<NDArray>>* out_aux,
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
std::vector<std::vector<NDArray>>* out_aux,
const std::vector<dgl_type_t>& ufeat_node_tids,
const std::vector<dgl_type_t>& out_node_tids);
template void SpMMCsrHetero<kDGLCPU, int32_t, double>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const std::vector<CSRMatrix>& csr,
const std::vector<NDArray>& ufeat, const std::vector<NDArray>& efeat,
std::vector<NDArray>* out, std::vector<std::vector<NDArray>>* out_aux,
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
std::vector<std::vector<NDArray>>* out_aux,
const std::vector<dgl_type_t>& ufeat_node_tids,
const std::vector<dgl_type_t>& out_node_tids);
template void SpMMCsrHetero<kDGLCPU, int64_t, double>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const std::vector<CSRMatrix>& csr,
const std::vector<NDArray>& ufeat, const std::vector<NDArray>& efeat,
std::vector<NDArray>* out, std::vector<std::vector<NDArray>>* out_aux,
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const std::vector<CSRMatrix>& csr, const std::vector<NDArray>& ufeat,
const std::vector<NDArray>& efeat, std::vector<NDArray>* out,
std::vector<std::vector<NDArray>>* out_aux,
const std::vector<dgl_type_t>& ufeat_node_tids,
const std::vector<dgl_type_t>& out_node_tids);
/** @brief Edge_softmax_csr forward op on Csr format. */
template <int XPU, typename IdType, typename DType>
void Edge_softmax_csr_forward(const std::string& op,
const BcastOff& bcast,
const CSRMatrix& csr,
NDArray ufeat,
NDArray efeat,
NDArray out) {
void Edge_softmax_csr_forward(
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out) {
SWITCH_OP(op, Op, {
cpu::Edge_softmax_csr_forward<IdType, DType, Op>(bcast, csr, ufeat, efeat, out);
cpu::Edge_softmax_csr_forward<IdType, DType, Op>(
bcast, csr, ufeat, efeat, out);
});
}
/** @brief Edge_softmax_csr backward op on Csr format. */
template <int XPU, typename IdType, typename DType>
void Edge_softmax_csr_backward(const std::string& op,
const BcastOff& bcast,
const CSRMatrix& csr,
NDArray out,
NDArray sds,
NDArray back_out) {
void Edge_softmax_csr_backward(
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray out, NDArray sds, NDArray back_out) {
SWITCH_OP(op, Op, {
cpu::Edge_softmax_csr_backward<IdType, DType, Op>(bcast, csr, out, sds, back_out);
cpu::Edge_softmax_csr_backward<IdType, DType, Op>(
bcast, csr, out, sds, back_out);
});
}
template void Edge_softmax_csr_forward<kDGLCPU, int32_t, float>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
template void Edge_softmax_csr_forward<kDGLCPU, int64_t, float>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
template void Edge_softmax_csr_forward<kDGLCPU, int32_t, double>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
template void Edge_softmax_csr_forward<kDGLCPU, int64_t, double>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
template void Edge_softmax_csr_backward<kDGLCPU, int32_t, float>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
template void Edge_softmax_csr_backward<kDGLCPU, int64_t, float>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
template void Edge_softmax_csr_backward<kDGLCPU, int32_t, double>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
template void Edge_softmax_csr_backward<kDGLCPU, int64_t, double>(
const std::string& op,
const BcastOff& bcast, const CSRMatrix& csr,
const std::string& op, const BcastOff& bcast, const CSRMatrix& csr,
NDArray ufeat, NDArray efeat, NDArray out);
/** @brief Generalized SpMM on Coo format. */
template <int XPU, typename IdType, typename DType>
void SpMMCoo(const std::string& op, const std::string& reduce,
const BcastOff& bcast,
const COOMatrix& coo,
NDArray ufeat,
NDArray efeat,
NDArray out,
void SpMMCoo(
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux) {
if (reduce == "sum") {
SWITCH_OP(op, Op, {
......@@ -247,21 +243,21 @@ void SpMMCoo(const std::string& op, const std::string& reduce,
}
template void SpMMCoo<kDGLCPU, int32_t, float>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const COOMatrix& coo,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
template void SpMMCoo<kDGLCPU, int64_t, float>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const COOMatrix& coo,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
template void SpMMCoo<kDGLCPU, int32_t, double>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const COOMatrix& coo,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
template void SpMMCoo<kDGLCPU, int64_t, double>(
const std::string& op, const std::string& reduce,
const BcastOff& bcast, const COOMatrix& coo,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
const std::string& op, const std::string& reduce, const BcastOff& bcast,
const COOMatrix& coo, NDArray ufeat, NDArray efeat, NDArray out,
std::vector<NDArray> out_aux);
} // namespace aten
} // namespace dgl
src/array/cpu/spmm_binary_ops.h
View file @ 8ac27dad
......@@ -7,6 +7,7 @@
#define DGL_ARRAY_CPU_SPMM_BINARY_OPS_H_
#include <dgl/array.h>
#include <dgl/bcast.h>
#include <limits>
namespace dgl {
namespace aten {
......
src/array/cuda/array_index_select.cu
View file @ 8ac27dad
......@@ -63,8 +63,10 @@ template NDArray IndexSelect<kDGLCUDA, int64_t, int64_t>(NDArray, IdArray);
template NDArray IndexSelect<kDGLCUDA, __half, int32_t>(NDArray, IdArray);
template NDArray IndexSelect<kDGLCUDA, __half, int64_t>(NDArray, IdArray);
#if BF16_ENABLED
template NDArray IndexSelect<kDGLCUDA, __nv_bfloat16, int32_t>(NDArray, IdArray);
template NDArray IndexSelect<kDGLCUDA, __nv_bfloat16, int64_t>(NDArray, IdArray);
template NDArray IndexSelect<kDGLCUDA, __nv_bfloat16, int32_t>(
NDArray, IdArray);
template NDArray IndexSelect<kDGLCUDA, __nv_bfloat16, int64_t>(
NDArray, IdArray);
#endif // BF16_ENABLED
template NDArray IndexSelect<kDGLCUDA, float, int32_t>(NDArray, IdArray);
template NDArray IndexSelect<kDGLCUDA, float, int64_t>(NDArray, IdArray);
......@@ -76,8 +78,8 @@ DType IndexSelect(NDArray array, int64_t index) {
auto device = runtime::DeviceAPI::Get(array->ctx);
DType ret = static_cast<DType>(0.0f);
device->CopyDataFromTo(
static_cast<DType*>(array->data) + index, 0, &ret, 0,
sizeof(DType), array->ctx, DGLContext{kDGLCPU, 0}, array->dtype);
static_cast<DType*>(array->data) + index, 0, &ret, 0, sizeof(DType),
array->ctx, DGLContext{kDGLCPU, 0}, array->dtype);
return ret;
}
......@@ -87,7 +89,8 @@ template uint32_t IndexSelect<kDGLCUDA, uint32_t>(NDArray array, int64_t index);
template uint64_t IndexSelect<kDGLCUDA, uint64_t>(NDArray array, int64_t index);
template __half IndexSelect<kDGLCUDA, __half>(NDArray array, int64_t index);
#if BF16_ENABLED
template __nv_bfloat16 IndexSelect<kDGLCUDA, __nv_bfloat16>(NDArray array, int64_t index);
template __nv_bfloat16 IndexSelect<kDGLCUDA, __nv_bfloat16>(
NDArray array, int64_t index);
#endif // BF16_ENABLED
template float IndexSelect<kDGLCUDA, float>(NDArray array, int64_t index);
template double IndexSelect<kDGLCUDA, double>(NDArray array, int64_t index);
......
src/array/cuda/array_op_impl.cu
View file @ 8ac27dad
This diff is collapsed.
src/array/cuda/array_scatter.cu
View file @ 8ac27dad
......@@ -4,6 +4,7 @@
* @brief Array scatter GPU implementation
*/
#include <dgl/array.h>
#include "../../runtime/cuda/cuda_common.h"
#include "./utils.h"
......@@ -13,8 +14,8 @@ namespace aten {
namespace impl {
template <typename DType, typename IdType>
__global__ void _ScatterKernel(const IdType* index, const DType* value,
int64_t length, DType* out) {
__global__ void _ScatterKernel(
const IdType* index, const DType* value, int64_t length, DType* out) {
int tx = blockIdx.x * blockDim.x + threadIdx.x;
int stride_x = gridDim.x * blockDim.x;
while (tx < length) {
......@@ -33,15 +34,15 @@ void Scatter_(IdArray index, NDArray value, NDArray out) {
cudaStream_t stream = runtime::getCurrentCUDAStream();
const int nt = cuda::FindNumThreads(len);
const int nb = (len + nt - 1) / nt;
CUDA_KERNEL_CALL(_ScatterKernel, nb, nt, 0, stream,
idx, val, len, outd);
CUDA_KERNEL_CALL(_ScatterKernel, nb, nt, 0, stream, idx, val, len, outd);
}
template void Scatter_<kDGLCUDA, int32_t, int32_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, int64_t, int32_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, __half, int32_t>(IdArray, NDArray, NDArray);
#if BF16_ENABLED
template void Scatter_<kDGLCUDA, __nv_bfloat16, int32_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, __nv_bfloat16, int32_t>(
IdArray, NDArray, NDArray);
#endif // BF16_ENABLED
template void Scatter_<kDGLCUDA, float, int32_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, double, int32_t>(IdArray, NDArray, NDArray);
......@@ -49,7 +50,8 @@ template void Scatter_<kDGLCUDA, int32_t, int64_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, int64_t, int64_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, __half, int64_t>(IdArray, NDArray, NDArray);
#if BF16_ENABLED
template void Scatter_<kDGLCUDA, __nv_bfloat16, int64_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, __nv_bfloat16, int64_t>(
IdArray, NDArray, NDArray);
#endif // BF16_ENABLED
template void Scatter_<kDGLCUDA, float, int64_t>(IdArray, NDArray, NDArray);
template void Scatter_<kDGLCUDA, double, int64_t>(IdArray, NDArray, NDArray);
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment