[REVIEW]Optimize cugraph-DGL csc codepath

[VibhuJawa]

This PR optimizes cugraph-DGL csc codepath and adds an end to end benchmark using cugraph-dgl.

sampled_dir = "/raid/vjawa/nov_1_bulksampling_benchmarks/ogbn_papers100M[2]_b512_f[10, 10, 10]"
dataset = HomogenousBulkSamplerDataset(meta_json_d["total_num_nodes"], edge_dir=edge_dir, sparse_format="csc", return_type="cugraph_dgl.nn.SparseGraph")
dataset.set_input_files(input_directory=sampled_dir+"/samples")
dataloader = torch.utils.data.DataLoader(dataset, collate_fn=lambda x:x, shuffle=False, num_workers=0, batch_size=None)

def run(dataloader):
    for input_nodes, output_nodes, blocks in dataloader: 
       pass

%%timeit
run(dataloader)

With PR :

2.48 s ± 14 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

MAIN:

%%timeit
9.52 s ± 151 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

E2E Benchmarks:

 python3 cugraph_dgl_benchmark.py

PR:

...
Epoch time = 70.41 seconds
Time to create MFG = 3.37 seconds
Time analysis for fanout = [10, 10, 10], batch_size = 512
mfg_creation_time_per_epoch = 3.37 seconds
feature_time_per_epoch = 44.09 seconds
m_fwd_time_per_epoch = 7.31 seconds
m_bkwd_time_per_epoch = 15.60 seconds

MAIN:

....
Epoch time = 84.72 seconds
Time to create MFG = 10.79 seconds
Time analysis for fanout = [10, 10, 10], batch_size = 512
mfg_creation_time_per_epoch = 10.79 seconds
feature_time_per_epoch = 47.09 seconds
m_fwd_time_per_epoch = 8.24 seconds
m_bkwd_time_per_epoch = 18.58 seconds

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.

[VibhuJawa]

/ok to test

[VibhuJawa]

@tingyu66 , This is the main optimization because transferring b/w CPU ->GPU 1 tensor at a time was slow.

[tingyu66]

Thanks for resolving the .item() overhead. 👍

[VibhuJawa]

@oorliu ,

The only dgl specific args for our benchmarking efforts are:

                --reverse_edges \
                --sampling_target_framework cugraph_dgl_csr

[rlratzel]

/ok to test

[alexbarghi-nv]

👍

[tingyu66]

Does this setting also work for cugraph-dgl COO code path?

[VibhuJawa]

I will have to test, I just focussed on CSC code path to ensure we have success there (Given it is the fastest one) but I dont see why it will not work.

[tingyu66]

But prior_sources_behavior needs to be True for dgl, right? I think COO path would be useful in the future for debugging purposes.

[VibhuJawa]

@tingyu66 , I have not spent time exploring COO code path, do you think I should focus on it , do we expect it to have equivalent speed (maybe after optimizations) ?

[VibhuJawa]

I will have to add something like sampling_target_framework=='cugraph_dgl_coo' with prior_sources_behavior =True . The above path is for cugraph-pyG .

[tingyu66]

I don't think getting COO up to speed is a priority, but it would be useful to include/document the parameter combinations needed for COO path, unless we no longer support it.

[VibhuJawa]

I will add that as a followup:

Filed #3981 to track.

[alexbarghi-nv]

/merge

[BradReesWork]

/ok to test

[alexbarghi-nv]

/ok to test