[Traversal] Fall back to CPU for Graph Traversal as a Workaround (#2115)

* workaround

* Allow Test for Traversal on GPU Graphs
Co-authored-by: Ubuntu <ubuntu@ip-172-31-1-5.us-west-2.compute.internal>
Co-authored-by: Quan (Andy) Gan <coin2028@hotmail.com>

python/dgl/traversal.py

@@ -44,7 +44,8 @@ def bfs_nodes_generator(graph, source, reverse=False):
         'DGLGraph is deprecated, Please use DGLHeteroGraph'
     assert len(graph.canonical_etypes) == 1, \
         'bfs_nodes_generator only support homogeneous graph'
-    gidx = graph._graph
+    # Workaround before support for GPU graph
+    gidx = graph._graph.copy_to(utils.to_dgl_context(F.cpu()))
     source = utils.toindex(source, dtype=graph._idtype_str)
     ret = _CAPI_DGLBFSNodes_v2(gidx, source.todgltensor(), reverse)
     all_nodes = utils.toindex(ret(0), dtype=graph._idtype_str).tousertensor()
@@ -88,7 +89,8 @@ def bfs_edges_generator(graph, source, reverse=False):
         'DGLGraph is deprecated, Please use DGLHeteroGraph'
     assert len(graph.canonical_etypes) == 1, \
         'bfs_edges_generator only support homogeneous graph'
-    gidx = graph._graph
+    # Workaround before support for GPU graph
+    gidx = graph._graph.copy_to(utils.to_dgl_context(F.cpu()))
     source = utils.toindex(source, dtype=graph._idtype_str)
     ret = _CAPI_DGLBFSEdges_v2(gidx, source.todgltensor(), reverse)
     all_edges = utils.toindex(ret(0), dtype=graph._idtype_str).tousertensor()
@@ -129,7 +131,8 @@ def topological_nodes_generator(graph, reverse=False):
         'DGLGraph is deprecated, Please use DGLHeteroGraph'
     assert len(graph.canonical_etypes) == 1, \
         'topological_nodes_generator only support homogeneous graph'
-    gidx = graph._graph
+    # Workaround before support for GPU graph
+    gidx = graph._graph.copy_to(utils.to_dgl_context(F.cpu()))
     ret = _CAPI_DGLTopologicalNodes_v2(gidx, reverse)
     all_nodes = utils.toindex(ret(0), dtype=graph._idtype_str).tousertensor()
     # TODO(minjie): how to support directly creating python list
@@ -177,7 +180,8 @@ def dfs_edges_generator(graph, source, reverse=False):
         'DGLGraph is deprecated, Please use DGLHeteroGraph'
     assert len(graph.canonical_etypes) == 1, \
         'dfs_edges_generator only support homogeneous graph'
-    gidx = graph._graph
+    # Workaround before support for GPU graph
+    gidx = graph._graph.copy_to(utils.to_dgl_context(F.cpu()))
     source = utils.toindex(source, dtype=graph._idtype_str)
     ret = _CAPI_DGLDFSEdges_v2(gidx, source.todgltensor(), reverse)
     all_edges = utils.toindex(ret(0), dtype=graph._idtype_str).tousertensor()
@@ -252,7 +256,8 @@ def dfs_labeled_edges_generator(
         'DGLGraph is deprecated, Please use DGLHeteroGraph'
     assert len(graph.canonical_etypes) == 1, \
         'dfs_labeled_edges_generator only support homogeneous graph'
-    gidx = graph._graph
+    # Workaround before support for GPU graph
+    gidx = graph._graph.copy_to(utils.to_dgl_context(F.cpu()))
     source = utils.toindex(source, dtype=graph._idtype_str)
     ret = _CAPI_DGLDFSLabeledEdges_v2(
         gidx,

tests/compute/test_traversal.py

@@ -18,7 +18,6 @@ def toset(x):
     # F.zerocopy_to_numpy may return a int
     return set(F.zerocopy_to_numpy(x).tolist())
 
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
 @parametrize_dtype
 def test_bfs(idtype, n=100):
     def _bfs_nx(g_nx, src):
@@ -60,7 +59,6 @@ def test_bfs(idtype, n=100):
     assert len(edges_dgl) == len(edges_nx)
     assert all(toset(x) == y for x, y in zip(edges_dgl, edges_nx))
 
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
 @parametrize_dtype
 def test_topological_nodes(idtype, n=100):
     a = sp.random(n, n, 3 / n, data_rvs=lambda n: np.ones(n))
@@ -88,7 +86,6 @@ def test_topological_nodes(idtype, n=100):
     assert all(toset(x) == toset(y) for x, y in zip(layers_dgl, layers_spmv))
 
 DFS_LABEL_NAMES = ['forward', 'reverse', 'nontree']
-@unittest.skipIf(F._default_context_str == 'gpu', reason="GPU not implemented")
 @parametrize_dtype
 def test_dfs_labeled_edges(idtype, example=False):
     dgl_g = dgl.DGLGraph().astype(idtype)