wehub-resource-sync
879 lines
28 KiB
C++
879 lines
28 KiB
C++
|
|
/**
|
|
* 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_
|