Repository: dmlc/dgl

Issue:
[GraphBolt][CUDA] Eliminate unnecessary synchronization inside `unique_and_compact`
## 🔨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

https://github.com/dmlc/dgl/blob/22a2513d0a85f973ce414755c1f26f5f8f12ef3b/python/dgl/graphbolt/internal/sample_utils.py#L361
introduces a CPU GPU synchronization causing the following gaps in my profile (marked with red drawing), costing around %10 slowdowns of the overall sampling stage (5.5 ms duration, around 0.5ms total gap):
![Unique_profile](https://github.com/dmlc/dgl/assets/16190101/1c5fcf77-efef-4e11-b0ed-66689e643e54)
Link to the image if now showing above: [Image](https://github.com/dmlc/dgl/assets/16190101/1c5fcf77-efef-4e11-b0ed-66689e643e54)

I suggest moving this assertion to the constructor of `CSCFormatBase` and disabling it for CUDA (CPU access is okay as there is no data transfer). We need to ensure that the code is bug-free in the first place.

Profile with the fix in #6863, notice that the marked gaps disappeared, significantly speeding up the sampling stage and the last `unique_and_compact` by %30 (920us to 690us).
![image](https://github.com/dmlc/dgl/assets/16190101/55f863a5-160b-4941-ba5c-e3dfb30d1168)
Link to the image if now showing above: [Image](https://github.com/dmlc/dgl/assets/16190101/55f863a5-160b-4941-ba5c-e3dfb30d1168)

If the images are not showing, then click the second link on the line below.

@frozenbugs @Rhett-Ying @peizhou001 


Discussion:

Resolving patch:
diff --git a/python/dgl/graphbolt/base.py b/python/dgl/graphbolt/base.py
--- a/python/dgl/graphbolt/base.py
+++ b/python/dgl/graphbolt/base.py
@@ -164,6 +164,14 @@ class CSCFormatBase:
     indptr: torch.Tensor = None
     indices: torch.Tensor = None
 
+    def __init__(self, indptr: torch.Tensor, indices: torch.Tensor):
+        self.indptr = indptr
+        self.indices = indices
+        if not indptr.is_cuda:
+            assert self.indptr[-1] == len(
+                self.indices
+            ), "The last element of indptr should be the same as the length of indices."
+
     def __repr__(self) -> str:
         return _csc_format_base_str(self)
 
diff --git a/python/dgl/graphbolt/internal/sample_utils.py b/python/dgl/graphbolt/internal/sample_utils.py
--- a/python/dgl/graphbolt/internal/sample_utils.py
+++ b/python/dgl/graphbolt/internal/sample_utils.py
@@ -254,9 +254,6 @@ def unique_and_compact_csc_formats(
     for etype, csc_format in csc_formats.items():
         if device is None:
             device = csc_format.indices.device
-        assert csc_format.indptr[-1] == len(
-            csc_format.indices
-        ), "The last element of indptr should be the same as the length of indices."
         src_type, _, dst_type = etype_str_to_tuple(etype)
         assert len(unique_dst_nodes.get(dst_type, [])) + 1 == len(
             csc_format.indptr
@@ -358,9 +355,6 @@ def compact_csc_format(
             assert isinstance(
                 dst_nodes, torch.Tensor
             ), "Edge type not supported in homogeneous graph."
-            assert csc_formats.indptr[-1] == len(
-                csc_formats.indices
-            ), "The last element of indptr should be the same as the length of indices."
             assert len(dst_nodes) + 1 == len(
                 csc_formats.indptr
             ), "The seed nodes should correspond to indptr."
@@ -381,9 +375,6 @@ def compact_csc_format(
         compacted_csc_formats = {}
         original_row_ids = copy.deepcopy(dst_nodes)
         for etype, csc_format in csc_formats.items():
-            assert csc_format.indptr[-1] == len(
-                csc_format.indices
-            ), "The last element of indptr should be the same as the length of indices."
             src_type, _, dst_type = etype_str_to_tuple(etype)
             assert len(dst_nodes.get(dst_type, [])) + 1 == len(
                 csc_format.indptr
diff --git a/python/dgl/graphbolt/minibatch.py b/python/dgl/graphbolt/minibatch.py
--- a/python/dgl/graphbolt/minibatch.py
+++ b/python/dgl/graphbolt/minibatch.py
@@ -202,7 +202,7 @@ def blocks(self):
                             v.indices,
                             torch.arange(
                                 0,
-                                v.indptr[-1],
+                                len(v.indices),
                                 device=v.indptr.device,
                                 dtype=v.indptr.dtype,
                             ),
@@ -227,7 +227,7 @@ def blocks(self):
                         sampled_csc.indices,
                         torch.arange(
                             0,
-                            sampled_csc.indptr[-1],
+                            len(sampled_csc.indices),
                             device=sampled_csc.indptr.device,
                             dtype=sampled_csc.indptr.dtype,
                         ),


Test patch:
diff --git a/tests/python/pytorch/graphbolt/test_base.py b/tests/python/pytorch/graphbolt/test_base.py
--- a/tests/python/pytorch/graphbolt/test_base.py
+++ b/tests/python/pytorch/graphbolt/test_base.py
@@ -244,3 +244,11 @@ def test_csc_format_base_representation():
 )"""
     )
     assert str(csc_format_base) == expected_result, print(csc_format_base)
+
+
+def test_csc_format_base_incorrect_indptr():
+    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
+    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4])
+    with pytest.raises(AssertionError):
+        # The value of last element in indptr is not corresponding to indices.
+        csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
diff --git a/tests/python/pytorch/graphbolt/test_graphbolt_utils.py b/tests/python/pytorch/graphbolt/test_graphbolt_utils.py
--- a/tests/python/pytorch/graphbolt/test_graphbolt_utils.py
+++ b/tests/python/pytorch/graphbolt/test_graphbolt_utils.py
@@ -350,14 +350,6 @@ def test_unique_and_compact_incorrect_indptr():
     with pytest.raises(AssertionError):
         gb.unique_and_compact_csc_formats(csc_formats, seeds)
 
-    seeds = torch.tensor([1, 3, 5, 2, 6])
-    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
-    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4])
-    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
-    # The value of last element in indptr is not corresponding to indices.
-    with pytest.raises(AssertionError):
-        gb.unique_and_compact_csc_formats(csc_formats, seeds)
-
 
 def test_compact_csc_format_hetero():
     dst_nodes = {
@@ -449,11 +441,3 @@ def test_compact_incorrect_indptr():
     # The number of seeds is not corresponding to indptr.
     with pytest.raises(AssertionError):
         gb.compact_csc_format(csc_formats, seeds)
-
-    seeds = torch.tensor([1, 3, 5, 2, 6])
-    indptr = torch.tensor([0, 2, 4, 6, 7, 11])
-    indices = torch.tensor([2, 3, 1, 4, 5, 2, 5, 1, 4, 4])
-    csc_formats = gb.CSCFormatBase(indptr=indptr, indices=indices)
-    # The value of last element in indptr is not corresponding to indices.
-    with pytest.raises(AssertionError):
-        gb.compact_csc_format(csc_formats, seeds)