Repository: dmlc/dgl
Issue category: Performance Issue

Issue:
[Bug][Performance] Significant overhead when `record_stream` for subgraph in dataloading
## 🐛 Bug

It seems there are many overheads observed in dataloader during feature prefetching when `use_alternate_streams=True` . When it calls `_record_stream` for every subgraph (batch), like the following: 
https://github.com/dmlc/dgl/blob/829ce109f1604ee59ac039331a3609896301e8de/python/dgl/dataloading/dataloader.py#L339
many unnecessary h2d copies (or gpu feature slicing if use_uva) are invoked from L339, due to materializing the column data(feature) within its `record_stream` method, i.e., within its `data` method, for example from here (when uva is disabled)
https://github.com/dmlc/dgl/blob/829ce109f1604ee59ac039331a3609896301e8de/python/dgl/frame.py#L246-L251
called by `record_stream` method for each column:
https://github.com/dmlc/dgl/blob/829ce109f1604ee59ac039331a3609896301e8de/python/dgl/frame.py#L551

It poses significant slow-down for both cases: `use_uva=True` and `use_uva=False`, for example, when running `example/graphsage/node_classification.py`, with uva on, `use_alternate_streams=True` causes **1.56x** more time for each epoch, while with uva off, `use_alternate_streams=True` (default) causes **1.55x** more time for each epoch.


## To Reproduce

set `use_alternate_streams=True` and run

```
python node_classification.py
```

## Environment

 - DGL Version (e.g., 1.0): 1.0
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): Pytorch 1.13
 - How you installed DGL (`conda`, `pip`, source): source
 - GPU models and configuration (e.g. V100): A5000




Hints:

Patch:
diff --git a/python/dgl/frame.py b/python/dgl/frame.py
index 46c52a20a99b..707964c5abc6 100644
--- a/python/dgl/frame.py
+++ b/python/dgl/frame.py
@@ -5,7 +5,7 @@
 from collections.abc import MutableMapping
 
 from . import backend as F
-from .base import DGLError, dgl_warning
+from .base import dgl_warning, DGLError
 from .init import zero_initializer
 from .storages import TensorStorage
 from .utils import gather_pinned_tensor_rows, pin_memory_inplace
@@ -41,6 +41,17 @@ def flatten(self):
             flat_index = F.gather_row(flat_index, index)
         return flat_index
 
+    def record_stream(self, stream):
+        """Record stream for index.
+
+        Parameters
+        ----------
+        stream : torch.cuda.Stream.
+        """
+        for index in self._indices:
+            if F.context(index) != F.cpu():
+                index.record_stream(stream)
+
 
 class LazyFeature(object):
     """Placeholder for feature prefetching.
@@ -548,7 +559,13 @@ def record_stream(self, stream):
         """
         if F.get_preferred_backend() != "pytorch":
             raise DGLError("record_stream only supports the PyTorch backend.")
-        self.data.record_stream(stream)
+        if self.index is not None and (
+            isinstance(self.index, _LazyIndex)
+            or F.context(self.index) != F.cpu()
+        ):
+            self.index.record_stream(stream)
+        if F.context(self.storage) != F.cpu():
+            self.storage.record_stream(stream)
 
 
 class Frame(MutableMapping):


Test patch: