2-hop sampling on GPU

[brovatten]

Hi,

I want to offline sample 2-hop neighbours for every node of my Wikidata subset (5M nodes, 100M edges). This has to be without batching. Ultimately I want a dictionary with each node and its 2-hop neighbourhood.

It’s important that all 1-hop neighbours share all its possible edges with eachother (including 2-hop edges hence requires 2-hop node sampling I think). Remaining 2-hop nodes are not as important, but would optimally also be contained.

The problem is dgl.khop(g, nodes,2) is too slow. I would need to run it on GPU or use something custom. I tried doing it manually with tensors but it also took too much time.

Thesis is closing in so appreciate all advice

[peizhou001]

Hi, @brovatten, unfortunately there is no existing GPU implementation of khop_graph. Post the source code here for reference or hack. It would also be helpful if you can provide the performance metrics so that it can be a baseline for future improvement.