chore: import upstream snapshot with attribution

2026-07-13 13:35:51 +08:00
commit c36a561cd8
2172 changed files with 455595 additions and 0 deletions
@@ -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_