dmlc--dgl
Israt Nisa
ab2bd1f13d
* init cuda support * cuSPARSE err * passed unittest for csr_mm/SpGEMM. int64 not supported * Debugging cuSPARSE error 3 * csrgeam only supports int32? * disabling int64 for cuda * refactor and add CSRMask * lint * oops * remove todo * rewrite CSRMask with CSRGetData * lint * fix test * address comments * lint * fix * addresses comments and rename BUG_ON Co-authored-by: Israt Nisa <nisisrat@amazon.com> Co-authored-by: Ubuntu <ubuntu@ip-172-31-30-71.ec2.internal> Co-authored-by: Quan Gan <coin2028@hotmail.com> Co-authored-by: Jinjing Zhou <VoVAllen@users.noreply.github.com> Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>
81 行
2.2 KiB
C++
81 行
2.2 KiB
C++
/*!
|
|
* 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 <vector>
|
|
#include <utility>
|
|
|
|
#include "array.h"
|
|
#include "./bcast.h"
|
|
#include "./base_heterograph.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_
|