pub_soft/dgl
3
0
Fork 0
mirror of https://github.com/dmlc/dgl.git synced 2026-06-04 19:44:23 +08:00
Code Issues Packages Projects Releases Wiki Activity

[GraphBolt][CUDA] Optimize CopyTo performance. (#7634)

Browse Source
This commit is contained in:
Muhammed Fatih BALIN
2024-08-01 13:10:43 -04:00
committed by GitHub
parent 5c902cd59e
commit dfd491568d
2 changed files with 4 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
python/dgl/graphbolt/base.py
View File

@@ -369,9 +369,6 @@ class CopyTo(IterDataPipe):
def __iter__(self):
for data in self.datapipe:
if self.non_blocking:
# The copy is non blocking only if contents of data are pinned.
assert data.is_pinned(), f"{data} should be pinned."
yield recursive_apply(
data, apply_to, self.device, self.non_blocking
)

9
python/dgl/graphbolt/dataloader.py
View File

@@ -231,11 +231,10 @@ class DataLoader(torch_data.DataLoader):
datapipe_graph = dp_utils.replace_dp(
datapipe_graph,
copier,
copier.datapipe.transform(
lambda x: x.pin_memory()
).prefetch(2)
# After the data gets pinned, we can copy non_blocking.
.copy_to(copier.device, non_blocking=True),
# Add prefetch so that CPU and GPU can run concurrently.
copier.datapipe.prefetch(2).copy_to(
copier.device, non_blocking=True
),
)
# The stages after feature fetching is still done in the main process.