[GraphBolt] sample_neighbors() on CPU with prob/mask is 14x slower than w/o prob/mask
- Dominant language
- Python
- Stars
- 14.3k
- Forks
- 3.1k
- PR merge metrics
- No merged PRs in 30d
Description
## 🔨Work Item
**IMPORTANT:**
* This template is only for dev team to track project progress. For feature request or bug report, please use the corresponding issue templates.
* DO NOT create a new work item if the purpose is to fix an existing issue or feature request. We will directly use the issue in the project tracker.
Project tracker: https://github.com/orgs/dmlc/projects/2
## Description
Below numbers are from [node classification example](https://github.com/dmlc/dgl/blob/master/examples/sampling/graphbolt/node_classification.py) on latest master branch(2024.06.14) with **CPU sampling**.
### prob data
```
# Add prob data.
num_edges = dataset.graph.total_num_edges
prob_data = torch.rand(num_edges)
prob_data[torch.randperm(len(prob_data))[: int(len(prob_data) * 0.5)]] = 0.0
dataset.graph.add_edge_attribute("prob", prob_data)
```
non-dist, gb, homo, nc, no prob
Training...
Training: 3it [00:05, 1.55s/it]---- Average time for sampling: 0.082120824418962
Training: 10it [00:13, 1.19s/it]---- Average time for sampling: 0.0836272995453328
Training: 16it [00:19, 1.12s/it]---- Average time for sampling: 0.08516605570912361
Training: 23it [00:27, 1.14s/it]---- Average time for sampling: 0.0846009589266032
Training: 30it [00:35, 1.13s/it]---- Average time for sampling: 0.0846672392077744
Training: 36it [00:42, 1.12s/it]---- Average time for sampling: 0.08583860701570908
non-dist, gb, homo, nc, prob
Training: 4it [00:22, 5.03s/it]---- Average time for sampling: 1.1615502193570137
Training: 11it [00:52, 4.32s/it]---- Average time for sampling: 1.228120240289718
Training: 18it [01:22, 4.19s/it]---- Average time for sampling: 1.2553682390290002
Training: 24it [01:47, 4.19s/it]---- Average time for sampling: 1.2268523721955717
Training: 31it [02:16, 4.18s/it]---- Average time for sampling: 1.2365567354112863
Training: 38it [02:46, 4.21s/it]---- Average time for sampling: 1.2496669557721665
### mask data
```
# Add prob data.
num_edges = dataset.graph.total_num_edges
prob_data = torch.rand(num_edges)
prob_data[torch.randperm(len(prob_data))[: int(len(prob_data) * 0.5)]] = 0.0
dataset.graph.add_edge_attribute("prob", prob_data)
mask_data = (prob_data > 0.2).to(torch.float32)
dataset.graph.add_edge_attribute("mask", mask_data)
```
GB + NoMask
Training...
Training: 3it [00:00, 3.37it/s]---- Average time for sampling: 0.0064875221811234954
Training: 10it [00:02, 4.05it/s]---- Average time for sampling: 0.006722455704584717
Training: 16it [00:04, 3.92it/s]---- Average time for sampling: 0.008144200344880422
Training: 23it [00:06, 3.68it/s]---- Average time for sampling: 0.008010426000691951
Training: 30it [00:07, 4.15it/s]---- Average time for sampling: 0.008361106105148793
Training: 36it [00:09, 4.31it/s]---- Average time for sampling: 0.00815806492852668
GB + Mask
Training...
Training: 3it [00:01, 2.61it/s]---- Average time for sampling: 0.04655262678861618
Training: 10it [00:03, 3.52it/s]---- Average time for sampling: 0.05128640588372946
Training: 16it [00:04, 3.61it/s]---- Average time for sampling: 0.054167306640495856
Training: 23it [00:06, 3.64it/s]---- Average time for sampling: 0.052452172990888356
Training: 30it [00:08, 3.74it/s]---- Average time for sampling: 0.05273009695112705
Training: 36it [00:10, 3.49it/s]---- Average time for sampling: 0.05279933324394127
### critical call
https://github.com/dmlc/dgl/blob/ed50c170dda9627730cb8ee4c7110205b6ea09de/graphbolt/src/fused_csc_sampling_graph.cc#L1322-L1328
## Depending work items or issues
Contributor guide
No contributing guide indexed for this repository
Assessment
This issue has not been assessed yet.