Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

[RFC] cugraph-based neighborhood sampling #50

Open
rusty1s opened this issue May 20, 2022 · 0 comments
Open

[RFC] cugraph-based neighborhood sampling #50

rusty1s opened this issue May 20, 2022 · 0 comments

Comments

@rusty1s
Copy link
Member

rusty1s commented May 20, 2022

🚀 The feature, motivation and pitch

GPU-based neighborhood sampling can accelerate mini-batch creation for graphs that fit into GPU memory.

Currently, the (solely) CPU-based sampling interface inside PyG looks as follows:

template <bool replace, bool directed>
tuple<torch::Tensor, torch::Tensor, torch::Tensor, torch::Tensor>
sample(const torch::Tensor &colptr, const torch::Tensor &row,
       const torch::Tensor &input_node, const vector<int64_t> num_neighbors)

The PyG routine expects:

  • (colptr, row): CSC/CSR representation of the graph
  • input_node: The seed nodes for which to sample neighbors
  • num_neighbors: A list of neighbors to sample in each layer
  • replace: Sample without or with replacement
  • directed: Whether sampled edges are directed or not. If not, we extract the full subgraph of sampled nodes.

It returns (re-labeled) row and col vectors of the sampled subgraph/adjacency matrix, as well as output_node_id and output_edge_id of the sampled nodes/edges to perform feature fetching in a later stage.

On the other side, the sampling interface inside cugraph looks as follows:

template <typename graph_t>
std::tuple<rmm::device_uvector<typename graph_t::edge_type>,
           rmm::device_uvector<typename graph_t::vertex_type>>
sample_neighbors_adjacency_list(raft::handle_t const& handle,
                                raft::random::RngState& rng_state,
                                graph_t const& graph,
                                typename graph_t::vertex_type const* ptr_d_start,
                                size_t num_start_vertices,
                                size_t sampling_size,
                                ops::gnn::graph::SamplingAlgoT sampling_algo)

The major difference seems to be that cugraph performs sampling for 1-hop, while PyG supports multi-hop sampling (which can be fixed easily by just calling the cugraph routine multiple times) [to be confirmed by @pyg-team/nvidia-team].

Roadmap

For integrating GPU-based sampling inside PyG, we thus need to:

  • Integrate torch-sparse neighorhood sampling interface in pyg-lib
  • cugraph as a dependency inside pyg-lib
  • Call the cugraph-based sampling routine inside the GPU-based dispatcher
  • Integrate changes on PyG side
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Projects
None yet
Development

No branches or pull requests

1 participant