wehub/dmlc--dgl
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
include/dgl/array.h
+17
View File
@@ -0,0 +1,17 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/array.h
* @brief Common array operations required by DGL.
*
* Note that this is not meant for a full support of array library such as ATen.
* Only a limited set of operators required by DGL are implemented.
*/
#ifndef DGL_ARRAY_H_
#define DGL_ARRAY_H_
#include "./aten/array_ops.h"
#include "./aten/coo.h"
#include "./aten/csr.h"
#include "./aten/macro.h"
#include "./aten/spmat.h"
#include "./aten/types.h"
#endif // DGL_ARRAY_H_
include/dgl/array_iterator.h
+173
View File
@@ -0,0 +1,173 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/array_iterator.h
* @brief Various iterators.
*/
#ifndef DGL_ARRAY_ITERATOR_H_
#define DGL_ARRAY_ITERATOR_H_
#ifdef __CUDA_ARCH__
#define CUB_INLINE __host__ __device__ __forceinline__
#else
#define CUB_INLINE inline
#endif // __CUDA_ARCH__
#include <algorithm>
#include <iterator>
#include <utility>
namespace dgl {
namespace aten {
using std::swap;
// Make std::pair work on both host and device
template <typename DType>
struct Pair {
Pair() = default;
Pair(const Pair& other) = default;
Pair(Pair&& other) = default;
CUB_INLINE Pair(DType a, DType b) : first(a), second(b) {}
CUB_INLINE Pair& operator=(const Pair& other) {
first = other.first;
second = other.second;
return *this;
}
CUB_INLINE operator std::pair<DType, DType>() const {
return std::make_pair(first, second);
}
CUB_INLINE bool operator==(const Pair& other) const {
return (first == other.first) && (second == other.second);
}
CUB_INLINE void swap(const Pair& other) const {
std::swap(first, other.first);
std::swap(second, other.second);
}
DType first, second;
};
template <typename DType>
CUB_INLINE void swap(const Pair<DType>& r1, const Pair<DType>& r2) {
r1.swap(r2);
}
// PairRef and PairIterator that serves as an iterator over a pair of arrays in
// a zipped fashion like zip(a, b).
template <typename DType>
struct PairRef {
PairRef() = delete;
PairRef(const PairRef& other) = default;
PairRef(PairRef&& other) = default;
CUB_INLINE PairRef(DType* const r, DType* const c) : a(r), b(c) {}
CUB_INLINE PairRef& operator=(const PairRef& other) {
*a = *other.a;
*b = *other.b;
return *this;
}
CUB_INLINE PairRef& operator=(const Pair<DType>& val) {
*a = val.first;
*b = val.second;
return *this;
}
CUB_INLINE operator Pair<DType>() const { return Pair<DType>(*a, *b); }
CUB_INLINE operator std::pair<DType, DType>() const {
return std::make_pair(*a, *b);
}
CUB_INLINE bool operator==(const PairRef& other) const {
return (*a == *(other.a)) && (*b == *(other.b));
}
CUB_INLINE void swap(const PairRef& other) const {
std::swap(*a, *other.a);
std::swap(*b, *other.b);
}
DType *a, *b;
};
template <typename DType>
CUB_INLINE void swap(const PairRef<DType>& r1, const PairRef<DType>& r2) {
r1.swap(r2);
}
template <typename DType>
struct PairIterator : public std::iterator<
std::random_access_iterator_tag, Pair<DType>,
std::ptrdiff_t, Pair<DType*>, PairRef<DType>> {
PairIterator() = default;
PairIterator(const PairIterator& other) = default;
PairIterator(PairIterator&& other) = default;
CUB_INLINE PairIterator(DType* x, DType* y) : a(x), b(y) {}
PairIterator& operator=(const PairIterator& other) = default;
PairIterator& operator=(PairIterator&& other) = default;
~PairIterator() = default;
CUB_INLINE bool operator==(const PairIterator& other) const {
return a == other.a;
}
CUB_INLINE bool operator!=(const PairIterator& other) const {
return a != other.a;
}
CUB_INLINE bool operator<(const PairIterator& other) const {
return a < other.a;
}
CUB_INLINE bool operator>(const PairIterator& other) const {
return a > other.a;
}
CUB_INLINE bool operator<=(const PairIterator& other) const {
return a <= other.a;
}
CUB_INLINE bool operator>=(const PairIterator& other) const {
return a >= other.a;
}
CUB_INLINE PairIterator& operator+=(const std::ptrdiff_t& movement) {
a += movement;
b += movement;
return *this;
}
CUB_INLINE PairIterator& operator-=(const std::ptrdiff_t& movement) {
a -= movement;
b -= movement;
return *this;
}
CUB_INLINE PairIterator& operator++() {
++a;
++b;
return *this;
}
CUB_INLINE PairIterator& operator--() {
--a;
--b;
return *this;
}
CUB_INLINE PairIterator operator++(int) {
PairIterator ret(*this);
operator++();
return ret;
}
CUB_INLINE PairIterator operator--(int) {
PairIterator ret(*this);
operator--();
return ret;
}
CUB_INLINE PairIterator operator+(const std::ptrdiff_t& movement) const {
return PairIterator(a + movement, b + movement);
}
CUB_INLINE PairIterator operator-(const std::ptrdiff_t& movement) const {
return PairIterator(a - movement, b - movement);
}
CUB_INLINE std::ptrdiff_t operator-(const PairIterator& other) const {
return a - other.a;
}
CUB_INLINE PairRef<DType> operator*() const { return PairRef<DType>(a, b); }
CUB_INLINE PairRef<DType> operator*() { return PairRef<DType>(a, b); }
CUB_INLINE PairRef<DType> operator[](size_t offset) const {
return PairRef<DType>(a + offset, b + offset);
}
CUB_INLINE PairRef<DType> operator[](size_t offset) {
return PairRef<DType>(a + offset, b + offset);
}
DType *a, *b;
};
}; // namespace aten
}; // namespace dgl
#endif // DGL_ARRAY_ITERATOR_H_
include/dgl/aten/array_ops.h
+393
View File
@@ -0,0 +1,393 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/aten/array_ops.h
* @brief Common array operations required by DGL.
*
* Note that this is not meant for a full support of array library such as ATen.
* Only a limited set of operators required by DGL are implemented.
*/
#ifndef DGL_ATEN_ARRAY_OPS_H_
#define DGL_ATEN_ARRAY_OPS_H_
#include <algorithm>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "./types.h"
namespace dgl {
namespace aten {
//////////////////////////////////////////////////////////////////////
// ID array
//////////////////////////////////////////////////////////////////////
/** @return A special array to represent null. */
inline NDArray NullArray(
const DGLDataType& dtype = DGLDataType{kDGLInt, 64, 1},
const DGLContext& ctx = DGLContext{kDGLCPU, 0}) {
return NDArray::Empty({0}, dtype, ctx);
}
/**
* @return Whether the input array is a null array.
*/
inline bool IsNullArray(NDArray array) { return array->shape[0] == 0; }
/**
* @brief Create a new id array with given length
* @param length The array length
* @param ctx The array context
* @param nbits The number of integer bits
* @return id array
*/
IdArray NewIdArray(
int64_t length, DGLContext ctx = DGLContext{kDGLCPU, 0},
uint8_t nbits = 64);
/**
* @brief Create a new float array with given length
* @param length The array length
* @param ctx The array context
* @param nbits The number of integer bits
* @return float array
*/
FloatArray NewFloatArray(int64_t length,
DGLContext ctx = DGLContext{kDGLCPU, 0},
uint8_t nbits = 32);
/**
* @brief Create a new id array using the given vector data
* @param vec The vector data
* @param nbits The integer bits of the returned array
* @param ctx The array context
* @return the id array
*/
template <typename T>
IdArray VecToIdArray(
const std::vector<T>& vec, uint8_t nbits = 64,
DGLContext ctx = DGLContext{kDGLCPU, 0});
/**
* @brief Return an array representing a 1D range.
* @param low Lower bound (inclusive).
* @param high Higher bound (exclusive).
* @param nbits result array's bits (32 or 64)
* @param ctx Device context
* @return range array
*/
IdArray Range(int64_t low, int64_t high, uint8_t nbits, DGLContext ctx);
/**
* @brief Return an array full of the given value
* @param val The value to fill.
* @param length Number of elements.
* @param nbits result array's bits (32 or 64)
* @param ctx Device context
* @return the result array
*/
IdArray Full(int64_t val, int64_t length, uint8_t nbits, DGLContext ctx);
/**
* @brief Return an array full of the given value with the given type.
* @param val The value to fill.
* @param length Number of elements.
* @param ctx Device context
* @return the result array
*/
template <typename DType>
NDArray Full(DType val, int64_t length, DGLContext ctx);
/** @brief Create a deep copy of the given array */
IdArray Clone(IdArray arr);
/** @brief Convert the idarray to the given bit width */
IdArray AsNumBits(IdArray arr, uint8_t bits);
/** @brief Arithmetic functions */
IdArray Add(IdArray lhs, IdArray rhs);
IdArray Sub(IdArray lhs, IdArray rhs);
IdArray Mul(IdArray lhs, IdArray rhs);
IdArray Div(IdArray lhs, IdArray rhs);
IdArray Mod(IdArray lhs, IdArray rhs);
IdArray Add(IdArray lhs, int64_t rhs);
IdArray Sub(IdArray lhs, int64_t rhs);
IdArray Mul(IdArray lhs, int64_t rhs);
IdArray Div(IdArray lhs, int64_t rhs);
IdArray Mod(IdArray lhs, int64_t rhs);
IdArray Add(int64_t lhs, IdArray rhs);
IdArray Sub(int64_t lhs, IdArray rhs);
IdArray Mul(int64_t lhs, IdArray rhs);
IdArray Div(int64_t lhs, IdArray rhs);
IdArray Mod(int64_t lhs, IdArray rhs);
IdArray Neg(IdArray array);
// XXX(minjie): currently using integer array for bool type
IdArray GT(IdArray lhs, IdArray rhs);
IdArray LT(IdArray lhs, IdArray rhs);
IdArray GE(IdArray lhs, IdArray rhs);
IdArray LE(IdArray lhs, IdArray rhs);
IdArray EQ(IdArray lhs, IdArray rhs);
IdArray NE(IdArray lhs, IdArray rhs);
IdArray GT(IdArray lhs, int64_t rhs);
IdArray LT(IdArray lhs, int64_t rhs);
IdArray GE(IdArray lhs, int64_t rhs);
IdArray LE(IdArray lhs, int64_t rhs);
IdArray EQ(IdArray lhs, int64_t rhs);
IdArray NE(IdArray lhs, int64_t rhs);
IdArray GT(int64_t lhs, IdArray rhs);
IdArray LT(int64_t lhs, IdArray rhs);
IdArray GE(int64_t lhs, IdArray rhs);
IdArray LE(int64_t lhs, IdArray rhs);
IdArray EQ(int64_t lhs, IdArray rhs);
IdArray NE(int64_t lhs, IdArray rhs);
/** @brief Stack two arrays (of len L) into a 2*L length array */
IdArray HStack(IdArray arr1, IdArray arr2);
/** @brief Return the indices of the elements that are non-zero. */
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>
ValueType IndexSelect(NDArray array, int64_t index);
/**
* @brief Return the data under the index. In numpy notation, A[I]
*/
NDArray IndexSelect(NDArray array, IdArray index);
/**
* @brief Return the data from `start` (inclusive) to `end` (exclusive).
*/
NDArray IndexSelect(NDArray array, int64_t start, int64_t end);
/**
* @brief Permute the elements of an array according to given indices.
*
* Only support 1D arrays.
*
* Equivalent to:
*
* <code>
* result = np.zeros_like(array)
* result[indices] = array
* </code>
*/
NDArray Scatter(NDArray array, IdArray indices);
/**
* @brief Scatter data into the output array.
*
* Equivalent to:
*
* <code>
* out[index] = value
* </code>
*/
void Scatter_(IdArray index, NDArray value, NDArray out);
/**
* @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.
*/
NDArray Repeat(NDArray array, IdArray repeats);
/**
* @brief Relabel the given ids to consecutive ids.
*
* Relabeling is done inplace. The mapping is created from the union
* of the give arrays.
*
* Example:
*
* Given two IdArrays [2, 3, 10, 0, 2] and [4, 10, 5], one possible return
* mapping is [2, 3, 10, 4, 0, 5], meaning the new ID 0 maps to the old ID
* 2, 1 maps to 3, so on and so forth.
*
* @param arrays The id arrays to relabel.
* @return mapping array M from new id to old id.
*/
IdArray Relabel_(const std::vector<IdArray>& arrays);
/**
* @brief concatenate the given id arrays to one array
*
* Example:
*
* Given two IdArrays [2, 3, 10, 0, 2] and [4, 10, 5]
* Return [2, 3, 10, 0, 2, 4, 10, 5]
*
* @param arrays The id arrays to concatenate.
* @return concatenated array.
*/
NDArray Concat(const std::vector<IdArray>& arrays);
/** @brief Return whether the array is a valid 1D int array*/
inline bool IsValidIdArray(const dgl::runtime::NDArray& arr) {
return arr->ndim == 1 && arr->dtype.code == kDGLInt;
}
/**
* @brief Packs a tensor containing padded sequences of variable length.
*
* Similar to \c pack_padded_sequence in PyTorch, except that
*
* 1. The length for each sequence (before padding) is inferred as the number
* of elements before the first occurrence of \c pad_value.
* 2. It does not sort the sequences by length.
* 3. Along with the tensor containing the packed sequence, it returns both the
* length, as well as the offsets to the packed tensor, of each sequence.
*
* @param array The tensor containing sequences padded to the same length
* @param pad_value The padding value
* @return A triplet of packed tensor, the length tensor, and the offset tensor
*
* @note Example: consider the following array with padding value -1:
*
* <code>
* [[1, 2, -1, -1],
* [3, 4, 5, -1]]
* </code>
*
* 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 offset tensor would be [0, 2], i.e. the offset to the packed tensor for
* each sequence (before padding)
*/
template <typename ValueType>
std::tuple<NDArray, IdArray, IdArray> Pack(NDArray array, ValueType pad_value);
/**
* @brief Batch-slice a 1D or 2D array, and then pack the list of sliced arrays
* by concatenation.
*
* If a 2D array is given, then the function is equivalent to:
*
* <code>
* def ConcatSlices(array, lengths):
* slices = [array[i, :l] for i, l in enumerate(lengths)]
* packed = np.concatenate(slices)
* offsets = np.cumsum([0] + lengths[:-1])
* return packed, offsets
* </code>
*
* If a 1D array is given, then the function is equivalent to
*
* <code>
* def ConcatSlices(array, lengths):
* slices = [array[:l] for l in lengths]
* packed = np.concatenate(slices)
* offsets = np.cumsum([0] + lengths[:-1])
* return packed, offsets
* </code>
*
* @param array A 1D or 2D tensor for slicing
* @param lengths A 1D tensor indicating the number of elements to slice
* @return The tensor with packed slices along with the offsets.
*/
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.
*
* 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
* an indptr array over a count array.
*
* @param array The 1D input array.
* @return Array after cumsum.
*/
IdArray CumSum(IdArray array, bool prepend_zero = false);
/**
* @brief Return the nonzero index.
*
* Only support 1D array. The result index array is in int64.
*
* @param array The input array.
* @return A 1D index array storing the positions of the non zero values.
*/
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.
*
* @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.
*/
std::pair<IdArray, IdArray> Sort(IdArray array, int num_bits = 0);
/**
* @brief Return a string that prints out some debug information.
*/
std::string ToDebugString(NDArray array);
// inline implementations
template <typename T>
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));
} else if (nbits == 64) {
std::copy(vec.begin(), vec.end(), static_cast<int64_t*>(ret->data));
} else {
LOG(FATAL) << "Only int32 or int64 is supported.";
}
return ret.CopyTo(ctx);
}
/**
* @brief Get the context of the first array, and check if the non-null arrays'
* contexts are the same.
*/
inline DGLContext GetContextOf(const std::vector<IdArray>& arrays) {
bool first = true;
DGLContext result;
for (auto& array : arrays) {
if (first) {
first = false;
result = array->ctx;
} else {
CHECK_EQ(array->ctx, result)
<< "Context of the input arrays are different";
}
}
return result;
}
} // namespace aten
} // namespace dgl
#endif // DGL_ATEN_ARRAY_OPS_H_
include/dgl/aten/coo.h
+878
View File
@@ -0,0 +1,878 @@
/**
* Copyright (c) 2020-2022 by Contributors
* @file dgl/aten/coo.h
* @brief Common COO operations required by DGL.
*/
#ifndef DGL_ATEN_COO_H_
#define DGL_ATEN_COO_H_
#include <dmlc/io.h>
#include <dmlc/serializer.h>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "./array_ops.h"
#include "./macro.h"
#include "./spmat.h"
#include "./types.h"
namespace dgl {
namespace aten {
struct CSRMatrix;
/**
* @brief Plain COO structure
*
* The data array stores integer ids for reading edge features.
* Note that we do allow duplicate non-zero entries -- multiple non-zero entries
* that have the same row, col indices. It corresponds to multigraph in
* graph terminology.
*/
constexpr uint64_t kDGLSerialize_AtenCooMatrixMagic = 0xDD61ffd305dff127;
// TODO(BarclayII): Graph queries on COO formats should support the case where
// data ordered by rows/columns instead of EID.
struct COOMatrix {
/** @brief the dense shape of the matrix */
int64_t num_rows = 0, num_cols = 0;
/** @brief COO index arrays */
IdArray row, col;
/** @brief data index array. When is null, assume it is from 0 to NNZ - 1. */
IdArray data;
/** @brief whether the row indices are sorted */
bool row_sorted = false;
/** @brief whether the column indices per row are sorted */
bool col_sorted = false;
/** @brief whether the matrix is in pinned memory */
bool is_pinned = false;
/** @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)
: num_rows(nrows),
num_cols(ncols),
row(rarr),
col(carr),
data(darr),
row_sorted(rsorted),
col_sorted(csorted) {
CheckValidity();
}
/** @brief constructor from SparseMatrix object */
explicit COOMatrix(const SparseMatrix& spmat)
: num_rows(spmat.num_rows),
num_cols(spmat.num_cols),
row(spmat.indices[0]),
col(spmat.indices[1]),
data(spmat.indices[2]),
row_sorted(spmat.flags[0]),
col_sorted(spmat.flags[1]) {
CheckValidity();
}
// 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});
}
bool Load(dmlc::Stream* fs) {
uint64_t magicNum;
CHECK(fs->Read(&magicNum)) << "Invalid Magic Number";
CHECK_EQ(magicNum, kDGLSerialize_AtenCooMatrixMagic)
<< "Invalid COOMatrix Data";
CHECK(fs->Read(&num_cols)) << "Invalid num_cols";
CHECK(fs->Read(&num_rows)) << "Invalid num_rows";
CHECK(fs->Read(&row)) << "Invalid row";
CHECK(fs->Read(&col)) << "Invalid col";
CHECK(fs->Read(&data)) << "Invalid data";
CHECK(fs->Read(&row_sorted)) << "Invalid row_sorted";
CHECK(fs->Read(&col_sorted)) << "Invalid col_sorted";
CheckValidity();
return true;
}
void Save(dmlc::Stream* fs) const {
fs->Write(kDGLSerialize_AtenCooMatrixMagic);
fs->Write(num_cols);
fs->Write(num_rows);
fs->Write(row);
fs->Write(col);
fs->Write(data);
fs->Write(row_sorted);
fs->Write(col_sorted);
}
inline void CheckValidity() const {
CHECK_SAME_DTYPE(row, col);
CHECK_SAME_CONTEXT(row, col);
if (!aten::IsNullArray(data)) {
CHECK_SAME_DTYPE(row, data);
CHECK_SAME_CONTEXT(row, data);
}
CHECK_NO_OVERFLOW(row->dtype, num_rows);
CHECK_NO_OVERFLOW(row->dtype, num_cols);
}
inline bool IsEmpty() const {
return aten::IsNullArray(row) && aten::IsNullArray(col) &&
aten::IsNullArray(data);
}
// Check and update the internal flag is_pinned.
// This function will initialize a cuda context.
inline bool CheckIfPinnedInCUDA() {
is_pinned = (aten::IsNullArray(row) || row.IsPinned()) &&
(aten::IsNullArray(col) || col.IsPinned()) &&
(aten::IsNullArray(data) || data.IsPinned());
return is_pinned;
}
/** @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);
}
/** @brief Return a copy of this matrix in pinned (page-locked) memory. */
inline COOMatrix PinMemory() {
if (!IsEmpty()) {
if (is_pinned) return *this;
auto new_coo = COOMatrix(
num_rows, num_cols, row.PinMemory(), col.PinMemory(),
aten::IsNullArray(data) ? data : data.PinMemory(), row_sorted,
col_sorted);
CHECK(new_coo.CheckIfPinnedInCUDA())
<< "An internal DGL error has occured while trying to pin a COO "
"matrix. Please file a bug at "
"'https://github.com/dmlc/dgl/issues' "
"with the above stacktrace.";
return new_coo;
}
is_pinned = true;
return *this;
}
/**
* @brief Pin the row, col and data (if not Null) of the matrix.
* @note This is an in-place method. Behavior depends on the current context,
* kDGLCPU: will be pinned;
* IsPinned: directly return;
* kDGLCUDA: invalid, will throw an error.
* The context check is deferred to pinning the NDArray.
*/
inline void PinMemory_() {
if (!IsEmpty()) {
if (is_pinned) return;
row.PinMemory_();
col.PinMemory_();
if (!aten::IsNullArray(data)) {
data.PinMemory_();
}
is_pinned = true;
}
is_pinned = true;
return;
}
/**
* @brief Unpin the row, col and data (if not Null) of the matrix.
* @note This is an in-place method. Behavior depends on the current context,
* IsPinned: will be unpinned;
* others: directly return.
* The context check is deferred to unpinning the NDArray.
*/
inline void UnpinMemory_() {
if (!IsEmpty()) {
if (!is_pinned) return;
row.UnpinMemory_();
col.UnpinMemory_();
if (!aten::IsNullArray(data)) {
data.UnpinMemory_();
}
is_pinned = false;
}
is_pinned = false;
return;
}
/**
* @brief Record stream for the row, col and data (if not Null) of the matrix.
* @param stream The stream that is using the graph
*/
inline void RecordStream(DGLStreamHandle stream) const {
row.RecordStream(stream);
col.RecordStream(stream);
if (!aten::IsNullArray(data)) {
data.RecordStream(stream);
}
}
};
///////////////////////// COO routines //////////////////////////
/** @brief Return true if the value (row, col) is non-zero */
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).
*/
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);
/** @brief Return the data array of the given 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); }
/**
* @brief Check whether the COO is sorted.
*
* It returns two flags: one for whether the row is sorted;
* the other for whether the columns of each row is sorted
* if the first flag is true.
*
* Complexity: O(NNZ)
*/
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.
*
* @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 Three arrays {rows, cols, data}
*/
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.
*/
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);
const auto& rst = COOGetDataAndIndices(mat, rows, cols);
return rst[2];
}
/**
* @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.
*
* @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);
/** @brief Return a transposed COO matrix */
COOMatrix COOTranspose(COOMatrix coo);
/**
* @brief Convert COO matrix to CSR matrix.
*
* If the input COO matrix does not have data array, the data array of
* the result CSR matrix stores a shuffle index for how the entries
* will be reordered in CSR. The i^th entry in the result CSR corresponds
* to the CSR.data[i] th entry in the input COO.
*
* Conversion complexity: O(nnz)
*
* - The function first check whether the input COO matrix is sorted
* using a linear scan.
* - If the COO matrix is row sorted, the conversion can be done very
* efficiently in a sequential scan. The result indices and data arrays
* are directly equal to the column and data arrays from the input.
* - If the COO matrix is further column sorted, the result CSR is
* also column sorted.
* - Otherwise, the conversion is more costly but still is O(nnz).
*
* @param coo Input COO matrix.
* @return CSR matrix.
*/
CSRMatrix COOToCSR(COOMatrix coo);
/**
* @brief Slice rows of the given matrix and return.
* @param coo COO matrix
* @param start Start row id (inclusive)
* @param end End row id (exclusive)
*/
COOMatrix COOSliceRows(COOMatrix coo, int64_t start, int64_t end);
COOMatrix COOSliceRows(COOMatrix coo, runtime::NDArray rows);
/**
* @brief Get the submatrix specified by the row and col ids.
*
* In numpy notation, given matrix M, row index array I, col index array J
* This function returns the submatrix M[I, J].
*
* @param coo The input coo matrix
* @param rows The row index to select
* @param cols The col index to select
* @return submatrix
*/
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.
*/
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.
*
* 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.
*/
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.
*
* 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);
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.
*/
COOMatrix COORemove(COOMatrix coo, IdArray entries);
/**
* @brief Reorder the rows and colmns according to the new row and column order.
* @param csr The input coo matrix.
* @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);
/**
* @brief Randomly select a fixed number of non-zero entries along each given
* row using arXiv:2210.13339, Labor sampling.
*
* The picked indices are returned in the form of a COO matrix.
*
* The passed random_seed makes it so that for any seed vertex s and its
* neighbor t, the rolled random variate r_t is the same for any call to this
* function with the same random seed. When sampling as part of the same batch,
* one would want identical seeds so that LABOR can globally sample. One example
* is that for heterogenous graphs, there is a single random seed passed for
* each edge type. This will sample much fewer vertices compared to having
* unique random seeds for each edge type. If one called this function
* individually for each edge type for a heterogenous graph with different
* random seeds, then it would run LABOR locally for each edge type, resulting
* into a larger number of vertices being sampled.
*
* If this function is called without a random_seed, we get the random seed by
* getting a random number from DGL.
*
*
* Examples:
*
* // coo.num_rows = 4;
* // coo.num_cols = 4;
* // coo.rows = [0, 0, 1, 3, 3]
* // coo.cols = [0, 1, 1, 2, 3]
* // coo.data = [2, 3, 0, 1, 4]
* COOMatrix coo = ...;
* IdArray rows = ... ; // [1, 3]
* COOMatrix sampled = COOLaborSampling(coo, rows, 2, NullArray(), 0 \
* , NullArray(), NullArray());
* // possible sampled coo matrix:
* // sampled.num_rows = 4
* // sampled.num_cols = 4
* // sampled.rows = [1, 3, 3]
* // sampled.cols = [1, 2, 3]
* // sampled.data = [3, 0, 4]
*
* @param mat Input coo matrix.
* @param rows Rows to sample from.
* @param num_samples Number of samples using labor sampling
* @param prob Probability array for nonuniform sampling
* @param importance_sampling Whether to enable importance sampling
* @param random_seed The random seed for the sampler
* @param seed2_contribution The contribution of the second random seed, [0, 1)
* @param NIDs global nids if sampling from a subgraph
* @return A pair of COOMatrix storing the picked row and col indices and edge
* weights if importance_sampling != 0 or prob argument was passed.
* Its data field stores the the index of the picked elements in the
* value array.
*/
std::pair<COOMatrix, FloatArray> COOLaborSampling(
COOMatrix mat, IdArray rows, int64_t num_samples,
FloatArray prob = NullArray(), int importance_sampling = 0,
IdArray random_seed = NullArray(), float seed2_contribution = 0,
IdArray NIDs = NullArray());
/**
* @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.
*
* If replace is false and a row has fewer non-zero values than num_samples,
* all the values are picked.
*
* Examples:
*
* // coo.num_rows = 4;
* // coo.num_cols = 4;
* // coo.rows = [0, 0, 1, 3, 3]
* // coo.cols = [0, 1, 1, 2, 3]
* // coo.data = [2, 3, 0, 1, 4]
* COOMatrix coo = ...;
* IdArray rows = ... ; // [1, 3]
* COOMatrix sampled = COORowWiseSampling(coo, rows, 2, FloatArray(), false);
* // possible sampled coo matrix:
* // sampled.num_rows = 4
* // sampled.num_cols = 4
* // sampled.rows = [1, 3, 3]
* // sampled.cols = [1, 2, 3]
* // sampled.data = [3, 0, 4]
*
* @param mat Input coo matrix.
* @param rows Rows to sample from.
* @param num_samples Number of samples
* @param prob_or_mask Unnormalized probability array or mask array.
* 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.
*/
COOMatrix COORowWiseSampling(
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
* along each given row independently.
*
* The function performs random choices along each row independently.
* In each row, num_samples samples is picked for each edge type. (The edge
* type is stored in etypes)
* The picked indices are returned in the form of a COO matrix.
*
* If replace is false and a row has fewer non-zero values than num_samples,
* all the values are picked.
*
* Examples:
*
* // coo.num_rows = 4;
* // coo.num_cols = 4;
* // coo.rows = [0, 0, 0, 0, 3]
* // coo.cols = [0, 1, 3, 2, 3]
* // coo.data = [2, 3, 0, 1, 4]
* // eid2etype_offset = [0, 3, 4, 5]
* 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);
* // possible sampled coo matrix:
* // sampled.num_rows = 4
* // sampled.num_cols = 4
* // sampled.rows = [0, 0, 0, 3]
* // sampled.cols = [0, 3, 2, 3]
* // sampled.data = [2, 0, 1, 4]
*
* @param mat Input coo matrix.
* @param rows Rows to sample from.
* @param eid2etype_offset The offset to each edge type.
* @param num_samples Number of samples
* @param prob_or_mask Unnormalized probability array or mask array.
* 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.
* @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,
const std::vector<int64_t>& num_samples,
const std::vector<NDArray>& prob_or_mask, bool replace = true);
/**
* @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.
*
* If replace is false and a row has fewer non-zero values than k,
* all the values are picked.
*
* Examples:
*
* // coo.num_rows = 4;
* // coo.num_cols = 4;
* // coo.rows = [0, 0, 1, 3, 3]
* // coo.cols = [0, 1, 1, 2, 3]
* // coo.data = [2, 3, 0, 1, 4]
* COOMatrix coo = ...;
* IdArray rows = ... ; // [0, 1, 3]
* FloatArray weight = ... ; // [1., 0., -1., 10., 20.]
* COOMatrix sampled = COORowWiseTopk(coo, rows, 1, weight);
* // possible sampled coo matrix:
* // sampled.num_rows = 4
* // sampled.num_cols = 4
* // sampled.rows = [0, 1, 3]
* // sampled.cols = [1, 1, 2]
* // sampled.data = [3, 0, 1]
*
* @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.
*/
COOMatrix COORowWiseTopk(
COOMatrix mat, IdArray rows, int64_t k, NDArray weight,
bool ascending = false);
/**
* @brief Union two COOMatrix into one COOMatrix.
*
* Two Matrix must have the same shape.
*
* Example:
*
* A = [[0, 0, 1, 0],
* [1, 0, 1, 1],
* [0, 1, 0, 0]]
*
* B = [[0, 1, 1, 0],
* [0, 0, 0, 1],
* [0, 0, 1, 0]]
*
* COOMatrix_A.num_rows : 3
* COOMatrix_A.num_cols : 4
* COOMatrix_B.num_rows : 3
* COOMatrix_B.num_cols : 4
*
* C = UnionCoo({A, B});
*
* C = [[0, 1, 2, 0],
* [1, 0, 1, 2],
* [0, 1, 1, 0]]
*
* COOMatrix_C.num_rows : 3
* COOMatrix_C.num_cols : 4
*/
COOMatrix UnionCoo(const std::vector<COOMatrix>& coos);
/**
* @brief DisjointUnion a list COOMatrix into one COOMatrix.
*
* Examples:
*
* A = [[0, 0, 1],
* [1, 0, 1],
* [0, 1, 0]]
*
* B = [[0, 0],
* [1, 0]]
*
* COOMatrix_A.num_rows : 3
* COOMatrix_A.num_cols : 3
* COOMatrix_B.num_rows : 2
* COOMatrix_B.num_cols : 2
*
* C = DisjointUnionCoo({A, B});
*
* C = [[0, 0, 1, 0, 0],
* [1, 0, 1, 0, 0],
* [0, 1, 0, 0, 0],
* [0, 0, 0, 0, 0],
* [0, 0, 0, 1, 0]]
* COOMatrix_C.num_rows : 5
* 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
* @return The combined COOMatrix.
*/
COOMatrix DisjointUnionCoo(const std::vector<COOMatrix>& coos);
/**
* @brief COOMatrix toSimple.
*
* A = [[0, 0, 0],
* [3, 0, 2],
* [1, 1, 0],
* [0, 0, 4]]
*
* B, cnt, edge_map = COOToSimple(A)
*
* B = [[0, 0, 0],
* [1, 0, 1],
* [1, 1, 0],
* [0, 0, 1]]
* cnt = [3, 2, 1, 1, 4]
* 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
* returned graph.
*/
std::tuple<COOMatrix, IdArray, IdArray> COOToSimple(const COOMatrix& coo);
/**
* @brief Split a COOMatrix into multiple disjoin components.
*
* Examples:
*
* C = [[0, 0, 1, 0, 0],
* [1, 0, 1, 0, 0],
* [0, 1, 0, 0, 0],
* [0, 0, 0, 0, 0],
* [0, 0, 0, 1, 0],
* [0, 0, 0, 0, 1]]
* COOMatrix_C.num_rows : 6
* COOMatrix_C.num_cols : 5
*
* batch_size : 2
* edge_cumsum : [0, 4, 6]
* src_vertex_cumsum : [0, 3, 6]
* dst_vertex_cumsum : [0, 3, 5]
*
* ret = DisjointPartitionCooBySizes(C,
* batch_size,
* edge_cumsum,
* src_vertex_cumsum,
* dst_vertex_cumsum)
*
* A = [[0, 0, 1],
* [1, 0, 1],
* [0, 1, 0]]
* COOMatrix_A.num_rows : 3
* COOMatrix_A.num_cols : 3
*
* B = [[0, 0],
* [1, 0],
* [0, 1]]
* COOMatrix_B.num_rows : 3
* COOMatrix_B.num_cols : 2
*
* @param coo COOMatrix to split.
* @param batch_size Number of disjoin components (Sub COOMatrix)
* @param edge_cumsum Number of edges of each components
* @param src_vertex_cumsum Number of src vertices of each component.
* @param dst_vertex_cumsum Number of dst vertices of each component.
* @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);
/**
* @brief Slice a contiguous chunk from a COOMatrix
*
* Examples:
*
* C = [[0, 0, 1, 0, 0],
* [1, 0, 1, 0, 0],
* [0, 1, 0, 0, 0],
* [0, 0, 0, 0, 0],
* [0, 0, 0, 1, 0],
* [0, 0, 0, 0, 1]]
* COOMatrix_C.num_rows : 6
* COOMatrix_C.num_cols : 5
*
* edge_range : [4, 6]
* src_vertex_range : [3, 6]
* dst_vertex_range : [3, 5]
*
* ret = COOSliceContiguousChunk(C,
* edge_range,
* src_vertex_range,
* dst_vertex_range)
*
* ret = [[0, 0],
* [1, 0],
* [0, 1]]
* COOMatrix_ret.num_rows : 3
* COOMatrix_ret.num_cols : 2
*
* @param coo COOMatrix to slice.
* @param edge_range ID range of the edges in the chunk
* @param src_vertex_range ID range of the src vertices in the chunk.
* @param dst_vertex_range ID range of the dst vertices in the chunk.
* @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);
/**
* @brief Create a LineGraph of input coo
*
* A = [[0, 0, 1],
* [1, 0, 1],
* [1, 1, 0]]
* A.row = [0, 1, 1, 2, 2]
* A.col = [2, 0, 2, 0, 1]
* A.eid = [0, 1, 2, 3, 4]
*
* B = COOLineGraph(A, backtracking=False)
*
* B = [[0, 0, 0, 0, 1],
* [1, 0, 0, 0, 0],
* [0, 0, 0, 1, 0],
* [0, 0, 0, 0, 0],
* [0, 1, 0, 0, 0]]
*
* C = COOLineGraph(A, backtracking=True)
*
* C = [[0, 0, 0, 1, 1],
* [1, 0, 0, 0, 0],
* [0, 0, 0, 1, 1],
* [1, 0, 0, 0, 0],
* [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
* @return LineGraph in COO format
*/
COOMatrix COOLineGraph(const COOMatrix& coo, bool backtracking);
/**
* @brief Generalized Sparse Matrix-Matrix Multiplication on COO.
* @param op The binary operator, could be `add`, `sub', `mul`, 'div',
* `copy_u`, `copy_e'.
* @param op The reduce operator, could be `sum`, `min`, `max'.
* @param coo The COO we apply SpMM on.
* @param ufeat The source node feature.
* @param efeat The edge feature.
* @param out The output feature on destination nodes.
* @param out_aux A list of NDArray's that contains auxiliary information such
* as the argmax on source nodes and edges for reduce operators such as
* `min` and `max`.
*/
void COOSpMM(
const std::string& op, const std::string& reduce, const COOMatrix& coo,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
/** @brief COOSpMM C interface without std::string. */
void COOSpMM(
const char* op, const char* reduce, const COOMatrix& coo, NDArray ufeat,
NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
/**
* @brief Generalized Sampled Dense-Dense Matrix Multiplication on COO.
* @param op The binary operator, could be `add`, `sub', `mul`, 'div',
* `dot`, `copy_u`, `copy_e'.
* @param coo The COO we apply SpMM on.
* @param ufeat The source node feature.
* @param vfeat The destination node feature.
* @param out The output feature on edge.
* @param lhs_target Type of `ufeat` (0: source, 1: edge, 2: destination).
* @param rhs_target Type of `ufeat` (0: source, 1: edge, 2: destination).
*/
void COOSDDMM(
const std::string& op, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
NDArray out, int lhs_target, int rhs_target);
/** @brief COOSDDMM C interface without std::string. */
void COOSDDMM(
const char* op, const COOMatrix& coo, NDArray ufeat, NDArray efeat,
NDArray out, int lhs_target, int rhs_target);
} // namespace aten
} // namespace dgl
namespace dmlc {
DMLC_DECLARE_TRAITS(has_saveload, dgl::aten::COOMatrix, true);
} // namespace dmlc
#endif // DGL_ATEN_COO_H_
include/dgl/aten/csr.h
+1082
View File
File diff suppressed because it is too large Load Diff
include/dgl/aten/macro.h
+434
View File
@@ -0,0 +1,434 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/aten/macro.h
* @brief Common macros for aten package.
*/
#ifndef DGL_ATEN_MACRO_H_
#define DGL_ATEN_MACRO_H_
///////////////////////// Dispatchers //////////////////////////
/**
* Dispatch according to device:
*
* ATEN_XPU_SWITCH(array->ctx.device_type, XPU, {
* // Now XPU is a placeholder for array->ctx.device_type
* DeviceSpecificImplementation<XPU>(...);
* });
*/
#define ATEN_XPU_SWITCH(val, XPU, op, ...) \
do { \
if ((val) == kDGLCPU) { \
constexpr auto XPU = kDGLCPU; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << "Operator " << (op) << " does not support " \
<< dgl::runtime::DeviceTypeCode2Str(val) << " device."; \
} \
} while (0)
/**
* Dispatch according to device:
*
* XXX(minjie): temporary macro that allows CUDA operator
*
* ATEN_XPU_SWITCH(array->ctx.device_type, XPU, {
* // Now XPU is a placeholder for array->ctx.device_type
* DeviceSpecificImplementation<XPU>(...);
* });
*
* We treat pinned memory as normal host memory if we don't want
* to enable CUDA UVA access for this operator
*/
#ifdef DGL_USE_CUDA
#define ATEN_XPU_SWITCH_CUDA(val, XPU, op, ...) \
do { \
if ((val) == kDGLCPU) { \
constexpr auto XPU = kDGLCPU; \
{ __VA_ARGS__ } \
} else if ((val) == kDGLCUDA) { \
constexpr auto XPU = kDGLCUDA; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << "Operator " << (op) << " does not support " \
<< dgl::runtime::DeviceTypeCode2Str(val) << " device."; \
} \
} while (0)
#else // DGL_USE_CUDA
#define ATEN_XPU_SWITCH_CUDA ATEN_XPU_SWITCH
#endif // DGL_USE_CUDA
/**
* Dispatch according to integral type (either int32 or int64):
*
* ATEN_ID_TYPE_SWITCH(array->dtype, IdType, {
* // Now IdType is the type corresponding to data type in array.
* // For instance, one can do this for a CPU array:
* DType *data = static_cast<DType *>(array->data);
* });
*/
#define ATEN_ID_TYPE_SWITCH(val, IdType, ...) \
do { \
CHECK_EQ((val).code, kDGLInt) << "ID must be integer type"; \
if ((val).bits == 32) { \
typedef int32_t IdType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 64) { \
typedef int64_t IdType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << "ID can only be int32 or int64"; \
} \
} while (0)
/**
* Dispatch according to bits (either int32 or int64):
*
* ATEN_ID_BITS_SWITCH(bits, IdType, {
* // Now IdType is the type corresponding to data type in array.
* // For instance, one can do this for a CPU array:
* DType *data = static_cast<DType *>(array->data);
* });
*/
#define ATEN_ID_BITS_SWITCH(bits, IdType, ...) \
do { \
CHECK((bits) == 32 || (bits) == 64) << "bits must be 32 or 64"; \
if ((bits) == 32) { \
typedef int32_t IdType; \
{ __VA_ARGS__ } \
} else if ((bits) == 64) { \
typedef int64_t IdType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << "ID can only be int32 or int64"; \
} \
} while (0)
/**
* Dispatch according to float type (either float32 or float64):
*
* ATEN_FLOAT_TYPE_SWITCH(array->dtype, FloatType, {
* // Now FloatType is the type corresponding to data type in array.
* // For instance, one can do this for a CPU array:
* FloatType *data = static_cast<FloatType *>(array->data);
* });
*/
#define ATEN_FLOAT_TYPE_SWITCH(val, FloatType, val_name, ...) \
do { \
CHECK_EQ((val).code, kDGLFloat) << (val_name) << " must be float type"; \
if ((val).bits == 32) { \
typedef float FloatType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 64) { \
typedef double FloatType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << (val_name) << " can only be float32 or float64"; \
} \
} while (0)
/**
* Dispatch according to float type, including 16bits
* (float16/bfloat16/float32/float64).
*/
#ifdef DGL_USE_CUDA
#if BF16_ENABLED
#define ATEN_FLOAT_TYPE_SWITCH_16BITS(val, FloatType, XPU, val_name, ...) \
do { \
CHECK((val).code == kDGLFloat || (val.code == kDGLBfloat)) \
<< (val_name) << " must be float type"; \
if ((val).bits == 32) { \
typedef float FloatType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 64) { \
typedef double FloatType; \
{ __VA_ARGS__ } \
} else if ( \
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLFloat) { \
typedef __half FloatType; \
{ __VA_ARGS__ } \
} else if ( \
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLBfloat) { \
typedef __nv_bfloat16 FloatType; \
{ __VA_ARGS__ } \
} else if ( \
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLFloat) { \
LOG(FATAL) << (val_name) << " can't be float16 on CPU"; \
} else if ( \
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLBfloat) { \
typedef BFloat16 FloatType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << (val_name) \
<< " can only be float16/bfloat16/float32/float64 on GPU"; \
} \
} while (0)
#else // BF16_ENABLED
#define ATEN_FLOAT_TYPE_SWITCH_16BITS(val, FloatType, XPU, val_name, ...) \
do { \
CHECK((val).code == kDGLFloat || (val.code == kDGLBfloat)) \
<< (val_name) << " must be float type"; \
if ((val).bits == 32) { \
typedef float FloatType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 64) { \
typedef double FloatType; \
{ __VA_ARGS__ } \
} else if ( \
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLFloat) { \
typedef __half FloatType; \
{ __VA_ARGS__ } \
} else if ( \
XPU == kDGLCUDA && (val).bits == 16 && (val).code == kDGLBfloat) { \
LOG(FATAL) << "bfloat16 requires CUDA >= 11.0"; \
} else if ( \
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLFloat) { \
LOG(FATAL) << (val_name) << " can't be float16 on CPU"; \
} else if ( \
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLBfloat) { \
typedef BFloat16 FloatType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << (val_name) \
<< " can only be float16/float32/float64 on GPU"; \
} \
} while (0)
#endif // BF16_ENABLED
#else // DGL_USE_CUDA
#define ATEN_FLOAT_TYPE_SWITCH_16BITS(val, FloatType, XPU, val_name, ...) \
do { \
CHECK((val).code == kDGLFloat || (val.code == kDGLBfloat)) \
<< (val_name) << " must be float type"; \
if ((val).bits == 32) { \
typedef float FloatType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 64) { \
typedef double FloatType; \
{ __VA_ARGS__ } \
} else if ( \
XPU == kDGLCPU && (val).bits == 16 && (val).code == kDGLBfloat) { \
typedef BFloat16 FloatType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << (val_name) \
<< " can only be bfloat16/float32/float64 on CPU"; \
} \
} while (0)
#endif // DGL_USE_CUDA
/**
* Dispatch according to data type (int32, int64, float32 or float64):
*
* ATEN_DTYPE_SWITCH(array->dtype, DType, {
* // Now DType is the type corresponding to data type in array.
* // For instance, one can do this for a CPU array:
* DType *data = static_cast<DType *>(array->data);
* });
*/
#define ATEN_DTYPE_SWITCH(val, DType, val_name, ...) \
do { \
if ((val).code == kDGLInt && (val).bits == 32) { \
typedef int32_t DType; \
{ __VA_ARGS__ } \
} else if ((val).code == kDGLInt && (val).bits == 64) { \
typedef int64_t DType; \
{ __VA_ARGS__ } \
} else if ((val).code == kDGLFloat && (val).bits == 32) { \
typedef float DType; \
{ __VA_ARGS__ } \
} else if ((val).code == kDGLFloat && (val).bits == 64) { \
typedef double DType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << (val_name) \
<< " can only be int32, int64, float32 or float64"; \
} \
} while (0)
/**
* Dispatch according to data type (int8, uint8, float32 or float64):
*
* ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(array->dtype, DType, {
* // Now DType is the type corresponding to data type in array.
* // For instance, one can do this for a CPU array:
* DType *data = static_cast<DType *>(array->data);
* });
*/
#define ATEN_FLOAT_INT8_UINT8_TYPE_SWITCH(val, DType, val_name, ...) \
do { \
if ((val).code == kDGLInt && (val).bits == 8) { \
typedef int8_t DType; \
{ __VA_ARGS__ } \
} else if ((val).code == kDGLUInt && (val).bits == 8) { \
typedef uint8_t DType; \
{ __VA_ARGS__ } \
} else if ((val).code == kDGLFloat && (val).bits == 32) { \
typedef float DType; \
{ __VA_ARGS__ } \
} else if ((val).code == kDGLFloat && (val).bits == 64) { \
typedef double DType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << (val_name) \
<< " can only be int8, uint8, float32 or float64"; \
} \
} while (0)
/**
* Dispatch data type only based on bit-width (8-bit, 16-bit, 32-bit, 64-bit):
*
* ATEN_DTYPE_BITS_ONLY_SWITCH(array->dtype, DType, {
* // Now DType is the type which has the same bit-width with the
* // data type in array.
* // Do not use for computation, but only for read and write.
* // For instance, one can do this for a CPU array:
* DType *data = static_cast<DType *>(array->data);
* });
*/
#define ATEN_DTYPE_BITS_ONLY_SWITCH(val, DType, val_name, ...) \
do { \
if ((val).bits == 8) { \
typedef int8_t DType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 16) { \
typedef int16_t DType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 32) { \
typedef int32_t DType; \
{ __VA_ARGS__ } \
} else if ((val).bits == 64) { \
typedef int64_t DType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << (val_name) \
<< " can only be 8-bit, 16-bit, 32-bit, or 64-bit"; \
} \
} while (0)
/**
* Dispatch according to integral type of CSR graphs.
* Identical to ATEN_ID_TYPE_SWITCH except for a different error message.
*/
#define ATEN_CSR_DTYPE_SWITCH(val, DType, ...) \
do { \
if ((val).code == kDGLInt && (val).bits == 32) { \
typedef int32_t DType; \
{ __VA_ARGS__ } \
} else if ((val).code == kDGLInt && (val).bits == 64) { \
typedef int64_t DType; \
{ __VA_ARGS__ } \
} else { \
LOG(FATAL) << "CSR matrix data can only be int32 or int64"; \
} \
} while (0)
// Macro to dispatch according to device context and index type.
#define ATEN_CSR_SWITCH(csr, XPU, IdType, op, ...) \
ATEN_XPU_SWITCH((csr).indptr->ctx.device_type, XPU, op, { \
ATEN_ID_TYPE_SWITCH((csr).indptr->dtype, IdType, {{__VA_ARGS__}}); \
});
// Macro to dispatch according to device context and index type.
#define ATEN_COO_SWITCH(coo, XPU, IdType, op, ...) \
ATEN_XPU_SWITCH((coo).row->ctx.device_type, XPU, op, { \
ATEN_ID_TYPE_SWITCH((coo).row->dtype, IdType, {{__VA_ARGS__}}); \
});
#define CHECK_VALID_CONTEXT(VAR1, VAR2) \
CHECK( \
((VAR1)->ctx == (VAR2)->ctx) || (VAR1).IsPinned() || \
((VAR1).NumElements() == 0)) /* Let empty arrays pass */ \
<< "Expected " << (#VAR2) << "(" << (VAR2)->ctx << ")" \
<< " to have the same device " \
<< "context as " << (#VAR1) << "(" << (VAR1)->ctx << "). " \
<< "Or " << (#VAR1) << "(" << (VAR1)->ctx << ")" \
<< " is pinned";
/**
* Macro to dispatch according to the context of array and dtype of csr
* to enable CUDA UVA ops.
* Context check is covered here to avoid confusion with CHECK_SAME_CONTEXT.
* If csr has the same context with array, same behivor as ATEN_CSR_SWITCH_CUDA.
* If csr is pinned, array's context will conduct the actual operation.
*/
#define ATEN_CSR_SWITCH_CUDA_UVA(csr, array, XPU, IdType, op, ...) \
do { \
CHECK_VALID_CONTEXT(csr.indices, array); \
ATEN_XPU_SWITCH_CUDA(array->ctx.device_type, XPU, op, { \
ATEN_ID_TYPE_SWITCH((csr).indptr->dtype, IdType, {{__VA_ARGS__}}); \
}); \
} while (0)
// Macro to dispatch according to device context (allowing cuda)
#ifdef DGL_USE_CUDA
#define ATEN_CSR_SWITCH_CUDA(csr, XPU, IdType, op, ...) \
ATEN_XPU_SWITCH_CUDA((csr).indptr->ctx.device_type, XPU, op, { \
ATEN_ID_TYPE_SWITCH((csr).indptr->dtype, IdType, {{__VA_ARGS__}}); \
});
// Macro to dispatch according to device context and index type.
#define ATEN_COO_SWITCH_CUDA(coo, XPU, IdType, op, ...) \
ATEN_XPU_SWITCH_CUDA((coo).row->ctx.device_type, XPU, op, { \
ATEN_ID_TYPE_SWITCH((coo).row->dtype, IdType, {{__VA_ARGS__}}); \
});
#else // DGL_USE_CUDA
#define ATEN_CSR_SWITCH_CUDA ATEN_CSR_SWITCH
#define ATEN_COO_SWITCH_CUDA ATEN_COO_SWITCH
#endif // DGL_USE_CUDA
///////////////////////// Array checks //////////////////////////
#define IS_INT32(a) ((a)->dtype.code == kDGLInt && (a)->dtype.bits == 32)
#define IS_INT64(a) ((a)->dtype.code == kDGLInt && (a)->dtype.bits == 64)
#define IS_FLOAT32(a) ((a)->dtype.code == kDGLFloat && (a)->dtype.bits == 32)
#define IS_FLOAT64(a) ((a)->dtype.code == kDGLFloat && (a)->dtype.bits == 64)
#define CHECK_IF(cond, prop, value_name, dtype_name) \
CHECK(cond) << "Expecting " << (prop) << " of " << (value_name) << " to be " \
<< (dtype_name)
#define CHECK_INT32(value, value_name) \
CHECK_IF(IS_INT32(value), "dtype", value_name, "int32")
#define CHECK_INT64(value, value_name) \
CHECK_IF(IS_INT64(value), "dtype", value_name, "int64")
#define CHECK_INT(value, value_name) \
CHECK_IF( \
IS_INT32(value) || IS_INT64(value), "dtype", value_name, \
"int32 or int64")
#define CHECK_FLOAT32(value, value_name) \
CHECK_IF(IS_FLOAT32(value), "dtype", value_name, "float32")
#define CHECK_FLOAT64(value, value_name) \
CHECK_IF(IS_FLOAT64(value), "dtype", value_name, "float64")
#define CHECK_FLOAT(value, value_name) \
CHECK_IF( \
IS_FLOAT32(value) || IS_FLOAT64(value), "dtype", value_name, \
"float32 or float64")
#define CHECK_NDIM(value, _ndim, value_name) \
CHECK_IF((value)->ndim == (_ndim), "ndim", value_name, _ndim)
#define CHECK_SAME_DTYPE(VAR1, VAR2) \
CHECK((VAR1)->dtype == (VAR2)->dtype) \
<< "Expected " << (#VAR2) << " to be the same type as " << (#VAR1) \
<< "(" << (VAR1)->dtype << ")" \
<< ". But got " << (VAR2)->dtype << ".";
#define CHECK_SAME_CONTEXT(VAR1, VAR2) \
CHECK((VAR1)->ctx == (VAR2)->ctx) \
<< "Expected " << (#VAR2) << " to have the same device context as " \
<< (#VAR1) << "(" << (VAR1)->ctx << ")" \
<< ". But got " << (VAR2)->ctx << ".";
#define CHECK_NO_OVERFLOW(dtype, val) \
do { \
if (sizeof(val) == 8 && (dtype).bits == 32) \
CHECK_LE((val), 0x7FFFFFFFL) \
<< "int32 overflow for argument " << (#val) << "."; \
} while (0);
#define CHECK_IS_ID_ARRAY(VAR) \
CHECK((VAR)->ndim == 1 && (IS_INT32(VAR) || IS_INT64(VAR))) \
<< "Expected argument " << (#VAR) << " to be an 1D integer array.";
#endif // DGL_ATEN_MACRO_H_
include/dgl/aten/spmat.h
+137
View File
@@ -0,0 +1,137 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/aten/spmat.h
* @brief Sparse matrix definitions
*/
#ifndef DGL_ATEN_SPMAT_H_
#define DGL_ATEN_SPMAT_H_
#include <string>
#include <vector>
#include "../runtime/object.h"
#include "./types.h"
namespace dgl {
/**
* @brief Sparse format.
*/
enum class SparseFormat {
kCOO = 1,
kCSR = 2,
kCSC = 3,
};
/**
* @brief Sparse format codes
*/
const dgl_format_code_t ALL_CODE = 0x7;
const dgl_format_code_t ANY_CODE = 0x0;
const dgl_format_code_t COO_CODE = 0x1;
const dgl_format_code_t CSR_CODE = 0x2;
const dgl_format_code_t CSC_CODE = 0x4;
// Parse sparse format from string.
inline SparseFormat ParseSparseFormat(const std::string& name) {
if (name == "coo")
return SparseFormat::kCOO;
else if (name == "csr")
return SparseFormat::kCSR;
else if (name == "csc")
return SparseFormat::kCSC;
else
LOG(FATAL) << "Sparse format not recognized";
return SparseFormat::kCOO;
}
// Create string from sparse format.
inline std::string ToStringSparseFormat(SparseFormat sparse_format) {
if (sparse_format == SparseFormat::kCOO)
return std::string("coo");
else if (sparse_format == SparseFormat::kCSR)
return std::string("csr");
else
return std::string("csc");
}
inline std::vector<SparseFormat> CodeToSparseFormats(dgl_format_code_t code) {
std::vector<SparseFormat> ret;
if (code & COO_CODE) ret.push_back(SparseFormat::kCOO);
if (code & CSR_CODE) ret.push_back(SparseFormat::kCSR);
if (code & CSC_CODE) ret.push_back(SparseFormat::kCSC);
return ret;
}
inline dgl_format_code_t SparseFormatsToCode(
const std::vector<SparseFormat>& formats) {
dgl_format_code_t ret = 0;
for (auto format : formats) {
switch (format) {
case SparseFormat::kCOO:
ret |= COO_CODE;
break;
case SparseFormat::kCSR:
ret |= CSR_CODE;
break;
case SparseFormat::kCSC:
ret |= CSC_CODE;
break;
default:
LOG(FATAL) << "Only support COO/CSR/CSC formats.";
}
}
return ret;
}
inline std::string CodeToStr(dgl_format_code_t code) {
std::string ret = "";
if (code & COO_CODE) ret += "coo ";
if (code & CSR_CODE) ret += "csr ";
if (code & CSC_CODE) ret += "csc ";
return ret;
}
inline SparseFormat DecodeFormat(dgl_format_code_t code) {
if (code & COO_CODE) return SparseFormat::kCOO;
if (code & CSC_CODE) return SparseFormat::kCSC;
return SparseFormat::kCSR;
}
// Sparse matrix object that is exposed to python API.
struct SparseMatrix : public runtime::Object {
// Sparse format.
int32_t format = 0;
// Shape of this matrix.
int64_t num_rows = 0, num_cols = 0;
// Index arrays. For CSR, it is {indptr, indices, data}. For COO, it is {row,
// col, data}.
std::vector<IdArray> indices;
// Boolean flags.
// TODO(minjie): We might revisit this later to provide a more general
// solution. Currently, we only consider aten::COOMatrix and aten::CSRMatrix.
std::vector<bool> flags;
SparseMatrix() {}
SparseMatrix(
int32_t fmt, int64_t nrows, int64_t ncols,
const std::vector<IdArray>& idx, const std::vector<bool>& flg)
: format(fmt),
num_rows(nrows),
num_cols(ncols),
indices(idx),
flags(flg) {}
static constexpr const char* _type_key = "aten.SparseMatrix";
DGL_DECLARE_OBJECT_TYPE_INFO(SparseMatrix, runtime::Object);
};
// Define SparseMatrixRef
DGL_DEFINE_OBJECT_REF(SparseMatrixRef, SparseMatrix);
} // namespace dgl
#endif // DGL_ATEN_SPMAT_H_
include/dgl/aten/types.h
+43
View File
@@ -0,0 +1,43 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/aten/types.h
* @brief Array and ID types
*/
#ifndef DGL_ATEN_TYPES_H_
#define DGL_ATEN_TYPES_H_
#include <cstdint>
#include "../runtime/ndarray.h"
namespace dgl {
typedef uint64_t dgl_id_t;
typedef uint64_t dgl_type_t;
/** @brief Type for dgl fomrat code, whose binary representation indices
* which sparse format is in use and which is not.
*
* Suppose the binary representation is xyz, then
* - x indicates whether csc is in use (1 for true and 0 for false).
* - y indicates whether csr is in use.
* - z indicates whether coo is in use.
*/
typedef uint8_t dgl_format_code_t;
using dgl::runtime::NDArray;
typedef NDArray IdArray;
typedef NDArray DegreeArray;
typedef NDArray BoolArray;
typedef NDArray IntArray;
typedef NDArray FloatArray;
typedef NDArray TypeArray;
namespace aten {
static const DGLContext CPU{kDGLCPU, 0};
} // namespace aten
} // namespace dgl
#endif // DGL_ATEN_TYPES_H_
include/dgl/base_heterograph.h
+918
View File
@@ -0,0 +1,918 @@
/**
* Copyright (c) 2019 by Contributors
* @file dgl/heterograph_interface.h
* @brief DGL heterogeneous graph index class.
*/
#ifndef DGL_BASE_HETEROGRAPH_H_
#define DGL_BASE_HETEROGRAPH_H_
#include <algorithm>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "./runtime/object.h"
#include "array.h"
#include "aten/spmat.h"
#include "aten/types.h"
#include "graph_interface.h"
namespace dgl {
// Forward declaration
class BaseHeteroGraph;
class HeteroPickleStates;
typedef std::shared_ptr<BaseHeteroGraph> HeteroGraphPtr;
struct FlattenedHeteroGraph;
typedef std::shared_ptr<FlattenedHeteroGraph> FlattenedHeteroGraphPtr;
struct HeteroSubgraph;
/** @brief Enum class for edge direction */
enum class EdgeDir {
kIn, // in edge direction
kOut // out edge direction
};
/**
* @brief Base heterogenous graph.
*
* In heterograph, nodes represent entities and edges represent relations.
* Nodes and edges are associated with types. The same pair of entity types
* can have multiple relation types between them, but relation type **uniquely**
* identifies the source and destination entity types.
*
* In a high-level, a heterograph is a data structure composed of:
* - A meta-graph that stores the entity-entity relation graph.
* - A dictionary of relation type to the bipartite graph representing the
* actual connections among entity nodes.
*/
class BaseHeteroGraph : public runtime::Object {
public:
explicit BaseHeteroGraph(GraphPtr meta_graph) : meta_graph_(meta_graph) {}
virtual ~BaseHeteroGraph() = default;
////////////////////// query/operations on meta graph ///////////////////////
/** @return the number of vertex types */
virtual uint64_t NumVertexTypes() const { return meta_graph_->NumVertices(); }
/** @return the number of edge types */
virtual uint64_t NumEdgeTypes() const { return meta_graph_->NumEdges(); }
/** @return given the edge type, find the source type */
virtual std::pair<dgl_type_t, dgl_type_t> GetEndpointTypes(
dgl_type_t etype) const {
return meta_graph_->FindEdge(etype);
}
/** @return the meta graph */
virtual GraphPtr meta_graph() const { return meta_graph_; }
/**
* @brief Return the bipartite graph of the given edge type.
* @param etype The edge type.
* @return The bipartite graph.
*/
virtual HeteroGraphPtr GetRelationGraph(dgl_type_t etype) const = 0;
///////////////////// query/operations on realized graph /////////////////////
/** @brief Add vertices to the given vertex type */
virtual void AddVertices(dgl_type_t vtype, uint64_t num_vertices) = 0;
/** @brief Add one edge to the given edge type */
virtual void AddEdge(dgl_type_t etype, dgl_id_t src, dgl_id_t dst) = 0;
/** @brief Add edges to the given edge type */
virtual void AddEdges(dgl_type_t etype, IdArray src_ids, IdArray dst_ids) = 0;
/**
* @brief Clear the graph. Remove all vertices/edges.
*/
virtual void Clear() = 0;
/**
* @brief Get the data type of node and edge IDs of this graph.
*/
virtual DGLDataType DataType() const = 0;
/**
* @brief Get the device context of this graph.
*/
virtual DGLContext Context() const = 0;
/**
* @brief Pin graph.
*/
virtual void PinMemory_() = 0;
/**
* @brief Check if this graph is pinned.
*/
virtual bool IsPinned() const = 0;
/**
* @brief Record stream for this graph.
* @param stream The stream that is using the graph
*/
virtual void RecordStream(DGLStreamHandle stream) = 0;
/**
* @brief Get the number of integer bits used to store node/edge ids (32 or
* 64).
*/
// TODO(BarclayII) replace NumBits() calls to DataType() calls
virtual uint8_t NumBits() const = 0;
/**
* @return whether the graph is a multigraph
*/
virtual bool IsMultigraph() const = 0;
/** @return whether the graph is read-only */
virtual bool IsReadonly() const = 0;
/** @return the number of vertices in the graph.*/
virtual uint64_t NumVertices(dgl_type_t vtype) const = 0;
/** @return the number of vertices for each type in the graph as a vector */
inline virtual std::vector<int64_t> NumVerticesPerType() const {
LOG(FATAL) << "[BUG] NumVerticesPerType() not supported on this object.";
return {};
}
/** @return the number of edges in the graph.*/
virtual uint64_t NumEdges(dgl_type_t etype) const = 0;
/** @return true if the given vertex is in the graph.*/
virtual bool HasVertex(dgl_type_t vtype, dgl_id_t vid) const = 0;
/** @return a 0-1 array indicating whether the given vertices are in the
* graph.
*/
virtual BoolArray HasVertices(dgl_type_t vtype, IdArray vids) const = 0;
/** @return true if the given edge is in the graph.*/
virtual bool HasEdgeBetween(
dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const = 0;
/** @return a 0-1 array indicating whether the given edges are in the graph.*/
virtual BoolArray HasEdgesBetween(
dgl_type_t etype, IdArray src_ids, IdArray dst_ids) const = 0;
/**
* @brief Find the predecessors of a vertex.
* @note The given vertex should belong to the source vertex type
* of the given edge type.
* @param etype The edge type
* @param vid The vertex id.
* @return the predecessor id array.
*/
virtual IdArray Predecessors(dgl_type_t etype, dgl_id_t dst) const = 0;
/**
* @brief Find the successors of a vertex.
* @note The given vertex should belong to the dest vertex type
* of the given edge type.
* @param etype The edge type
* @param vid The vertex id.
* @return the successor id array.
*/
virtual IdArray Successors(dgl_type_t etype, dgl_id_t src) const = 0;
/**
* @brief Get all edge ids between the two given endpoints
* @note The given src and dst vertices should belong to the source vertex
* type and the dest vertex type of the given edge type, respectively.
* @param etype The edge type
* @param src The source vertex.
* @param dst The destination vertex.
* @return the edge id array.
*/
virtual IdArray EdgeId(
dgl_type_t etype, dgl_id_t src, dgl_id_t dst) const = 0;
/**
* @brief Get all edge ids between the given endpoint pairs.
*
* @param etype The edge type
* @param src The src vertex ids.
* @param dst The dst vertex ids.
* @return EdgeArray containing all edges between all pairs.
*/
virtual EdgeArray EdgeIdsAll(
dgl_type_t etype, IdArray src, IdArray dst) const = 0;
/**
* @brief Get edge ids between the given endpoint pairs.
*
* Only find one matched edge Ids even if there are multiple matches due to
* parallel edges. The i^th Id in the returned array is for edge (src[i],
* dst[i]).
*
* @param etype The edge type
* @param src The src vertex ids.
* @param dst The dst vertex ids.
* @return EdgeArray containing all edges between all pairs.
*/
virtual IdArray EdgeIdsOne(
dgl_type_t etype, IdArray src, IdArray dst) const = 0;
/**
* @brief Find the edge ID and return the pair of endpoints
* @param etype The edge type
* @param eid The edge ID
* @return a pair whose first element is the source and the second the
* destination.
*/
virtual std::pair<dgl_id_t, dgl_id_t> FindEdge(
dgl_type_t etype, dgl_id_t eid) const = 0;
/**
* @brief Find the edge IDs and return their source and target node IDs.
* @param etype The edge type
* @param eids The edge ID array.
* @return EdgeArray containing all edges with id in eid. The order is
* preserved.
*/
virtual EdgeArray FindEdges(dgl_type_t etype, IdArray eids) const = 0;
/**
* @brief Get the in edges of the vertex.
* @note The given vertex should belong to the dest vertex type
* of the given edge type.
* @param etype The edge type
* @param vid The vertex id.
* @return the edges
*/
virtual EdgeArray InEdges(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Get the in edges of the vertices.
* @note The given vertex should belong to the dest vertex type
* of the given edge type.
* @param etype The edge type
* @param vids The vertex id array.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray InEdges(dgl_type_t etype, IdArray vids) const = 0;
/**
* @brief Get the out edges of the vertex.
* @note The given vertex should belong to the source vertex type
* of the given edge type.
* @param etype The edge type
* @param vid The vertex id.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray OutEdges(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Get the out edges of the vertices.
* @note The given vertex should belong to the source vertex type
* of the given edge type.
* @param etype The edge type
* @param vids The vertex id array.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray OutEdges(dgl_type_t etype, IdArray vids) const = 0;
/**
* @brief Get all the edges in the graph.
* @note If order is "srcdst", the returned edges list is sorted by their src
* and dst ids. If order is "eid", they are in their edge id order. Otherwise,
* in the arbitrary order.
* @param etype The edge type
* @param order The order of the returned edge list.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray Edges(
dgl_type_t etype, const std::string& order = "") const = 0;
/**
* @brief Get the in degree of the given vertex.
* @note The given vertex should belong to the dest vertex type of the given
* edge type.
* @param etype The edge type
* @param vid The vertex id.
* @return the in degree
*/
virtual uint64_t InDegree(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Get the in degrees of the given vertices.
* @note The given vertex should belong to the dest vertex type of the given
* edge type.
* @param etype The edge type
* @param vid The vertex id array.
* @return the in degree array
*/
virtual DegreeArray InDegrees(dgl_type_t etype, IdArray vids) const = 0;
/**
* @brief Get the out degree of the given vertex.
* @note The given vertex should belong to the source vertex type of the given
* edge type.
* @param etype The edge type
* @param vid The vertex id.
* @return the out degree
*/
virtual uint64_t OutDegree(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Get the out degrees of the given vertices.
* @note The given vertex should belong to the source vertex type of the given
* edge type.
* @param etype The edge type
* @param vid The vertex id array.
* @return the out degree array
*/
virtual DegreeArray OutDegrees(dgl_type_t etype, IdArray vids) const = 0;
/**
* @brief Return the successor vector
* @note The given vertex should belong to the source vertex type of the given
* edge type.
* @param vid The vertex id.
* @return the successor vector iterator pair.
*/
virtual DGLIdIters SuccVec(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Return the out edge id vector
* @note The given vertex should belong to the source vertex type of the given
* edge type.
* @param vid The vertex id.
* @return the out edge id vector iterator pair.
*/
virtual DGLIdIters OutEdgeVec(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Return the predecessor vector
* @note The given vertex should belong to the dest vertex type of the given
* edge type.
* @param vid The vertex id.
* @return the predecessor vector iterator pair.
*/
virtual DGLIdIters PredVec(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Return the in edge id vector
* @note The given vertex should belong to the dest vertex type of the given
* edge type.
* @param vid The vertex id.
* @return the in edge id vector iterator pair.
*/
virtual DGLIdIters InEdgeVec(dgl_type_t etype, dgl_id_t vid) const = 0;
/**
* @brief Get the adjacency matrix of the graph.
*
* TODO(minjie): deprecate this interface; replace it with GetXXXMatrix.
*
* By default, a row of returned adjacency matrix represents the destination
* of an edge and the column represents the source.
*
* If the fmt is 'csr', the function should return three arrays, representing
* indptr, indices and edge ids
*
* If the fmt is 'coo', the function should return one array of shape (2,
* nnz), representing a horitonzal stack of row and col indices.
*
* @param transpose A flag to transpose the returned adjacency matrix.
* @param fmt the format of the returned adjacency matrix.
* @return a vector of IdArrays.
*/
virtual std::vector<IdArray> GetAdj(
dgl_type_t etype, bool transpose, const std::string& fmt) const = 0;
/**
* @brief Determine which format to use with a preference.
*
* Otherwise, it will return whatever DGL thinks is the most appropriate given
* the arguments.
*
* @param etype Edge type.
* @param preferred_formats Preferred sparse formats.
* @return Available sparse format.
*/
virtual SparseFormat SelectFormat(
dgl_type_t etype, dgl_format_code_t preferred_formats) const = 0;
/**
* @brief Return sparse formats already created for the graph.
*
* @return a number of type dgl_format_code_t.
*/
virtual dgl_format_code_t GetCreatedFormats() const = 0;
/**
* @brief Return allowed sparse formats for the graph.
*
* @return a number of type dgl_format_code_t.
*/
virtual dgl_format_code_t GetAllowedFormats() const = 0;
/**
* @brief Return the graph in specified available formats.
*
* @return The new graph.
*/
virtual HeteroGraphPtr GetGraphInFormat(dgl_format_code_t formats) const = 0;
/**
* @brief Get adjacency matrix in COO format.
* @param etype Edge type.
* @return COO matrix.
*/
virtual aten::COOMatrix GetCOOMatrix(dgl_type_t etype) const = 0;
/**
* @brief Get adjacency matrix in CSR format.
*
* The row and column sizes are equal to the number of dsttype and srctype
* nodes, respectively.
*
* @param etype Edge type.
* @return CSR matrix.
*/
virtual aten::CSRMatrix GetCSRMatrix(dgl_type_t etype) const = 0;
/**
* @brief Get adjacency matrix in CSC format.
*
* A CSC matrix is equivalent to the transpose of a CSR matrix.
* We reuse the CSRMatrix data structure as return value. The row and column
* sizes are equal to the number of dsttype and srctype nodes, respectively.
*
* @param etype Edge type.
* @return A CSR matrix.
*/
virtual aten::CSRMatrix GetCSCMatrix(dgl_type_t etype) const = 0;
/**
* @brief Extract the induced subgraph by the given vertices.
*
* The length of the given vector should be equal to the number of vertex
* types. Empty arrays can be provided if no vertex is needed for the type.
* The result subgraph has the same meta graph with the parent, but some types
* can have no node/edge.
*
* @param vids the induced vertices per type.
* @return the subgraph.
*/
virtual HeteroSubgraph VertexSubgraph(
const std::vector<IdArray>& vids) const = 0;
/**
* @brief Extract the induced subgraph by the given edges.
*
* The length of the given vector should be equal to the number of edge types.
* Empty arrays can be provided if no edge is needed for the type. The result
* subgraph has the same meta graph with the parent, but some types can have
* no node/edge.
*
* @param eids The edges in the subgraph.
* @param preserve_nodes If true, the vertices will not be relabeled, so some
* vertices may have no incident edges.
* @return the subgraph.
*/
virtual HeteroSubgraph EdgeSubgraph(
const std::vector<IdArray>& eids, bool preserve_nodes = false) const = 0;
/**
* @brief Convert the list of requested unitgraph graphs into a single
* unitgraph graph.
*
* @param etypes The list of edge type IDs.
* @return The flattened graph, with induced source/edge/destination
* types/IDs.
*/
virtual FlattenedHeteroGraphPtr Flatten(
const std::vector<dgl_type_t>& etypes) const {
LOG(FATAL) << "Flatten operation unsupported";
return nullptr;
}
/** @brief Cast this graph to immutable graph */
virtual GraphPtr AsImmutableGraph() const {
LOG(FATAL) << "AsImmutableGraph not supported.";
return nullptr;
}
static constexpr const char* _type_key = "graph.HeteroGraph";
DGL_DECLARE_OBJECT_TYPE_INFO(BaseHeteroGraph, runtime::Object);
protected:
/** @brief meta graph */
GraphPtr meta_graph_;
// empty constructor
BaseHeteroGraph() {}
};
// Define HeteroGraphRef
DGL_DEFINE_OBJECT_REF(HeteroGraphRef, BaseHeteroGraph);
/**
* @brief Hetero-subgraph data structure.
*
* This class can be used as arguments and return values of a C API.
*
* <code>
* DGL_REGISTER_GLOBAL("some_c_api")
* .set_body([] (DGLArgs args, DGLRetValue* rv) {
* HeteroSubgraphRef subg = args[0];
* std::shared_ptr<HeteroSubgraph> ret = do_something( ... );
* *rv = HeteroSubgraphRef(ret);
* });
* </code>
*/
struct HeteroSubgraph : public runtime::Object {
/** @brief The heterograph. */
HeteroGraphPtr graph;
/**
* @brief The induced vertex ids of each entity type.
* The vector length is equal to the number of vertex types in the parent
* graph. Each array i has the same length as the number of vertices in type
* i. Empty array is allowed if the mapping is identity.
*/
std::vector<IdArray> induced_vertices;
/**
* @brief The induced edge ids of each relation type.
* The vector length is equal to the number of edge types in the parent graph.
* Each array i has the same length as the number of edges in type i.
* Empty array is allowed if the mapping is identity.
*/
std::vector<IdArray> induced_edges;
static constexpr const char* _type_key = "graph.HeteroSubgraph";
DGL_DECLARE_OBJECT_TYPE_INFO(HeteroSubgraph, runtime::Object);
};
// Define HeteroSubgraphRef
DGL_DEFINE_OBJECT_REF(HeteroSubgraphRef, HeteroSubgraph);
/** @brief The flattened heterograph */
struct FlattenedHeteroGraph : public runtime::Object {
/** @brief The graph */
HeteroGraphRef graph;
/**
* @brief Mapping from source node ID to node type in parent graph
* @note The induced type array guarantees that the same type always appear
* contiguously.
*/
IdArray induced_srctype;
/**
* @brief The set of node types in parent graph appearing in source nodes.
*/
IdArray induced_srctype_set;
/** @brief Mapping from source node ID to local node ID in parent graph */
IdArray induced_srcid;
/**
* @brief Mapping from edge ID to edge type in parent graph
* @note The induced type array guarantees that the same type always appear
* contiguously.
*/
IdArray induced_etype;
/**
* @brief The set of edge types in parent graph appearing in edges.
*/
IdArray induced_etype_set;
/** @brief Mapping from edge ID to local edge ID in parent graph */
IdArray induced_eid;
/**
* @brief Mapping from destination node ID to node type in parent graph
* @note The induced type array guarantees that the same type always appear
* contiguously.
*/
IdArray induced_dsttype;
/**
* @brief The set of node types in parent graph appearing in destination
* nodes.
*/
IdArray induced_dsttype_set;
/** @brief Mapping from destination node ID to local node ID in parent graph
*/
IdArray induced_dstid;
void VisitAttrs(runtime::AttrVisitor* v) final {
v->Visit("graph", &graph);
v->Visit("induced_srctype", &induced_srctype);
v->Visit("induced_srctype_set", &induced_srctype_set);
v->Visit("induced_srcid", &induced_srcid);
v->Visit("induced_etype", &induced_etype);
v->Visit("induced_etype_set", &induced_etype_set);
v->Visit("induced_eid", &induced_eid);
v->Visit("induced_dsttype", &induced_dsttype);
v->Visit("induced_dsttype_set", &induced_dsttype_set);
v->Visit("induced_dstid", &induced_dstid);
}
static constexpr const char* _type_key = "graph.FlattenedHeteroGraph";
DGL_DECLARE_OBJECT_TYPE_INFO(FlattenedHeteroGraph, runtime::Object);
};
DGL_DEFINE_OBJECT_REF(FlattenedHeteroGraphRef, FlattenedHeteroGraph);
// Declarations of functions and algorithms
/**
* @brief Create a heterograph from meta graph and a list of bipartite graph,
* additionally specifying number of nodes per type.
*/
HeteroGraphPtr CreateHeteroGraph(
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& rel_graphs,
const std::vector<int64_t>& num_nodes_per_type = {});
/**
* @brief Create a heterograph from COO input.
* @param num_vtypes Number of vertex types. Must be 1 or 2.
* @param num_src Number of nodes in the source type.
* @param num_dst Number of nodes in the destination type.
* @param row Src node ids of the edges.
* @param col Dst node ids of the edges.
* @param row_sorted Whether the `row` array is in sorted ascending order.
* @param col_sorted When `row_sorted` is true, whether the columns within each
* row are also sorted. When `row_sorted` is false, this flag must also be
* false.
* @param formats Sparse formats used for storing this graph.
* @return A heterograph pointer.
*/
HeteroGraphPtr CreateFromCOO(
int64_t num_vtypes, int64_t num_src, int64_t num_dst, IdArray row,
IdArray col, bool row_sorted = false, bool col_sorted = false,
dgl_format_code_t formats = ALL_CODE);
/**
* @brief Create a heterograph from COO input.
* @param num_vtypes Number of vertex types. Must be 1 or 2.
* @param mat The COO matrix
* @param formats Sparse formats used for storing this graph.
* @return A heterograph pointer.
*/
HeteroGraphPtr CreateFromCOO(
int64_t num_vtypes, const aten::COOMatrix& mat,
dgl_format_code_t formats = ALL_CODE);
/**
* @brief Create a heterograph from CSR input.
* @param num_vtypes Number of vertex types. Must be 1 or 2.
* @param num_src Number of nodes in the source type.
* @param num_dst Number of nodes in the destination type.
* @param indptr Indptr array
* @param indices Indices array
* @param edge_ids Edge ids
* @param formats Sparse formats for storing this graph.
* @return A heterograph pointer.
*/
HeteroGraphPtr CreateFromCSR(
int64_t num_vtypes, int64_t num_src, int64_t num_dst, IdArray indptr,
IdArray indices, IdArray edge_ids, dgl_format_code_t formats = ALL_CODE);
/**
* @brief Create a heterograph from CSR input.
* @param num_vtypes Number of vertex types. Must be 1 or 2.
* @param mat The CSR matrix
* @param formats Sparse formats for storing this graph.
* @return A heterograph pointer.
*/
HeteroGraphPtr CreateFromCSR(
int64_t num_vtypes, const aten::CSRMatrix& mat,
dgl_format_code_t formats = ALL_CODE);
/**
* @brief Create a heterograph from CSC input.
* @param num_vtypes Number of vertex types. Must be 1 or 2.
* @param num_src Number of nodes in the source type.
* @param num_dst Number of nodes in the destination type.
* @param indptr Indptr array
* @param indices Indices array
* @param edge_ids Edge ids
* @param formats Sparse formats used for storing this graph.
* @return A heterograph pointer.
*/
HeteroGraphPtr CreateFromCSC(
int64_t num_vtypes, int64_t num_src, int64_t num_dst, IdArray indptr,
IdArray indices, IdArray edge_ids, dgl_format_code_t formats = ALL_CODE);
/**
* @brief Create a heterograph from CSC input.
* @param num_vtypes Number of vertex types. Must be 1 or 2.
* @param mat The CSC matrix
* @param formats Sparse formats available for storing this graph.
* @return A heterograph pointer.
*/
HeteroGraphPtr CreateFromCSC(
int64_t num_vtypes, const aten::CSRMatrix& mat,
dgl_format_code_t formats = ALL_CODE);
/**
* @brief Extract the subgraph of the in edges of the given nodes.
* @param graph Graph
* @param nodes Node IDs of each type
* @param relabel_nodes Whether to remove isolated nodes and relabel the rest
* ones
* @return Subgraph containing only the in edges. The returned graph has
* the same schema as the original one.
*/
HeteroSubgraph InEdgeGraph(
const HeteroGraphPtr graph, const std::vector<IdArray>& nodes,
bool relabel_nodes = false);
/**
* @brief Extract the subgraph of the out edges of the given nodes.
* @param graph Graph
* @param nodes Node IDs of each type
* @param relabel_nodes Whether to remove isolated nodes and relabel the rest
* ones
* @return Subgraph containing only the out edges. The returned graph has
* the same schema as the original one.
*/
HeteroSubgraph OutEdgeGraph(
const HeteroGraphPtr graph, const std::vector<IdArray>& nodes,
bool relabel_nodes = false);
/**
* @brief Joint union multiple graphs into one graph.
*
* All input graphs should have the same metagraph.
*
* TODO(xiangsx): remove the meta_graph argument
*
* @param meta_graph Metagraph of the inputs and result.
* @param component_graphs Input graphs
* @return One graph that unions all the components
*/
HeteroGraphPtr JointUnionHeteroGraph(
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& component_graphs);
/**
* @brief Union multiple graphs into one with each input graph as one disjoint
* component.
*
* All input graphs should have the same metagraph.
*
* TODO(minjie): remove the meta_graph argument
*
* @tparam IdType Graph's index data type, can be int32_t or int64_t
* @param meta_graph Metagraph of the inputs and result.
* @param component_graphs Input graphs
* @return One graph that unions all the components
*/
template <class IdType>
HeteroGraphPtr DisjointUnionHeteroGraph(
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& component_graphs);
HeteroGraphPtr DisjointUnionHeteroGraph2(
GraphPtr meta_graph, const std::vector<HeteroGraphPtr>& component_graphs);
/**
* @brief Slice a contiguous subgraph, e.g. retrieve a component graph from a
* batched graph.
*
* TODO(mufei): remove the meta_graph argument
*
* @param meta_graph Metagraph of the input and result.
* @param batched_graph Input graph.
* @param num_nodes_per_type Number of vertices of each type in the result.
* @param start_nid_per_type Start vertex ID of each type to slice.
* @param num_edges_per_type Number of edges of each type in the result.
* @param start_eid_per_type Start edge ID of each type to slice.
* @return Sliced graph
*/
HeteroGraphPtr SliceHeteroGraph(
GraphPtr meta_graph, HeteroGraphPtr batched_graph,
IdArray num_nodes_per_type, IdArray start_nid_per_type,
IdArray num_edges_per_type, IdArray start_eid_per_type);
/**
* @brief Split a graph into multiple disjoin components.
*
* Edges across different components are ignored. All the result graphs have the
* same metagraph as the input one.
*
* The `vertex_sizes` and `edge_sizes` arrays the concatenation of arrays of
* each node/edge type. Suppose there are N vertex types, then the array length
* should be B*N, where B is the number of components to split.
*
* TODO(minjie): remove the meta_graph argument; use vector<IdArray> for
* vertex_sizes and edge_sizes.
*
* @tparam IdType Graph's index data type, can be int32_t or int64_t
* @param meta_graph Metagraph.
* @param batched_graph Input graph.
* @param vertex_sizes Number of vertices of each component.
* @param edge_sizes Number of vertices of each component.
* @return A list of graphs representing each disjoint components.
*/
template <class IdType>
std::vector<HeteroGraphPtr> DisjointPartitionHeteroBySizes(
GraphPtr meta_graph, HeteroGraphPtr batched_graph, IdArray vertex_sizes,
IdArray edge_sizes);
std::vector<HeteroGraphPtr> DisjointPartitionHeteroBySizes2(
GraphPtr meta_graph, HeteroGraphPtr batched_graph, IdArray vertex_sizes,
IdArray edge_sizes);
/**
* @brief Structure for pickle/unpickle.
*
* The design principle is to leverage the NDArray class as much as possible so
* that when they are converted to backend-specific tensors, we could leverage
* the efficient pickle/unpickle solutions from the backend framework.
*
* NOTE(minjie): This is a temporary solution before we support shared memory
* storage ourselves.
*
* This class can be used as arguments and return values of a C API.
*/
struct HeteroPickleStates : public runtime::Object {
/** @brief version number */
int64_t version = 0;
/** @brief Metainformation
*
* metagraph, number of nodes per type, format, flags
*/
std::string meta;
/** @brief Arrays representing graph structure (coo or csr) */
std::vector<IdArray> arrays;
/* To support backward compatibility, we have to retain fields in the old
* version of HeteroPickleStates
*/
/** @brief Metagraph(64bits ImmutableGraph) */
GraphPtr metagraph;
/** @brief Number of nodes per type */
std::vector<int64_t> num_nodes_per_type;
/** @brief adjacency matrices of each relation graph */
std::vector<std::shared_ptr<SparseMatrix> > adjs;
static constexpr const char* _type_key = "graph.HeteroPickleStates";
DGL_DECLARE_OBJECT_TYPE_INFO(HeteroPickleStates, runtime::Object);
};
// Define HeteroPickleStatesRef
DGL_DEFINE_OBJECT_REF(HeteroPickleStatesRef, HeteroPickleStates);
/**
* @brief Create a heterograph from pickling states.
*
* @param states Pickle states
* @return A heterograph pointer
*/
HeteroGraphPtr HeteroUnpickle(const HeteroPickleStates& states);
/**
* @brief Get the pickling state of the relation graph structure in backend
* tensors.
*
* @return a HeteroPickleStates object
*/
HeteroPickleStates HeteroPickle(HeteroGraphPtr graph);
/**
* @brief Old version of HeteroUnpickle, for backward compatibility
*
* @param states Pickle states
* @return A heterograph pointer
*/
HeteroGraphPtr HeteroUnpickleOld(const HeteroPickleStates& states);
/**
* @brief Create heterograph from pickling states pickled by ForkingPickler.
*
* This is different from HeteroUnpickle where
* (1) Backward compatibility is not required,
* (2) All graph formats are pickled instead of only one.
*/
HeteroGraphPtr HeteroForkingUnpickle(const HeteroPickleStates& states);
/**
* @brief Get the pickling states of the relation graph structure in backend
* tensors for ForkingPickler.
*
* This is different from HeteroPickle where
* (1) Backward compatibility is not required,
* (2) All graph formats are pickled instead of only one.
*/
HeteroPickleStates HeteroForkingPickle(HeteroGraphPtr graph);
#define FORMAT_HAS_CSC(format) ((format)&CSC_CODE)
#define FORMAT_HAS_CSR(format) ((format)&CSR_CODE)
#define FORMAT_HAS_COO(format) ((format)&COO_CODE)
} // namespace dgl
#endif // DGL_BASE_HETEROGRAPH_H_
include/dgl/bcast.h
+68
View File
@@ -0,0 +1,68 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/aten/bcast.h
* @brief Broadcast related function C++ header.
*/
#ifndef DGL_BCAST_H_
#define DGL_BCAST_H_
#include <string>
#include <vector>
#include "./runtime/ndarray.h"
using namespace dgl::runtime;
namespace dgl {
/**
* @brief Broadcast offsets and auxiliary information.
*/
struct BcastOff {
/**
* @brief offset vector of lhs operand and rhs operand.
* @note lhs_offset[i] indicates the start position of the scalar
* in lhs operand that required to compute the i-th element
* in the output, likewise for rhs_offset.
*
* For example, when lhs array has shape (1, 3) and rhs array
* has shape (5, 1), the resulting array would have shape (5, 3),
* then both lhs_offset and rhs_offset would contain 15 elements.
*
* lhs_offset: 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2, 0, 1, 2
* rhs_offset: 0, 0, 0, 1, 1, 1, 2, 2, 2, 3, 3, 3, 4, 4, 4
*
* in order to compute the 7-th (row 2, column 0) element in the output,
* we need the 0-th element in the lhs array and the 2-th element in the
* rhs array.
*/
std::vector<int64_t> lhs_offset, rhs_offset;
/** @brief Whether broadcast is required or not. */
bool use_bcast;
/**
* @brief Auxiliary information for kernel computation
* @note lhs_len refers to the left hand side operand length.
* e.g. 15 for shape (1, 3, 5)
* rhs_len refers to the right hand side operand length.
* e.g. 15 for shape (3, 1, 5)
* out_len refers to the output length.
* e.g. 45 for shape (3, 3, 5)
* reduce_size refers to the reduction size (for op like dot).
* e.g. 1 for add, 5 for dot and lhs_shape,rhs_shape=(3,5)
*/
int64_t lhs_len, rhs_len, out_len, reduce_size;
};
/**
* @brief: Compute broadcast and auxiliary information given operator
* and operands for kernel computation.
* @param op: a string indicates the operator, could be `add`, `sub`,
* `mul`, `div`, `dot`, 'copy_u`, `copy_e`.
* @param lhs The left hand side operand of NDArray class.
* @param rhs The right hand side operand of NDArray class.
* @return the broadcast information of BcastOff class.
*/
BcastOff CalcBcastOff(const std::string& op, NDArray lhs, NDArray rhs);
} // namespace dgl
#endif // DGL_BCAST_H_
include/dgl/env_variable.h
+18
View File
@@ -0,0 +1,18 @@
/**
* Copyright (c) 2023 by Contributors
* @file dgl/env_variable.h
* @brief Class about envrionment variables.
*/
#ifndef DGL_ENV_VARIABLE_H_
#define DGL_ENV_VARIABLE_H_
#include <cstdlib>
namespace dgl {
static const char* kDGLParallelForGrainSize =
std::getenv("DGL_PARALLEL_FOR_GRAIN_SIZE");
} // namespace dgl
#endif // DGL_ENV_VARIABLE_H_
include/dgl/graph.h
+388
View File
@@ -0,0 +1,388 @@
/**
* Copyright (c) 2018 by Contributors
* @file dgl/graph.h
* @brief DGL graph index class.
*/
#ifndef DGL_GRAPH_H_
#define DGL_GRAPH_H_
#include <cstdint>
#include <memory>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "graph_interface.h"
namespace dgl {
class Graph;
class GraphOp;
typedef std::shared_ptr<Graph> MutableGraphPtr;
/** @brief Mutable graph based on adjacency list. */
class Graph : public GraphInterface {
public:
/** @brief default constructor */
Graph() {}
/** @brief construct a graph from the coo format. */
Graph(IdArray src_ids, IdArray dst_ids, size_t num_nodes);
/** @brief default copy constructor */
Graph(const Graph& other) = default;
#ifndef _MSC_VER
/** @brief default move constructor */
Graph(Graph&& other) = default;
#else
Graph(Graph&& other) {
adjlist_ = other.adjlist_;
reverse_adjlist_ = other.reverse_adjlist_;
all_edges_src_ = other.all_edges_src_;
all_edges_dst_ = other.all_edges_dst_;
read_only_ = other.read_only_;
num_edges_ = other.num_edges_;
other.Clear();
}
#endif // _MSC_VER
/** @brief default assign constructor */
Graph& operator=(const Graph& other) = default;
/** @brief default destructor */
~Graph() = default;
/**
* @brief Add vertices to the graph.
* @note Since vertices are integers enumerated from zero, only the number of
* vertices to be added needs to be specified.
* @param num_vertices The number of vertices to be added.
*/
void AddVertices(uint64_t num_vertices) override;
/**
* @brief Add one edge to the graph.
* @param src The source vertex.
* @param dst The destination vertex.
*/
void AddEdge(dgl_id_t src, dgl_id_t dst) override;
/**
* @brief Add edges to the graph.
* @param src_ids The source vertex id array.
* @param dst_ids The destination vertex id array.
*/
void AddEdges(IdArray src_ids, IdArray dst_ids) override;
/**
* @brief Clear the graph. Remove all vertices/edges.
*/
void Clear() override {
adjlist_.clear();
reverse_adjlist_.clear();
all_edges_src_.clear();
all_edges_dst_.clear();
read_only_ = false;
num_edges_ = 0;
}
DGLContext Context() const override { return DGLContext{kDGLCPU, 0}; }
uint8_t NumBits() const override { return 64; }
/**
* @note not const since we have caches
* @return whether the graph is a multigraph
*/
bool IsMultigraph() const override;
/**
* @return whether the graph is read-only
*/
bool IsReadonly() const override { return false; }
/** @return the number of vertices in the graph.*/
uint64_t NumVertices() const override { return adjlist_.size(); }
/** @return the number of edges in the graph.*/
uint64_t NumEdges() const override { return num_edges_; }
/** @return a 0-1 array indicating whether the given vertices are in the
* graph.
*/
BoolArray HasVertices(IdArray vids) const override;
/** @return true if the given edge is in the graph.*/
bool HasEdgeBetween(dgl_id_t src, dgl_id_t dst) const override;
/** @return a 0-1 array indicating whether the given edges are in the graph.*/
BoolArray HasEdgesBetween(IdArray src_ids, IdArray dst_ids) const override;
/**
* @brief Find the predecessors of a vertex.
* @param vid The vertex id.
* @param radius The radius of the neighborhood. Default is immediate neighbor
* (radius=1).
* @return the predecessor id array.
*/
IdArray Predecessors(dgl_id_t vid, uint64_t radius = 1) const override;
/**
* @brief Find the successors of a vertex.
* @param vid The vertex id.
* @param radius The radius of the neighborhood. Default is immediate neighbor
* (radius=1).
* @return the successor id array.
*/
IdArray Successors(dgl_id_t vid, uint64_t radius = 1) const override;
/**
* @brief Get all edge ids between the two given endpoints
* @note Edges are associated with an integer id start from zero.
* The id is assigned when the edge is being added to the graph.
* @param src The source vertex.
* @param dst The destination vertex.
* @return the edge id array.
*/
IdArray EdgeId(dgl_id_t src, dgl_id_t dst) const override;
/**
* @brief Get all edge ids between the given endpoint pairs.
* @note Edges are associated with an integer id start from zero.
* The id is assigned when the edge is being added to the graph.
* If duplicate pairs exist, the returned edge IDs will also duplicate.
* The order of returned edge IDs will follow the order of src-dst pairs
* first, and ties are broken by the order of edge ID.
* @return EdgeArray containing all edges between all pairs.
*/
EdgeArray EdgeIds(IdArray src, IdArray dst) const override;
/**
* @brief Find the edge ID and return the pair of endpoints
* @param eid The edge ID
* @return a pair whose first element is the source and the second the
* destination.
*/
std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_id_t eid) const override {
return std::make_pair(all_edges_src_[eid], all_edges_dst_[eid]);
}
/**
* @brief Find the edge IDs and return their source and target node IDs.
* @param eids The edge ID array.
* @return EdgeArray containing all edges with id in eid. The order is
* preserved.
*/
EdgeArray FindEdges(IdArray eids) const override;
/**
* @brief Get the in edges of the vertex.
* @note The returned dst id array is filled with vid.
* @param vid The vertex id.
* @return the edges
*/
EdgeArray InEdges(dgl_id_t vid) const override;
/**
* @brief Get the in edges of the vertices.
* @param vids The vertex id array.
* @return the id arrays of the two endpoints of the edges.
*/
EdgeArray InEdges(IdArray vids) const override;
/**
* @brief Get the out edges of the vertex.
* @note The returned src id array is filled with vid.
* @param vid The vertex id.
* @return the id arrays of the two endpoints of the edges.
*/
EdgeArray OutEdges(dgl_id_t vid) const override;
/**
* @brief Get the out edges of the vertices.
* @param vids The vertex id array.
* @return the id arrays of the two endpoints of the edges.
*/
EdgeArray OutEdges(IdArray vids) const override;
/**
* @brief Get all the edges in the graph.
* @note If sorted is true, the returned edges list is sorted by their src and
* dst ids. Otherwise, they are in their edge id order.
* @param sorted Whether the returned edge list is sorted by their src and dst
* ids.
* @return the id arrays of the two endpoints of the edges.
*/
EdgeArray Edges(const std::string& order = "") const override;
/**
* @brief Get the in degree of the given vertex.
* @param vid The vertex id.
* @return the in degree
*/
uint64_t InDegree(dgl_id_t vid) const override {
CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
return reverse_adjlist_[vid].succ.size();
}
/**
* @brief Get the in degrees of the given vertices.
* @param vid The vertex id array.
* @return the in degree array
*/
DegreeArray InDegrees(IdArray vids) const override;
/**
* @brief Get the out degree of the given vertex.
* @param vid The vertex id.
* @return the out degree
*/
uint64_t OutDegree(dgl_id_t vid) const override {
CHECK(HasVertex(vid)) << "invalid vertex: " << vid;
return adjlist_[vid].succ.size();
}
/**
* @brief Get the out degrees of the given vertices.
* @param vid The vertex id array.
* @return the out degree array
*/
DegreeArray OutDegrees(IdArray vids) const override;
/**
* @brief Construct the induced subgraph of the given vertices.
*
* The induced subgraph is a subgraph formed by specifying a set of vertices
* V' and then selecting all of the edges from the original graph that connect
* two vertices in V'.
*
* Vertices and edges in the original graph will be "reindexed" to local
* index. The local index of the vertices preserve the order of the given id
* array, while the local index of the edges preserve the index order in the
* original graph. Vertices not in the original graph are ignored.
*
* The result subgraph is read-only.
*
* @param vids The vertices in the subgraph.
* @return the induced subgraph
*/
Subgraph VertexSubgraph(IdArray vids) const override;
/**
* @brief Construct the induced edge subgraph of the given edges.
*
* The induced edges subgraph is a subgraph formed by specifying a set of
* edges E' and then selecting all of the nodes from the original graph that
* are endpoints in E'.
*
* Vertices and edges in the original graph will be "reindexed" to local
* index. The local index of the edges preserve the order of the given id
* array, while the local index of the vertices preserve the index order in
* the original graph. Edges not in the original graph are ignored.
*
* The result subgraph is read-only.
*
* @param eids The edges in the subgraph.
* @return the induced edge subgraph
*/
Subgraph EdgeSubgraph(
IdArray eids, bool preserve_nodes = false) const override;
/**
* @brief Return the successor vector
* @param vid The vertex id.
* @return the successor vector
*/
DGLIdIters SuccVec(dgl_id_t vid) const override {
auto data = adjlist_[vid].succ.data();
auto size = adjlist_[vid].succ.size();
return DGLIdIters(data, data + size);
}
/**
* @brief Return the out edge id vector
* @param vid The vertex id.
* @return the out edge id vector
*/
DGLIdIters OutEdgeVec(dgl_id_t vid) const override {
auto data = adjlist_[vid].edge_id.data();
auto size = adjlist_[vid].edge_id.size();
return DGLIdIters(data, data + size);
}
/**
* @brief Return the predecessor vector
* @param vid The vertex id.
* @return the predecessor vector
*/
DGLIdIters PredVec(dgl_id_t vid) const override {
auto data = reverse_adjlist_[vid].succ.data();
auto size = reverse_adjlist_[vid].succ.size();
return DGLIdIters(data, data + size);
}
/**
* @brief Return the in edge id vector
* @param vid The vertex id.
* @return the in edge id vector
*/
DGLIdIters InEdgeVec(dgl_id_t vid) const override {
auto data = reverse_adjlist_[vid].edge_id.data();
auto size = reverse_adjlist_[vid].edge_id.size();
return DGLIdIters(data, data + size);
}
/**
* @brief Get the adjacency matrix of the graph.
*
* By default, a row of returned adjacency matrix represents the destination
* of an edge and the column represents the source.
* @param transpose A flag to transpose the returned adjacency matrix.
* @param fmt the format of the returned adjacency matrix.
* @return a vector of three IdArray.
*/
std::vector<IdArray> GetAdj(
bool transpose, const std::string& fmt) const override;
/** @brief Create an empty graph */
static MutableGraphPtr Create() { return std::make_shared<Graph>(); }
/** @brief Create from coo */
static MutableGraphPtr CreateFromCOO(
int64_t num_nodes, IdArray src_ids, IdArray dst_ids) {
return std::make_shared<Graph>(src_ids, dst_ids, num_nodes);
}
protected:
friend class GraphOp;
/** @brief Internal edge list type */
struct EdgeList {
/** @brief successor vertex list */
std::vector<dgl_id_t> succ;
/** @brief out edge list */
std::vector<dgl_id_t> edge_id;
};
typedef std::vector<EdgeList> AdjacencyList;
/** @brief adjacency list using vector storage */
AdjacencyList adjlist_;
/** @brief reverse adjacency list using vector storage */
AdjacencyList reverse_adjlist_;
/** @brief all edges' src endpoints in their edge id order */
std::vector<dgl_id_t> all_edges_src_;
/** @brief all edges' dst endpoints in their edge id order */
std::vector<dgl_id_t> all_edges_dst_;
/** @brief read only flag */
bool read_only_ = false;
/** @brief number of edges */
uint64_t num_edges_ = 0;
};
} // namespace dgl
#endif // DGL_GRAPH_H_
include/dgl/graph_interface.h
+465
View File
@@ -0,0 +1,465 @@
/**
* Copyright (c) 2018 by Contributors
* @file dgl/graph_interface.h
* @brief DGL graph index class.
*/
#ifndef DGL_GRAPH_INTERFACE_H_
#define DGL_GRAPH_INTERFACE_H_
#include <algorithm>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "./runtime/object.h"
#include "array.h"
namespace dgl {
const dgl_id_t DGL_INVALID_ID = static_cast<dgl_id_t>(-1);
/**
* @brief This class references data in std::vector.
*
* This isn't a STL-style iterator. It provides a STL data container interface.
* but it doesn't own data itself. instead, it only references data in
* std::vector.
*/
class DGLIdIters {
public:
/** @brief default constructor to create an empty range */
DGLIdIters() {}
/** @brief constructor with given begin and end */
DGLIdIters(const dgl_id_t *begin, const dgl_id_t *end) {
this->begin_ = begin;
this->end_ = end;
}
const dgl_id_t *begin() const { return this->begin_; }
const dgl_id_t *end() const { return this->end_; }
dgl_id_t operator[](int64_t i) const { return *(this->begin_ + i); }
size_t size() const { return this->end_ - this->begin_; }
private:
const dgl_id_t *begin_{nullptr}, *end_{nullptr};
};
/**
* @brief int32 version for DGLIdIters
*
*/
class DGLIdIters32 {
public:
/** @brief default constructor to create an empty range */
DGLIdIters32() {}
/** @brief constructor with given begin and end */
DGLIdIters32(const int32_t *begin, const int32_t *end) {
this->begin_ = begin;
this->end_ = end;
}
const int32_t *begin() const { return this->begin_; }
const int32_t *end() const { return this->end_; }
int32_t operator[](int32_t i) const { return *(this->begin_ + i); }
size_t size() const { return this->end_ - this->begin_; }
private:
const int32_t *begin_{nullptr}, *end_{nullptr};
};
/* @brief structure used to represent a list of edges */
typedef struct {
/* @brief the two endpoints and the id of the edge */
IdArray src, dst, id;
} EdgeArray;
// forward declaration
struct Subgraph;
class GraphRef;
class GraphInterface;
typedef std::shared_ptr<GraphInterface> GraphPtr;
/**
* @brief dgl graph index interface.
*
* DGL's graph is directed. Vertices are integers enumerated from zero.
*
* When calling functions supporing multiple edges (e.g. AddEdges, HasEdges),
* the input edges are represented by two id arrays for source and destination
* vertex ids. In the general case, the two arrays should have the same length.
* If the length of src id array is one, it represents one-many connections.
* If the length of dst id array is one, it represents many-one connections.
*/
class GraphInterface : public runtime::Object {
public:
virtual ~GraphInterface() = default;
/**
* @brief Add vertices to the graph.
* @note Since vertices are integers enumerated from zero, only the number of
* vertices to be added needs to be specified.
* @param num_vertices The number of vertices to be added.
*/
virtual void AddVertices(uint64_t num_vertices) = 0;
/**
* @brief Add one edge to the graph.
* @param src The source vertex.
* @param dst The destination vertex.
*/
virtual void AddEdge(dgl_id_t src, dgl_id_t dst) = 0;
/**
* @brief Add edges to the graph.
* @param src_ids The source vertex id array.
* @param dst_ids The destination vertex id array.
*/
virtual void AddEdges(IdArray src_ids, IdArray dst_ids) = 0;
/**
* @brief Clear the graph. Remove all vertices/edges.
*/
virtual void Clear() = 0;
/**
* @brief Get the device context of this graph.
*/
virtual DGLContext Context() const = 0;
/**
* @brief Get the number of integer bits used to store node/edge ids
* (32 or 64).
*/
virtual uint8_t NumBits() const = 0;
/**
* @return whether the graph is a multigraph
*/
virtual bool IsMultigraph() const = 0;
/**
* @return whether the graph is unibipartite
*/
virtual bool IsUniBipartite() const {
EdgeArray edges = Edges();
IdArray src = edges.src;
IdArray dst = edges.dst;
bool is_unibipartite = true;
const size_t n = edges.src.NumElements();
ATEN_ID_TYPE_SWITCH(src->dtype, IdType, {
auto src_v = src.ToVector<IdType>();
std::sort(src_v.begin(), src_v.end());
auto dst_v = dst.ToVector<IdType>();
std::sort(dst_v.begin(), dst_v.end());
// std::set_intersection() requires output, so this is better
for (size_t i = 0, j = 0; i < n && j < n;) {
if (src_v[i] < dst_v[j]) {
++i;
} else if (src_v[i] == dst_v[j]) {
is_unibipartite = false;
break;
} else {
++j;
}
}
});
return is_unibipartite;
}
/**
* @return whether the graph is read-only
*/
virtual bool IsReadonly() const = 0;
/** @return the number of vertices in the graph.*/
virtual uint64_t NumVertices() const = 0;
/** @return the number of edges in the graph.*/
virtual uint64_t NumEdges() const = 0;
/** @return true if the given vertex is in the graph.*/
virtual bool HasVertex(dgl_id_t vid) const { return vid < NumVertices(); }
/** @return a 0-1 array indicating whether the given vertices are in the
* graph.
*/
virtual BoolArray HasVertices(IdArray vids) const = 0;
/** @return true if the given edge is in the graph.*/
virtual bool HasEdgeBetween(dgl_id_t src, dgl_id_t dst) const = 0;
/** @return a 0-1 array indicating whether the given edges are in the graph.*/
virtual BoolArray HasEdgesBetween(IdArray src_ids, IdArray dst_ids) const = 0;
/**
* @brief Find the predecessors of a vertex.
* @param vid The vertex id.
* @param radius The radius of the neighborhood. Default is immediate neighbor
* (radius=1).
* @return the predecessor id array.
*/
virtual IdArray Predecessors(dgl_id_t vid, uint64_t radius = 1) const = 0;
/**
* @brief Find the successors of a vertex.
* @param vid The vertex id.
* @param radius The radius of the neighborhood. Default is immediate neighbor
* (radius=1).
* @return the successor id array.
*/
virtual IdArray Successors(dgl_id_t vid, uint64_t radius = 1) const = 0;
/**
* @brief Get all edge ids between the two given endpoints
* @note Edges are associated with an integer id start from zero.
* The id is assigned when the edge is being added to the graph.
* @param src The source vertex.
* @param dst The destination vertex.
* @return the edge id array.
*/
virtual IdArray EdgeId(dgl_id_t src, dgl_id_t dst) const = 0;
/**
* @brief Get all edge ids between the given endpoint pairs.
* @note Edges are associated with an integer id start from zero.
* The id is assigned when the edge is being added to the graph.
* If duplicate pairs exist, the returned edge IDs will also duplicate.
* The order of returned edge IDs will follow the order of src-dst pairs
* first, and ties are broken by the order of edge ID.
* @return EdgeArray containing all edges between all pairs.
*/
virtual EdgeArray EdgeIds(IdArray src, IdArray dst) const = 0;
/**
* @brief Find the edge ID and return the pair of endpoints
* @param eid The edge ID
* @return a pair whose first element is the source and the second the
* destination.
*/
virtual std::pair<dgl_id_t, dgl_id_t> FindEdge(dgl_id_t eid) const = 0;
/**
* @brief Find the edge IDs and return their source and target node IDs.
* @param eids The edge ID array.
* @return EdgeArray containing all edges with id in eid. The order is
* preserved.
*/
virtual EdgeArray FindEdges(IdArray eids) const = 0;
/**
* @brief Get the in edges of the vertex.
* @note The returned dst id array is filled with vid.
* @param vid The vertex id.
* @return the edges
*/
virtual EdgeArray InEdges(dgl_id_t vid) const = 0;
/**
* @brief Get the in edges of the vertices.
* @param vids The vertex id array.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray InEdges(IdArray vids) const = 0;
/**
* @brief Get the out edges of the vertex.
* @note The returned src id array is filled with vid.
* @param vid The vertex id.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray OutEdges(dgl_id_t vid) const = 0;
/**
* @brief Get the out edges of the vertices.
* @param vids The vertex id array.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray OutEdges(IdArray vids) const = 0;
/**
* @brief Get all the edges in the graph.
* @note If order is "srcdst", the returned edges list is sorted by their src
* and dst ids. If order is "eid", they are in their edge id order.
* Otherwise, in the arbitrary order.
* @param order The order of the returned edge list.
* @return the id arrays of the two endpoints of the edges.
*/
virtual EdgeArray Edges(const std::string &order = "") const = 0;
/**
* @brief Get the in degree of the given vertex.
* @param vid The vertex id.
* @return the in degree
*/
virtual uint64_t InDegree(dgl_id_t vid) const = 0;
/**
* @brief Get the in degrees of the given vertices.
* @param vid The vertex id array.
* @return the in degree array
*/
virtual DegreeArray InDegrees(IdArray vids) const = 0;
/**
* @brief Get the out degree of the given vertex.
* @param vid The vertex id.
* @return the out degree
*/
virtual uint64_t OutDegree(dgl_id_t vid) const = 0;
/**
* @brief Get the out degrees of the given vertices.
* @param vid The vertex id array.
* @return the out degree array
*/
virtual DegreeArray OutDegrees(IdArray vids) const = 0;
/**
* @brief Construct the induced subgraph of the given vertices.
*
* The induced subgraph is a subgraph formed by specifying a set of vertices
* V' and then selecting all of the edges from the original graph that connect
* two vertices in V'.
*
* Vertices and edges in the original graph will be "reindexed" to local
* index. The local index of the vertices preserve the order of the given id
* array, while the local index of the edges preserve the index order in the
* original graph. Vertices not in the original graph are ignored.
*
* The result subgraph is read-only.
*
* @param vids The vertices in the subgraph.
* @return the induced subgraph
*/
virtual Subgraph VertexSubgraph(IdArray vids) const = 0;
/**
* @brief Construct the induced edge subgraph of the given edges.
*
* The induced edges subgraph is a subgraph formed by specifying a set of
* edges E' and then selecting all of the nodes from the original graph that
* are endpoints in E'.
*
* Vertices and edges in the original graph will be "reindexed" to local
* index. The local index of the edges preserve the order of the given id
* array, while the local index of the vertices preserve the index order in
* the original graph. Edges not in the original graph are ignored.
*
* The result subgraph is read-only.
*
* @param eids The edges in the subgraph.
* @param preserve_nodes If true, the vertices will not be relabeled, so some
* vertices may have no incident edges.
* @return the induced edge subgraph
*/
virtual Subgraph EdgeSubgraph(
IdArray eids, bool preserve_nodes = false) const = 0;
/**
* @brief Return the successor vector
* @param vid The vertex id.
* @return the successor vector iterator pair.
*/
virtual DGLIdIters SuccVec(dgl_id_t vid) const = 0;
/**
* @brief Return the out edge id vector
* @param vid The vertex id.
* @return the out edge id vector iterator pair.
*/
virtual DGLIdIters OutEdgeVec(dgl_id_t vid) const = 0;
/**
* @brief Return the predecessor vector
* @param vid The vertex id.
* @return the predecessor vector iterator pair.
*/
virtual DGLIdIters PredVec(dgl_id_t vid) const = 0;
/**
* @brief Return the in edge id vector
* @param vid The vertex id.
* @return the in edge id vector iterator pair.
*/
virtual DGLIdIters InEdgeVec(dgl_id_t vid) const = 0;
/**
* @brief Get the adjacency matrix of the graph.
*
* By default, a row of returned adjacency matrix represents the destination
* of an edge and the column represents the source.
*
* If the fmt is 'csr', the function should return three arrays, representing
* indptr, indices and edge ids
*
* If the fmt is 'coo', the function should return one array of shape (2,
* nnz), representing a horitonzal stack of row and col indices.
*
* @param transpose A flag to transpose the returned adjacency matrix.
* @param fmt the format of the returned adjacency matrix.
* @return a vector of IdArrays.
*/
virtual std::vector<IdArray> GetAdj(
bool transpose, const std::string &fmt) const = 0;
/**
* @brief Sort the columns in CSR.
*
* This sorts the columns in each row based on the column Ids.
* The edge ids should be sorted accordingly.
*/
virtual void SortCSR() {}
static constexpr const char *_type_key = "graph.Graph";
DGL_DECLARE_OBJECT_TYPE_INFO(GraphInterface, runtime::Object);
};
// Define GraphRef
DGL_DEFINE_OBJECT_REF(GraphRef, GraphInterface);
/** @brief Subgraph data structure */
struct Subgraph : public runtime::Object {
/** @brief The graph. */
GraphPtr graph;
/**
* @brief The induced vertex ids.
* @note This is also a map from the new vertex id to the vertex id in the
* parent graph.
*/
IdArray induced_vertices;
/**
* @brief The induced edge ids.
* @note This is also a map from the new edge id to the edge id in the parent
* graph.
*/
IdArray induced_edges;
static constexpr const char *_type_key = "graph.Subgraph";
DGL_DECLARE_OBJECT_TYPE_INFO(Subgraph, runtime::Object);
};
/** @brief Subgraph data structure for negative subgraph */
struct NegSubgraph : public Subgraph {
/** @brief The existence of the negative edges in the parent graph. */
IdArray exist;
/** @brief The Ids of head nodes */
IdArray head_nid;
/** @brief The Ids of tail nodes */
IdArray tail_nid;
};
/** @brief Subgraph data structure for halo subgraph */
struct HaloSubgraph : public Subgraph {
/** @brief Indicate if a node belongs to the partition. */
IdArray inner_nodes;
};
// Define SubgraphRef
DGL_DEFINE_OBJECT_REF(SubgraphRef, Subgraph);
} // namespace dgl
#endif // DGL_GRAPH_INTERFACE_H_
include/dgl/graph_op.h
+184
View File
@@ -0,0 +1,184 @@
/**
* Copyright (c) 2018 by Contributors
* @file dgl/graph_op.h
* @brief Operations on graph index.
*/
#ifndef DGL_GRAPH_OP_H_
#define DGL_GRAPH_OP_H_
#include <vector>
#include "graph.h"
#include "immutable_graph.h"
namespace dgl {
class GraphOp {
public:
/**
* @brief Return a new graph with all the edges reversed.
*
* The returned graph preserves the vertex and edge index in the original
* graph.
*
* @return the reversed graph
*/
static GraphPtr Reverse(GraphPtr graph);
/**
* @brief Return the line graph.
*
* If i~j and j~i are two edges in original graph G, then
* (i,j)~(j,i) and (j,i)~(i,j) are the "backtracking" edges on
* the line graph.
*
* @param graph The input graph.
* @param backtracking Whether the backtracking edges are included or not
* @return the line graph
*/
static GraphPtr LineGraph(GraphPtr graph, bool backtracking);
/**
* @brief Return a disjoint union of the input graphs.
*
* The new graph will include all the nodes/edges in the given graphs.
* Nodes/Edges will be relabled by adding the cumsum of the previous graph
* sizes in the given sequence order. For example, giving input [g1, g2, g3],
* where they have 5, 6, 7 nodes respectively. Then node#2 of g2 will become
* node#7 in the result graph. Edge ids are re-assigned similarly.
*
* The input list must be either ALL mutable graphs or ALL immutable graphs.
* The returned graph type is also determined by the input graph type.
*
* @param graphs A list of input graphs to be unioned.
* @return the disjoint union of the graphs
*/
static GraphPtr DisjointUnion(std::vector<GraphPtr> graphs);
/**
* @brief Partition the graph into several subgraphs.
*
* This is a reverse operation of DisjointUnion. The graph will be partitioned
* into num graphs. This requires the given number of partitions to evenly
* divides the number of nodes in the graph.
*
* If the input graph is mutable, the result graphs are mutable.
* If the input graph is immutable, the result graphs are immutable.
*
* @param graph The graph to be partitioned.
* @param num The number of partitions.
* @return a list of partitioned graphs
*/
static std::vector<GraphPtr> DisjointPartitionByNum(
GraphPtr graph, int64_t num);
/**
* @brief Partition the graph into several subgraphs.
*
* This is a reverse operation of DisjointUnion. The graph will be partitioned
* based on the given sizes. This requires the sum of the given sizes is equal
* to the number of nodes in the graph.
*
* If the input graph is mutable, the result graphs are mutable.
* If the input graph is immutable, the result graphs are immutable.
*
* @param graph The graph to be partitioned.
* @param sizes The number of partitions.
* @return a list of partitioned graphs
*/
static std::vector<GraphPtr> DisjointPartitionBySizes(
GraphPtr graph, IdArray sizes);
/**
* @brief Map vids in the parent graph to the vids in the subgraph.
*
* If the Id doesn't exist in the subgraph, -1 will be used.
*
* @param parent_vid_map An array that maps the vids in the parent graph to
* the subgraph. The elements store the vertex Ids in the parent graph, and
* the indices indicate the vertex Ids in the subgraph.
* @param query The vertex Ids in the parent graph.
* @return an Id array that contains the subgraph node Ids.
*/
static IdArray MapParentIdToSubgraphId(IdArray parent_vid_map, IdArray query);
/**
* @brief Expand an Id array based on the offset array.
*
* For example,
* ids: [0, 1, 2, 3, 4],
* offset: [0, 2, 2, 5, 6, 7],
* result: [0, 0, 2, 2, 2, 3, 4].
* The offset array has one more element than the ids array.
* (offset[i], offset[i+1]) shows the location of ids[i] in the result array.
*
* @param ids An array that contains the node or edge Ids.
* @param offset An array that contains the offset after expansion.
* @return a expanded Id array.
*/
static IdArray ExpandIds(IdArray ids, IdArray offset);
/**
* @brief Convert the graph to a simple graph.
* @param graph The input graph.
* @return a new immutable simple graph with no multi-edge.
*/
static GraphPtr ToSimpleGraph(GraphPtr graph);
/**
* @brief Convert the graph to a mutable bidirected graph.
*
* If the original graph has m edges for i -> j and n edges for
* j -> i, the new graph will have max(m, n) edges for both
* i -> j and j -> i.
*
* @param graph The input graph.
* @return a new mutable bidirected graph.
*/
static GraphPtr ToBidirectedMutableGraph(GraphPtr graph);
/**
* @brief Same as BidirectedMutableGraph except that the returned graph is
* immutable.
* @param graph The input graph.
* @return a new immutable bidirected
* graph.
*/
static GraphPtr ToBidirectedImmutableGraph(GraphPtr graph);
/**
* @brief Same as BidirectedMutableGraph except that the returned graph is
* immutable and call gk_csr_MakeSymmetric in GKlib. This is more efficient
* than ToBidirectedImmutableGraph. It return a null pointer if the conversion
* fails.
*
* @param graph The input graph.
* @return a new immutable bidirected graph.
*/
static GraphPtr ToBidirectedSimpleImmutableGraph(ImmutableGraphPtr ig);
/**
* @brief Get a induced subgraph with HALO nodes.
* The HALO nodes are the ones that can be reached from `nodes` within
* `num_hops`.
* @param graph The input graph.
* @param nodes The input nodes that form the core of the induced subgraph.
* @param num_hops The number of hops to reach.
* @return the induced subgraph with HALO nodes.
*/
static HaloSubgraph GetSubgraphWithHalo(
GraphPtr graph, IdArray nodes, int num_hops);
/**
* @brief Reorder the nodes in the immutable graph.
* @param graph The input graph.
* @param new_order The node Ids in the new graph. The index in `new_order` is
* old node Ids.
* @return the graph with reordered node Ids
*/
static GraphPtr ReorderImmutableGraph(
ImmutableGraphPtr ig, IdArray new_order);
};
} // namespace dgl
#endif // DGL_GRAPH_OP_H_
include/dgl/graph_serializer.h
+29
View File
@@ -0,0 +1,29 @@
/**
* Copyright (c) 2018 by Contributors
* @file graph/graph_serializer.cc
* @brief DGL serializer APIs
*/
#ifndef DGL_GRAPH_SERIALIZER_H_
#define DGL_GRAPH_SERIALIZER_H_
#include <memory>
namespace dgl {
// Util class to call the private/public empty constructor, which is needed for
// serialization
class Serializer {
public:
template <typename T>
static T* new_object() {
return new T();
}
template <typename T>
static std::shared_ptr<T> make_shared() {
return std::shared_ptr<T>(new T());
}
};
} // namespace dgl
#endif // DGL_GRAPH_SERIALIZER_H_
include/dgl/graph_traversal.h
+101
View File
@@ -0,0 +1,101 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/graph_traversal.h
* @brief common graph traversal operations
*/
#ifndef DGL_GRAPH_TRAVERSAL_H_
#define DGL_GRAPH_TRAVERSAL_H_
#include "array.h"
#include "base_heterograph.h"
namespace dgl {
///////////////////////// Graph Traverse routines //////////////////////////
/**
* @brief Class for representing frontiers.
*
* Each frontier is a list of nodes/edges (specified by their ids).
* An optional tag can be specified on each node/edge (represented by an int
* value).
*/
struct Frontiers {
/** @brief a vector store for the nodes/edges in all the frontiers */
IdArray ids;
/**
* @brief a vector store for node/edge tags. Dtype is int64.
* Empty if no tags are requested
*/
IdArray tags;
/** @brief a section vector to indicate each frontier Dtype is int64. */
IdArray sections;
};
namespace aten {
/**
* @brief Traverse the graph in a breadth-first-search (BFS) order.
*
* @param csr The input csr matrix.
* @param sources Source nodes.
* @return A Frontiers object containing the search result
*/
Frontiers BFSNodesFrontiers(const CSRMatrix& csr, IdArray source);
/**
* @brief Traverse the graph in a breadth-first-search (BFS) order, returning
* the edges of the BFS tree.
*
* @param csr The input csr matrix.
* @param sources Source nodes.
* @return A Frontiers object containing the search result
*/
Frontiers BFSEdgesFrontiers(const CSRMatrix& csr, IdArray source);
/**
* @brief Traverse the graph in topological order.
*
* @param csr The input csr matrix.
* @return A Frontiers object containing the search result
*/
Frontiers TopologicalNodesFrontiers(const CSRMatrix& csr);
/**
* @brief Traverse the graph in a depth-first-search (DFS) order.
*
* @param csr The input csr matrix.
* @param sources Source nodes.
* @return A Frontiers object containing the search result
*/
Frontiers DGLDFSEdges(const CSRMatrix& csr, IdArray source);
/**
* @brief Traverse the graph in a depth-first-search (DFS) order and return the
* recorded edge tag if return_labels is specified.
*
* The traversal visit edges in its DFS order. Edges have three tags:
* FORWARD(0), REVERSE(1), NONTREE(2)
*
* A FORWARD edge is one in which `u` has been visisted but `v` has not.
* A REVERSE edge is one in which both `u` and `v` have been visisted and the
* edge is in the DFS tree.
* A NONTREE edge is one in which both `u` and `v` have been visisted but the
* edge is NOT in the DFS tree.
*
* @param csr The input csr matrix.
* @param sources Source nodes.
* @param has_reverse_edge If true, REVERSE edges are included
* @param has_nontree_edge If true, NONTREE edges are included
* @param return_labels If true, return the recorded edge tags.
* @return A Frontiers object containing the search result
*/
Frontiers DGLDFSLabeledEdges(
const CSRMatrix& csr, IdArray source, const bool has_reverse_edge,
const bool has_nontree_edge, const bool return_labels);
} // namespace aten
} // namespace dgl
#endif // DGL_GRAPH_TRAVERSAL_H_
include/dgl/immutable_graph.h
+1021
View File
File diff suppressed because it is too large Load Diff
include/dgl/kernel.h
+71
View File
@@ -0,0 +1,71 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/aten/kernel.h
* @brief Sparse matrix operators.
*/
#ifndef DGL_KERNEL_H_
#define DGL_KERNEL_H_
#include <string>
#include <utility>
#include <vector>
#include "./base_heterograph.h"
#include "./bcast.h"
#include "array.h"
namespace dgl {
namespace aten {
/**
* @brief Generalized Sparse Matrix-Matrix Multiplication.
* @param op The binary operator, could be `add`, `sub', `mul`, 'div',
* `copy_u`, `copy_e'.
* @param op The reduce operator, could be `sum`, `min`, `max'.
* @param graph The graph we apply SpMM on.
* @param ufeat The source node feature.
* @param efeat The edge feature.
* @param out The output feature on destination nodes.
* @param out_aux A list of NDArray's that contains auxiliary information such
* as the argmax on source nodes and edges for reduce operators such as
* `min` and `max`.
*/
void SpMM(
const std::string& op, const std::string& reduce, HeteroGraphPtr graph,
NDArray ufeat, NDArray efeat, NDArray out, std::vector<NDArray> out_aux);
/**
* @brief Generalized Sampled Dense-Dense Matrix Multiplication.
* @param op The binary operator, could be `add`, `sub', `mul`, 'div',
* `dot`, `copy_u`, `copy_e'.
* @param graph The graph we apply SpMM on.
* @param ufeat The source node feature.
* @param vfeat The destination node feature.
* @param out The output feature on edge.
*/
void SDDMM(
const std::string& op, HeteroGraphPtr graph, NDArray ufeat, NDArray efeat,
NDArray out);
/**
* @brief Sparse-sparse matrix multiplication.
*
* The sparse matrices must have scalar weights (i.e. \a A_weights and \a
* B_weights are 1D vectors.)
*/
std::pair<CSRMatrix, NDArray> CSRMM(
CSRMatrix A, NDArray A_weights, CSRMatrix B, NDArray B_weights);
/**
* @brief Summing up a list of sparse matrices.
*
* The sparse matrices must have scalar weights (i.e. the arrays in \a A_weights
* are 1D vectors.)
*/
std::pair<CSRMatrix, NDArray> CSRSum(
const std::vector<CSRMatrix>& A, const std::vector<NDArray>& A_weights);
} // namespace aten
} // namespace dgl
#endif // DGL_KERNEL_H_
include/dgl/lazy.h
+54
View File
@@ -0,0 +1,54 @@
/**
* Copyright (c) 2019 by Contributors
* @file dgl/lazy.h
* @brief Lazy object that will be materialized only when being queried.
*/
#ifndef DGL_LAZY_H_
#define DGL_LAZY_H_
#include <memory>
namespace dgl {
/**
* @brief Lazy object that will be materialized only when being queried.
*
* The object should be immutable -- no mutation once materialized.
* The object is currently not threaad safe.
*/
template <typename T>
class Lazy {
public:
/** @brief default constructor to construct a lazy object */
Lazy() {}
/**
* @brief constructor to construct an object with given value (non-lazy case)
*/
explicit Lazy(const T& val) : ptr_(new T(val)) {}
/** @brief destructor */
~Lazy() = default;
/**
* @brief Get the value of this object. If the object has not been
* instantiated, using the provided function to create it.
* @param fn The creator function.
* @return the object value.
*/
template <typename Fn>
const T& Get(Fn fn) {
if (!ptr_) {
ptr_.reset(new T(fn()));
}
return *ptr_;
}
private:
/** @brief the internal data pointer */
std::shared_ptr<T> ptr_{nullptr};
};
} // namespace dgl
#endif // DGL_LAZY_H_
include/dgl/nodeflow.h
+101
View File
@@ -0,0 +1,101 @@
/**
* Copyright (c) 2019 by Contributors
* @file dgl/nodeflow.h
* @brief DGL NodeFlow class.
*/
#ifndef DGL_NODEFLOW_H_
#define DGL_NODEFLOW_H_
#include <memory>
#include <string>
#include <vector>
#include "./runtime/object.h"
#include "graph_interface.h"
namespace dgl {
class ImmutableGraph;
/**
* @brief A NodeFlow graph stores the sampling results for a sampler that
* samples nodes/edges in layers.
*
* We store multiple layers of the sampling results in a single graph, which
* results in a more compact format. We store extra information, such as the
* node and edge mapping from the NodeFlow graph to the parent graph.
*/
struct NodeFlowObject : public runtime::Object {
/** @brief The graph. */
GraphPtr graph;
/**
* @brief the offsets of each layer.
*/
IdArray layer_offsets;
/**
* @brief the offsets of each flow.
*/
IdArray flow_offsets;
/**
* @brief The node mapping from the NodeFlow graph to the parent graph.
*/
IdArray node_mapping;
/**
* @brief The edge mapping from the NodeFlow graph to the parent graph.
*/
IdArray edge_mapping;
static constexpr const char *_type_key = "graph.NodeFlow";
DGL_DECLARE_OBJECT_TYPE_INFO(NodeFlowObject, runtime::Object);
};
// Define NodeFlow as the reference class of NodeFlowObject
class NodeFlow : public runtime::ObjectRef {
public:
DGL_DEFINE_OBJECT_REF_METHODS(NodeFlow, runtime::ObjectRef, NodeFlowObject);
/** @brief create a new nodeflow reference */
static NodeFlow Create() {
return NodeFlow(std::make_shared<NodeFlowObject>());
}
};
/**
* @brief Get a slice on a graph that represents a NodeFlow.
*
* The entire block has to be taken as a slice. Users have to specify the
* correct starting and ending location of a layer.
*
* If remap is false, the returned arrays can be viewed as a sub-matrix slice
* of the adjmat of the input graph. Let the adjmat of the input graph be A,
* then the slice is equal to (in numpy syntax):
* A[layer1_start:layer1_end, layer0_start:layer0_end]
*
* If remap is true, the returned arrays represents an adjacency matrix
* of shape NxM, where N is the number of nodes in layer1 and M is
* the number of nodes in layer0. Nodes in layer0 will be remapped to
* [0, M) and nodes in layer1 will be remapped to [0, N).
*
* A row of the returned adjacency matrix represents the destination
* of an edge and the column represents the source.
*
* If fmt == "csr", the function returns three arrays: indptr, indices, eid.
* If fmt == "coo", the function returns two arrays: idx, eid. Here, the idx
* array is the concatenation of src and dst node id arrays.
*
* @param graph An immutable graph.
* @param fmt the format of the returned adjacency matrix.
* @param layer0_size the size of the first layer in the block.
* @param layer1_start the location where the second layer starts.
* @param layer1_end the location where the secnd layer ends.
* @param remap Indicates to remap all vertex ids and edge Ids to local Id
* space.
* @return a vector of IdArrays.
*/
std::vector<IdArray> GetNodeFlowSlice(
const ImmutableGraph &graph, const std::string &fmt, size_t layer0_size,
size_t layer1_start, size_t layer1_end, bool remap);
} // namespace dgl
#endif // DGL_NODEFLOW_H_
include/dgl/packed_func_ext.h
+176
View File
@@ -0,0 +1,176 @@
/**
* Copyright (c) 2019 by Contributors
* @file packed_func_ext.h
* @brief Extension package to PackedFunc
* This enables pass ObjectRef types into/from PackedFunc.
*/
#ifndef DGL_PACKED_FUNC_EXT_H_
#define DGL_PACKED_FUNC_EXT_H_
#include <memory>
#include <sstream>
#include <string>
#include <type_traits>
#include "./runtime/container.h"
#include "./runtime/object.h"
#include "./runtime/packed_func.h"
namespace dgl {
namespace runtime {
/**
* @brief Runtime type checker for node type.
* @tparam T the type to be checked.
*/
template <typename T>
struct ObjectTypeChecker {
static inline bool Check(Object* sptr) {
// This is the only place in the project where RTTI is used
// It can be turned off, but will make non strict checking.
// TODO(tqchen) possibly find alternative to turn of RTTI
using ContainerType = typename T::ContainerType;
return sptr->derived_from<ContainerType>();
}
static inline void PrintName(std::ostringstream& os) { // NOLINT(*)
using ContainerType = typename T::ContainerType;
os << ContainerType::_type_key;
}
};
template <typename T>
struct ObjectTypeChecker<List<T> > {
static inline bool Check(Object* sptr) {
if (sptr == nullptr) return false;
if (!sptr->is_type<ListObject>()) return false;
ListObject* n = static_cast<ListObject*>(sptr);
for (const auto& p : n->data) {
if (!ObjectTypeChecker<T>::Check(p.get())) return false;
}
return true;
}
static inline void PrintName(std::ostringstream& os) { // NOLINT(*)
os << "list<";
ObjectTypeChecker<T>::PrintName(os);
os << ">";
}
};
template <typename V>
struct ObjectTypeChecker<Map<std::string, V> > {
static inline bool Check(Object* sptr) {
if (sptr == nullptr) return false;
if (!sptr->is_type<StrMapObject>()) return false;
StrMapObject* n = static_cast<StrMapObject*>(sptr);
for (const auto& kv : n->data) {
if (!ObjectTypeChecker<V>::Check(kv.second.get())) return false;
}
return true;
}
static inline void PrintName(std::ostringstream& os) { // NOLINT(*)
os << "map<string";
os << ',';
ObjectTypeChecker<V>::PrintName(os);
os << '>';
}
};
template <typename K, typename V>
struct ObjectTypeChecker<Map<K, V> > {
static inline bool Check(Object* sptr) {
if (sptr == nullptr) return false;
if (!sptr->is_type<MapObject>()) return false;
MapObject* n = static_cast<MapObject*>(sptr);
for (const auto& kv : n->data) {
if (!ObjectTypeChecker<K>::Check(kv.first.get())) return false;
if (!ObjectTypeChecker<V>::Check(kv.second.get())) return false;
}
return true;
}
static inline void PrintName(std::ostringstream& os) { // NOLINT(*)
os << "map<";
ObjectTypeChecker<K>::PrintName(os);
os << ',';
ObjectTypeChecker<V>::PrintName(os);
os << '>';
}
};
template <typename T>
inline std::string NodeTypeName() {
std::ostringstream os;
ObjectTypeChecker<T>::PrintName(os);
return os.str();
}
// extensions for DGLArgValue
template <typename TObjectRef>
inline TObjectRef DGLArgValue::AsObjectRef() const {
static_assert(
std::is_base_of<ObjectRef, TObjectRef>::value,
"Conversion only works for ObjectRef derived class");
if (type_code_ == kNull) return TObjectRef();
DGL_CHECK_TYPE_CODE(type_code_, kObjectHandle);
std::shared_ptr<Object>& sptr = *ptr<std::shared_ptr<Object> >();
CHECK(ObjectTypeChecker<TObjectRef>::Check(sptr.get()))
<< "Expected type " << NodeTypeName<TObjectRef>() << " but get "
<< sptr->type_key();
return TObjectRef(sptr);
}
inline std::shared_ptr<Object>& DGLArgValue::obj_sptr() {
DGL_CHECK_TYPE_CODE(type_code_, kObjectHandle);
return *ptr<std::shared_ptr<Object> >();
}
template <typename TObjectRef, typename>
inline bool DGLArgValue::IsObjectType() const {
DGL_CHECK_TYPE_CODE(type_code_, kObjectHandle);
std::shared_ptr<Object>& sptr = *ptr<std::shared_ptr<Object> >();
return ObjectTypeChecker<TObjectRef>::Check(sptr.get());
}
// extensions for DGLRetValue
inline DGLRetValue& DGLRetValue::operator=(
const std::shared_ptr<Object>& other) {
if (other.get() == nullptr) {
SwitchToPOD(kNull);
} else {
SwitchToClass<std::shared_ptr<Object> >(kObjectHandle, other);
}
return *this;
}
inline DGLRetValue& DGLRetValue::operator=(const ObjectRef& other) {
if (!other.defined()) {
SwitchToPOD(kNull);
} else {
SwitchToClass<std::shared_ptr<Object> >(kObjectHandle, other.obj_);
}
return *this;
}
template <typename TObjectRef>
inline TObjectRef DGLRetValue::AsObjectRef() const {
static_assert(
std::is_base_of<ObjectRef, TObjectRef>::value,
"Conversion only works for ObjectRef");
if (type_code_ == kNull) return TObjectRef();
DGL_CHECK_TYPE_CODE(type_code_, kObjectHandle);
return TObjectRef(*ptr<std::shared_ptr<Object> >());
}
inline void DGLArgsSetter::operator()(
size_t i, const ObjectRef& other) const { // NOLINT(*)
if (other.defined()) {
values_[i].v_handle = const_cast<std::shared_ptr<Object>*>(&(other.obj_));
type_codes_[i] = kObjectHandle;
} else {
type_codes_[i] = kNull;
}
}
} // namespace runtime
} // namespace dgl
#endif // DGL_PACKED_FUNC_EXT_H_
include/dgl/random.h
+260
View File
@@ -0,0 +1,260 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/random.h
* @brief Random number generators
*/
#ifndef DGL_RANDOM_H_
#define DGL_RANDOM_H_
#include <dgl/array.h>
#include <dmlc/logging.h>
#include <dmlc/thread_local.h>
#include <random>
#include <thread>
#include <vector>
#include <pcg_random.hpp>
namespace dgl {
namespace {
// Get a unique integer ID representing this thread.
inline uint32_t GetThreadId() {
static int num_threads = 0;
static std::mutex mutex;
static thread_local int id = -1;
if (id == -1) {
std::lock_guard<std::mutex> guard(mutex);
id = num_threads;
num_threads++;
}
return id;
}
}; // namespace
/**
* @brief Thread-local Random Number Generator class
*/
class RandomEngine {
public:
/** @brief Constructor with default seed */
RandomEngine() {
std::random_device rd;
SetSeed(rd());
}
/** @brief Constructor with given seed */
explicit RandomEngine(uint64_t seed, uint64_t stream = GetThreadId()) {
SetSeed(seed, stream);
}
/** @brief Get the thread-local random number generator instance */
static RandomEngine* ThreadLocal() {
return dmlc::ThreadLocalStore<RandomEngine>::Get();
}
/**
* @brief Set the seed of this random number generator
*/
void SetSeed(uint64_t seed, uint64_t stream = GetThreadId()) {
rng_.seed(seed, stream);
}
/**
* @brief Generate an arbitrary random 32-bit integer.
*/
int32_t RandInt32() { return static_cast<int32_t>(rng_()); }
/**
* @brief Generate a uniform random integer in [0, upper)
*/
template <typename T>
T RandInt(T upper) {
return RandInt<T>(0, upper);
}
/**
* @brief Generate a uniform random integer in [lower, upper)
*/
template <typename T>
T RandInt(T lower, T upper) {
CHECK_LT(lower, upper);
std::uniform_int_distribution<T> dist(lower, upper - 1);
return dist(rng_);
}
/**
* @brief Generate a uniform random float in [0, 1)
*/
template <typename T>
T Uniform() {
return Uniform<T>(0., 1.);
}
/**
* @brief Generate a uniform random float in [lower, upper)
*/
template <typename T>
T Uniform(T lower, T upper) {
// Although the result is in [lower, upper), we allow lower == upper as in
// www.cplusplus.com/reference/random/uniform_real_distribution/uniform_real_distribution/
CHECK_LE(lower, upper);
std::uniform_real_distribution<T> dist(lower, upper);
return dist(rng_);
}
/**
* @brief Pick a random integer between 0 to N-1 according to given
* probabilities.
* @tparam IdxType Return integer type.
* @param prob Array of N unnormalized probability of each element. Must be
* non-negative.
* @return An integer randomly picked from 0 to N-1.
*/
template <typename IdxType>
IdxType Choice(FloatArray prob);
/**
* @brief Pick random integers between 0 to N-1 according to given
* probabilities
*
* If replace is false, the number of picked integers must not larger than N.
*
* @tparam IdxType Id type
* @tparam FloatType Probability value type
* @param num Number of integers to choose
* @param prob Array of N unnormalized probability of each element. Must be
* non-negative.
* @param out The output buffer to write selected indices.
* @param replace If true, choose with replacement.
*/
template <typename IdxType, typename FloatType>
void Choice(IdxType num, FloatArray prob, IdxType* out, bool replace = true);
/**
* @brief Pick random integers between 0 to N-1 according to given
* probabilities
*
* If replace is false, the number of picked integers must not larger than N.
*
* @tparam IdxType Id type
* @tparam FloatType Probability value type
* @param num Number of integers to choose
* @param prob Array of N unnormalized probability of each element. Must be
* non-negative.
* @param replace If true, choose with replacement.
* @return Picked indices
*/
template <typename IdxType, typename FloatType>
IdArray Choice(IdxType num, FloatArray prob, bool replace = true) {
const DGLDataType dtype{kDGLInt, sizeof(IdxType) * 8, 1};
IdArray ret = IdArray::Empty({num}, dtype, prob->ctx);
Choice<IdxType, FloatType>(
num, prob, static_cast<IdxType*>(ret->data), replace);
return ret;
}
/**
* @brief Pick random integers from population by uniform distribution.
*
* If replace is false, num must not be larger than population.
*
* @tparam IdxType Return integer type
* @param num Number of integers to choose
* @param population Total number of elements to choose from.
* @param out The output buffer to write selected indices.
* @param replace If true, choose with replacement.
*/
template <typename IdxType>
void UniformChoice(
IdxType num, IdxType population, IdxType* out, bool replace = true);
/**
* @brief Pick random integers from population by uniform distribution.
*
* If replace is false, num must not be larger than population.
*
* @tparam IdxType Return integer type
* @param num Number of integers to choose
* @param population Total number of elements to choose from.
* @param replace If true, choose with replacement.
* @return Picked indices
*/
template <typename IdxType>
IdArray UniformChoice(IdxType num, IdxType population, bool replace = true) {
const DGLDataType dtype{kDGLInt, sizeof(IdxType) * 8, 1};
// TODO(minjie): only CPU implementation right now
IdArray ret = IdArray::Empty({num}, dtype, DGLContext{kDGLCPU, 0});
UniformChoice<IdxType>(
num, population, static_cast<IdxType*>(ret->data), replace);
return ret;
}
/**
* @brief Pick random integers with different probability for different
* segments.
*
* For example, if split=[0, 4, 10] and bias=[1.5, 1], it means to pick some
* integers from 0 to 9, which is divided into two segments. 0-3 are in the
* first segment and the rest belongs to the second. The weight(bias) of each
* candidate in the first segment is upweighted to 1.5.
*
* candidate | 0 1 2 3 | 4 5 6 7 8 9 |
* split ^ ^ ^
* bias | 1.5 | 1 |
*
*
* The complexity of this operator is O(k * log(T)) where k is the number of
* integers we want to pick, and T is the number of segments. It is much
* faster compared with assigning probability for each candidate, of which the
* complexity is O(k * log(N)) where N is the number of all candidates.
*
* If replace is false, num must not be larger than population.
*
* @tparam IdxType Return integer type
* @param num Number of integers to choose
* @param split Array of T+1 split positions of different segments(including
* start and end)
* @param bias Array of T weight of each segments.
* @param out The output buffer to write selected indices.
* @param replace If true, choose with replacement.
*/
template <typename IdxType, typename FloatType>
void BiasedChoice(
IdxType num, const IdxType* split, FloatArray bias, IdxType* out,
bool replace = true);
/**
* @brief Pick random integers with different probability for different
* segments.
*
* If replace is false, num must not be larger than population.
*
* @tparam IdxType Return integer type
* @param num Number of integers to choose
* @param split Split positions of different segments
* @param bias Weights of different segments
* @param replace If true, choose with replacement.
*/
template <typename IdxType, typename FloatType>
IdArray BiasedChoice(
IdxType num, const IdxType* split, FloatArray bias, bool replace = true) {
const DGLDataType dtype{kDGLInt, sizeof(IdxType) * 8, 1};
IdArray ret = IdArray::Empty({num}, dtype, DGLContext{kDGLCPU, 0});
BiasedChoice<IdxType, FloatType>(
num, split, bias, static_cast<IdxType*>(ret->data), replace);
return ret;
}
private:
pcg32 rng_;
};
}; // namespace dgl
#endif // DGL_RANDOM_H_
include/dgl/runtime/bfloat16.h
+68
View File
@@ -0,0 +1,68 @@
/**
* Copyright (c) 2023 by Contributors
* @file dgl/runtime/ndarray.h
* @brief BFloat16 CPU header
*/
#ifndef DGL_RUNTIME_BFLOAT16_H_
#define DGL_RUNTIME_BFLOAT16_H_
#include <cmath>
class BFloat16 {
uint16_t val;
public:
constexpr BFloat16() : val(0) {}
// Disable lint "explicit" warning, since implicit usage on constructor is
// expected.
BFloat16(float f) { // NOLINT
if (std::isnan(f)) {
val = 0x7FC0;
} else {
union {
uint16_t iraw16[2];
uint32_t iraw32;
float f32;
};
f32 = f;
const uint32_t rounding_bias = 0x00007FFF + (iraw16[1] & 0x1);
val = static_cast<uint16_t>((iraw32 + rounding_bias) >> 16);
}
}
static constexpr BFloat16 Min() {
BFloat16 min;
min.val = 0xFF80;
return min;
}
static constexpr BFloat16 Max() {
BFloat16 max;
max.val = 0x7F80;
return max;
}
BFloat16& operator-=(const float& rhs) {
float lhs = (*this);
(*this) = lhs - rhs;
return *this;
}
BFloat16& operator+=(const float& rhs) {
float lhs = (*this);
(*this) = lhs + rhs;
return *this;
}
operator float() const {
union {
float f;
uint16_t raw[2];
};
raw[0] = 0;
raw[1] = val;
return f;
}
};
#endif // DGL_RUNTIME_BFLOAT16_H_
include/dgl/runtime/c_backend_api.h
+131
View File
@@ -0,0 +1,131 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/runtime/c_backend_api.h
* @brief DGL runtime backend API.
*
* The functions defined in this header are intended to be
* used by compiled dgl operators, usually user do not need to use these
* function directly.
*/
#ifndef DGL_RUNTIME_C_BACKEND_API_H_
#define DGL_RUNTIME_C_BACKEND_API_H_
#include "c_runtime_api.h"
#ifdef __cplusplus
extern "C" {
#endif
// Backend related functions.
/**
* @brief Backend function for modules to get function
* from its environment mod_node (its imports and global function).
* The user do should not call DGLFuncFree on func.
*
* @param mod_node The module handle.
* @param func_name The name of the function.
* @param out The result function.
* @return 0 when no error is thrown, -1 when failure happens
*/
DGL_DLL int DGLBackendGetFuncFromEnv(
void* mod_node, const char* func_name, DGLFunctionHandle* out);
/**
* @brief Backend function to register system-wide library symbol.
*
* @param name The name of the symbol
* @param ptr The symbol address.
* @return 0 when no error is thrown, -1 when failure happens
*/
DGL_DLL int DGLBackendRegisterSystemLibSymbol(const char* name, void* ptr);
/**
* @brief Backend function to allocate temporal workspace.
*
* @note The result allocate spaced is ensured to be aligned to
* kTempAllocaAlignment.
*
* @param nbytes The size of the space requested.
* @param device_type The device type which the space will be allocated.
* @param device_id The device id which the space will be allocated.
* @param dtype_code_hint The type code of the array elements. Only used in
* certain backends such as OpenGL.
* @param dtype_bits_hint The type bits of the array elements. Only used in
* certain backends such as OpenGL.
* @return nullptr when error is thrown, a valid ptr if success
*/
DGL_DLL void* DGLBackendAllocWorkspace(
int device_type, int device_id, uint64_t nbytes, int dtype_code_hint,
int dtype_bits_hint);
/**
* @brief Backend function to free temporal workspace.
*
* @param ptr The result allocated space pointer.
* @param device_type The device type which the space will be allocated.
* @param device_id The device id which the space will be allocated.
* @return 0 when no error is thrown, -1 when failure happens
*
* @sa DGLBackendAllocWorkspace
*/
DGL_DLL int DGLBackendFreeWorkspace(int device_type, int device_id, void* ptr);
/**
* @brief Environment for DGL parallel task.
*/
typedef struct {
/**
* @brief Auxiliary used for synchronization
*/
void* sync_handle;
/** @brief total amount of task */
int32_t num_task;
} DGLParallelGroupEnv;
/**
* @brief The callback function to execute a parallel lambda
* @param task_id the task id of the function.
* @param penv The parallel environment backs the execution.
* @param cdata The supporting closure data.
*/
typedef int (*FDGLParallelLambda)(
int task_id, DGLParallelGroupEnv* penv, void* cdata);
/**
* @brief Backend function for running parallel jobs.
*
* @param flambda The parallel function to be launched.
* @param cdata The closure data.
* @param num_task Number of tasks to launch, can be 0, means launch
* with all available threads.
*
* @return 0 when no error is thrown, -1 when failure happens
*/
DGL_DLL int DGLBackendParallelLaunch(
FDGLParallelLambda flambda, void* cdata, int num_task);
/**
* @brief BSP barrrier between parallel threads
* @param task_id the task id of the function.
* @param penv The parallel environment backs the execution.
* @return 0 when no error is thrown, -1 when failure happens
*/
DGL_DLL int DGLBackendParallelBarrier(int task_id, DGLParallelGroupEnv* penv);
/**
* @brief Simple static initialization fucntion.
* Run f once and set handle to be not null.
* This function is mainly used for test purpose.
*
* @param handle An global address to indicate f
* @param f The function to be ran
* @param cdata The closure data to pass to the function.
* @param nbytes Number of bytes in the closure data.
* @return 0 when no error is thrown, -1 when failure happens
*/
DGL_DLL int DGLBackendRunOnce(
void** handle, int (*f)(void*), void* cdata, int nbytes);
#ifdef __cplusplus
} // DGL_EXTERN_C
#endif
#endif // DGL_RUNTIME_C_BACKEND_API_H_
include/dgl/runtime/c_object_api.h
+71
View File
@@ -0,0 +1,71 @@
/**
* Copyright (c) 2019 by Contributors
* @file dgl/runtime/c_object_api.h
*
* @brief DGL Object C API, used to extend and prototype new CAPIs.
*
* @note Most API functions are registerd as PackedFunc and
* can be grabbed via DGLFuncGetGlobal
*/
#ifndef DGL_RUNTIME_C_OBJECT_API_H_
#define DGL_RUNTIME_C_OBJECT_API_H_
#include "./c_runtime_api.h"
#ifdef __cplusplus
extern "C" {
#endif
/** @brief handle to object */
typedef void* ObjectHandle;
/**
* @brief free the object handle
* @param handle The object handle to be freed.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLObjectFree(ObjectHandle handle);
/**
* @brief Convert type key to type index.
* @param type_key The key of the type.
* @param out_index the corresponding type index.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLObjectTypeKey2Index(const char* type_key, int* out_index);
/**
* @brief Get runtime type index of the object.
* @param handle the object handle.
* @param out_index the corresponding type index.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLObjectGetTypeIndex(ObjectHandle handle, int* out_index);
/**
* @brief get attributes given key
* @param handle The object handle
* @param key The attribute name
* @param out_value The attribute value
* @param out_type_code The type code of the attribute.
* @param out_success Whether get is successful.
* @return 0 when success, -1 when failure happens
* @note API calls always exchanges with type bits=64, lanes=1
*/
DGL_DLL int DGLObjectGetAttr(
ObjectHandle handle, const char* key, DGLValue* out_value,
int* out_type_code, int* out_success);
/**
* @brief get attributes names in the object.
* @param handle The object handle
* @param out_size The number of functions
* @param out_array The array of function names.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLObjectListAttrNames(
ObjectHandle handle, int* out_size, const char*** out_array);
#ifdef __cplusplus
} // DGL_EXTERN_C
#endif
#endif // DGL_RUNTIME_C_OBJECT_API_H_
include/dgl/runtime/c_runtime_api.h
+629
View File
@@ -0,0 +1,629 @@
/**
* Copyright (c) 2016-2022 by Contributors
* @file dgl/runtime/c_runtime_api.h
* @brief DGL runtime library.
*
* This runtime is adapted from TVM project (commit: 2ce5277)
*/
#ifndef DGL_RUNTIME_C_RUNTIME_API_H_
#define DGL_RUNTIME_C_RUNTIME_API_H_
// Macros to do weak linking
#ifdef _MSC_VER
#define DGL_WEAK __declspec(selectany)
#else
#define DGL_WEAK __attribute__((weak))
#endif
#ifdef __EMSCRIPTEN__
#include <emscripten/emscripten.h>
#define DGL_DLL EMSCRIPTEN_KEEPALIVE
#endif
#ifndef DGL_DLL
#ifdef _WIN32
#ifdef DGL_EXPORTS
#define DGL_DLL __declspec(dllexport)
#else
#define DGL_DLL __declspec(dllimport)
#endif
#else
#define DGL_DLL
#endif
#endif
// DGL version
#define DGL_VERSION "2.5"
#ifdef __cplusplus
extern "C" {
#endif
#include <stddef.h>
#include <stdint.h>
/** @brief type of array index. */
typedef int64_t dgl_index_t;
/**
* @brief The device type in DGLContext.
*/
#ifdef __cplusplus
typedef enum : int32_t {
#else
typedef enum {
#endif
/** @brief CPU device */
kDGLCPU = 1,
/** @brief CUDA GPU device */
kDGLCUDA = 2,
// add more devices once supported
} DGLDeviceType;
/**
* @brief The object type code is used in DGL FFI to indicate the types of
* objects passed between C and Python.
*/
typedef enum {
kObjectInt = 0U,
kObjectUInt = 1U,
kObjectFloat = 2U,
kHandle = 3U,
kNull = 4U,
kDGLDataType = 5U,
kDGLContext = 6U,
kArrayHandle = 7U,
kObjectHandle = 8U,
kModuleHandle = 9U,
kFuncHandle = 10U,
kStr = 11U,
kBytes = 12U,
kNDArrayContainer = 13U,
// Extension codes for other frameworks to integrate DGL PackedFunc.
// To make sure each framework's id do not conflict, use first and
// last sections to mark ranges.
// Open an issue at the repo if you need a section of code.
kExtBegin = 15U,
kNNVMFirst = 16U,
kNNVMLast = 20U,
// The following section of code is used for non-reserved types.
kExtReserveEnd = 64U,
kExtEnd = 128U
} DGLObjectTypeCode;
/**
* @brief The type code options DGLDataType.
*/
typedef enum {
/** @brief signed integer */
kDGLInt = 0U,
/** @brief unsigned integer */
kDGLUInt = 1U,
/** @brief IEEE floating point */
kDGLFloat = 2U,
/** @brief bfloat16 */
kDGLBfloat = 4U,
// add more data types if we are going to support them
} DGLDataTypeCode;
/**
* @brief The data type the tensor can hold. The data type is assumed to follow
* the native endian-ness. An explicit error message should be raised when
* attempting to export an array with non-native endianness
*
* Examples
* - float: type_code = 2, bits = 32, lanes=1
* - float4(vectorized 4 float): type_code = 2, bits = 32, lanes=4
* - int8: type_code = 0, bits = 8, lanes=1
*/
typedef struct {
/**
* @brief Type code of base types.
* We keep it uint8_t instead of DGLDataTypeCode for minimal memory
* footprint, but the value should be one of DGLDataTypeCode enum values.
* */
uint8_t code;
/**
* @brief Number of bits, common choices are 8, 16, 32.
*/
uint8_t bits;
/** @brief Number of lanes in the type, used for vector types. */
uint16_t lanes;
} DGLDataType;
/**
* @brief The Device information, abstract away common device types.
*/
typedef struct {
/** @brief The device type used in the device. */
DGLDeviceType device_type;
/**
* @brief The device index.
* For vanilla CPU memory, pinned memory, or managed memory, this is set to 0.
*/
int32_t device_id;
} DGLContext;
/**
* @brief The tensor array stucture to DGL API.
* The structure is heavily inspired by DLTensor from DLPack.
*/
typedef struct {
/**
* @brief The data pointer points to the allocated data.
*
* Depending on the device context, it can be a CPU pointer, or a CUDA
* device pointer or acl_mem handle in OpenCL.
* This pointer is always aligned to 256 bytes as in CUDA. Use the
* `byte_offset` field to mark the beginning of the actual data (if the
* address is not 256 byte aligned).
*
* Note that as of Nov 2021, multiply libraries (CuPy, PyTorch, TensorFlow,
* TVM, perhaps others) do not adhere to this 256 byte alignment requirement
* on CPU/CUDA/ROCm, and always use `byte_offset=0`. This is likely to be
* fixed in the future; at the moment it is recommended
* to not rely on the data pointer being correctly aligned.
*
* For a DGLArray, the size of memory required to store the contents of
* data can be calculated as follows:
*
* @code{.c}
* static inline size_t GetDataSize(const DGLArray* t) {
* size_t size = 1;
* for (int32_t i = 0; i < t->ndim; ++i) {
* size *= t->shape[i];
* }
* size *= (t->dtype.bits * t->dtype.lanes + 7) / 8;
* return size;
* }
* @endcode
*/
void* data;
/** @brief The device of the tensor */
DGLContext ctx;
/** @brief Number of dimensions */
int32_t ndim;
/** @brief The data type of the pointer*/
DGLDataType dtype;
/** @brief The shape of the tensor */
int64_t* shape;
/**
* @brief strides of the tensor (in number of elements, not bytes)
* can be NULL, indicating tensor is compact and row-majored.
*/
int64_t* strides;
/** @brief The offset in bytes to the beginning pointer to data */
uint64_t byte_offset;
} DGLArray;
/** @brief the array handle */
typedef DGLArray* DGLArrayHandle;
/**
* @brief Union type of values
* being passed through API and function calls.
*/
typedef union {
int64_t v_int64;
double v_float64;
void* v_handle;
const char* v_str;
DGLDataType v_type;
DGLContext v_ctx;
} DGLValue;
/**
* @brief Byte array type used to pass in byte array
* When kBytes is used as data type.
*/
typedef struct {
const char* data;
size_t size;
} DGLByteArray;
/** @brief Handle to DGL runtime modules. */
typedef void* DGLModuleHandle;
/** @brief Handle to packed function handle. */
typedef void* DGLFunctionHandle;
/** @brief Handle to hold return value. */
typedef void* DGLRetValueHandle;
/**
* @brief The stream that is specific to device
* can be NULL, which indicates the default one.
*/
typedef void* DGLStreamHandle;
/**
* @brief Used for implementing C API function.
* Set last error message before return.
* @param msg The error message to be set.
*/
DGL_DLL void DGLAPISetLastError(const char* msg);
/**
* @brief return str message of the last error
* all function in this file will return 0 when success
* and -1 when an error occured,
* DGLGetLastError can be called to retrieve the error
*
* this function is threadsafe and can be called by different thread
*
* @return error info
*/
DGL_DLL const char* DGLGetLastError(void);
/**
* @brief Load module from file.
* @param file_name The file name to load the module from.
* @param format The format of the module.
* @param out The result module
*
* @return 0 when success, -1 when failure happens
* @note The resulting module do not contain import relation.
* It can be reconstructed by DGLModImport.
*/
DGL_DLL int DGLModLoadFromFile(
const char* file_name, const char* format, DGLModuleHandle* out);
/**
* @brief Add dep to mod's dependency.
* This allows functions in this module to use modules.
*
* @param mod The module handle.
* @param dep The dependent module to be imported.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLModImport(DGLModuleHandle mod, DGLModuleHandle dep);
/**
* @brief Get function from the module.
* @param mod The module handle.
* @param func_name The name of the function.
* @param query_imports Whether to query imported modules
* @param out The result function, can be NULL if it is not available.
* @return 0 when no error is thrown, -1 when failure happens
*/
DGL_DLL int DGLModGetFunction(
DGLModuleHandle mod, const char* func_name, int query_imports,
DGLFunctionHandle* out);
/**
* @brief Free front-end extension type resource.
* @param handle The extension handle.
* @param type_code The type of of the extension type.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLExtTypeFree(void* handle, int type_code);
/**
* @brief Free the Module
* @param mod The module to be freed.
*
* @note This may not free up the module's resources.
* If there is active DGLFunctionHandle uses the module
* Or if this module is imported by another active module.
*
* The all functions remains valid until DGLFuncFree is called.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLModFree(DGLModuleHandle mod);
/**
* @brief Free the function when it is no longer needed.
* @param func The function handle
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLFuncFree(DGLFunctionHandle func);
/**
* @brief Call a Packed DGL Function.
*
* @param func node handle of the function.
* @param arg_values The arguments
* @param type_codes The type codes of the arguments
* @param num_args Number of arguments.
*
* @param ret_val The return value.
* @param ret_type_code the type code of return value.
*
* @return 0 when success, -1 when failure happens
* @note DGL calls always exchanges with type bits=64, lanes=1
*
* @note API calls always exchanges with type bits=64, lanes=1
* If API call returns container handles (e.g. FunctionHandle)
* these handles should be managed by the front-end.
* The front-end need to call free function (e.g. DGLFuncFree)
* to free these handles.
*/
DGL_DLL int DGLFuncCall(
DGLFunctionHandle func, DGLValue* arg_values, int* type_codes, int num_args,
DGLValue* ret_val, int* ret_type_code);
/**
* @brief Set the return value of DGLPackedCFunc.
*
* This function is called by DGLPackedCFunc to set the return value.
* When this function is not called, the function returns null by default.
*
* @param ret The return value handle, pass by ret in DGLPackedCFunc
* @param value The value to be returned.
* @param type_code The type of the value to be returned.
* @param num_ret Number of return values, for now only 1 is supported.
*/
DGL_DLL int DGLCFuncSetReturn(
DGLRetValueHandle ret, DGLValue* value, int* type_code, int num_ret);
/**
* @brief Inplace translate callback argument value to return value.
* This is only needed for non-POD arguments.
*
* @param value The value to be translated.
* @param code The type code to be translated.
* @note This function will do a shallow copy when necessary.
*
* @return 0 when success, -1 when failure happens.
*/
DGL_DLL int DGLCbArgToReturn(DGLValue* value, int code);
/**
* @brief C type of packed function.
*
* @param args The arguments
* @param type_codes The type codes of the arguments
* @param num_args Number of arguments.
* @param ret The return value handle.
* @param resource_handle The handle additional resouce handle from fron-end.
* @return 0 if success, -1 if failure happens, set error via
* DGLAPISetLastError.
* @sa DGLCFuncSetReturn
*/
typedef int (*DGLPackedCFunc)(
DGLValue* args, int* type_codes, int num_args, DGLRetValueHandle ret,
void* resource_handle);
/**
* @brief C callback to free the resource handle in C packed function.
* @param resource_handle The handle additional resouce handle from fron-end.
*/
typedef void (*DGLPackedCFuncFinalizer)(void* resource_handle);
/**
* @brief Signature for extension function declarer.
*
* DGL call this function to get the extension functions
* The declarer will call register_func to register function and their name.
*
* @param register_func_handle The register function
* @return 0 if success, -1 if failure happens
*/
typedef int (*DGLExtensionFuncDeclarer)(DGLFunctionHandle register_func_handle);
/**
* @brief Wrap a DGLPackedCFunc to become a FunctionHandle.
*
* The resource_handle will be managed by DGL API, until the function is no
* longer used.
*
* @param func The packed C function.
* @param resource_handle The resource handle from front-end, can be NULL.
* @param fin The finalizer on resource handle when the FunctionHandle get
* freed, can be NULL.
* @param out the result function handle.
* @return 0 when success, -1 when failure happens.
*/
DGL_DLL int DGLFuncCreateFromCFunc(
DGLPackedCFunc func, void* resource_handle, DGLPackedCFuncFinalizer fin,
DGLFunctionHandle* out);
/**
* @brief Register the function to runtime's global table.
*
* The registered function then can be pulled by the backend by the name.
*
* @param name The name of the function.
* @param f The function to be registered.
* @param override Whether allow override already registered function.
*/
DGL_DLL int DGLFuncRegisterGlobal(
const char* name, DGLFunctionHandle f, int override);
/**
* @brief Get a global function.
*
* @param name The name of the function.
* @param out the result function pointer, NULL if it does not exist.
*
* @note The function handle of global function is managed by DGL runtime,
* So DGLFuncFree is should not be called when it get deleted.
*/
DGL_DLL int DGLFuncGetGlobal(const char* name, DGLFunctionHandle* out);
/**
* @brief List all the globally registered function name
* @param out_size The number of functions
* @param out_array The array of function names.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLFuncListGlobalNames(int* out_size, const char*** out_array);
// Array related apis for quick proptyping
/**
* @brief Allocate a nd-array's memory,
* including space of shape, of given spec.
*
* @param shape The shape of the array, the data content will be copied to out
* @param ndim The number of dimension of the array.
* @param dtype_code The type code of the dtype
* @param dtype_bits The number of bits of dtype
* @param dtype_lanes The number of lanes in the dtype.
* @param device_type The device type of context
* @param device_id The device id of context.
* @param out The output handle.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLArrayAlloc(
const dgl_index_t* shape, int ndim, int dtype_code, int dtype_bits,
int dtype_lanes, int device_type, int device_id, DGLArrayHandle* out);
/**
* @brief Allocate a nd-array's with shared memory,
* including space of shape, of given spec.
*
* @param the name of the shared memory
* @param shape The shape of the array, the data content will be copied to out
* @param ndim The number of dimension of the array.
* @param dtype_code The type code of the dtype
* @param dtype_bits The number of bits of dtype
* @param dtype_lanes The number of lanes in the dtype.
* @param is_create whether the shared memory is created
* @param out The output handle.
* @return 0 when success, -1 when failure happens
*/
int DGLArrayAllocSharedMem(
const char* mem_name, const dgl_index_t* shape, int ndim, int dtype_code,
int dtype_bits, int dtype_lanes, bool is_create, DGLArrayHandle* out);
/**
* @brief Free the DGL Array.
* @param handle The array handle to be freed.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLArrayFree(DGLArrayHandle handle);
/**
* @brief Copy array data from CPU byte array.
* @param handle The array handle.
* @param data the data pointer
* @param nbytes The number of bytes to copy.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLArrayCopyFromBytes(
DGLArrayHandle handle, void* data, size_t nbytes);
/**
* @brief Copy array data to CPU byte array.
* @param handle The array handle.
* @param data the data pointer
* @param nbytes The number of bytes to copy.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLArrayCopyToBytes(
DGLArrayHandle handle, void* data, size_t nbytes);
/**
* @brief Copy the array, both from and to must be valid during the copy.
* @param from The array to be copied from.
* @param to The target space.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLArrayCopyFromTo(DGLArrayHandle from, DGLArrayHandle to);
/**
* @brief Create a new runtime stream.
*
* @param device_type The device type of context
* @param device_id The device id of context
* @param out The new stream handle
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLStreamCreate(
int device_type, int device_id, DGLStreamHandle* out);
/**
* @brief Free a created stream handle.
*
* @param device_type The device type of context
* @param device_id The device id of context
* @param stream The stream to be freed
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLStreamFree(
int device_type, int device_id, DGLStreamHandle stream);
/**
* @brief Set the runtime stream of current thread to be stream.
* The subsequent calls to the same device_type
* will use the setted stream handle.
* The specific type of stream is runtime device dependent.
*
* @param device_type The device type of context
* @param device_id The device id of context.
* @param handle The stream handle.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLSetStream(
int device_type, int device_id, DGLStreamHandle handle);
/**
* @brief Get the runtime stream of current thread.
*
* @param device_type The device type of context
* @param device_id The device id of context.
* @param handle The stream handle.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLGetStream(
int device_type, int device_id, DGLStreamHandle* handle);
/**
* @brief Wait until all computations on stream completes.
*
* @param device_type The device type of context
* @param device_id The device id of context.
* @param stream The stream to be synchronized.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLSynchronize(
int device_type, int device_id, DGLStreamHandle stream);
/**
* @brief Synchronize two streams of execution.
*
* @param device_type The device type of context
* @param device_id The device id of context
* @param src The source stream to synchronize.
* @param dst The destination stream to synchronize.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLStreamStreamSynchronize(
int device_type, int device_id, DGLStreamHandle src, DGLStreamHandle dst);
/**
* @brief Load tensor adapter.
* @return 0 when success, -1 when failure happens.
*/
DGL_DLL int DGLLoadTensorAdapter(const char* path);
/**
* @brief Pin host memory.
*/
int DGLArrayPinData(DGLArrayHandle handle, DGLContext ctx);
/**
* @brief Unpin host memory.
*/
int DGLArrayUnpinData(DGLArrayHandle handle, DGLContext ctx);
/**
* @brief Record the stream that's using this tensor.
*/
int DGLArrayRecordStream(DGLArrayHandle handle, DGLStreamHandle stream);
/**
* @brief Bug report macro.
*
* This serves as a sanity check on system side to make sure the code is correct
* by checking whether a condition always holds for complex reasons. Failing
* the condition signifies a system bug instead of users giving invalid inputs
* or using the functionality incorrectly.
*
* Hints the user to file a bug report if the condition fails.
*/
#define BUG_IF_FAIL(cond) \
CHECK(cond) \
<< "A bug has been occurred. " \
"Please file a bug report at https://github.com/dmlc/dgl/issues. " \
"Message: "
#ifdef __cplusplus
} // DGL_EXTERN_C
#endif
#endif // DGL_RUNTIME_C_RUNTIME_API_H_
include/dgl/runtime/config.h
+32
View File
@@ -0,0 +1,32 @@
/**
* Copyright (c) 2019 by Contributors
* @file runtime/config.h
* @brief DGL runtime config
*/
#ifndef DGL_RUNTIME_CONFIG_H_
#define DGL_RUNTIME_CONFIG_H_
namespace dgl {
namespace runtime {
class Config {
public:
static Config* Global() {
static Config config;
return &config;
}
// Enabling or disable use libxsmm for Spmm
void EnableLibxsmm(bool);
bool IsLibxsmmAvailable() const;
private:
Config();
bool libxsmm_;
};
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_CONFIG_H_
include/dgl/runtime/container.h
+685
View File
@@ -0,0 +1,685 @@
/**
* Copyright (c) 2019 by Contributors
* @file runtime/container.h
* @brief Defines the container object data structures.
*/
#ifndef DGL_RUNTIME_CONTAINER_H_
#define DGL_RUNTIME_CONTAINER_H_
#include <memory>
#include <string>
#include <unordered_map>
#include <utility>
#include <vector>
#include "object.h"
#include "packed_func.h"
namespace dgl {
namespace runtime {
/**
* @brief value object.
*
* It is typically used to wrap a non-Object type to Object type.
* Any type that is supported by DGLRetValue is supported by this.
*/
class ValueObject : public Object {
public:
/** @brief the value data */
DGLRetValue data;
static constexpr const char* _type_key = "Value";
DGL_DECLARE_OBJECT_TYPE_INFO(ValueObject, Object);
};
/** @brief Construct a value object. */
template <typename T>
inline std::shared_ptr<ValueObject> MakeValue(T&& val) {
auto obj = std::make_shared<ValueObject>();
obj->data = val;
return obj;
}
/** @brief Vallue reference type */
class Value : public ObjectRef {
public:
Value() {}
explicit Value(std::shared_ptr<Object> o) : ObjectRef(o) {}
const ValueObject* operator->() const {
return static_cast<const ValueObject*>(obj_.get());
}
using ContainerType = ValueObject;
};
/** @brief list obj content in list */
class ListObject : public Object {
public:
/** @brief the data content */
std::vector<std::shared_ptr<Object> > data;
void VisitAttrs(AttrVisitor* visitor) final {
// Visitor to list have no effect.
}
static constexpr const char* _type_key = "List";
DGL_DECLARE_OBJECT_TYPE_INFO(ListObject, Object);
};
/** @brief map obj content */
class MapObject : public Object {
public:
void VisitAttrs(AttrVisitor* visitor) final {
// Visitor to map have no effect.
}
// hash function
struct Hash {
size_t operator()(const std::shared_ptr<Object>& n) const {
return std::hash<Object*>()(n.get());
}
};
// comparator
struct Equal {
bool operator()(
const std::shared_ptr<Object>& a,
const std::shared_ptr<Object>& b) const {
return a.get() == b.get();
}
};
/** @brief The corresponding conatiner type */
using ContainerType = std::unordered_map<
std::shared_ptr<Object>, std::shared_ptr<Object>, Hash, Equal>;
/** @brief the data content */
ContainerType data;
static constexpr const char* _type_key = "Map";
DGL_DECLARE_OBJECT_TYPE_INFO(MapObject, Object);
};
/** @brief specialized map obj with string as key */
class StrMapObject : public Object {
public:
void VisitAttrs(AttrVisitor* visitor) final {
// Visitor to map have no effect.
}
/** @brief The corresponding conatiner type */
using ContainerType =
std::unordered_map<std::string, std::shared_ptr<Object> >;
/** @brief the data content */
ContainerType data;
static constexpr const char* _type_key = "StrMap";
DGL_DECLARE_OBJECT_TYPE_INFO(StrMapObject, Object);
};
/**
* @brief iterator adapter that adapts TIter to return another type.
* @tparam Converter a struct that contains converting function
* @tparam TIter the content iterator type.
*/
template <typename Converter, typename TIter>
class IterAdapter {
public:
explicit IterAdapter(TIter iter) : iter_(iter) {}
inline IterAdapter& operator++() { // NOLINT(*)
++iter_;
return *this;
}
inline IterAdapter& operator++(int) { // NOLINT(*)
++iter_;
return *this;
}
inline IterAdapter operator+(int offset) const { // NOLINT(*)
return IterAdapter(iter_ + offset);
}
inline bool operator==(IterAdapter other) const {
return iter_ == other.iter_;
}
inline bool operator!=(IterAdapter other) const { return !(*this == other); }
inline const typename Converter::ResultType operator*() const {
return Converter::convert(*iter_);
}
private:
TIter iter_;
};
/**
* @brief List container of ObjectRef.
*
* List implements copy on write semantics, which means list is mutable
* but copy will happen when list is referenced in more than two places.
*
* That is said when using this container for runtime arguments or return
* values, try use the constructor to create the list at once (for example
* from an existing vector).
*
* operator[] only provide const access, use Set to mutate the content.
*
* @tparam T The content ObjectRef type.
*
* @note The element type must subclass \c ObjectRef. Otherwise, the
* compiler would throw an error:
*
* <code>
* error: no type named 'type' in 'struct std::enable_if<false, void>'
* </code>
*
* Example:
*
* <code>
* // List<int> list; // fails
* // List<NDArray> list2; // fails
* List<Value> list; // works
* list.push_back(Value(MakeValue(1))); // works
* list.push_back(Value(MakeValue(NDArray::Empty(shape, dtype, ctx)))); //
* works
* </code>
*/
template <
typename T,
typename =
typename std::enable_if<std::is_base_of<ObjectRef, T>::value>::type>
class List : public ObjectRef {
public:
/**
* @brief default constructor
*/
List() { obj_ = std::make_shared<ListObject>(); }
/**
* @brief move constructor
* @param other source
*/
List(List<T>&& other) { // NOLINT(*)
obj_ = std::move(other.obj_);
}
/**
* @brief copy constructor
* @param other source
*/
List(const List<T>& other) : ObjectRef(other.obj_) { // NOLINT(*)
}
/**
* @brief constructor from pointer
* @param n the container pointer
*/
explicit List(std::shared_ptr<Object> n) : ObjectRef(n) {}
/**
* @brief constructor from iterator
* @param begin begin of iterator
* @param end end of iterator
* @tparam IterType The type of iterator
*/
template <typename IterType>
List(IterType begin, IterType end) {
assign(begin, end);
}
/**
* @brief constructor from initializer list
* @param init The initalizer list
*/
List(std::initializer_list<T> init) { // NOLINT(*)
assign(init.begin(), init.end());
}
/**
* @brief constructor from vector
* @param init The vector
*/
List(const std::vector<T>& init) { // NOLINT(*)
assign(init.begin(), init.end());
}
/**
* @brief Constructs a container with n elements. Each element is a copy of
* val
* @param n The size of the container
* @param val The init value
*/
explicit List(size_t n, const T& val) {
auto tmp_obj = std::make_shared<ListObject>();
for (size_t i = 0; i < n; ++i) {
tmp_obj->data.push_back(val.obj_);
}
obj_ = std::move(tmp_obj);
}
/**
* @brief move assign operator
* @param other The source of assignment
* @return reference to self.
*/
List<T>& operator=(List<T>&& other) {
obj_ = std::move(other.obj_);
return *this;
}
/**
* @brief copy assign operator
* @param other The source of assignment
* @return reference to self.
*/
List<T>& operator=(const List<T>& other) {
obj_ = other.obj_;
return *this;
}
/**
* @brief reset the list to content from iterator.
* @param begin begin of iterator
* @param end end of iterator
* @tparam IterType The type of iterator
*/
template <typename IterType>
void assign(IterType begin, IterType end) {
auto n = std::make_shared<ListObject>();
for (IterType it = begin; it != end; ++it) {
n->data.push_back((*it).obj_);
}
obj_ = std::move(n);
}
/**
* @brief Read i-th element from list.
* @param i The index
* @return the i-th element.
*/
inline const T operator[](size_t i) const {
return T(static_cast<const ListObject*>(obj_.get())->data[i]);
}
/** @return The size of the list */
inline size_t size() const {
if (obj_.get() == nullptr) return 0;
return static_cast<const ListObject*>(obj_.get())->data.size();
}
/**
* @brief copy on write semantics
* Do nothing if current handle is the unique copy of the list.
* Otherwise make a new copy of the list to ensure the current handle
* hold a unique copy.
*
* @return Handle to the internal obj container(which ganrantees to be unique)
*/
inline ListObject* CopyOnWrite() {
if (obj_.get() == nullptr || !obj_.unique()) {
obj_ = std::make_shared<ListObject>(
*static_cast<const ListObject*>(obj_.get()));
}
return static_cast<ListObject*>(obj_.get());
}
/**
* @brief push a new item to the back of the list
* @param item The item to be pushed.
*/
inline void push_back(const T& item) {
ListObject* n = this->CopyOnWrite();
n->data.push_back(item.obj_);
}
/**
* @brief set i-th element of the list.
* @param i The index
* @param value The value to be setted.
*/
inline void Set(size_t i, const T& value) {
ListObject* n = this->CopyOnWrite();
n->data[i] = value.obj_;
}
/** @return whether list is empty */
inline bool empty() const { return size() == 0; }
/** @brief Copy the content to a vector */
inline std::vector<T> ToVector() const {
return std::vector<T>(begin(), end());
}
/** @brief specify container obj */
using ContainerType = ListObject;
struct Ptr2ObjectRef {
using ResultType = T;
static inline T convert(const std::shared_ptr<Object>& n) { return T(n); }
};
using iterator = IterAdapter<
Ptr2ObjectRef, std::vector<std::shared_ptr<Object> >::const_iterator>;
using reverse_iterator = IterAdapter<
Ptr2ObjectRef,
std::vector<std::shared_ptr<Object> >::const_reverse_iterator>;
/** @return begin iterator */
inline iterator begin() const {
return iterator(static_cast<const ListObject*>(obj_.get())->data.begin());
}
/** @return end iterator */
inline iterator end() const {
return iterator(static_cast<const ListObject*>(obj_.get())->data.end());
}
/** @return rbegin iterator */
inline reverse_iterator rbegin() const {
return reverse_iterator(
static_cast<const ListObject*>(obj_.get())->data.rbegin());
}
/** @return rend iterator */
inline reverse_iterator rend() const {
return reverse_iterator(
static_cast<const ListObject*>(obj_.get())->data.rend());
}
};
/**
* @brief Map container of ObjectRef->ObjectRef.
*
* Map implements copy on write semantics, which means map is mutable
* but copy will happen when list is referenced in more than two places.
*
* That is said when using this container for runtime arguments or return
* values, try use the constructor to create it at once (for example
* from an existing std::map).
*
* operator[] only provide const acces, use Set to mutate the content.
*
* @tparam K The key ObjectRef type.
* @tparam V The value ObjectRef type.
*
* @note The element type must subclass \c ObjectRef. Otherwise, the
* compiler would throw an error:
*
* <code>
* error: no type named 'type' in 'struct std::enable_if<false, void>'
* </code>
*
* Example:
*
* <code>
* // Map<std::string, int> map; // fails
* // Map<std::string, NDArray> map2; // fails
* Map<std::string, Value> map; // works
* map.Set("key1", Value(MakeValue(1))); // works
* map.Set("key2", Value(MakeValue(NDArray::Empty(shape, dtype, ctx)))); //
* works
* </code>
*/
template <
typename K, typename V,
typename = typename std::enable_if<
std::is_base_of<ObjectRef, K>::value ||
std::is_base_of<std::string, K>::value>::type,
typename =
typename std::enable_if<std::is_base_of<ObjectRef, V>::value>::type>
class Map : public ObjectRef {
public:
/**
* @brief default constructor
*/
Map() { obj_ = std::make_shared<MapObject>(); }
/**
* @brief move constructor
* @param other source
*/
Map(Map<K, V>&& other) { // NOLINT(*)
obj_ = std::move(other.obj_);
}
/**
* @brief copy constructor
* @param other source
*/
Map(const Map<K, V>& other) : ObjectRef(other.obj_) { // NOLINT(*)
}
/**
* @brief constructor from pointer
* @param n the container pointer
*/
explicit Map(std::shared_ptr<Object> n) : ObjectRef(n) {}
/**
* @brief constructor from iterator
* @param begin begin of iterator
* @param end end of iterator
* @tparam IterType The type of iterator
*/
template <typename IterType>
Map(IterType begin, IterType end) {
assign(begin, end);
}
/**
* @brief constructor from initializer list
* @param init The initalizer list
*/
Map(std::initializer_list<std::pair<K, V> > init) { // NOLINT(*)
assign(init.begin(), init.end());
}
/**
* @brief constructor from vector
* @param init The vector
*/
template <typename Hash, typename Equal>
Map(const std::unordered_map<K, V, Hash, Equal>& init) { // NOLINT(*)
assign(init.begin(), init.end());
}
/**
* @brief move assign operator
* @param other The source of assignment
* @return reference to self.
*/
Map<K, V>& operator=(Map<K, V>&& other) {
obj_ = std::move(other.obj_);
return *this;
}
/**
* @brief copy assign operator
* @param other The source of assignment
* @return reference to self.
*/
Map<K, V>& operator=(const Map<K, V>& other) {
obj_ = other.obj_;
return *this;
}
/**
* @brief reset the list to content from iterator.
* @param begin begin of iterator
* @param end end of iterator
* @tparam IterType The type of iterator
*/
template <typename IterType>
void assign(IterType begin, IterType end) {
auto n = std::shared_ptr<MapObject>();
for (IterType i = begin; i != end; ++i) {
n->data.emplace(std::make_pair(i->first.obj_, i->second.obj_));
}
obj_ = std::move(n);
}
/**
* @brief Read element from map.
* @param key The key
* @return the corresonding element.
*/
inline const V operator[](const K& key) const {
return V(static_cast<const MapObject*>(obj_.get())->data.at(key.obj_));
}
/**
* @brief Read element from map.
* @param key The key
* @return the corresonding element.
*/
inline const V at(const K& key) const {
return V(static_cast<const MapObject*>(obj_.get())->data.at(key.obj_));
}
/** @return The size of the list */
inline size_t size() const {
if (obj_.get() == nullptr) return 0;
return static_cast<const MapObject*>(obj_.get())->data.size();
}
/** @return The size of the list */
inline size_t count(const K& key) const {
if (obj_.get() == nullptr) return 0;
return static_cast<const MapObject*>(obj_.get())->data.count(key.obj_);
}
/**
* @brief copy on write semantics
* Do nothing if current handle is the unique copy of the list.
* Otherwise make a new copy of the list to ensure the current handle
* hold a unique copy.
*
* @return Handle to the internal obj container(which ganrantees to be unique)
*/
inline MapObject* CopyOnWrite() {
if (obj_.get() == nullptr || !obj_.unique()) {
obj_ = std::make_shared<MapObject>(
*static_cast<const MapObject*>(obj_.get()));
}
return static_cast<MapObject*>(obj_.get());
}
/**
* @brief set the Map.
* @param key The index key.
* @param value The value to be setted.
*/
inline void Set(const K& key, const V& value) {
MapObject* n = this->CopyOnWrite();
n->data[key.obj_] = value.obj_;
}
/** @return whether list is empty */
inline bool empty() const { return size() == 0; }
/** @brief specify container obj */
using ContainerType = MapObject;
struct Ptr2ObjectRef {
using ResultType = std::pair<K, V>;
static inline ResultType convert(
const std::pair<std::shared_ptr<Object>, std::shared_ptr<Object> >& n) {
return std::make_pair(K(n.first), V(n.second));
}
};
using iterator =
IterAdapter<Ptr2ObjectRef, MapObject::ContainerType::const_iterator>;
/** @return begin iterator */
inline iterator begin() const {
return iterator(static_cast<const MapObject*>(obj_.get())->data.begin());
}
/** @return end iterator */
inline iterator end() const {
return iterator(static_cast<const MapObject*>(obj_.get())->data.end());
}
/** @return begin iterator */
inline iterator find(const K& key) const {
return iterator(
static_cast<const MapObject*>(obj_.get())->data.find(key.obj_));
}
};
// specialize of string map
template <typename V, typename T1, typename T2>
class Map<std::string, V, T1, T2> : public ObjectRef {
public:
// for code reuse
Map() { obj_ = std::make_shared<StrMapObject>(); }
Map(Map<std::string, V>&& other) { // NOLINT(*)
obj_ = std::move(other.obj_);
}
Map(const Map<std::string, V>& other) : ObjectRef(other.obj_) { // NOLINT(*)
}
explicit Map(std::shared_ptr<Object> n) : ObjectRef(n) {}
template <typename IterType>
Map(IterType begin, IterType end) {
assign(begin, end);
}
Map(std::initializer_list<std::pair<std::string, V> > init) { // NOLINT(*)
assign(init.begin(), init.end());
}
template <typename Hash, typename Equal>
Map(const std::unordered_map<std::string, V, Hash, Equal>&
init) { // NOLINT(*)
assign(init.begin(), init.end());
}
Map<std::string, V>& operator=(Map<std::string, V>&& other) {
obj_ = std::move(other.obj_);
return *this;
}
Map<std::string, V>& operator=(const Map<std::string, V>& other) {
obj_ = other.obj_;
return *this;
}
template <typename IterType>
void assign(IterType begin, IterType end) {
auto n = std::make_shared<StrMapObject>();
for (IterType i = begin; i != end; ++i) {
n->data.emplace(std::make_pair(i->first, i->second.obj_));
}
obj_ = std::move(n);
}
inline const V operator[](const std::string& key) const {
return V(static_cast<const StrMapObject*>(obj_.get())->data.at(key));
}
inline const V at(const std::string& key) const {
return V(static_cast<const StrMapObject*>(obj_.get())->data.at(key));
}
inline size_t size() const {
if (obj_.get() == nullptr) return 0;
return static_cast<const StrMapObject*>(obj_.get())->data.size();
}
inline size_t count(const std::string& key) const {
if (obj_.get() == nullptr) return 0;
return static_cast<const StrMapObject*>(obj_.get())->data.count(key);
}
inline StrMapObject* CopyOnWrite() {
if (obj_.get() == nullptr || !obj_.unique()) {
obj_ = std::make_shared<MapObject>(
*static_cast<const MapObject*>(obj_.get()));
}
return static_cast<StrMapObject*>(obj_.get());
}
inline void Set(const std::string& key, const V& value) {
StrMapObject* n = this->CopyOnWrite();
n->data[key] = value.obj_;
}
inline bool empty() const { return size() == 0; }
using ContainerType = StrMapObject;
struct Ptr2ObjectRef {
using ResultType = std::pair<std::string, V>;
static inline ResultType convert(
const std::pair<std::string, std::shared_ptr<Object> >& n) {
return std::make_pair(n.first, V(n.second));
}
};
using iterator =
IterAdapter<Ptr2ObjectRef, StrMapObject::ContainerType::const_iterator>;
/** @return begin iterator */
inline iterator begin() const {
return iterator(static_cast<const StrMapObject*>(obj_.get())->data.begin());
}
/** @return end iterator */
inline iterator end() const {
return iterator(static_cast<const StrMapObject*>(obj_.get())->data.end());
}
/** @return begin iterator */
inline iterator find(const std::string& key) const {
return iterator(
static_cast<const StrMapObject*>(obj_.get())->data.find(key));
}
};
/**
* @brief Helper function to convert a List<Value> object to a vector.
* @tparam T element type
* @param list Input list object.
* @return std vector
*/
template <typename T>
inline std::vector<T> ListValueToVector(const List<Value>& list) {
std::vector<T> ret;
ret.reserve(list.size());
for (Value val : list)
// (BarclayII) apparently MSVC 2017 CL 19.10 had trouble parsing
// ret.push_back(val->data)
// So I kindly tell it how to properly parse it.
ret.push_back(val->data.operator T());
return ret;
}
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_CONTAINER_H_
include/dgl/runtime/device_api.h
+268
View File
@@ -0,0 +1,268 @@
/**
* Copyright (c) 2016 by Contributors
* @file dgl/runtime/device_api.h
* @brief Abstract device memory management API
*/
#ifndef DGL_RUNTIME_DEVICE_API_H_
#define DGL_RUNTIME_DEVICE_API_H_
#include <string>
#include "c_runtime_api.h"
#include "packed_func.h"
namespace dgl {
namespace runtime {
/**
* @brief the query type into GetAttr
*/
enum DeviceAttrKind : int {
kExist = 0,
kMaxThreadsPerBlock = 1,
kWarpSize = 2,
kMaxSharedMemoryPerBlock = 3,
kComputeVersion = 4,
kDeviceName = 5,
kMaxClockRate = 6,
kMultiProcessorCount = 7,
kMaxThreadDimensions = 8
};
/** @brief Number of bytes each allocation must align to */
constexpr int kAllocAlignment = 64;
/** @brief Number of bytes each allocation must align to in temporary allocation
*/
constexpr int kTempAllocaAlignment = 64;
/** @brief Maximum size that can be allocated on stack */
constexpr int kMaxStackAlloca = 1024;
/**
* @brief DGL Runtime Device API, abstracts the device
* specific interface for memory management.
*/
class DeviceAPI {
public:
/** @brief virtual destructor */
virtual ~DeviceAPI() {}
/**
* @brief Check whether the device is available.
*/
virtual bool IsAvailable() { return true; }
/**
* @brief Set the environment device id to ctx
* @param ctx The context to be set.
*/
virtual void SetDevice(DGLContext ctx) = 0;
/**
* @brief Get attribute of specified device.
* @param ctx The device context
* @param kind The result kind
* @param rv The return value.
* @sa DeviceAttrKind
*/
virtual void GetAttr(
DGLContext ctx, DeviceAttrKind kind, DGLRetValue* rv) = 0;
/**
* @brief Allocate a data space on device.
* @param ctx The device context to perform operation.
* @param nbytes The number of bytes in memory.
* @param alignment The alignment of the memory.
* @param type_hint The type of elements. Only needed by certain backends such
* as OpenGL, as nbytes & alignment are sufficient for most backends.
* @return The allocated device pointer.
*/
virtual void* AllocDataSpace(
DGLContext ctx, size_t nbytes, size_t alignment,
DGLDataType type_hint) = 0;
/**
* @brief Free a data space on device.
* @param ctx The device context to perform operation.
* @param ptr The data space.
*/
virtual void FreeDataSpace(DGLContext ctx, void* ptr) = 0;
/**
* @brief copy data from one place to another
* @param from The source array.
* @param from_offset The byte offeset in the from.
* @param to The target array.
* @param to_offset The byte offset in the to.
* @param num_bytes The size of the memory in bytes.
* @param ctx_from The source context.
* @param ctx_to The target context.
* @param type_hint The type of elements, only needed by certain backends,
* can be useful for cross device endian converison.
*/
virtual void CopyDataFromTo(
const void* from, size_t from_offset, void* to, size_t to_offset,
size_t num_bytes, DGLContext ctx_from, DGLContext ctx_to,
DGLDataType type_hint) = 0;
/**
* @brief copy data between device and CPU while recording the event.
* @param from The source array.
* @param from_offset The byte offeset in the from.
* @param to The target array.
* @param to_offset The byte offset in the to.
* @param num_bytes The size of the memory in bytes.
* @param ctx_from The source context.
* @param ctx_to The target context.
* @param type_hint The type of elements, only needed by certain backends,
* can be useful for cross device endian converison.
* @param pytorch_ctx The context pointer from PyTorch's CachingHostAllocator.
* @note This function only works when PyTorch CachingHostAllocator is
* available.
*/
virtual void RecordedCopyDataFromTo(
void* from, size_t from_offset, void* to, size_t to_offset,
size_t num_bytes, DGLContext ctx_from, DGLContext ctx_to,
DGLDataType type_hint, void* pytorch_ctx) = 0;
/**
* @brief Create a new stream of execution.
*
* @param ctx The context of allocation.
*/
DGL_DLL virtual DGLStreamHandle CreateStream(DGLContext ctx);
/**
* @brief Free a stream of execution
*
* @param ctx The context of the stream
* @param stream The pointer to be freed.
*/
DGL_DLL virtual void FreeStream(DGLContext ctx, DGLStreamHandle stream);
/**
* @brief Synchronize the stream
* @param ctx The context to perform operation.
* @param stream The stream to be sync.
*/
virtual void StreamSync(DGLContext ctx, DGLStreamHandle stream) = 0;
/**
* @brief Set the stream
* @param ctx The context to set stream.
* @param stream The stream to be set.
*/
virtual void SetStream(DGLContext ctx, DGLStreamHandle stream) {}
/**
* @brief Get the stream
*/
virtual DGLStreamHandle GetStream() const { return nullptr; }
/**
* @brief Synchronize 2 streams of execution.
*
* An event is created in event_src stream that the second then
* stream waits on. Neither event_src or event_dst need to be of
* the same device ID as the context, but they must be of the same
* device type.
*
* @param ctx The context of the streams.
* @param event_src The source stream to synchronize.
* @param event_dst The destination stream to synchronize.
*/
DGL_DLL virtual void SyncStreamFromTo(
DGLContext ctx, DGLStreamHandle event_src, DGLStreamHandle event_dst);
/**
* @brief Pin host memory using cudaHostRegister().
*
* @param ptr The host memory pointer to be pinned.
* @param nbytes The size to be pinned.
* @return false when pinning an empty tensor. true otherwise.
*/
DGL_DLL virtual bool PinData(void* ptr, size_t nbytes);
/**
* @brief Unpin host memory using cudaHostUnregister().
*
* @param ptr The host memory pointer to be unpinned.
*/
DGL_DLL virtual void UnpinData(void* ptr);
/**
* @brief Allocate the pinned memory using PyTorch CachingHostAllocator.
*
* @param nbytes The size to be pinned.
* @param ctx Pointer to the context pointer from PyTorch's
* CachingHostAllocator.
* @param deleter Pointer to the deleter function from PyTorch's
* CachingHostAllocator.
*/
DGL_DLL virtual void* AllocPinnedDataSpace(
size_t nbytes, void** ctx, void** deleter);
/**
* @brief 'Deallocate' the pinned memory from PyTorch CachingHostAllocator.
* @note It avoids unnecessary cudaFreeHost calls and puts the memory
* block into CachingHostAllocator's free list.
* @param deleter Pointer to the deleter function from PyTorch's
* CachingHostAllocator.
*/
DGL_DLL virtual void FreePinnedDataSpace(void** deleter);
/**
* @brief Check whether the memory is in pinned memory.
*/
DGL_DLL virtual bool IsPinned(const void* ptr) { return false; }
/**
* @brief Allocate temporal workspace for backend execution.
*
* \note We have the following assumption about backend temporal
* workspace allocation, and backend will optimize for such assumption:
*
* - Only a few allocation will happen, and space will be released after use.
* - The release order is usually in reverse order of allocate (stack style).
* - Repeative pattern of same allocations over different runs.
* - Workspace should not overlap between different threads(i.e. be
* threadlocal)
*
* @param ctx The context of allocation.
* @param nbytes The size to be allocated.
* @param type_hint The type of elements. Only needed by certain backends such
* as OpenGL, as nbytes is sufficient for most backends.
*/
DGL_DLL virtual void* AllocWorkspace(
DGLContext ctx, size_t nbytes, DGLDataType type_hint = {});
/**
* @brief Free temporal workspace in backend execution.
*
* @param ctx The context of allocation.
* @param ptr The pointer to be freed.
*/
DGL_DLL virtual void FreeWorkspace(DGLContext ctx, void* ptr);
/**
* @brief Get device API based on context.
* @param ctx The context
* @param allow_missing Whether allow missing
* @return The corresponding device API.
*/
DGL_DLL static DeviceAPI* Get(DGLContext ctx, bool allow_missing = false);
/**
* @brief Get device API based on device type.
* @param dev_type The device type
* @param allow_missing Whether allow missing
* @return The corresponding device API.
*/
DGL_DLL static DeviceAPI* Get(
DGLDeviceType dev_type, bool allow_missing = false);
};
/** @brief The device type bigger than this is RPC device */
constexpr int kRPCSessMask = 128;
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_DEVICE_API_H_
include/dgl/runtime/dlpack_convert.h
+84
View File
@@ -0,0 +1,84 @@
/**
* Copyright (c) 2022 by Contributors
* @file include/dgl/runtime/dlpack_convert.h
* @brief Conversion between NDArray and DLPack.
*/
#ifndef DGL_RUNTIME_DLPACK_CONVERT_H_
#define DGL_RUNTIME_DLPACK_CONVERT_H_
#include "c_runtime_api.h"
#include "ndarray.h"
struct DLManagedTensor;
namespace dgl {
namespace runtime {
struct DLPackConvert {
/**
* @brief Create a DGL NDArray from a DLPack tensor.
*
* This allows us to create a NDArray using the memory
* allocated by an external deep learning framework
* that is DLPack compatible.
*
* The memory is retained until the NDArray went out of scope.
* @param tensor The DLPack tensor to copy from.
* @return The created NDArray view.
*/
static NDArray FromDLPack(DLManagedTensor* tensor);
/**
* @brief Deleter for NDArray converted from DLPack.
*
* This is used from data which is passed from external
* DLPack(DLManagedTensor) that are not allocated inside of DGL. This enables
* us to create NDArray from memory allocated by other frameworks that are
* DLPack compatible
*/
static void DLPackDeleter(NDArray::Container* ptr);
/** @brief Convert a DGL NDArray to a DLPack tensor.
*
* @param from The DGL NDArray.
* @return A DLPack tensor.
*/
static DLManagedTensor* ToDLPack(const NDArray& from);
};
} // namespace runtime
} // namespace dgl
#ifdef __cplusplus
extern "C" {
#endif
/**
* @brief Delete (free) a DLManagedTensor's data.
* @param dltensor Pointer to the DLManagedTensor.
*/
DGL_DLL void DGLDLManagedTensorCallDeleter(DLManagedTensor* dltensor);
/**
* @brief Produce an array from the DLManagedTensor that shares data memory
* with the DLManagedTensor.
* @param from The source DLManagedTensor.
* @param out The output array handle.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLArrayFromDLPack(DLManagedTensor* from, DGLArrayHandle* out);
/**
* @brief Produce a DLMangedTensor from the array that shares data memory with
* the array.
* @param from The source array.
* @param out The DLManagedTensor handle.
* @return 0 when success, -1 when failure happens
*/
DGL_DLL int DGLArrayToDLPack(
DGLArrayHandle from, DLManagedTensor** out, int alignment = 0);
#ifdef __cplusplus
} // DGL_EXTERN_C
#endif
#endif // DGL_RUNTIME_DLPACK_CONVERT_H_
include/dgl/runtime/module.h
+175
View File
@@ -0,0 +1,175 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/runtime/module.h
* @brief Runtime container of the functions generated by DGL,
* This is used to support dynamically link, load and save
* functions from different convention under unified API.
*/
#ifndef DGL_RUNTIME_MODULE_H_
#define DGL_RUNTIME_MODULE_H_
#include <dmlc/io.h>
#include <memory>
#include <string>
#include <unordered_map>
#include <vector>
#include "c_runtime_api.h"
namespace dgl {
namespace runtime {
// The internal container of module.
class ModuleNode;
class PackedFunc;
/**
* @brief Module container of DGL.
*/
class Module {
public:
Module() {}
// constructor from container.
explicit Module(std::shared_ptr<ModuleNode> n) : node_(n) {}
/**
* @brief Get packed function from current module by name.
*
* @param name The name of the function.
* @param query_imports Whether also query dependency modules.
* @return The result function.
* This function will return PackedFunc(nullptr) if function do not exist.
* @note Implemented in packed_func.cc
*/
inline PackedFunc GetFunction(
const std::string& name, bool query_imports = false);
/** @return internal container */
inline ModuleNode* operator->();
/** @return internal container */
inline const ModuleNode* operator->() const;
// The following functions requires link with runtime.
/**
* @brief Import another module into this module.
* @param other The module to be imported.
*
* @note Cyclic dependency is not allowed among modules,
* An error will be thrown when cyclic dependency is detected.
*/
DGL_DLL void Import(Module other);
/**
* @brief Load a module from file.
* @param file_name The name of the host function module.
* @param format The format of the file.
* @note This function won't load the import relationship.
* Re-create import relationship by calling Import.
*/
DGL_DLL static Module LoadFromFile(
const std::string& file_name, const std::string& format = "");
private:
std::shared_ptr<ModuleNode> node_;
};
/**
* @brief Base node container of module.
* Do not create this directly, instead use Module.
*/
class ModuleNode {
public:
/** @brief virtual destructor */
virtual ~ModuleNode() {}
/** @return The module type key */
virtual const char* type_key() const = 0;
/**
* @brief Get a PackedFunc from module.
*
* The PackedFunc may not be fully initialized,
* there might still be first time running overhead when
* executing the function on certain devices.
* For benchmarking, use prepare to eliminate
*
* @param name the name of the function.
* @param sptr_to_self The shared_ptr that points to this module node.
*
* @return PackedFunc(nullptr) when it is not available.
*
* @note The function will always remain valid.
* If the function need resource from the module(e.g. late linking),
* it should capture sptr_to_self.
*/
virtual PackedFunc GetFunction(
const std::string& name,
const std::shared_ptr<ModuleNode>& sptr_to_self) = 0;
/**
* @brief Save the module to file.
* @param file_name The file to be saved to.
* @param format The format of the file.
*/
virtual void SaveToFile(
const std::string& file_name, const std::string& format);
/**
* @brief Save the module to binary stream.
* @param stream The binary stream to save to.
* @note It is recommended to implement this for device modules,
* but not necessarily host modules.
* We can use this to do AOT loading of bundled device functions.
*/
DGL_DLL virtual void SaveToBinary(dmlc::Stream* stream);
/**
* @brief Get the source code of module, when available.
* @param format Format of the source code, can be empty by default.
* @return Possible source code when available.
*/
DGL_DLL virtual std::string GetSource(const std::string& format = "");
/**
* @brief Get a function from current environment
* The environment includes all the imports as well as Global functions.
*
* @param name name of the function.
* @return The corresponding function.
*/
DGL_DLL const PackedFunc* GetFuncFromEnv(const std::string& name);
/** @return The module it imports from */
const std::vector<Module>& imports() const { return imports_; }
protected:
friend class Module;
/** @brief The modules this module depend on */
std::vector<Module> imports_;
private:
/** @brief Cache used by GetImport */
std::unordered_map<std::string, std::unique_ptr<PackedFunc> > import_cache_;
};
/** @brief namespace for constant symbols */
namespace symbol {
/** @brief Global variable to store module context. */
constexpr const char* dgl_module_ctx = "__dgl_module_ctx";
/** @brief Global variable to store device module blob */
constexpr const char* dgl_dev_mblob = "__dgl_dev_mblob";
/** @brief Number of bytes of device module blob. */
constexpr const char* dgl_dev_mblob_nbytes = "__dgl_dev_mblob_nbytes";
/** @brief global function to set device */
constexpr const char* dgl_set_device = "__dgl_set_device";
/** @brief Auxiliary counter to global barrier. */
constexpr const char* dgl_global_barrier_state = "__dgl_global_barrier_state";
/**
* @brief Prepare the global barrier before kernels that uses global barrier.
*/
constexpr const char* dgl_prepare_global_barrier =
"__dgl_prepare_global_barrier";
/** @brief Placeholder for the module's entry function. */
constexpr const char* dgl_module_main = "__dgl_main__";
} // namespace symbol
// implementations of inline functions.
inline ModuleNode* Module::operator->() { return node_.get(); }
inline const ModuleNode* Module::operator->() const { return node_.get(); }
} // namespace runtime
} // namespace dgl
#include "packed_func.h"
#endif // DGL_RUNTIME_MODULE_H_
include/dgl/runtime/ndarray.h
+890
View File
@@ -0,0 +1,890 @@
/**
* Copyright (c) 2017-2022 by Contributors
* @file dgl/runtime/ndarray.h
* @brief Abstract device memory management API
*/
#ifndef DGL_RUNTIME_NDARRAY_H_
#define DGL_RUNTIME_NDARRAY_H_
#include <atomic>
#include <memory>
#include <string>
#include <utility>
#include <vector>
#include "bfloat16.h"
#include "c_runtime_api.h"
#include "serializer.h"
#include "shared_mem.h"
#ifdef DGL_USE_CUDA
#include <cuda_runtime.h>
#define BF16_ENABLED (defined(CUDART_VERSION) && CUDART_VERSION >= 11000)
#include <cuda_fp16.h>
#if BF16_ENABLED
#include <cuda_bf16.h>
#endif // BF16_ENABLED
#endif // DGL_USE_CUDA
// forward declaration
inline std::ostream& operator<<(std::ostream& os, DGLDataType t);
namespace dgl {
/**
* @brief Type traits that converts a C type to a DGLDataType.
*
* Usage:
* DGLDataTypeTraits<int>::dtype == dtype
*/
template <typename T>
struct DGLDataTypeTraits {
static constexpr DGLDataType dtype{0, 0, 0}; // dummy
};
#define GEN_DGLDATATYPETRAITS_FOR(T, code, bits) \
template <> \
struct DGLDataTypeTraits<T> { \
static constexpr DGLDataType dtype{code, bits, 1}; \
}
GEN_DGLDATATYPETRAITS_FOR(int8_t, kDGLInt, 8);
GEN_DGLDATATYPETRAITS_FOR(uint8_t, kDGLUInt, 8);
GEN_DGLDATATYPETRAITS_FOR(int16_t, kDGLInt, 16);
GEN_DGLDATATYPETRAITS_FOR(int32_t, kDGLInt, 32);
GEN_DGLDATATYPETRAITS_FOR(int64_t, kDGLInt, 64);
// XXX(BarclayII) most DL frameworks do not support unsigned int and long
// arrays, so I'm just converting uints to signed DTypes.
GEN_DGLDATATYPETRAITS_FOR(uint32_t, kDGLInt, 32);
GEN_DGLDATATYPETRAITS_FOR(uint64_t, kDGLInt, 64);
#ifdef DGL_USE_CUDA
GEN_DGLDATATYPETRAITS_FOR(__half, kDGLFloat, 16);
#if BF16_ENABLED
GEN_DGLDATATYPETRAITS_FOR(__nv_bfloat16, kDGLBfloat, 16);
#endif // BF16_ENABLED
#endif // DGL_USE_CUDA
GEN_DGLDATATYPETRAITS_FOR(float, kDGLFloat, 32);
GEN_DGLDATATYPETRAITS_FOR(double, kDGLFloat, 64);
#undef GEN_DGLDATATYPETRAITS_FOR
namespace runtime {
/**
* @brief DLPack converter.
*/
struct DLPackConvert;
/**
* @brief Managed NDArray.
* The array is backed by reference counted blocks.
*/
class NDArray {
public:
// internal container type
struct Container;
/** @brief default constructor */
NDArray() {}
/**
* @brief cosntruct a NDArray that refers to data
* @param data The data this NDArray refers to
*/
explicit inline NDArray(Container* data);
/**
* @brief copy constructor
* @param other The value to be copied
*/
inline NDArray(const NDArray& other); // NOLINT(*)
/**
* @brief move constructor
* @param other The value to be moved
*/
NDArray(NDArray&& other) // NOLINT(*)
: data_(other.data_) {
other.data_ = nullptr;
}
/** @brief destructor */
~NDArray() { this->reset(); }
/**
* @brief Swap this array with another NDArray
* @param other The other NDArray
*/
void swap(NDArray& other) { // NOLINT(*)
std::swap(data_, other.data_);
}
/**
* @brief copy assignmemt
* @param other The value to be assigned.
* @return reference to self.
*/
NDArray& operator=(const NDArray& other) { // NOLINT(*)
// copy-and-swap idiom
NDArray(other).swap(*this); // NOLINT(*)
return *this;
}
/**
* @brief move assignmemt
* @param other The value to be assigned.
* @return reference to self.
*/
NDArray& operator=(NDArray&& other) { // NOLINT(*)
// copy-and-swap idiom
NDArray(std::move(other)).swap(*this); // NOLINT(*)
return *this;
}
/** @return If NDArray is defined */
bool defined() const { return data_ != nullptr; }
/** @return If both NDArray reference the same container */
bool same_as(const NDArray& other) const { return data_ == other.data_; }
/** @brief reset the content of NDArray to be nullptr */
inline void reset();
/**
* @return the reference counter
* @note this number is approximate in multi-threaded setting.
*/
inline int use_count() const;
/** @return Pointer to content of DGLArray */
inline const DGLArray* operator->() const;
/** @return True if the ndarray is contiguous. */
bool IsContiguous() const;
/** @return the data pointer with type. */
template <typename T>
inline T* Ptr() const {
if (!defined())
return nullptr;
else
return static_cast<T*>(operator->()->data);
}
/**
* @brief Copy data content from/into another array.
* @param other The source array to be copied from.
* @note The copy runs on the dgl internal stream if it involves a GPU
* context.
*/
inline void CopyFrom(DGLArray* other);
inline void CopyFrom(const NDArray& other);
inline void CopyTo(DGLArray* other) const;
inline void CopyTo(const NDArray& other) const;
/**
* @brief Copy the data to another context.
* @param ctx The target context.
* @return The array under another context.
*/
inline NDArray CopyTo(const DGLContext& ctx) const;
/**
* @brief Return a new array with a copy of the content.
*/
inline NDArray Clone() const;
/**
* @brief Return a copy of the current instance of NDArray in pinned
* (page-locked) memory.
* @note This is an out-of-place method, which utilizes PyTorch's
* CachingHostAllocator for allocating pinned memory and copying data
* from the current NDAarray. As a result, PyTorch is responsible for
* managing the lifecycle of the returned NDArray, including deciding
* when to flush the data for reuse or call cudaFreeHost. The current
* context must be kDGLCPU, otherwise, an error will be thrown.
*/
inline NDArray PinMemory();
/**
* @brief In-place method to pin the current array by calling PinContainer
* on the underlying NDArray:Container.
* @note This is an in-place method that flags the memory as page-locked by
* utilizing cudaHostRegister at the underlying level to pin the current
* instance of NDArray. The current context must be kDGLCPU, otherwise,
* an error will be thrown.
*/
inline void PinMemory_();
/**
* @brief In-place method to unpin the current array by calling UnpinContainer
* on the underlying NDArray:Container.
* @note This is an in-place method. Behavior depends on the current context,
* IsPinned: will be unpinned;
* others: directly return.
*/
inline void UnpinMemory_();
/**
* @brief Check if the array is pinned.
*/
inline bool IsPinned() const;
/**
* @brief Record streams that are using the underlying tensor.
* @param stream The stream that is using the underlying tensor.
*/
inline void RecordStream(DGLStreamHandle stream) const;
/**
* @brief Load NDArray from stream
* @param stream The input data stream
* @return Whether load is successful
*/
bool Load(dmlc::Stream* stream);
/**
* @brief Save NDArray to stream
* @param stream The output data stream
*/
void Save(dmlc::Stream* stream) const;
/**
* @brief Create a NDArray that shares the data memory with the current one.
* @param shape The shape of the new array.
* @param dtype The data type of the new array.
* @param offset The offset (in bytes) of the starting pointer.
* @note The memory size of new array must be smaller than the current one.
*/
DGL_DLL NDArray
CreateView(std::vector<int64_t> shape, DGLDataType dtype, int64_t offset = 0);
/**
* @brief Create an empty NDArray.
* @param shape The shape of the new array.
* @param dtype The data type of the new array.
* @param ctx The context of the array.
* @return The created Array
*/
DGL_DLL static NDArray Empty(
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
/**
* @brief Create an empty NDArray in pinned memory.
* @param shape The shape of the new array.
* @param dtype The data type of the new array.
* @param ctx The context of the array.
* @return The created array.
*/
DGL_DLL static NDArray PinnedEmpty(
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
/**
* @brief Create an empty NDArray with shared memory.
* @param name The name of shared memory.
* @param shape The shape of the new array.
* @param dtype The data type of the new array.
* @param ctx The context of the array.
* @param is_create whether to create shared memory.
* @return The created Array
*/
DGL_DLL static NDArray EmptyShared(
const std::string& name, std::vector<int64_t> shape, DGLDataType dtype,
DGLContext ctx, bool is_create);
/**
* @brief Get the size of the array in the number of bytes.
*/
size_t GetSize() const;
/**
* @brief Get the number of elements in this array.
*/
int64_t NumElements() const;
/**
* @brief Create a NDArray by copying from std::vector.
* @tparam T Type of vector data. Determines the dtype of returned array.
*/
template <typename T>
DGL_DLL static NDArray FromVector(
const std::vector<T>& vec, DGLContext ctx = DGLContext{kDGLCPU, 0});
/**
* @brief Create a NDArray from a raw pointer.
*/
DGL_DLL static NDArray CreateFromRaw(
const std::vector<int64_t>& shape, DGLDataType dtype, DGLContext ctx,
void* raw, bool auto_free);
/**
* @brief Create a std::vector from a 1D NDArray.
* @tparam T Type of vector data.
* @note Type casting is NOT performed. The caller has to make sure that the
* vector type matches the dtype of NDArray.
*/
template <typename T>
std::vector<T> ToVector() const;
std::shared_ptr<SharedMemory> GetSharedMem() const;
/**
* @brief Function to copy data from one array to another.
* @param from The source array.
* @param to The target array.
* @param (optional) stream The stream used in copy.
*/
DGL_DLL static void CopyFromTo(DGLArray* from, DGLArray* to);
DGL_DLL static void CopyFromTo(
DGLArray* from, DGLArray* to, DGLStreamHandle stream);
/**
* @brief Function to copy data between device and CPU while recording the
* event.
* @param from The source array.
* @param to The target array.
* @param pytorch_ctx The context pointer from PyTorch's CachingHostAllocator.
* @note This function fuses data-copy and event recording to ensure
* CachingHostAllocator works properly.
*/
DGL_DLL static void RecordedCopyFromTo(
DGLArray* from, DGLArray* to, void* pytorch_ctx);
/**
* @brief Function to pin the DGLArray of a Container.
* @param ptr The container to be pinned.
* @note Data of the given array will be pinned inplace.
* Behavior depends on the current context,
* kDGLCPU: will be pinned;
* IsPinned: directly return;
* kDGLCUDA: invalid, will throw an error.
*/
DGL_DLL static void PinContainer(Container* ptr);
/**
* @brief Function to unpin the DGLArray of a Container.
* @param ptr The container to be unpinned.
* @note Data of the given array will be unpinned inplace.
* Behavior depends on the current context,
* IsPinned: will be unpinned;
* others: directly return.
*/
DGL_DLL static void UnpinContainer(Container* ptr);
/**
* @brief Function check if the DGLArray of a Container is pinned.
* @param ptr The container to be checked.
* @return true if pinned.
*/
DGL_DLL static bool IsContainerPinned(Container* ptr);
/**
* @brief Record streams that are using this tensor.
* @param ptr Pointer of the tensor to be recorded.
* @param stream The stream that is using this tensor.
*/
DGL_DLL static void RecordStream(DGLArray* tensor, DGLStreamHandle stream);
// internal namespace
struct Internal {
// Default deleter for the container
static void DefaultDeleter(NDArray::Container* ptr);
// Local create function which allocates tensor metadata
// but does not allocate space for the data.
static NDArray Create(
std::vector<int64_t> shape, DGLDataType dtype, DGLContext ctx);
// Implementation of API function
static DGLArray* MoveAsDGLArray(NDArray arr);
};
private:
/** @brief Internal Data content */
Container* data_{nullptr};
// enable internal functions
friend struct Internal;
friend struct DLPackConvert;
friend class DGLRetValue;
friend class DGLArgsSetter;
};
/**
* @brief Save a DGLArray to stream
* @param strm The outpu stream
* @param tensor The tensor to be saved.
*/
inline bool SaveDGLArray(dmlc::Stream* strm, const DGLArray* tensor);
/**
* @brief Reference counted Container object used to back NDArray.
*
* This object is DGLArray compatible:
* the pointer to the NDArrayContainer can be directly
* interpreted as a DGLArray*
*
* @note: do not use this function directly, use NDArray.
*/
struct NDArray::Container {
public:
/** NOTE: the first part of this structure is the same as
* DLManagedTensor, note that, however, the deleter
* is only called when the reference counter goes to 0
*/
/**
* @brief Tensor structure.
* @note it is important that the first field is DGLArray
* So that this data structure is DGLArray compatible.
* The head ptr of this struct can be viewed as DGLArray*.
*/
DGLArray dl_tensor;
/**
* @brief addtional context, reserved for recycling
* @note We can attach additional content here
* which the current container depend on
* (e.g. reference to original memory when creating views).
*/
void* manager_ctx{nullptr};
/**
* @brief Customized deleter
*
* @note The customized deleter is helpful to enable
* different ways of memory allocator that are not
* currently defined by the system.
*/
void (*deleter)(Container* self) = nullptr;
/** @brief default constructor */
Container() {
dl_tensor.data = nullptr;
dl_tensor.ndim = 0;
dl_tensor.shape = nullptr;
dl_tensor.strides = nullptr;
dl_tensor.byte_offset = 0;
}
/** @brief pointer to shared memory */
std::shared_ptr<SharedMemory> mem;
/** @brief developer function, increases reference counter */
void IncRef() { ref_counter_.fetch_add(1, std::memory_order_relaxed); }
/** @brief developer function, decrease reference counter */
void DecRef() {
if (ref_counter_.fetch_sub(1, std::memory_order_release) == 1) {
std::atomic_thread_fence(std::memory_order_acquire);
if (this->deleter != nullptr) {
(*this->deleter)(this);
}
}
}
private:
friend struct DLPackConvert;
friend class NDArray;
friend class RPCWrappedFunc;
/**
* @brief The shape container,
* can be used for shape data.
*/
std::vector<int64_t> shape_;
/**
* @brief The stride container,
* can be used for stride data.
*/
std::vector<int64_t> stride_;
/** @brief The internal array object */
std::atomic<int> ref_counter_{0};
/** @brief Whether underlying dl_tensor is pinned by DGL. */
bool pinned_by_dgl_{false};
/** @brief Whether underlying dl_tensor is pinned by PyTorch
* (CachingHostAllocator). */
bool pinned_by_pytorch_{false};
/** @brief The PyTorch storage ctx ptr if pinned_by_pytorch_ = True. */
void* pytorch_ctx_{nullptr};
/** @brief Pointer to the corresp. PyTorch deleter if pinned_by_pytorch_ =
* True.
*/
void* pytorch_raw_deleter_{nullptr};
};
// implementations of inline functions
// the usages of functions are documented in place.
inline NDArray::NDArray(Container* data) : data_(data) {
if (data_) data_->IncRef();
}
inline NDArray::NDArray(const NDArray& other) : data_(other.data_) {
if (data_) data_->IncRef();
}
inline void NDArray::reset() {
if (data_) {
data_->DecRef();
data_ = nullptr;
}
}
inline void NDArray::CopyFrom(DGLArray* other) {
CHECK(data_ != nullptr);
CopyFromTo(other, &(data_->dl_tensor));
}
inline void NDArray::CopyFrom(const NDArray& other) {
CHECK(other.data_ != nullptr);
// Copy between two devices
if (data_->dl_tensor.ctx.device_type !=
other.data_->dl_tensor.ctx.device_type) {
CHECK(data_ != nullptr);
auto to_ctx_type = data_->dl_tensor.ctx.device_type;
auto cpu_data = (to_ctx_type == kDGLCPU ? data_ : other.data_);
// Pinned by PyTorch
if (cpu_data->pinned_by_pytorch_) {
// To ensure correct behavior, the event must be recorded after
// cudaMemcpyAsync as long as the memory is pinned by PyTorch.
void* pytorch_ctx = cpu_data->pytorch_ctx_;
RecordedCopyFromTo(
&(other.data_->dl_tensor), &(data_->dl_tensor), pytorch_ctx);
return;
}
}
CopyFrom(&(other.data_->dl_tensor));
}
inline void NDArray::CopyTo(DGLArray* other) const {
CHECK(data_ != nullptr);
CopyFromTo(&(data_->dl_tensor), other);
}
inline void NDArray::CopyTo(const NDArray& other) const {
CHECK(other.data_ != nullptr);
// copy between two devices
if (data_->dl_tensor.ctx.device_type !=
other.data_->dl_tensor.ctx.device_type) {
CHECK(data_ != nullptr);
auto from_ctx_type = data_->dl_tensor.ctx.device_type;
auto cpu_data = (from_ctx_type == kDGLCPU ? data_ : other.data_);
// pinned by PyTorch
if (cpu_data->pinned_by_pytorch_) {
// To ensure correct behavior, the event must be recorded after
// cudaMemcpyAsync as long as the memory is pinned by PyTorch.
void* pytorch_ctx = cpu_data->pytorch_ctx_;
RecordedCopyFromTo(
&(data_->dl_tensor), &(other.data_->dl_tensor), pytorch_ctx);
return;
}
}
CopyTo(&(other.data_->dl_tensor));
}
inline NDArray NDArray::CopyTo(const DGLContext& ctx) const {
CHECK(data_ != nullptr);
const DGLArray* array = operator->();
NDArray ret = Empty(
std::vector<int64_t>(array->shape, array->shape + array->ndim),
array->dtype, ctx);
this->CopyTo(ret);
return ret;
}
inline NDArray NDArray::Clone() const {
CHECK(data_ != nullptr);
const DGLArray* array = operator->();
return this->CopyTo(array->ctx);
}
inline NDArray NDArray::PinMemory() {
CHECK(data_ != nullptr);
const DGLArray* array = operator->();
auto ctx = array->ctx;
NDArray ret = PinnedEmpty(
std::vector<int64_t>(array->shape, array->shape + array->ndim),
array->dtype, ctx);
this->CopyTo(ret);
return ret;
}
inline void NDArray::PinMemory_() {
CHECK(data_ != nullptr);
PinContainer(data_);
}
inline void NDArray::UnpinMemory_() {
CHECK(data_ != nullptr);
UnpinContainer(data_);
}
inline bool NDArray::IsPinned() const {
CHECK(data_ != nullptr);
return IsContainerPinned(data_);
}
inline void NDArray::RecordStream(DGLStreamHandle stream) const {
CHECK(data_ != nullptr);
RecordStream(&(data_->dl_tensor), stream);
}
inline int NDArray::use_count() const {
if (data_ == nullptr) return 0;
return data_->ref_counter_.load(std::memory_order_relaxed);
}
inline const DGLArray* NDArray::operator->() const {
return &(data_->dl_tensor);
}
/** @brief Magic number for NDArray file */
constexpr uint64_t kDGLNDArrayMagic = 0xDD5E40F096B4A13F;
inline bool SaveDGLArray(dmlc::Stream* strm, DGLArray* tensor) {
uint64_t header = kDGLNDArrayMagic, reserved = 0;
strm->Write(header);
strm->Write(reserved);
// Always save data as CPU context
//
// Parameters that get serialized should be in CPU by default.
// So even the array's context is GPU, it will be stored as CPU array.
// This is used to prevent case when another user loads the parameters
// back on machine that do not have GPU or related context.
//
// We can always do array.CopyTo(target_ctx) to get a corresponding
// array in the target context.
DGLContext cpu_ctx;
cpu_ctx.device_type = kDGLCPU;
cpu_ctx.device_id = 0;
strm->Write(cpu_ctx);
strm->Write(tensor->ndim);
strm->Write(tensor->dtype);
int ndim = tensor->ndim;
strm->WriteArray(tensor->shape, ndim);
int type_bytes = tensor->dtype.bits / 8;
int64_t num_elems = 1;
for (int i = 0; i < ndim; ++i) {
num_elems *= tensor->shape[i];
}
int64_t data_byte_size = type_bytes * num_elems;
strm->Write(data_byte_size);
if (DMLC_IO_NO_ENDIAN_SWAP && tensor->ctx.device_type == kDGLCPU &&
tensor->strides == nullptr && tensor->byte_offset == 0) {
// quick path
strm->Write(tensor->data, data_byte_size);
} else {
std::vector<uint8_t> bytes(data_byte_size);
CHECK_EQ(
DGLArrayCopyToBytes(tensor, dmlc::BeginPtr(bytes), data_byte_size), 0)
<< DGLGetLastError();
if (!DMLC_IO_NO_ENDIAN_SWAP) {
dmlc::ByteSwap(dmlc::BeginPtr(bytes), type_bytes, num_elems);
}
strm->Write(dmlc::BeginPtr(bytes), data_byte_size);
}
return true;
}
/**
* @brief Convert type code to its name
* @param type_code The type code .
* @return The name of type code.
*/
inline const char* TypeCode2Str(int type_code) {
switch (type_code) {
case kDGLInt:
return "int";
case kDGLUInt:
return "uint";
case kDGLFloat:
return "float";
case kStr:
return "str";
case kBytes:
return "bytes";
case kHandle:
return "handle";
case kNull:
return "NULL";
case kObjectHandle:
return "ObjectHandle";
case kArrayHandle:
return "ArrayHandle";
case kDGLDataType:
return "DGLDataType";
case kDGLContext:
return "DGLContext";
case kFuncHandle:
return "FunctionHandle";
case kModuleHandle:
return "ModuleHandle";
case kNDArrayContainer:
return "NDArrayContainer";
default:
LOG(FATAL) << "unknown type_code=" << static_cast<int>(type_code);
return "";
}
}
/**
* @brief Convert device type code to its name
* @param device_type The device type code.
* @return The name of the device.
*/
inline const char* DeviceTypeCode2Str(DGLDeviceType device_type) {
switch (device_type) {
case kDGLCPU:
return "cpu";
case kDGLCUDA:
return "cuda";
default:
LOG(FATAL) << "Unsupported device type code="
<< static_cast<int>(device_type);
return "";
}
}
/**
* @brief convert a string to DGL type.
* @param s The string to be converted.
* @return The corresponding dgl type.
*/
inline DGLDataType String2DGLDataType(std::string s) {
DGLDataType t;
t.bits = 32;
t.lanes = 1;
const char* scan;
if (s.substr(0, 3) == "int") {
t.code = kDGLInt;
scan = s.c_str() + 3;
} else if (s.substr(0, 4) == "uint") {
t.code = kDGLUInt;
scan = s.c_str() + 4;
} else if (s.substr(0, 5) == "float") {
t.code = kDGLFloat;
scan = s.c_str() + 5;
} else if (s.substr(0, 6) == "handle") {
t.code = kHandle;
t.bits = 64; // handle uses 64 bit by default.
scan = s.c_str() + 6;
} else {
scan = s.c_str();
LOG(FATAL) << "unknown type " << s;
}
char* xdelim; // emulate sscanf("%ux%u", bits, lanes)
uint8_t bits = static_cast<uint8_t>(strtoul(scan, &xdelim, 10));
if (bits != 0) t.bits = bits;
if (*xdelim == 'x') {
t.lanes = static_cast<uint16_t>(strtoul(xdelim + 1, nullptr, 10));
}
return t;
}
/**
* @brief convert a DGL type to string.
* @param t The type to be converted.
* @return The corresponding dgl type in string.
*/
inline std::string DGLDataType2String(DGLDataType t) {
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
std::ostringstream os;
os << t;
return os.str();
#else
std::string repr = "";
repr += TypeCode2Str(t.code);
if (t.code == kHandle) return repr;
repr += std::to_string(static_cast<int>(t.bits));
if (t.lanes != 1) {
repr += "x" + std::to_string(static_cast<int>(t.lanes));
}
return repr;
#endif
}
// macro to check type code.
#define DGL_CHECK_TYPE_CODE(CODE, T) \
CHECK_EQ(CODE, T) << " expected " << TypeCode2Str(T) << " but get " \
<< TypeCode2Str(CODE)
} // namespace runtime
} // namespace dgl
namespace dmlc {
DMLC_DECLARE_TRAITS(has_saveload, dgl::runtime::NDArray, true);
} // namespace dmlc
///////////////// Operator overloading for NDArray /////////////////
dgl::runtime::NDArray operator+(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator-(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator*(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator/(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator%(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator+(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator-(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator*(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator/(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator%(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator+(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator-(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator*(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator/(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator%(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator-(const dgl::runtime::NDArray& array);
dgl::runtime::NDArray operator>(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator<(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator>=(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator<=(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator==(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator!=(
const dgl::runtime::NDArray& a1, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator>(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator<(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator>=(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator<=(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator==(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator!=(const dgl::runtime::NDArray& a1, int64_t rhs);
dgl::runtime::NDArray operator>(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator<(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator>=(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator<=(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator==(int64_t lhs, const dgl::runtime::NDArray& a2);
dgl::runtime::NDArray operator!=(int64_t lhs, const dgl::runtime::NDArray& a2);
std::ostream& operator<<(std::ostream& os, dgl::runtime::NDArray array);
///////////////// Operator overloading for DGLDataType /////////////////
/** @brief Check whether two data types are the same.*/
inline bool operator==(const DGLDataType& ty1, const DGLDataType& ty2) {
return ty1.code == ty2.code && ty1.bits == ty2.bits && ty1.lanes == ty2.lanes;
}
/** @brief Check whether two data types are different.*/
inline bool operator!=(const DGLDataType& ty1, const DGLDataType& ty2) {
return !(ty1 == ty2);
}
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
inline std::ostream& operator<<(std::ostream& os, DGLDataType t) {
os << dgl::runtime::TypeCode2Str(t.code);
if (t.code == kHandle) return os;
os << static_cast<int>(t.bits);
if (t.lanes != 1) {
os << 'x' << static_cast<int>(t.lanes);
}
return os;
}
#endif
///////////////// Operator overloading for DGLContext /////////////////
/** @brief Check whether two device contexts are the same.*/
inline bool operator==(const DGLContext& ctx1, const DGLContext& ctx2) {
return ctx1.device_type == ctx2.device_type &&
ctx1.device_id == ctx2.device_id;
}
/** @brief Check whether two device contexts are different.*/
inline bool operator!=(const DGLContext& ctx1, const DGLContext& ctx2) {
return !(ctx1 == ctx2);
}
#ifndef _LIBCPP_SGX_NO_IOSTREAMS
inline std::ostream& operator<<(std::ostream& os, const DGLContext& ctx) {
return os << dgl::runtime::DeviceTypeCode2Str(ctx.device_type) << ":"
<< ctx.device_id;
}
#endif
#endif // DGL_RUNTIME_NDARRAY_H_
include/dgl/runtime/object.h
+328
View File
@@ -0,0 +1,328 @@
/**
* Copyright (c) 2019 by Contributors
* @file runtime/object.h
* @brief Defines the Object data structures.
*/
#ifndef DGL_RUNTIME_OBJECT_H_
#define DGL_RUNTIME_OBJECT_H_
#include <dmlc/logging.h>
#include <memory>
#include <string>
#include <type_traits>
#include <vector>
namespace dgl {
namespace runtime {
// forward declaration
class Object;
class ObjectRef;
class NDArray;
/**
* @brief Visitor class to each object attribute.
* The content is going to be called for each field.
*/
class AttrVisitor {
public:
//! \cond Doxygen_Suppress
virtual void Visit(const char* key, double* value) = 0;
virtual void Visit(const char* key, int64_t* value) = 0;
virtual void Visit(const char* key, uint64_t* value) = 0;
virtual void Visit(const char* key, int* value) = 0;
virtual void Visit(const char* key, bool* value) = 0;
virtual void Visit(const char* key, std::string* value) = 0;
virtual void Visit(const char* key, ObjectRef* value) = 0;
virtual void Visit(const char* key, NDArray* value) = 0;
template <
typename ENum,
typename = typename std::enable_if<std::is_enum<ENum>::value>::type>
void Visit(const char* key, ENum* ptr) {
static_assert(
std::is_same<int, typename std::underlying_type<ENum>::type>::value,
"declare enum to be enum int to use visitor");
this->Visit(key, reinterpret_cast<int*>(ptr));
}
//! \endcond
};
/**
* @brief base class of object container.
* All object's internal is stored as std::shared_ptr<Object>
*/
class Object {
public:
/** @brief virtual destructor */
virtual ~Object() {}
/** @return The unique type key of the object */
virtual const char* type_key() const = 0;
/**
* @brief Apply visitor to each field of the Object
* Visitor could mutate the content of the object.
* override if Object contains attribute fields.
* @param visitor The visitor
*/
virtual void VisitAttrs(AttrVisitor* visitor) {}
/** @return the type index of the object */
virtual uint32_t type_index() const = 0;
/**
* @brief Whether this object derives from object with type_index=tid.
* Implemented by DGL_DECLARE_OBJECT_TYPE_INFO
*
* @param tid The type index.
* @return the check result.
*/
virtual bool _DerivedFrom(uint32_t tid) const;
/**
* @brief get a runtime unique type index given a type key
* @param type_key Type key of a type.
* @return the corresponding type index.
*/
static uint32_t TypeKey2Index(const char* type_key);
/**
* @brief get type key from type index.
* @param index The type index
* @return the corresponding type key.
*/
static const char* TypeIndex2Key(uint32_t index);
/**
* @return whether the type is derived from
*/
template <typename T>
inline bool derived_from() const;
/**
* @return whether the object is of type T
* @tparam The type to be checked.
*/
template <typename T>
inline bool is_type() const;
// object ref can see this
friend class ObjectRef;
static constexpr const char* _type_key = "Object";
};
/** @brief base class of all reference object */
class ObjectRef {
public:
/** @brief type indicate the container type */
using ContainerType = Object;
/**
* @brief Comparator
*
* Compare with the two are referencing to the same object (compare by
* address).
*
* @param other Another object ref.
* @return the compare result.
* @sa same_as
*/
inline bool operator==(const ObjectRef& other) const;
/**
* @brief Comparator
*
* Compare with the two are referencing to the same object (compare by
* address).
*
* @param other Another object ref.
* @return the compare result.
*/
inline bool same_as(const ObjectRef& other) const;
/**
* @brief Comparator
*
* The operator overload allows ObjectRef be used in std::map.
*
* @param other Another object ref.
* @return the compare result.
*/
inline bool operator<(const ObjectRef& other) const;
/**
* @brief Comparator
* @param other Another object ref.
* @return the compare result.
* @sa same_as
*/
inline bool operator!=(const ObjectRef& other) const;
/** @return the hash function for ObjectRef */
inline size_t hash() const;
/** @return whether the expression is null */
inline bool defined() const;
/** @return the internal type index of Object */
inline uint32_t type_index() const;
/** @return the internal object pointer */
inline const Object* get() const;
/** @return the internal object pointer */
inline const Object* operator->() const;
/**
* @brief Downcast this object to its actual type.
* This returns nullptr if the object is not of the requested type.
* Example usage:
*
* if (const Banana *banana = obj->as<Banana>()) {
* // This is a Banana!
* }
* @tparam T the target type, must be subtype of Object
*/
template <typename T>
inline const T* as() const;
/** @brief default constructor */
ObjectRef() = default;
explicit ObjectRef(std::shared_ptr<Object> obj) : obj_(obj) {}
/** @brief the internal object, do not touch */
std::shared_ptr<Object> obj_;
};
/**
* @brief helper macro to declare type information in a base object.
*
* This is macro should be used in abstract base class definition
* because it does not define type_key and type_index.
*/
#define DGL_DECLARE_BASE_OBJECT_INFO(TypeName, Parent) \
const bool _DerivedFrom(uint32_t tid) const override { \
static uint32_t tidx = TypeKey2Index(TypeName::_type_key); \
if (tidx == tid) return true; \
return Parent::_DerivedFrom(tid); \
}
/**
* @brief helper macro to declare type information in a terminal class
*
* This is macro should be used in terminal class definition.
*
* For example:
*
* // This class is an abstract class and cannot create instances
* class SomeBaseClass : public Object {
* public:
* static constexpr const char* _type_key = "some_base";
* DGL_DECLARE_BASE_OBJECT_INFO(SomeBaseClass, Object);
* };
*
* // Child class that allows instantiation
* class SomeChildClass : public SomeBaseClass {
* public:
* static constexpr const char* _type_key = "some_child";
* DGL_DECLARE_OBJECT_TYPE_INFO(SomeChildClass, SomeBaseClass);
* };
*/
#define DGL_DECLARE_OBJECT_TYPE_INFO(TypeName, Parent) \
const char* type_key() const final { return TypeName::_type_key; } \
uint32_t type_index() const final { \
static uint32_t tidx = TypeKey2Index(TypeName::_type_key); \
return tidx; \
} \
bool _DerivedFrom(uint32_t tid) const final { \
static uint32_t tidx = TypeKey2Index(TypeName::_type_key); \
if (tidx == tid) return true; \
return Parent::_DerivedFrom(tid); \
}
/** @brief Macro to generate common object reference class method definition */
#define DGL_DEFINE_OBJECT_REF_METHODS(TypeName, BaseTypeName, ObjectName) \
TypeName() {} \
explicit TypeName(std::shared_ptr<runtime::Object> obj) \
: BaseTypeName(obj) {} \
const ObjectName* operator->() const { \
return static_cast<const ObjectName*>(obj_.get()); \
} \
ObjectName* operator->() { return static_cast<ObjectName*>(obj_.get()); } \
std::shared_ptr<ObjectName> sptr() const { \
return CHECK_NOTNULL(std::dynamic_pointer_cast<ObjectName>(obj_)); \
} \
operator bool() const { return this->defined(); } \
using ContainerType = ObjectName
/** @brief Macro to generate object reference class definition */
#define DGL_DEFINE_OBJECT_REF(TypeName, ObjectName) \
class TypeName : public ::dgl::runtime::ObjectRef { \
public: \
DGL_DEFINE_OBJECT_REF_METHODS( \
TypeName, ::dgl::runtime::ObjectRef, ObjectName); \
}
// implementations of inline functions after this
template <typename T>
inline bool Object::is_type() const {
// use static field so query only happens once.
static uint32_t type_id = Object::TypeKey2Index(T::_type_key);
return type_id == this->type_index();
}
template <typename T>
inline bool Object::derived_from() const {
// use static field so query only happens once.
static uint32_t type_id = Object::TypeKey2Index(T::_type_key);
return this->_DerivedFrom(type_id);
}
inline const Object* ObjectRef::get() const { return obj_.get(); }
inline const Object* ObjectRef::operator->() const { return obj_.get(); }
inline bool ObjectRef::defined() const { return obj_.get() != nullptr; }
inline bool ObjectRef::operator==(const ObjectRef& other) const {
return obj_.get() == other.obj_.get();
}
inline bool ObjectRef::same_as(const ObjectRef& other) const {
return obj_.get() == other.obj_.get();
}
inline bool ObjectRef::operator<(const ObjectRef& other) const {
return obj_.get() < other.obj_.get();
}
inline bool ObjectRef::operator!=(const ObjectRef& other) const {
return obj_.get() != other.obj_.get();
}
inline size_t ObjectRef::hash() const {
return std::hash<Object*>()(obj_.get());
}
inline uint32_t ObjectRef::type_index() const {
CHECK(obj_.get() != nullptr) << "null type";
return get()->type_index();
}
template <typename T>
inline const T* ObjectRef::as() const {
const Object* ptr = get();
if (ptr && ptr->is_type<T>()) {
return static_cast<const T*>(ptr);
}
return nullptr;
}
/** @brief The hash function for nodes */
struct ObjectHash {
size_t operator()(const ObjectRef& a) const { return a.hash(); }
};
/** @brief The equal comparator for nodes */
struct ObjectEqual {
bool operator()(const ObjectRef& a, const ObjectRef& b) const {
return a.get() == b.get();
}
};
} // namespace runtime
} // namespace dgl
namespace std {
template <>
struct hash<::dgl::runtime::ObjectRef> {
std::size_t operator()(const ::dgl::runtime::ObjectRef& k) const {
return k.hash();
}
};
} // namespace std
#endif // DGL_RUNTIME_OBJECT_H_
include/dgl/runtime/packed_func.h
+1088
View File
File diff suppressed because it is too large Load Diff
include/dgl/runtime/parallel_for.h
+182
View File
@@ -0,0 +1,182 @@
/**
* Copyright (c) 2021 by Contributors
* @file runtime/container.h
* @brief Defines the container object data structures.
*/
#ifndef DGL_RUNTIME_PARALLEL_FOR_H_
#define DGL_RUNTIME_PARALLEL_FOR_H_
#include <dgl/env_variable.h>
#include <dmlc/omp.h>
#include <algorithm>
#include <atomic>
#include <cstdlib>
#include <exception>
#include <string>
#include <utility>
#include <vector>
namespace {
int64_t divup(int64_t x, int64_t y) { return (x + y - 1) / y; }
} // namespace
namespace dgl {
namespace runtime {
namespace {
struct DefaultGrainSizeT {
size_t grain_size;
DefaultGrainSizeT() : DefaultGrainSizeT(1) {}
explicit DefaultGrainSizeT(size_t default_grain_size) {
auto var = dgl::kDGLParallelForGrainSize;
if (var) {
grain_size = std::stoul(var);
} else {
grain_size = default_grain_size;
}
}
size_t operator()() { return grain_size; }
};
} // namespace
inline size_t compute_num_threads(size_t begin, size_t end, size_t grain_size) {
#ifdef _OPENMP
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));
#else
return 1;
#endif
}
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.
*/
template <typename F>
void parallel_for(
const size_t begin, const size_t end, const size_t grain_size, F&& f) {
if (begin >= end) {
return;
}
#ifdef _OPENMP
auto num_threads = compute_num_threads(begin, end, grain_size);
// (BarclayII) the exception code is borrowed from PyTorch.
std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
std::exception_ptr eptr;
#pragma omp parallel num_threads(num_threads)
{
auto tid = omp_get_thread_num();
auto chunk_size = divup((end - begin), num_threads);
auto begin_tid = begin + tid * chunk_size;
if (begin_tid < end) {
auto end_tid = std::min(end, static_cast<size_t>(chunk_size + begin_tid));
try {
f(begin_tid, end_tid);
} catch (...) {
if (!err_flag.test_and_set()) eptr = std::current_exception();
}
}
}
if (eptr) std::rethrow_exception(eptr);
#else
f(begin, end);
#endif
}
/**
* @brief OpenMP-based parallel for loop with default grain size.
*
* parallel_for with grain size to default value, either 1 or controlled through
* environment variable DGL_PARALLEL_FOR_GRAIN_SIZE.
* If grain size is set to 1, the function behaves the same way as OpenMP
* parallel for pragma with static scheduling.
*/
template <typename 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 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:
*
* parallel_reduce(
* 0, // starting index
* 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;
* },
* [] (float result, float partial_result) {
* return std::max(result, partial_result);
* });
*/
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) {
if (begin >= end) {
return ident;
}
int num_threads = compute_num_threads(begin, end, grain_size);
if (num_threads == 1) {
return f(begin, end, ident);
}
std::vector<DType> results(num_threads, ident);
std::atomic_flag err_flag = ATOMIC_FLAG_INIT;
std::exception_ptr eptr;
#pragma omp parallel num_threads(num_threads)
{
auto tid = omp_get_thread_num();
auto chunk_size = divup((end - begin), num_threads);
auto begin_tid = begin + tid * chunk_size;
if (begin_tid < end) {
auto end_tid = std::min(end, static_cast<size_t>(chunk_size + begin_tid));
try {
results[tid] = f(begin_tid, end_tid, ident);
} catch (...) {
if (!err_flag.test_and_set()) eptr = std::current_exception();
}
}
}
if (eptr) std::rethrow_exception(eptr);
DType out = ident;
for (int64_t i = 0; i < num_threads; ++i) out = sf(out, results[i]);
return out;
}
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_PARALLEL_FOR_H_
include/dgl/runtime/registry.h
+146
View File
@@ -0,0 +1,146 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/runtime/registry.h
* @brief This file defines the DGL global function registry.
*
* The registered functions will be made available to front-end
* as well as backend users.
*
* The registry stores type-erased functions.
* Each registered function is automatically exposed
* to front-end language(e.g. python).
*
* Front-end can also pass callbacks as PackedFunc, or register
* then into the same global registry in C++.
* The goal is to mix the front-end language and the DGL back-end.
*
* @code
* // register the function as MyAPIFuncName
* DGL_REGISTER_GLOBAL(MyAPIFuncName)
* .set_body([](DGLArgs args, DGLRetValue* rv) {
* // my code.
* });
* @endcode
*/
#ifndef DGL_RUNTIME_REGISTRY_H_
#define DGL_RUNTIME_REGISTRY_H_
#include <string>
#include <vector>
#include "packed_func.h"
namespace dgl {
namespace runtime {
/** @brief Registry for global function */
class Registry {
public:
/**
* @brief set the body of the function to be f
* @param f The body of the function.
*/
DGL_DLL Registry& set_body(PackedFunc f); // NOLINT(*)
/**
* @brief set the body of the function to be f
* @param f The body of the function.
*/
Registry& set_body(PackedFunc::FType f) { // NOLINT(*)
return set_body(PackedFunc(f));
}
/**
* @brief set the body of the function to be TypedPackedFunc.
*
* @code
*
* DGL_REGISTER_API("addone")
* .set_body_typed<int(int)>([](int x) { return x + 1; });
*
* @endcode
*
* @param f The body of the function.
* @tparam FType the signature of the function.
* @tparam FLambda The type of f.
*/
template <typename FType, typename FLambda>
Registry& set_body_typed(FLambda f) {
return set_body(TypedPackedFunc<FType>(f).packed());
}
/**
* @brief Register a function with given name
* @param name The name of the function.
* @param override Whether allow oveeride existing function.
* @return Reference to theregistry.
*/
DGL_DLL static Registry& Register(
const std::string& name, bool override = false); // NOLINT(*)
/**
* @brief Erase global function from registry, if exist.
* @param name The name of the function.
* @return Whether function exist.
*/
DGL_DLL static bool Remove(const std::string& name);
/**
* @brief Get the global function by name.
* @param name The name of the function.
* @return pointer to the registered function,
* nullptr if it does not exist.
*/
DGL_DLL static const PackedFunc* Get(const std::string& name); // NOLINT(*)
/**
* @brief Get the names of currently registered global function.
* @return The names
*/
DGL_DLL static std::vector<std::string> ListNames();
// Internal class.
struct Manager;
protected:
/** @brief name of the function */
std::string name_;
/** @brief internal packed function */
PackedFunc func_;
friend struct Manager;
};
/** @brief helper macro to supress unused warning */
#if defined(__GNUC__)
#define DGL_ATTRIBUTE_UNUSED __attribute__((unused))
#else
#define DGL_ATTRIBUTE_UNUSED
#endif
#define DGL_STR_CONCAT_(__x, __y) __x##__y
#define DGL_STR_CONCAT(__x, __y) DGL_STR_CONCAT_(__x, __y)
#define DGL_FUNC_REG_VAR_DEF \
static DGL_ATTRIBUTE_UNUSED ::dgl::runtime::Registry& __mk_##DGL
#define DGL_TYPE_REG_VAR_DEF \
static DGL_ATTRIBUTE_UNUSED ::dgl::runtime::ExtTypeVTable* __mk_##DGLT
/**
* @brief Register a function globally.
* @code
* DGL_REGISTER_GLOBAL("MyPrint")
* .set_body([](DGLArgs args, DGLRetValue* rv) {
* });
* @endcode
*/
#define DGL_REGISTER_GLOBAL(OpName) \
DGL_STR_CONCAT(DGL_FUNC_REG_VAR_DEF, __COUNTER__) = \
::dgl::runtime::Registry::Register(OpName)
/**
* @brief Macro to register extension type.
* This must be registered in a cc file
* after the trait extension_class_info is defined.
*/
#define DGL_REGISTER_EXT_TYPE(T) \
DGL_STR_CONCAT(DGL_TYPE_REG_VAR_DEF, __COUNTER__) = \
::dgl::runtime::ExtTypeVTable::Register_<T>()
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_REGISTRY_H_
include/dgl/runtime/serializer.h
+52
View File
@@ -0,0 +1,52 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/runtime/serializer.h
* @brief Serializer extension to support DGL data types
* Include this file to enable serialization of DGLDataType, DGLContext
*/
#ifndef DGL_RUNTIME_SERIALIZER_H_
#define DGL_RUNTIME_SERIALIZER_H_
#include <dmlc/io.h>
#include <dmlc/serializer.h>
#include "c_runtime_api.h"
#include "smart_ptr_serializer.h"
namespace dmlc {
namespace serializer {
template <>
struct Handler<DGLDataType> {
inline static void Write(Stream *strm, const DGLDataType &dtype) {
Handler<uint8_t>::Write(strm, dtype.code);
Handler<uint8_t>::Write(strm, dtype.bits);
Handler<uint16_t>::Write(strm, dtype.lanes);
}
inline static bool Read(Stream *strm, DGLDataType *dtype) {
if (!Handler<uint8_t>::Read(strm, &(dtype->code))) return false;
if (!Handler<uint8_t>::Read(strm, &(dtype->bits))) return false;
if (!Handler<uint16_t>::Read(strm, &(dtype->lanes))) return false;
return true;
}
};
template <>
struct Handler<DGLContext> {
inline static void Write(Stream *strm, const DGLContext &ctx) {
int32_t device_type = static_cast<int32_t>(ctx.device_type);
Handler<int32_t>::Write(strm, device_type);
Handler<int32_t>::Write(strm, ctx.device_id);
}
inline static bool Read(Stream *strm, DGLContext *ctx) {
int32_t device_type = 0;
if (!Handler<int32_t>::Read(strm, &(device_type))) return false;
ctx->device_type = static_cast<DGLDeviceType>(device_type);
if (!Handler<int32_t>::Read(strm, &(ctx->device_id))) return false;
return true;
}
};
} // namespace serializer
} // namespace dmlc
#endif // DGL_RUNTIME_SERIALIZER_H_
include/dgl/runtime/shared_mem.h
+91
View File
@@ -0,0 +1,91 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/runtime/ndarray.h
* @brief shared memory management.
*/
#ifndef DGL_RUNTIME_SHARED_MEM_H_
#define DGL_RUNTIME_SHARED_MEM_H_
#ifdef _WIN32
#include <windows.h>
#endif // _WIN32
#include <string>
namespace dgl {
namespace runtime {
/**
* @brief This class owns shared memory.
*
* When the object is gone, the shared memory will also be destroyed.
* When the shared memory is destroyed, the file corresponding to
* the shared memory is removed.
*/
class SharedMemory {
/**
* @brief whether the shared memory is owned by the object.
*
* If shared memory is created in the object, it'll be owned by the object
* and will be responsible for deleting it when the object is destroyed.
*/
bool own_;
/* @brief the file descripter of the shared memory. */
#ifndef _WIN32
int fd_;
#else // !_WIN32
HANDLE handle_;
#endif // _WIN32
/* @brief the address of the shared memory. */
void *ptr_;
/* @brief the size of the shared memory. */
size_t size_;
/**
* @brief the name of the object.
*
* In Unix, shared memory is identified by a file. Thus, `name` is actually
* the file name that identifies the shared memory.
*/
std::string name;
public:
/* @brief Get the filename of shared memory file
*/
std::string GetName() const { return name; }
/**
* @brief constructor of the shared memory.
* @param name The file corresponding to the shared memory.
*/
explicit SharedMemory(const std::string &name);
/**
* @brief destructor of the shared memory.
* It deallocates the shared memory and removes the corresponding file.
*/
~SharedMemory();
/**
* @brief create shared memory.
* It creates the file and shared memory.
* @param sz the size of the shared memory.
* @return the address of the shared memory
*/
void *CreateNew(size_t sz);
/**
* @brief allocate shared memory that has been created.
* @param sz the size of the shared memory.
* @return the address of the shared memory
*/
void *Open(size_t sz);
/**
* @brief check if the shared memory exist.
* @param name the name of the shared memory.
* @return a boolean value to indicate if the shared memory exists.
*/
static bool Exist(const std::string &name);
};
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_SHARED_MEM_H_
include/dgl/runtime/smart_ptr_serializer.h
+54
View File
@@ -0,0 +1,54 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/runtime/serializer.h
* @brief Serializer extension to support DGL data types
* Include this file to enable serialization of DGLDataType, DGLContext
*/
#ifndef DGL_RUNTIME_SMART_PTR_SERIALIZER_H_
#define DGL_RUNTIME_SMART_PTR_SERIALIZER_H_
#include <dgl/graph_serializer.h>
#include <dmlc/io.h>
#include <dmlc/serializer.h>
#include <memory>
namespace dmlc {
namespace serializer {
//! \cond Doxygen_Suppress
template <typename T>
struct Handler<std::shared_ptr<T>> {
inline static void Write(Stream *strm, const std::shared_ptr<T> &data) {
Handler<T>::Write(strm, *data.get());
}
inline static bool Read(Stream *strm, std::shared_ptr<T> *data) {
// When read, the default initialization behavior of shared_ptr is
// shared_ptr<T>(), which is holding a nullptr. Here we need to manually
// reset to a real object for further loading
if (!(*data)) {
data->reset(dgl::Serializer::new_object<T>());
}
return Handler<T>::Read(strm, data->get());
}
};
template <typename T>
struct Handler<std::unique_ptr<T>> {
inline static void Write(Stream *strm, const std::unique_ptr<T> &data) {
Handler<T>::Write(strm, *data.get());
}
inline static bool Read(Stream *strm, std::unique_ptr<T> *data) {
// When read, the default initialization behavior of unique_ptr is
// unique_ptr<T>(), which is holding a nullptr. Here we need to manually
// reset to a real object for further loading
if (!(*data)) {
data->reset(dgl::Serializer::new_object<T>());
}
return Handler<T>::Read(strm, data->get());
}
};
} // namespace serializer
} // namespace dmlc
#endif // DGL_RUNTIME_SMART_PTR_SERIALIZER_H_
include/dgl/runtime/tensordispatch.h
+281
View File
@@ -0,0 +1,281 @@
/**
* Copyright (c) 2020-2022 by Contributors
* @file array/tensordispatch.h
* @brief This file defines the dispatcher of tensor operators to
* framework-specific implementations.
*
* The dispatcher consists of a TensorDispatcher singleton in DGL C library and
* one separately-built shared library per supported backend.
*
* Those shared libraries contain wrappers of the framework-specific operators.
* The wrappers are defined with extern "C", meaning that the C++ compiler will
* not do name mangling for those functions so that DGL can conveniently locate
* them using dlsym(3) (or GetProcAddress in Windows).
*
* The TensorDispatcher singleton maintains a mapping from an array operator to
* the address of the corresponding symbol in the shared library. During
* initialization, the TensorDispatcher checks which backend DGL is using.
* It then locates and opens the corresponding shared library using dlopen(3)
* (or LoadLibrary in Windows), and populates the said mapping above with
* dlsym(3) (or GetProcAddress in Windows).
*
* A tensor operator in TensorDispatcher first checks whether the corresponding
* symbol address is found in the mapping. If so, it calls the function located
* at the symbol address instead, allocate/free pieces of memory on CPU/GPU. If
* not, it falls back to DeviceAPI::AllocWorkspace/FreeWorkspace.
*/
#ifndef DGL_RUNTIME_TENSORDISPATCH_H_
#define DGL_RUNTIME_TENSORDISPATCH_H_
#include <stddef.h>
#include <tensoradapter.h>
#if defined(WIN32) || defined(_WIN32)
#include <windows.h>
#endif // WIN32
#ifdef DGL_USE_CUDA
#include <cuda_runtime.h>
#endif // DGL_USE_CUDA
#include "ndarray.h"
/**
* @brief Casts a pointer \c entry to a function pointer with signature of \c
* func.
*/
#define FUNCCAST(func, entry) (*reinterpret_cast<decltype(&(func))>(entry))
namespace dgl {
namespace runtime {
/**
* @brief Dispatcher that delegates the function calls to framework-specific C++
* APIs.
*
* This class is not thread-safe.
*/
class TensorDispatcher {
public:
/** @brief Get the singleton instance. */
static TensorDispatcher* Global() {
static TensorDispatcher inst;
return &inst;
}
/** @brief Whether an adapter library is available. */
inline bool IsAvailable() { return available_; }
/** @brief Load symbols from the given tensor adapter library path. */
bool Load(const char* path_cstr);
/**
* @brief Allocate a piece of CPU memory via PyTorch's CPUAllocator.
* Used in CPUDeviceAPI::AllocWorkspace().
*
* @param nbytes The size to be allocated.
* @return Pointer to the allocated memory.
*/
inline void* CPUAllocWorkspace(size_t nbytes) {
auto entry = entrypoints_[Op::kCPURawAlloc];
return FUNCCAST(tensoradapter::CPURawAlloc, entry)(nbytes);
}
/**
* @brief Free the CPU memory.
* Used in CPUDeviceAPI::FreeWorkspace().
*
* @param ptr Pointer to the memory to be freed.
*/
inline void CPUFreeWorkspace(void* ptr) {
auto entry = entrypoints_[Op::kCPURawDelete];
FUNCCAST(tensoradapter::CPURawDelete, entry)(ptr);
}
#ifdef DGL_USE_CUDA
/**
* @brief Allocate a piece of GPU memory via
* PyTorch's THCCachingAllocator.
* Used in CUDADeviceAPI::AllocWorkspace().
*
* @note THCCachingAllocator specify the device to allocate on
* via cudaGetDevice(). Make sure to call cudaSetDevice()
* before invoking this function.
*
* @param nbytes The size to be allocated.
* @param stream The stream to be allocated on.
* @return Pointer to the allocated memory.
*/
inline void* CUDAAllocWorkspace(size_t nbytes, cudaStream_t stream) {
auto entry = entrypoints_[Op::kCUDARawAlloc];
return FUNCCAST(tensoradapter::CUDARawAlloc, entry)(nbytes, stream);
}
/**
* @brief Free the GPU memory.
* Used in CUDADeviceAPI::FreeWorkspace().
*
* @param ptr Pointer to the memory to be freed.
*/
inline void CUDAFreeWorkspace(void* ptr) {
auto entry = entrypoints_[Op::kCUDARawDelete];
FUNCCAST(tensoradapter::CUDARawDelete, entry)(ptr);
}
/**
* @brief Find the current PyTorch CUDA stream
* Used in runtime::getCurrentCUDAStream().
*
* @note PyTorch pre-allocates/sets the current CUDA stream
* on current device via cudaGetDevice(). Make sure to call cudaSetDevice()
* before invoking this function.
*
* @return cudaStream_t stream handle
*/
inline cudaStream_t CUDAGetCurrentStream() {
auto entry = entrypoints_[Op::kCUDACurrentStream];
return FUNCCAST(tensoradapter::CUDACurrentStream, entry)();
}
/**
* @brief Allocate a piece of pinned CPU memory via PyTorch
* CachingHostAllocator.
* @note Used in CUDADeviceAPI::AllocPinnedDataSpace().
* @param nbytes The size to be allocated.
* @param ctx Pointer to the PyTorch storage ctx ptr returned from the
* allocator.
* @param deleter Pointer to the delete function ptr returned from the
* allocator.
* @return Raw pointer to the allocated memory.
*/
inline void* CUDAAllocHostWorkspace(
size_t nbytes, void** ctx, void** deleter) {
auto entry = entrypoints_[Op::kCUDARawHostAlloc];
auto alloc_func = FUNCCAST(tensoradapter::CUDARawHostAlloc, entry);
return alloc_func(nbytes, ctx, deleter);
}
/**
* @brief Insert the pinned memory block (allocated via PyTorch
* CachingHostAllocator) back to the free list for future usage.(ref:
* pytorch/pytorch/blob/master/aten/src/ATen/cuda/CachingHostAllocator.cpp).
* @note Used in CUDADeviceAPI::FreePinnedDataSpace().
* @param deleter Pointer to the delete function ptr returned from the
* allocator.
*/
inline void CUDAFreeHostWorkspace(void** deleter) {
auto entry = entrypoints_[Op::kCUDARawHostDelete];
FUNCCAST(tensoradapter::CUDARawHostDelete, entry)(deleter);
}
/**
* @brief Invoke the record_event function call from PyTorch
* CachingHostAllocator.
* @note This function assoicates a CUDA stream (used by a copy kernel) to the
* pinned data. In the free path of this data, which is achieved by
* calling CUDAFreeHostWorkspace, the set of associated streams is then
* consumed to ensure proper functionlity. (ref:
* pytorch/pytorch/blob/master/aten/src/ATen/cuda/CachingHostAllocator.cpp).
* Used in CUDADeviceAPI::RecordedCopyDataFromTo().
*
* @param data Pointer of the tensor to be recorded.
* @param ctx PyTorch storage ctx ptr returned from the allocator.
* @param stream The stream that currently consumes this tensor.
* @param device_id Device of the tensor.
*/
inline void CUDARecordHostAlloc(
void* data, void* ctx, cudaStream_t stream, int device_id) {
auto entry = entrypoints_[Op::kCUDARecordHostAlloc];
auto recorded_alloc = FUNCCAST(tensoradapter::CUDARecordHostAlloc, entry);
recorded_alloc(data, ctx, stream, device_id);
}
/**
* @brief Release cached pinned memory allocations via cudaHostFree.
* @note Used in CUDADeviceAPI::PinData() before pinning any host memory by
* DGL.
*/
inline void CUDAHostAllocatorEmptyCache() {
auto entry = entrypoints_[Op::kCUDAHostAllocatorEmptyCache];
FUNCCAST(tensoradapter::CUDAHostAllocatorEmptyCache, entry)();
}
#endif // DGL_USE_CUDA
/**
* @brief Record streams that are using this tensor.
* Used in NDArray::RecordStream().
*
* @param ptr Pointer of the tensor to be recorded.
* @param stream The stream that is using this tensor.
* @param device_id Device of the tensor.
*/
inline void RecordStream(void* ptr, DGLStreamHandle stream, int device_id) {
#ifdef DGL_USE_CUDA
auto entry = entrypoints_[Op::kRecordStream];
FUNCCAST(tensoradapter::RecordStream, entry)
(ptr, static_cast<cudaStream_t>(stream), device_id);
#endif
}
private:
/** @brief ctor */
TensorDispatcher() = default;
/** @brief dtor */
~TensorDispatcher();
/**
* @brief List of symbols in the adapter library.
*
* Must match the functions in tensoradapter/include/tensoradapter.h.
*/
static constexpr const char* names_[] = {
"CPURawAlloc", "CPURawDelete",
#ifdef DGL_USE_CUDA
"CUDARawAlloc", "CUDARawDelete",
"CUDACurrentStream", "RecordStream",
"CUDARawHostAlloc", "CUDARawHostDelete",
"CUDARecordHostAlloc", "CUDAHostAllocatorEmptyCache",
#endif // DGL_USE_CUDA
};
/** @brief Index of each function to the symbol list */
class Op {
public:
static constexpr int kCPURawAlloc = 0;
static constexpr int kCPURawDelete = 1;
#ifdef DGL_USE_CUDA
static constexpr int kCUDARawAlloc = 2;
static constexpr int kCUDARawDelete = 3;
static constexpr int kCUDACurrentStream = 4;
static constexpr int kRecordStream = 5;
static constexpr int kCUDARawHostAlloc = 6;
static constexpr int kCUDARawHostDelete = 7;
static constexpr int kCUDARecordHostAlloc = 8;
static constexpr int kCUDAHostAllocatorEmptyCache = 9;
#endif // DGL_USE_CUDA
};
/** @brief Number of functions */
static constexpr int num_entries_ = sizeof(names_) / sizeof(names_[0]);
/** @brief Entrypoints of each function */
void* entrypoints_[num_entries_] = {
nullptr, nullptr,
#ifdef DGL_USE_CUDA
nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr, nullptr,
#endif // DGL_USE_CUDA
};
bool available_ = false;
#if defined(WIN32) || defined(_WIN32)
HINSTANCE handle_;
#else // !WIN32
void* handle_;
#endif // WIN32
};
}; // namespace runtime
}; // namespace dgl
#undef FUNCCAST
#endif // DGL_RUNTIME_TENSORDISPATCH_H_
include/dgl/runtime/threading_backend.h
+86
View File
@@ -0,0 +1,86 @@
/**
* Copyright (c) 2018 by Contributors
* @file dgl/runtime/threading_backend.h
* @brief Utilities for manipulating thread pool threads.
*/
#ifndef DGL_RUNTIME_THREADING_BACKEND_H_
#define DGL_RUNTIME_THREADING_BACKEND_H_
#include <functional>
#include <memory>
#include <vector>
namespace dgl {
namespace runtime {
namespace threading {
/**
* @brief A platform-agnostic abstraction for managing a collection of
* thread pool threads.
*/
class ThreadGroup {
public:
class Impl;
/**
* @brief Creates a collection of threads which run a provided function.
*
* @param num_workers The total number of worker threads in this group.
Includes main thread if `exclude_worker0 = true`
* @param worker_callback A callback which is run in its own thread.
Receives the worker_id as an argument.
* @param exclude_worker0 Whether to use the main thread as a worker.
* If `true`, worker0 will not be launched in a new thread and
* `worker_callback` will only be called for values >= 1. This
* allows use of the main thread as a worker.
*/
ThreadGroup(
int num_workers, std::function<void(int)> worker_callback,
bool exclude_worker0 = false);
~ThreadGroup();
/**
* @brief Blocks until all non-main threads in the pool finish.
*/
void Join();
enum AffinityMode : int {
kBig = 1,
kLittle = -1,
};
/**
* @brief configure the CPU id affinity
*
* @param mode The preferred CPU type (1 = big, -1 = little).
* @param nthreads The number of threads to use (0 = use all).
* @param exclude_worker0 Whether to use the main thread as a worker.
* If `true`, worker0 will not be launched in a new thread and
* `worker_callback` will only be called for values >= 1. This
* allows use of the main thread as a worker.
*
* @return The number of workers to use.
*/
int Configure(AffinityMode mode, int nthreads, bool exclude_worker0);
private:
Impl* impl_;
};
/**
* @brief Platform-agnostic no-op.
*/
// This used to be Yield(), renaming to YieldThread() because windows.h defined
// it as a macro in later SDKs.
void YieldThread();
/**
* @return the maximum number of effective workers for this system.
*/
int MaxConcurrency();
} // namespace threading
} // namespace runtime
} // namespace dgl
#endif // DGL_RUNTIME_THREADING_BACKEND_H_
include/dgl/runtime/util.h
+53
View File
@@ -0,0 +1,53 @@
/**
* Copyright (c) 2017 by Contributors
* @file dgl/runtime/util.h
* @brief Useful runtime util.
*/
#ifndef DGL_RUNTIME_UTIL_H_
#define DGL_RUNTIME_UTIL_H_
#include "c_runtime_api.h"
namespace dgl {
namespace runtime {
/**
* @brief Check whether type matches the given spec.
* @param t The type
* @param code The type code.
* @param bits The number of bits to be matched.
* @param lanes The number of lanes sin the type.
*/
inline bool TypeMatch(DGLDataType t, int code, int bits, int lanes = 1) {
return t.code == code && t.bits == bits && t.lanes == lanes;
}
} // namespace runtime
} // namespace dgl
// Forward declare the intrinsic id we need
// in structure fetch to enable stackvm in runtime
namespace dgl {
namespace ir {
namespace intrinsic {
/** @brief The kind of structure field info used in intrinsic */
enum DGLStructFieldKind : int {
// array head address
kArrAddr,
kArrData,
kArrShape,
kArrStrides,
kArrNDim,
kArrTypeCode,
kArrTypeBits,
kArrTypeLanes,
kArrByteOffset,
kArrDeviceId,
kArrDeviceType,
kArrKindBound_,
// DGLValue field
kDGLValueContent,
kDGLValueKindBound_
};
} // namespace intrinsic
} // namespace ir
} // namespace dgl
#endif // DGL_RUNTIME_UTIL_H_
include/dgl/sampler.h
+59
View File
@@ -0,0 +1,59 @@
/**
* Copyright (c) 2018 by Contributors
* @file dgl/sampler.h
* @brief DGL sampler header.
*/
#ifndef DGL_SAMPLER_H_
#define DGL_SAMPLER_H_
#include <cstdlib>
#include <ctime>
#include <string>
#include <vector>
#include "graph_interface.h"
#include "nodeflow.h"
namespace dgl {
class ImmutableGraph;
class SamplerOp {
public:
/**
* @brief Sample a graph from the seed vertices with neighbor sampling.
* The neighbors are sampled with a uniform distribution.
*
* @param graph A graph for sampling.
* @param seeds the nodes where we should start to sample.
* @param edge_type the type of edges we should sample neighbors.
* @param num_hops the number of hops to sample neighbors.
* @param expand_factor the max number of neighbors to sample.
* @param add_self_loop whether to add self loop to the sampled subgraph
* @param probability the transition probability (float/double).
* @return a NodeFlow graph.
*/
template <typename ValueType>
static NodeFlow NeighborSample(
const ImmutableGraph *graph, const std::vector<dgl_id_t> &seeds,
const std::string &edge_type, int num_hops, int expand_factor,
const bool add_self_loop, const ValueType *probability);
/**
* @brief Sample a graph from the seed vertices with layer sampling.
* The layers are sampled with a uniform distribution.
*
* @param graph A graph for sampling.
* @param seeds the nodes where we should start to sample.
* @param edge_type the type of edges we should sample neighbors.
* @param layer_sizes The size of layers.
* @return a NodeFlow graph.
*/
static NodeFlow LayerUniformSample(
const ImmutableGraph *graph, const std::vector<dgl_id_t> &seeds,
const std::string &neigh_type, IdArray layer_sizes);
};
} // namespace dgl
#endif // DGL_SAMPLER_H_
include/dgl/sampling/negative.h
+44
View File
@@ -0,0 +1,44 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/sampling/negative.h
* @brief Negative sampling.
*/
#ifndef DGL_SAMPLING_NEGATIVE_H_
#define DGL_SAMPLING_NEGATIVE_H_
#include <dgl/array.h>
#include <dgl/base_heterograph.h>
#include <utility>
namespace dgl {
namespace sampling {
/**
* @brief Given an edge type, uniformly sample source-destination pairs that do
* not have an edge in between using rejection sampling.
*
* @note This function may not return the same number of elements as the given
* number of samples.
* @note This function requires sorting the CSR or CSC matrix of the graph
* in-place. It prefers CSC over CSR.
*
* @param hg The graph.
* @param etype The edge type.
* @param num_samples The number of negative examples to sample.
* @param num_trials The number of rejection sampling trials.
* @param exclude_self_loops Do not include the examples where the source equals
* the destination.
* @param replace Whether to sample with replacement.
* @param redundancy How much redundant negative examples to take in case of
* duplicate examples.
* @return The pair of source and destination tensors.
*/
std::pair<IdArray, IdArray> GlobalUniformNegativeSampling(
HeteroGraphPtr hg, dgl_type_t etype, int64_t num_samples, int num_trials,
bool exclude_self_loops, bool replace, double redundancy);
}; // namespace sampling
}; // namespace dgl
#endif // DGL_SAMPLING_NEGATIVE_H_
include/dgl/sampling/neighbor.h
+131
View File
@@ -0,0 +1,131 @@
/**
* Copyright (c) 2020 by Contributors
* @file dgl/sampling/neighbor.h
* @brief Neighborhood-based sampling.
*/
#ifndef DGL_SAMPLING_NEIGHBOR_H_
#define DGL_SAMPLING_NEIGHBOR_H_
#include <dgl/array.h>
#include <dgl/base_heterograph.h>
#include <tuple>
#include <vector>
namespace dgl {
namespace sampling {
/**
* @brief Sample from the neighbors of the given nodes and return the sampled
* edges as a graph.
*
* When sampling with replacement, the sampled subgraph could have parallel
* edges.
*
* For sampling without replace, if fanout > the number of neighbors, all the
* neighbors will be sampled.
*
* @param hg The input graph.
* @param nodes Node IDs of each type. The vector length must be equal to the
* number of node types. Empty array is allowed.
* @param fanouts Number of sampled neighbors for each edge type. The vector
* length should be equal to the number of edge types, or one if they all have
* the same fanout.
* @param dir Edge direction.
* @param probability A vector of 1D float arrays, indicating the transition
* probability of each edge by edge type. An empty float array assumes uniform
* transition.
* @param exclude_edges Edges IDs of each type which will be excluded during
* sampling. The vector length must be equal to the number of edges types. Empty
* array is allowed.
* @param replace If true, sample with replacement.
* @return Sampled neighborhoods as a graph. The return graph has the same
* schema as the original one.
*/
HeteroSubgraph SampleNeighbors(
const HeteroGraphPtr hg, const std::vector<IdArray>& nodes,
const std::vector<int64_t>& fanouts, EdgeDir dir,
const std::vector<FloatArray>& probability,
const std::vector<IdArray>& exclude_edges, bool replace = true);
/**
* @brief Sample from the neighbors of the given nodes and convert a graph into
* a bipartite-structured graph for message passing.
*
* Specifically, we create one node type \c ntype_l on the "left" side and
* another node type \c ntype_r on the "right" side for each node type \c ntype.
* The nodes of type \c ntype_r would contain the nodes designated by the
* caller, and node type \c ntype_l would contain the nodes that has an edge
* connecting to one of the designated nodes.
*
* The nodes of \c ntype_l would also contain the nodes in node type \c ntype_r.
* When sampling with replacement, the sampled subgraph could have parallel
* edges.
*
* For sampling without replace, if fanout > the number of neighbors, all the
* neighbors will be sampled.
*
* Non-deterministic algorithm, requires nodes parameter to store unique Node
* IDs.
*
* @tparam IdType Graph's index data type, can be int32_t or int64_t
* @param hg The input graph.
* @param nodes Node IDs of each type. The vector length must be equal to the
* number of node types. Empty array is allowed.
* @param mapping External parameter that should be set to a vector of IdArrays
* filled with -1, required for mapping of nodes in returned
* graph
* @param fanouts Number of sampled neighbors for each edge type. The vector
* length should be equal to the number of edge types, or one if they all have
* the same fanout.
* @param dir Edge direction.
* @param probability A vector of 1D float arrays, indicating the transition
* probability of each edge by edge type. An empty float array assumes uniform
* transition.
* @param exclude_edges Edges IDs of each type which will be excluded during
* sampling. The vector length must be equal to the number of edges types. Empty
* array is allowed.
* @param replace If true, sample with replacement.
* @return Sampled neighborhoods as a graph. The return graph has the same
* schema as the original one.
*/
template <typename IdType>
std::tuple<HeteroGraphPtr, std::vector<IdArray>, std::vector<IdArray>>
SampleNeighborsFused(
const HeteroGraphPtr hg, const std::vector<IdArray>& nodes,
const std::vector<IdArray>& mapping, const std::vector<int64_t>& fanouts,
EdgeDir dir, const std::vector<NDArray>& prob_or_mask,
const std::vector<IdArray>& exclude_edges, bool replace = true);
/**
* Select the neighbors with k-largest weights on the connecting edges for each
* given node.
*
* If k > the number of neighbors, all the neighbors are sampled.
*
* @param hg The input graph.
* @param nodes Node IDs of each type. The vector length must be equal to the
* number of node types. Empty array is allowed.
* @param k The k value for each edge type. The vector length should be equal to
* the number of edge types, or one if they all have the same fanout.
* @param dir Edge direction.
* @param weight A vector of 1D float arrays, indicating the weights associated
* witheach edge.
* @param ascending If true, elements are sorted by ascending order, equivalent
* to find the K smallest values. Otherwise, find K largest values.
* @return Sampled neighborhoods as a graph. The return graph has the same
* schema as the original one.
*/
HeteroSubgraph SampleNeighborsTopk(
const HeteroGraphPtr hg, const std::vector<IdArray>& nodes,
const std::vector<int64_t>& k, EdgeDir dir,
const std::vector<FloatArray>& weight, bool ascending = false);
HeteroSubgraph SampleNeighborsBiased(
const HeteroGraphPtr hg, const IdArray& nodes, const int64_t fanouts,
const NDArray& bias, const NDArray& tag_offset, const EdgeDir dir,
const bool replace);
} // namespace sampling
} // namespace dgl
#endif // DGL_SAMPLING_NEIGHBOR_H_
include/dgl/sampling/randomwalks.h
+89
View File
@@ -0,0 +1,89 @@
/**
* Copyright (c) 2019 by Contributors
* @file dgl/samplinig/randomwalks.h
* @brief Random walk functions.
*/
#ifndef DGL_SAMPLING_RANDOMWALKS_H_
#define DGL_SAMPLING_RANDOMWALKS_H_
#include <dgl/array.h>
#include <dgl/base_heterograph.h>
#include <tuple>
#include <utility>
#include <vector>
namespace dgl {
namespace sampling {
/**
* @brief Metapath-based random walk.
* @param hg The heterograph.
* @param seeds A 1D array of seed nodes, with the type the source type of the
* first edge type in the metapath.
* @param metapath A 1D array of edge types representing the metapath.
* @param prob A vector of 1D float arrays, indicating the transition
* probability of each edge by edge type. An empty float array assumes uniform
* transition.
* @return A pair of
* 1. One 2D array of shape (len(seeds), len(metapath) + 1) with node
* IDs. The paths that terminated early are padded with -1.
* 2. One 2D array of shape (len(seeds), len(metapath)) with edge IDs.
* The paths that terminated early are padded with -1.
* 3. One 1D array of shape (len(metapath) + 1) with node type IDs.
*/
std::tuple<IdArray, IdArray, TypeArray> RandomWalk(
const HeteroGraphPtr hg, const IdArray seeds, const TypeArray metapath,
const std::vector<FloatArray> &prob);
/**
* @brief Metapath-based random walk with restart probability.
* @param hg The heterograph.
* @param seeds A 1D array of seed nodes, with the type the source type of the
* first edge type in the metapath.
* @param metapath A 1D array of edge types representing the metapath.
* @param prob A vector of 1D float arrays, indicating the transition
* probability of each edge by edge type. An empty float array assumes uniform
* transition.
* @param restart_prob Restart probability.
* @return A pair of
* 1. One 2D array of shape (len(seeds), len(metapath) + 1) with node
* IDs. The paths that terminated early are padded with -1.
* 2. One 2D array of shape (len(seeds), len(metapath)) with edge IDs.
* The paths that terminated early are padded with -1.
* 3. One 1D array of shape (len(metapath) + 1) with node type IDs.
*/
std::tuple<IdArray, IdArray, TypeArray> RandomWalkWithRestart(
const HeteroGraphPtr hg, const IdArray seeds, const TypeArray metapath,
const std::vector<FloatArray> &prob, double restart_prob);
/**
* @brief Metapath-based random walk with stepwise restart probability. Useful
* for PinSAGE-like models.
* @param hg The heterograph.
* @param seeds A 1D array of seed nodes, with the type the source type of the
* first edge type in the metapath.
* @param metapath A 1D array of edge types representing the metapath.
* @param prob A vector of 1D float arrays, indicating the transition
* probability of each edge by edge type. An empty float array assumes uniform
* transition.
* @param restart_prob Restart probability array which has the same number of
* elements as \c metapath, indicating the probability to terminate after
* transition.
* @return A pair of
* 1. One 2D array of shape (len(seeds), len(metapath) + 1) with node
* IDs. The paths that terminated early are padded with -1.
* 2. One 2D array of shape (len(seeds), len(metapath)) with edge IDs.
* The paths that terminated early are padded with -1.
* 3. One 1D array of shape (len(metapath) + 1) with node type IDs.
*/
std::tuple<IdArray, IdArray, TypeArray> RandomWalkWithStepwiseRestart(
const HeteroGraphPtr hg, const IdArray seeds, const TypeArray metapath,
const std::vector<FloatArray> &prob, FloatArray restart_prob);
}; // namespace sampling
}; // namespace dgl
#endif // DGL_SAMPLING_RANDOMWALKS_H_
include/dgl/scheduler.h
+64
View File
@@ -0,0 +1,64 @@
/**
* Copyright (c) 2018 by Contributors
* @file dgl/scheduler.h
* @brief Operations on graph index.
*/
#ifndef DGL_SCHEDULER_H_
#define DGL_SCHEDULER_H_
#include <vector>
#include "runtime/ndarray.h"
namespace dgl {
typedef dgl::runtime::NDArray IdArray;
namespace sched {
/**
* @brief Generate degree bucketing schedule
* @tparam IdType Graph's index data type, can be int32_t or int64_t
* @param msg_ids The edge id for each message
* @param vids The destination vertex for each message
* @param recv_ids The recv nodes (for checking zero degree nodes)
* @note If there are multiple messages going into the same destination vertex,
* then there will be multiple copies of the destination vertex in vids.
* @return a vector of 5 IdArrays for degree bucketing. The 5 arrays are:
* degrees: degrees for each bucket
* nids: destination node ids
* nid_section: number of nodes in each bucket (used to split nids)
* mids: message ids
* mid_section: number of messages in each bucket (used to split mids)
*/
template <class IdType>
std::vector<IdArray> DegreeBucketing(
const IdArray& msg_ids, const IdArray& vids, const IdArray& recv_ids);
/**
* @brief Generate degree bucketing schedule for group_apply edge
* @tparam IdType Graph's index data type, can be int32_t or int64_t
* @param uids One end vertex of edge by which edges are grouped
* @param vids The other end vertex of edge
* @param eids Edge ids
* @note This function always generate group_apply schedule based on degrees of
* nodes in uids. Therefore, if group_apply by source nodes, then uids
* should be source. If group_apply by destination nodes, then uids
* should be destination.
* @return a vector of 5 IdArrays for degree bucketing. The 5 arrays are:
* degrees: degrees for each bucket
* new_uids: uids reordered by degree bucket
* new_vids: vids reordered by degree bucket
* new_edis: eids reordered by degree bucket
* sections: number of edges in each degree bucket (used to partition
* new_uids, new_vids, and new_eids)
*/
template <class IdType>
std::vector<IdArray> GroupEdgeByNodeDegree(
const IdArray& uids, const IdArray& vids, const IdArray& eids);
} // namespace sched
} // namespace dgl
#endif // DGL_SCHEDULER_H_
include/dgl/transform.h
+135
View File
@@ -0,0 +1,135 @@
/**
* Copyright (c) 2019 by Contributors
* @file dgl/transform.h
* @brief DGL graph transformations
*/
#ifndef DGL_TRANSFORM_H_
#define DGL_TRANSFORM_H_
#include <tuple>
#include <utility>
#include <vector>
#include "array.h"
#include "base_heterograph.h"
namespace dgl {
namespace transform {
/**
* @brief Given a list of graphs, remove the common nodes that do not have
* inbound and outbound edges.
*
* The graphs should have identical node ID space (i.e. should have the same set
* of nodes, including types and IDs).
*
* @param graphs The list of graphs.
* @param always_preserve The list of nodes to preserve regardless of whether
* the inbound or outbound edges exist.
*
* @return A pair. The first element is the list of compacted graphs, and the
* second element is the mapping from the compacted graphs and the original
* graph.
*/
std::pair<std::vector<HeteroGraphPtr>, std::vector<IdArray>> CompactGraphs(
const std::vector<HeteroGraphPtr> &graphs,
const std::vector<IdArray> &always_preserve);
/**
* @brief Convert a graph into a bipartite-structured graph for message passing.
*
* Specifically, we create one node type \c ntype_l on the "left" side and
* another node type \c ntype_r on the "right" side for each node type \c ntype.
* The nodes of type \c ntype_r would contain the nodes designated by the
* caller, and node type \c ntype_l would contain the nodes that has an edge
* connecting to one of the designated nodes.
*
* The nodes of \c ntype_l would also contain the nodes in node type \c ntype_r.
*
* This function is often used for constructing a series of dependency graphs
* for multi-layer message passing, where we first construct a series of
* frontier graphs on the original node space, and run the following to get the
* bipartite graph needed for message passing with each GNN layer:
*
* <code>
* bipartites = [None] * len(num_layers)
* for l in reversed(range(len(layers))):
* bipartites[l], seeds = to_bipartite(frontier[l], seeds)
* x = graph.ndata["h"][seeds]
* for g, layer in zip(bipartites, layers):
* x_src = x
* x_dst = x[:len(g.dsttype)]
* x = sageconv(g, (x_src, x_dst))
* output = x
* </code>
*
* @param graph The graph.
* @param rhs_nodes Designated nodes that would appear on the right side.
* @param include_rhs_in_lhs If false, do not include the nodes of node type \c
* ntype_r in \c ntype_l.
*
* @return A triplet containing
* * The bipartite-structured graph,
* * The induced node from the left side for each graph,
* * The induced edges.
*
* @note If include_rhs_in_lhs is true, then for each node type \c ntype, the
* nodes in rhs_nodes[ntype] would always appear first in the nodes of type \c
* ntype_l in the new graph.
*/
std::tuple<HeteroGraphPtr, std::vector<IdArray>, std::vector<IdArray>> ToBlock(
HeteroGraphPtr graph, const std::vector<IdArray> &rhs_nodes,
bool include_rhs_in_lhs);
/**
* @brief Convert a multigraph to a simple graph.
*
* @return A triplet of
* * @c hg : The said simple graph.
* * @c count : The array of edge occurrences per edge type.
* * @c edge_map : The mapping from original edge IDs to new edge IDs per edge
* type.
*
* @note Example: consider a graph with the following edges
*
* [(0, 1), (1, 3), (2, 2), (1, 3), (1, 4), (1, 4)]
*
* Then ToSimpleGraph(g) would yield the following elements:
*
* * The first element would be the simple graph itself with the following edges
*
* [(0, 1), (1, 3), (1, 4), (2, 2)]
*
* * The second element is an array \c count. \c count[i] stands for the number
* of edges connecting simple_g.src[i] and simple_g.dst[i] in the original
* graph.
*
* count[0] = [1, 2, 2, 1]
*
* * One can find the mapping between edges from the original graph to the new
* simple graph.
*
* edge_map[0] = [0, 1, 3, 1, 2, 2]
*/
std::tuple<HeteroGraphPtr, std::vector<IdArray>, std::vector<IdArray>>
ToSimpleGraph(const HeteroGraphPtr graph);
/**
* @brief Remove edges from a graph.
*
* @param graph The graph.
* @param eids The edge IDs to remove per edge type.
*
* @return A pair of the graph with edges removed, as well as the edge ID
* mapping from the original graph to the new graph per edge type.
*/
std::pair<HeteroGraphPtr, std::vector<IdArray>> RemoveEdges(
const HeteroGraphPtr graph, const std::vector<IdArray> &eids);
}; // namespace transform
}; // namespace dgl
#endif // DGL_TRANSFORM_H_
include/dgl/zerocopy_serializer.h
+191
View File
@@ -0,0 +1,191 @@
/**
* Copyright (c) 2020 by Contributors
* @file rpc/shared_mem_serializer.h
* @brief headers for serializer.
*/
#ifndef DGL_ZEROCOPY_SERIALIZER_H_
#define DGL_ZEROCOPY_SERIALIZER_H_
#include <dgl/runtime/ndarray.h>
#include <dmlc/io.h>
#include <dmlc/memory_io.h>
#include <dmlc/serializer.h>
#include <deque>
#include <memory>
#include <queue>
#include <string>
#include <tuple>
#include <utility>
#include <vector>
#include "dmlc/logging.h"
namespace dgl {
/**
*
* StreamWithBuffer is backed up by dmlc::MemoryFixedSizeStream or
* dmlc::MemoryStringStream. This class supports serializing and deserializing
* NDArrays stored in shared memory. If the stream is created for
* sending/recving data through network, the data pointer of the NDArray will be
* transmitted directly without and copy. Otherwise, the stream is for
* sending/recving data to another process on the same machine, so if an NDArray
* is stored in shared memory, it will just record the shared memory name
* instead of the actual data buffer.
*
* For example:
*
* std::string blob;
* // Send to local
* StreamWithBuffer strm(&blob, false);
* // Send to remote
* StreamWithBuffer strm(&blob, true);
* // Receive from local
* StreamWithBuffer strm(&blob, false);
* // Receive from remote
* std::vector<void*> ptr_list
* StreamWithBuffer strm(&blob, ptr_list);
*/
class StreamWithBuffer : public dmlc::SeekStream {
public:
// Buffer type. Storing NDArray to maintain the reference counting to ensure
// the liveness of data pointer
struct Buffer {
dgl::runtime::NDArray tensor = dgl::runtime::NDArray();
void* data = nullptr;
int64_t size = 0;
Buffer(const dgl::runtime::NDArray& tensor, void* data, int64_t data_size)
: tensor(tensor), data(data), size(data_size) {}
explicit Buffer(void* data) : data(data) {}
};
/**
* @brief This constructor is for writing scenario or reading from local
* machine
* @param strm The backup stream to write/load from
* @param send_to_remote Whether this stream will be deserialized at remote
* machine or the local machine. If true, will record the data pointer into
* buffer list.
*/
StreamWithBuffer(std::unique_ptr<dmlc::SeekStream> strm, bool send_to_remote)
: strm_(std::move(strm)),
buffer_list_(),
send_to_remote_(send_to_remote) {}
/**
* @brief This constructor is for reading from remote
* @param strm The stream to write/load from zerocopy write/load
* @param data_ptr_list list of pointer to reconstruct NDArray
*
* For example:
* std::string blob;
* std::vector<void*> data_ptr_list;
* // Read from remote sended pointer list
* StreamWithBuffer buf_strm(&blob, data_ptr_list)
*/
StreamWithBuffer(
std::unique_ptr<dmlc::SeekStream> strm,
const std::vector<void*>& data_ptr_list)
: strm_(std::move(strm)), send_to_remote_(true) {
for (void* data : data_ptr_list) {
buffer_list_.emplace_back(data);
}
}
/**
* @brief Construct stream backed up by string
* @param blob The string to write/load from zerocopy write/load
* @param send_to_remote Whether this stream will be deserialized at remote
* machine or the local machine. If true, will record the data pointer into
* buffer list.
*/
StreamWithBuffer(std::string* blob, bool send_to_remote)
: strm_(new dmlc::MemoryStringStream(blob)),
send_to_remote_(send_to_remote) {}
/**
* @brief Construct stream backed up by string
* @param p_buffer buffer pointer
* @param size buffer size
* @param send_to_remote Whether this stream will be deserialized at remote
* machine or the local machine. If true, will record the data pointer into
* buffer list.
*/
StreamWithBuffer(char* p_buffer, size_t size, bool send_to_remote)
: strm_(new dmlc::MemoryFixedSizeStream(p_buffer, size)),
send_to_remote_(send_to_remote) {}
/**
* @brief Construct stream backed up by string, and reconstruct NDArray
* from data_ptr_list
* @param blob The string to write/load from zerocopy write/load
* @param data_ptr_list pointer list for NDArrays to deconstruct from
*/
StreamWithBuffer(std::string* blob, const std::vector<void*>& data_ptr_list)
: strm_(new dmlc::MemoryStringStream(blob)), send_to_remote_(true) {
for (void* data : data_ptr_list) {
buffer_list_.emplace_back(data);
}
}
/**
* @brief Construct stream backed up by string, and reconstruct NDArray
* from data_ptr_list
* @param p_buffer buffer pointer
* @param size buffer size
* @param data_ptr_list pointer list for NDArrays to deconstruct from
*/
StreamWithBuffer(
char* p_buffer, size_t size, const std::vector<void*>& data_ptr_list)
: strm_(new dmlc::MemoryFixedSizeStream(p_buffer, size)),
send_to_remote_(true) {
for (void* data : data_ptr_list) {
buffer_list_.emplace_back(data);
}
}
// delegate methods to strm_
virtual size_t Read(void* ptr, size_t size) { return strm_->Read(ptr, size); }
virtual void Write(const void* ptr, size_t size) { strm_->Write(ptr, size); }
virtual void Seek(size_t pos) { strm_->Seek(pos); }
virtual size_t Tell(void) { return strm_->Tell(); }
using dmlc::Stream::Read;
using dmlc::Stream::Write;
/**
* @brief push NDArray into stream
* If send_to_remote=true, the NDArray will be saved to the buffer list
* If send_to_remote=false, the NDArray will be saved to the backedup string
*/
void PushNDArray(const runtime::NDArray& tensor);
/**
* @brief pop NDArray from stream
* If send_to_remote=true, the NDArray will be reconstructed from buffer list
* If send_to_remote=false, the NDArray will be reconstructed from shared
* memory
*/
dgl::runtime::NDArray PopNDArray();
/**
* @brief Get whether this stream is for remote usage
*/
bool send_to_remote() { return send_to_remote_; }
/**
* @brief Get underlying buffer list
*/
const std::deque<Buffer>& buffer_list() const { return buffer_list_; }
private:
std::unique_ptr<dmlc::SeekStream> strm_;
std::deque<Buffer> buffer_list_;
bool send_to_remote_;
}; // namespace dgl
} // namespace dgl
#endif // DGL_ZEROCOPY_SERIALIZER_H_