dmlc--dgl
Quan (Andy) Gan
929d863447
* test * more stuff * add test * fixes * optimize algo * replace unordered_map with arrays * lint * lint x2 * oops * disable gpu csrmm tests * remove gpu invocation * optimize with openmp * remove python functions * add back with docstrings * lint * lint * update python interface * functionize * functionize * lint * lint
82 行
2.2 KiB
C++
82 行
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.
|
|
*
|
|
* \note B is transposed (i.e. in CSC format).
|
|
*/
|
|
std::pair<CSRMatrix, NDArray> CSRMM(
|
|
CSRMatrix A,
|
|
NDArray A_weights,
|
|
CSRMatrix B,
|
|
NDArray B_weights);
|
|
|
|
/*!
|
|
* \brief Sparse-sparse matrix summation.
|
|
*/
|
|
std::pair<CSRMatrix, NDArray> CSRSum(
|
|
const std::vector<CSRMatrix>& A,
|
|
const std::vector<NDArray>& A_weights);
|
|
|
|
/*!
|
|
* \brief Return a sparse matrix with the values of A but nonzero entry locations of B.
|
|
*/
|
|
NDArray CSRMask(const CSRMatrix& A, NDArray A_weights, const CSRMatrix& B);
|
|
|
|
} // namespace aten
|
|
} // namespace dgl
|
|
|
|
#endif // DGL_KERNEL_H_
|