Merge branch 'master'

graphbolt/CMakeLists.txt

@@ -1,5 +1,6 @@
 cmake_minimum_required(VERSION 3.5)
 project(graphbolt C CXX)
+set (CMAKE_CXX_STANDARD 17)
 
 # Find PyTorch cmake files and PyTorch versions with the python interpreter
 # $PYTHON_INTERP ("python3" or "python" if empty)

graphbolt/include/graphbolt/csc_sampling_graph.h

@@ -16,6 +16,21 @@
 namespace graphbolt {
 namespace sampling {
 
+enum SamplerType { NEIGHBOR, LABOR };
+
+template <SamplerType S>
+struct SamplerArgs;
+
+template <>
+struct SamplerArgs<SamplerType::NEIGHBOR> {};
+
+template <>
+struct SamplerArgs<SamplerType::LABOR> {
+  const torch::Tensor& indices;
+  int64_t random_seed;
+  int64_t num_nodes;
+};
+
 /**
  * @brief A sampling oriented csc format graph.
  *
@@ -143,6 +158,9 @@ class CSCSamplingGraph : public torch::CustomClassHolder {
    * @param replace Boolean indicating whether the sample is preformed with or
    * without replacement. If True, a value can be selected multiple times.
    * Otherwise, each value can be selected only once.
+   * @param layer Boolean indicating whether neighbors should be sampled in a
+   * layer sampling fashion. Uses the LABOR-0 algorithm to increase overlap of
+   * sampled edges, see arXiv:2210.13339.
    * @param return_eids Boolean indicating whether edge IDs need to be returned,
    * typically used when edge features are required.
    * @param probs_name An optional string specifying the name of an edge
@@ -156,7 +174,7 @@ class CSCSamplingGraph : public torch::CustomClassHolder {
    */
   c10::intrusive_ptr<SampledSubgraph> SampleNeighbors(
       const torch::Tensor& nodes, const std::vector<int64_t>& fanouts,
-      bool replace, bool return_eids,
+      bool replace, bool layer, bool return_eids,
       torch::optional<std::string> probs_name) const;
 
   /**
@@ -204,6 +222,13 @@ class CSCSamplingGraph : public torch::CustomClassHolder {
       const std::string& shared_memory_name);
 
  private:
+  template <SamplerType S>
+  c10::intrusive_ptr<SampledSubgraph> SampleNeighborsImpl(
+      const torch::Tensor& nodes, const std::vector<int64_t>& fanouts,
+      bool replace, bool return_eids,
+      const torch::optional<torch::Tensor>& probs_or_mask,
+      SamplerArgs<S> args) const;
+
   /**
    * @brief Build a CSCSamplingGraph from shared memory tensors.
    *
@@ -298,10 +323,11 @@ class CSCSamplingGraph : public torch::CustomClassHolder {
  *
  * @return A tensor containing the picked neighbors.
  */
+template <SamplerType S>
 torch::Tensor Pick(
     int64_t offset, int64_t num_neighbors, int64_t fanout, bool replace,
     const torch::TensorOptions& options,
-    const torch::optional<torch::Tensor>& probs_or_mask);
+    const torch::optional<torch::Tensor>& probs_or_mask, SamplerArgs<S> args);
 
 /**
  * @brief Picks a specified number of neighbors for a node per edge type,
@@ -330,11 +356,19 @@ torch::Tensor Pick(
  *
  * @return A tensor containing the picked neighbors.
  */
+template <SamplerType S>
 torch::Tensor PickByEtype(
     int64_t offset, int64_t num_neighbors, const std::vector<int64_t>& fanouts,
     bool replace, const torch::TensorOptions& options,
     const torch::Tensor& type_per_edge,
-    const torch::optional<torch::Tensor>& probs_or_mask);
+    const torch::optional<torch::Tensor>& probs_or_mask, SamplerArgs<S> args);
+
+template <bool NonUniform, bool Replace, typename T = float>
+torch::Tensor LaborPick(
+    int64_t offset, int64_t num_neighbors, int64_t fanout,
+    const torch::TensorOptions& options,
+    const torch::optional<torch::Tensor>& probs_or_mask,
+    SamplerArgs<SamplerType::LABOR> args);
 
 }  // namespace sampling
 }  // namespace graphbolt