Add a SAGEConv model using the primitives in cugraph-ops

[tingyu66]

This PR adds a GraphSAGE model torch_geometric.nn.SAGEConvCuGraph that leverages the optimized aggregation function in cugraph-ops, similar to #6278.

Fixes rapidsai/cugraph-ops#178.

[tingyu66]

@rusty1s This is now ready for review. Regarding the CI failure, should I disable the test for now as CI is cpu-only?

The interface of SAGEConvCuGraph largely resembles SAGEConv. It accepts edge_index in both torch.Tensor and SparseTensor type. As you will see in the code, we rely on the utilities in torch_sparse to generate a csc representation which is required by cugraph-ops. The conversion is actually the performance bottleneck, so it would be nice if the dataloader can create a csc layout in SparseStorage.

cugraph-ops provides aggregation primitives specialized for both full graphs and bipartite graphs, with the latter being more performant. That is why you see the extra argument max_num_neighbors, which allows the model to opt for the faster option when given bipartite graphs.

CC: @stadlmax

[rusty1s]

Thanks. You can use the @onlyCUDA decorator in torch_geometric.testing to only test this on GPU devices.

[codecov]

Codecov Report

Merging #6388 (b0a2d54) into master (5868159) will increase coverage by 0.00%.
The diff coverage is 88.88%.

@@           Coverage Diff           @@
##           master    #6388   +/-   ##
=======================================
  Coverage   86.23%   86.24%           
=======================================
  Files         413      416    +3     
  Lines       22751    22769   +18     
=======================================
+ Hits        19620    19636   +16     
- Misses       3131     3133    +2     

Impacted Files	Coverage Δ
torch_geometric/nn/conv/fused_gat_conv.py	100.00% <ø> (ø)
torch_geometric/nn/conv/cugraph/sage_conv.py	80.00% <80.00%> (ø)
torch_geometric/nn/conv/__init__.py	100.00% <100.00%> (ø)
torch_geometric/nn/conv/cugraph/__init__.py	100.00% <100.00%> (ø)
torch_geometric/nn/conv/cugraph/base.py	100.00% <100.00%> (ø)

📣 We’re building smart automated test selection to slash your CI/CD build times. Learn more

[rusty1s]

Hey @tingyu66 Can you give me write access to the PR?

[tingyu66]

Hey @tingyu66 Can you give me write access to the PR?

Invited.

[rusty1s]

@tingyu66 Thanks. I made a past. Can you confirm that everything runs through? I sadly still cannot test it :(

[tingyu66]

@tingyu66 Thanks. I made a past. Can you confirm that everything runs through? I sadly still cannot test it :(

Thank you for updating to_csc. The test has passed locally on my end.

[rusty1s]

Super, then we are good to merge. Can we update the other modules to look similar?

[tingyu66]

Super, then we are good to merge. Can we update the other modules to look similar?

Sure. I will make changes accordingly.

[tingyu66]

Shall we merge this in? Still need 1 approving review.