Repository: dmlc/dgl

Issue:
[GraphBolt][Cuda] An error with GPU sampling when graph is pinned.
## 🔨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

<!-- short description of the work item -->
When graph is pinned and there is `original_edge_ids`, GPU sampling will report error.
The following test code have not add to the test, it will be added after the problem is solved.
**Test Code**
```
@pytest.mark.parametrize("labor", [False, True])
@pytest.mark.parametrize("is_pinned", [False, True])
def test_sample_neighbors_return_eids_homo(labor, is_pinned):
    """Original graph in COO:
    1   0   1   0   1
    1   0   1   1   0
    0   1   0   1   0
    0   1   0   0   1
    1   0   0   0   1
    """
    # Initialize data.
    total_num_edges = 12
    indptr = torch.LongTensor([0, 3, 5, 7, 9, 12])
    indices = torch.LongTensor([0, 1, 4, 2, 3, 0, 1, 1, 2, 0, 3, 4])
    assert indptr[-1] == total_num_edges
    assert indptr[-1] == len(indices)

    # Add edge id mapping from CSC graph -> original graph.
    edge_attributes = {gb.ORIGINAL_EDGE_ID: torch.randperm(total_num_edges)}

    # Construct FusedCSCSamplingGraph.
    graph = gb.fused_csc_sampling_graph(
        indptr, indices, edge_attributes=edge_attributes
    )
    if F._default_context_str == "gpu":
        if is_pinned:
            graph.pin_memory_()
        else:
            graph = graph.to(F.ctx())

    # Generate subgraph via sample neighbors.
    nodes = torch.LongTensor([1, 3, 4]).to(F.ctx())
    sampler = graph.sample_layer_neighbors if labor else graph.sample_neighbors
    subgraph = sampler(nodes, fanouts=torch.LongTensor([-1]))

    # Verify in subgraph.
    expected_reverse_edge_ids = edge_attributes[gb.ORIGINAL_EDGE_ID][
        torch.tensor([3, 4, 7, 8, 9, 10, 11])
    ].to(F.ctx())
    assert torch.equal(
        torch.sort(expected_reverse_edge_ids)[0],
        torch.sort(subgraph.original_edge_ids)[0],
    )
    assert subgraph.original_column_node_ids is None
    assert subgraph.original_row_node_ids is None
```
**Error**: When running is_pinned=True
```
    def _convert_to_sampled_subgraph(
        self,
        C_sampled_subgraph: torch.ScriptObject,
    ) -> SampledSubgraphImpl:
        """An internal function used to convert a fused homogeneous sampled
        subgraph to general struct 'SampledSubgraphImpl'."""
        indptr = C_sampled_subgraph.indptr
        indices = C_sampled_subgraph.indices
        type_per_edge = C_sampled_subgraph.type_per_edge
        column = C_sampled_subgraph.original_column_node_ids
        original_edge_ids = C_sampled_subgraph.original_edge_ids
    
        has_original_eids = (
            self.edge_attributes is not None
            and ORIGINAL_EDGE_ID in self.edge_attributes
        )
        if has_original_eids:
>           original_edge_ids = torch.index_select(
                self.edge_attributes[ORIGINAL_EDGE_ID],
                dim=0,
                index=original_edge_ids,
            )
E           RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cpu and cuda:0! (when checking argument for argument index in method wrapper_CUDA__index_select)
```
@mfbalin 

## Depending work items or issues

<!-- what must be done before this -->



Discussion:
I was trying to see where the logic for edge-ids was located. Now, I know where. Will fix it in a quick PR.

Resolving patch:
diff --git a/python/dgl/graphbolt/impl/fused_csc_sampling_graph.py b/python/dgl/graphbolt/impl/fused_csc_sampling_graph.py
--- a/python/dgl/graphbolt/impl/fused_csc_sampling_graph.py
+++ b/python/dgl/graphbolt/impl/fused_csc_sampling_graph.py
@@ -417,10 +417,8 @@ def _convert_to_sampled_subgraph(
             and ORIGINAL_EDGE_ID in self.edge_attributes
         )
         if has_original_eids:
-            original_edge_ids = torch.index_select(
-                self.edge_attributes[ORIGINAL_EDGE_ID],
-                dim=0,
-                index=original_edge_ids,
+            original_edge_ids = torch.ops.graphbolt.index_select(
+                self.edge_attributes[ORIGINAL_EDGE_ID], original_edge_ids
             )
         if type_per_edge is None:
             # The sampled graph is already a homogeneous graph.


Test patch: