Repository: dmlc/dgl

Issue:
Cannot use UVA on a HeteroGraph
## 🐛 Bug
EdgePredictionSampler fails in `g.edge_subgraph(seed_edges, relabel_nodes=False, output_device=self.output_device)` when using uva on a Heterogeneous graph.


## To Reproduce

Steps to reproduce the behavior:

1. Create edge dataloader, with parameters `(g, {'etype': eids.to('cuda')}, sampler, device='cuda', use_uva=True)`, where `g` is of instance HeteroGraph.
2. Iterate graph loader and it will throw the exception shown in additional context.

Snippet of code used. It will work, when eids is on CPU and `use_uva=False`. 
```
        sampler = dgl.dataloading.MultiLayerFullNeighborSampler(3, prefetch_edge_feats={etype: ['norm'] for etype in g.etypes})
        neg_sampler = UniformRecommendableItemSampler(1)

        dataloader = EdgeDataLoader(
            g, {'ui': eids.to(self.device)}, sampler, device=self.device, batch_size=self.batch_size,
            shuffle=True, drop_last=False, exclude='reverse_types', use_uva=self.use_cuda,
            reverse_etypes={'ui': 'iu'}, 
            negative_sampler=neg_sampler)
```
NOTE: UniformRecommendableItemSampler is a custom sampler limiting negative nodes to a subset of possible nodes.

## Expected behavior

It would not throw errors. I assume it to be `_CAPI_DGLHeteroEdgeSubgraph` that cannot handle eids on the cuda device while the graph is on CPU.

## Environment

 - DGL Version (e.g., 1.0): 0.8.0.post1
 - Backend Library & Version (e.g., PyTorch 0.4.1, MXNet/Gluon 1.3): 1.9.0
 - OS (e.g., Linux): Ubuntu 20.04
 - How you installed DGL (`conda`, `pip`, source): Pip
 - Build command you used (if compiling from source):
 - Python version: 3.8.10
 - CUDA/cuDNN version (if applicable): 11.1
 - GPU models and configuration (e.g. V100): 2070 GTX

## Additional context

```
Traceback (most recent call last):
  File "/pathtovenv/.venv/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 514, in __next__
    self._next_non_threaded() if not self.use_thread else self._next_threaded()
  File "/pathtovenv/.venv/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 487, in _next_non_threaded
    batch = next(self.dataloader_it)
  File "/pathtovenv/.venv/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 521, in __next__
    data = self._next_data()
  File "/pathtovenv/.venv/lib/python3.8/site-packages/torch/utils/data/dataloader.py", line 561, in _next_data
    data = self._dataset_fetcher.fetch(index)  # may raise StopIteration
  File "/pathtovenv/.venv/lib/python3.8/site-packages/torch/utils/data/_utils/fetch.py", line 35, in fetch
    return self.collate_fn(data)
  File "/pathtovenv/.venv/lib/python3.8/site-packages/dgl/dataloading/dataloader.py", line 531, in __call__
    batch = self.sample_func(self.g, items)
  File "/pathtovenv/.venv/lib/python3.8/site-packages/dgl/dataloading/base.py", line 409, in sample
    pair_graph = g.edge_subgraph(
  File "/pathtovenv/.venv/lib/python3.8/site-packages/dgl/utils/internal.py", line 909, in _fn
    return func(*args, **kwargs)
  File "/pathtovenv/.venv/lib/python3.8/site-packages/dgl/subgraph.py", line 309, in edge_subgraph
    sgi = graph._graph.edge_subgraph(induced_edges, not relabel_nodes)
  File "/pathtovenv/.venv/lib/python3.8/site-packages/dgl/heterograph_index.py", line 924, in edge_subgraph
    return _CAPI_DGLHeteroEdgeSubgraph(self, eids, preserve_nodes)
  File "dgl/_ffi/_cython/./function.pxi", line 287, in dgl._ffi._cy3.core.FunctionBase.__call__
  File "dgl/_ffi/_cython/./function.pxi", line 222, in dgl._ffi._cy3.core.FuncCall
  File "dgl/_ffi/_cython/./function.pxi", line 211, in dgl._ffi._cy3.core.FuncCall3
  File "dgl/_ffi/_cython/./base.pxi", line 155, in dgl._ffi._cy3.core.CALL
dgl._ffi.base.DGLError: [13:51:57] /opt/dgl/src/graph/heterograph.cc:127: Check failed: rg->Context() == ctx (cuda:0 vs. cpu:0) : Each relation graph must have the same context.
Stack trace:
  [bt] (0) /pathtovenv/.venv/lib/python3.8/site-packages/dgl/libdgl.so(dmlc::LogMessageFatal::~LogMessageFatal()+0x4f) [0x7f682d99452f]
  [bt] (1) /pathtovenv/.venv/lib/python3.8/site-packages/dgl/libdgl.so(dgl::HeteroGraph::HeteroGraph(std::shared_ptr<dgl::GraphInterface>, std::vector<std::shared_ptr<dgl::BaseHeteroGraph>, std::allocator<std::shared_ptr<dgl::BaseHeteroGraph> > > const&, std::vector<long, std::allocator<long> > const&)+0x88e) [0x7f682dd213be]
  [bt] (2) /pathtovenv/.venv/lib/python3.8/site-packages/dgl/libdgl.so(dgl::HeteroGraph::EdgeSubgraph(std::vector<dgl::runtime::NDArray, std::allocator<dgl::runtime::NDArray> > const&, bool) const+0x95b) [0x7f682dd2604b]
  [bt] (3) /pathtovenv/.venv/lib/python3.8/site-packages/dgl/libdgl.so(+0x6b189b) [0x7f682dd3589b]
  [bt] (4) /pathtovenv/.venv/lib/python3.8/site-packages/dgl/libdgl.so(DGLFuncCall+0x48) [0x7f682dcb8e58]
  [bt] (5) /pathtovenv/.venv/lib/python3.8/site-packages/dgl/_ffi/_cy3/core.cpython-38-x86_64-linux-gnu.so(+0x16019) [0x7f6962d0e019]
  [bt] (6) /pathtovenv/.venv/lib/python3.8/site-packages/dgl/_ffi/_cy3/core.cpython-38-x86_64-linux-gnu.so(+0x162fb) [0x7f6962d0e2fb]
  [bt] (7) /pathtovenv/.venv/bin/python(_PyObject_MakeTpCall+0x296) [0x5f6a46]
  [bt] (8) /pathtovenv/.venv/bin/python(_PyEval_EvalFrameDefault+0x854c) [0x57322c]
```



Discussion:
The crash occurs whenever user only specifies a part of the edge types in `edge_subgraph`.  Similar errors are also observed in `node_subgraph`, `in_subgraph`, etc.
I need to investigate further to figure out where the fix goes, but it looks like the empty relations are created with a context of a CPU instead of GPU (as the ones that have edges go through UVA). There was an issue like this a while back for empty tensors defaulting to the CPU.
```
(cuda-gdb) print ctx
$5 = {device_type = kDLGPU, device_id = 0}
(cuda-gdb) print rg->NumEdgeTypes()
print rg->Num$6 = 1
(cuda-gdb) print rg->NumEdges()
Too few arguments in function call.
(cuda-gdb) print rg->NumEdges(0)
$7 = 0
(cuda-gdb) frame
#2  0x00007ffe8f946f25 in dgl::(anonymous namespace)::HeteroGraphSanityCheck (meta_graph=..., rel_graphs=...) at /home/dominique/src/dgl/src/graph/heterograph.cc:127
127         CHECK_EQ(rg->Context(), ctx) << "Each relation graph must have the same context.";
```
I think this is the fix the for `edge_subgraph`:
```
diff --git a/python/dgl/subgraph.py b/python/dgl/subgraph.py
index 9b1c9893..cde894e6 100644
--- a/python/dgl/subgraph.py
+++ b/python/dgl/subgraph.py
@@ -303,9 +303,10 @@ def edge_subgraph(graph, edges, *, relabel_nodes=True, store_ids=True, output_de
             return utils.prepare_tensor(graph, e, 'edges["{}"]'.format(etype))

     edges = {graph.to_canonical_etype(etype): e for etype, e in edges.items()}
+    edges_device = F.context(list(edges.values())[0]) if len(edges) > 0 else g.device
     induced_edges = []
     for cetype in graph.canonical_etypes:
-        eids = edges.get(cetype, F.copy_to(F.tensor([], graph.idtype), graph.device))
+        eids = edges.get(cetype, F.copy_to(F.tensor([], graph.idtype), edges_device))
         induced_edges.append(_process_edges(cetype, eids))
     device = context_of(induced_edges)
     sgi = graph._graph.edge_subgraph(induced_edges, not relabel_nodes)
```

Resolving patch:
diff --git a/python/dgl/subgraph.py b/python/dgl/subgraph.py
--- a/python/dgl/subgraph.py
+++ b/python/dgl/subgraph.py
@@ -10,7 +10,6 @@
 from . import backend as F
 from . import graph_index
 from . import heterograph_index
-from . import ndarray as nd
 from .heterograph import DGLHeteroGraph
 from . import utils
 from .utils import recursive_apply, context_of
@@ -142,12 +141,12 @@ def _process_nodes(ntype, v):
             return F.astype(F.nonzero_1d(F.copy_to(v, graph.device)), graph.idtype)
         else:
             return utils.prepare_tensor(graph, v, 'nodes["{}"]'.format(ntype))
+    nodes = {ntype: _process_nodes(ntype, v) for ntype, v in nodes.items()}
+    device = context_of(nodes)
 
-    induced_nodes = []
-    for ntype in graph.ntypes:
-        nids = nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), graph.device))
-        induced_nodes.append(_process_nodes(ntype, nids))
-    device = context_of(induced_nodes)
+    induced_nodes = [
+        nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), device))
+        for ntype in graph.ntypes]
     sgi = graph._graph.node_subgraph(induced_nodes, relabel_nodes)
     induced_edges = sgi.induced_edges
     # (BarclayII) should not write induced_nodes = sgi.induced_nodes due to the same
@@ -301,13 +300,13 @@ def _process_edges(etype, e):
             return F.astype(F.nonzero_1d(F.copy_to(e, graph.device)), graph.idtype)
         else:
             return utils.prepare_tensor(graph, e, 'edges["{}"]'.format(etype))
-
     edges = {graph.to_canonical_etype(etype): e for etype, e in edges.items()}
-    induced_edges = []
-    for cetype in graph.canonical_etypes:
-        eids = edges.get(cetype, F.copy_to(F.tensor([], graph.idtype), graph.device))
-        induced_edges.append(_process_edges(cetype, eids))
-    device = context_of(induced_edges)
+    edges = {etype: _process_edges(etype, e) for etype, e in edges.items()}
+    device = context_of(edges)
+    induced_edges = [
+        edges.get(cetype, F.copy_to(F.tensor([], graph.idtype), device))
+        for cetype in graph.canonical_etypes]
+
     sgi = graph._graph.edge_subgraph(induced_edges, not relabel_nodes)
     induced_nodes_or_device = sgi.induced_nodes if relabel_nodes else device
     subg = _create_hetero_subgraph(
@@ -430,12 +429,9 @@ def in_subgraph(graph, nodes, *, relabel_nodes=False, store_ids=True, output_dev
         nodes = {graph.ntypes[0] : nodes}
     nodes = utils.prepare_tensor_dict(graph, nodes, 'nodes')
     device = context_of(nodes)
-    nodes_all_types = []
-    for ntype in graph.ntypes:
-        if ntype in nodes:
-            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
-        else:
-            nodes_all_types.append(nd.NULL[graph._idtype_str])
+    nodes_all_types = [
+        F.to_dgl_nd(nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), device)))
+        for ntype in graph.ntypes]
 
     sgi = _CAPI_DGLInSubgraph(graph._graph, nodes_all_types, relabel_nodes)
     induced_nodes_or_device = sgi.induced_nodes if relabel_nodes else device
@@ -560,12 +556,9 @@ def out_subgraph(graph, nodes, *, relabel_nodes=False, store_ids=True, output_de
         nodes = {graph.ntypes[0] : nodes}
     nodes = utils.prepare_tensor_dict(graph, nodes, 'nodes')
     device = context_of(nodes)
-    nodes_all_types = []
-    for ntype in graph.ntypes:
-        if ntype in nodes:
-            nodes_all_types.append(F.to_dgl_nd(nodes[ntype]))
-        else:
-            nodes_all_types.append(nd.NULL[graph._idtype_str])
+    nodes_all_types = [
+        F.to_dgl_nd(nodes.get(ntype, F.copy_to(F.tensor([], graph.idtype), device)))
+        for ntype in graph.ntypes]
 
     sgi = _CAPI_DGLOutSubgraph(graph._graph, nodes_all_types, relabel_nodes)
     induced_nodes_or_device = sgi.induced_nodes if relabel_nodes else device
@@ -693,7 +686,8 @@ def khop_in_subgraph(graph, nodes, k, *, relabel_nodes=True, store_ids=True, out
 
     last_hop_nodes = nodes
     k_hop_nodes_ = [last_hop_nodes]
-    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), graph.device)
+    device = context_of(nodes)
+    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), device)
     for _ in range(k):
         current_hop_nodes = {nty: [] for nty in graph.ntypes}
         for cetype in graph.canonical_etypes:
@@ -853,7 +847,8 @@ def khop_out_subgraph(graph, nodes, k, *, relabel_nodes=True, store_ids=True, ou
 
     last_hop_nodes = nodes
     k_hop_nodes_ = [last_hop_nodes]
-    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), graph.device)
+    device = context_of(nodes)
+    place_holder = F.copy_to(F.tensor([], dtype=graph.idtype), device)
     for _ in range(k):
         current_hop_nodes = {nty: [] for nty in graph.ntypes}
         for cetype in graph.canonical_etypes:


Test patch:
diff --git a/tests/compute/test_subgraph.py b/tests/compute/test_subgraph.py
--- a/tests/compute/test_subgraph.py
+++ b/tests/compute/test_subgraph.py
@@ -100,6 +100,8 @@ def create_test_heterograph(idtype):
         ('user', 'wishes', 'game'): ([0, 2], [1, 0]),
         ('developer', 'develops', 'game'): ([0, 1], [0, 1])
     }, idtype=idtype, device=F.ctx())
+    for etype in g.etypes:
+        g.edges[etype].data['weight'] = F.randn((g.num_edges(etype),))
     assert g.idtype == idtype
     assert g.device == F.ctx()
     return g
@@ -629,6 +631,28 @@ def test_subframes(parent_idx_device, child_device):
     if parent_device == 'uva':
         g.unpin_memory_()
 
+@unittest.skipIf(F._default_context_str != "gpu", reason="UVA only available on GPU")
+@pytest.mark.parametrize('device', [F.cpu(), F.cuda()])
+@parametrize_dtype
+def test_uva_subgraph(idtype, device):
+    g = create_test_heterograph(idtype)
+    g = g.to(F.cpu())
+    g.create_formats_()
+    g.pin_memory_()
+    indices = {'user': F.copy_to(F.tensor([0], idtype), device)}
+    edge_indices = {'follows': F.copy_to(F.tensor([0], idtype), device)}
+    assert g.subgraph(indices).device == device
+    assert g.edge_subgraph(edge_indices).device == device
+    assert g.in_subgraph(indices).device == device
+    assert g.out_subgraph(indices).device == device
+    if dgl.backend.backend_name != 'tensorflow':
+        # (BarclayII) Most of Tensorflow functions somehow do not preserve device: a CPU tensor
+        # becomes a GPU tensor after operations such as concat(), unique() or even sin().
+        # Not sure what should be the best fix.
+        assert g.khop_in_subgraph(indices, 1)[0].device == device
+        assert g.khop_out_subgraph(indices, 1)[0].device == device
+    assert g.sample_neighbors(indices, 1).device == device
+    g.unpin_memory_()
+
 if __name__ == '__main__':
-    test_khop_out_subgraph(F.int64)
-    test_subframes(('cpu', F.cpu()), F.cuda())
+    test_uva_subgraph(F.int64, F.cpu())