[Feature] Add to_cugraph and from_cugraph to DGL (#4166)

* Added to_cugraph and from_cugraph functionality

* fix from_cugraph example

* Addressed Reviews

* Fix linting

* Apply suggestions to docstrings from code review
Co-authored-by: Minjie Wang <minjie.wang@nyu.edu>

* Add API docs and remove `from_cugraph` alias

* move cugraph tests to  tests/cugraph/test_basics.py

* remove ununsed imports from test_basics.py

* remove pytest.importorskip as no longer needed
Co-authored-by: Minjie Wang <minjie.wang@nyu.edu>
Co-authored-by: Minjie Wang <wmjlyjemaine@gmail.com>

docs/source/api/python/dgl.DGLGraph.rst

@@ -132,6 +132,7 @@ under the ``dgl`` namespace.
     DGLGraph.add_self_loop
     DGLGraph.remove_self_loop
     DGLGraph.to_simple
+    DGLGraph.to_cugraph
     DGLGraph.reorder_graph
 
 Adjacency and incidence matrix

docs/source/api/python/dgl.rst

@@ -18,6 +18,7 @@ Operators for constructing :class:`DGLGraph` from raw data formats.
 
     graph
     heterograph
+    from_cugraph
     from_scipy
     from_networkx
     bipartite_from_scipy
@@ -93,6 +94,7 @@ Operators for generating new graphs by manipulating the structure of the existin
     to_bidirected
     to_bidirected_stale
     to_block
+    to_cugraph
     to_double
     to_float
     to_half

python/dgl/convert.py

@@ -29,6 +29,8 @@ __all__ = [
     'from_networkx',
     'bipartite_from_networkx',
     'to_networkx',
+    'from_cugraph',
+    'to_cugraph'
 ]
 
 def graph(data,
@@ -1620,6 +1622,110 @@ def to_networkx(g, node_attrs=None, edge_attrs=None):
 
 DGLHeteroGraph.to_networkx = to_networkx
 
+def to_cugraph(g):
+    """Convert a DGL graph to a :class:`cugraph.Graph` and return.
+
+    Parameters
+    ----------
+    g : DGLGraph
+        A homogeneous graph.
+
+    Returns
+    -------
+    cugraph.Graph
+        The converted cugraph graph.
+
+    Notes
+    -----
+    The function only supports GPU graph input.
+
+    Examples
+    --------
+    The following example uses PyTorch backend.
+
+    >>> import dgl
+    >>> import cugraph
+    >>> import torch
+
+    >>> g = dgl.graph((torch.tensor([1, 2]), torch.tensor([1, 3]))).to('cuda')
+    >>> cugraph_g = g.to_cugraph()
+    >>> cugraph_g.edges()
+        src  dst
+    0    2    3
+    1    1    1
+    """
+
+    if g.device.type != 'cuda':
+        raise DGLError(f"Cannot convert a {g.device.type} graph to cugraph." +
+                        "Call g.to('cuda') first.")
+    if not g.is_homogeneous:
+        raise DGLError("dgl.to_cugraph only supports homogeneous graphs.")
+
+    try:
+        import cugraph
+        import cudf
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError("to_cugraph requires cugraph which could not be imported")
+
+    edgelist = g.edges()
+    src_ser = cudf.from_dlpack(F.zerocopy_to_dlpack(edgelist[0]))
+    dst_ser = cudf.from_dlpack(F.zerocopy_to_dlpack(edgelist[1]))
+    cudf_data = cudf.DataFrame({'source':src_ser, 'destination':dst_ser})
+    g_cugraph = cugraph.Graph(directed=True)
+    g_cugraph.from_cudf_edgelist(cudf_data,
+                                 source='source',
+                                 destination='destination')
+    return g_cugraph
+
+DGLHeteroGraph.to_cugraph = to_cugraph
+
+def from_cugraph(cugraph_graph):
+    """Create a graph from a :class:`cugraph.Graph` object.
+
+    Parameters
+    ----------
+    cugraph_graph : cugraph.Graph
+        The cugraph graph object holding the graph structure. Node and edge attributes are
+        dropped.
+
+        If the input graph is undirected, DGL converts it to a directed graph
+        by :func:`cugraph.Graph.to_directed`.
+
+    Returns
+    -------
+    DGLGraph
+        The created graph.
+
+    Examples
+    --------
+
+    The following example uses PyTorch backend.
+
+    >>> import dgl
+    >>> import cugraph
+    >>> import cudf
+
+    Create a cugraph graph.
+    >>> cugraph_g = cugraph.Graph(directed=True)
+    >>> df = cudf.DataFrame({"source":[0, 1, 2, 3],
+                     "destination":[1, 2, 3, 0]})
+    >>> cugraph_g.from_cudf_edgelist(df)
+
+    Convert it into a DGLGraph
+    >>> g = dgl.from_cugraph(cugraph_g)
+    >>> g.edges()
+    (tensor([1, 2, 3, 0], device='cuda:0'), tensor([2, 3, 0, 1], device='cuda:0'))
+    """
+    if not cugraph_graph.is_directed():
+        cugraph_graph = cugraph_graph.to_directed()
+
+    edges = cugraph_graph.edges()
+    src_t = F.zerocopy_from_dlpack(edges['src'].to_dlpack())
+    dst_t = F.zerocopy_from_dlpack(edges['dst'].to_dlpack())
+    g = graph((src_t,dst_t))
+
+    return g
+
 ############################################################
 # Internal APIs
 ############################################################

tests/cugraph/test_basics.py

@@ -9,3 +9,56 @@ import cugraph
 def test_dummy():
     cg = cugraph.Graph()
     assert cg is not None
+
+def test_to_cugraph_conversion():
+    g = dgl.graph((F.tensor([0, 1, 2, 3]), F.tensor([1, 0, 3, 2]))).to('cuda')
+    cugraph_g = g.to_cugraph()
+
+    assert cugraph_g.number_of_nodes()==g.number_of_nodes()
+    assert cugraph_g.number_of_edges()==g.number_of_edges()
+
+    assert cugraph_g.has_edge(0, 1)
+    assert cugraph_g.has_edge(1, 0)
+    assert cugraph_g.has_edge(3, 2)
+
+def test_from_cugraph_conversion():
+    # cudf is a dependency of cugraph
+    import cudf
+
+    # directed graph conversion test
+    cugraph_g = cugraph.Graph(directed=True)
+    df = cudf.DataFrame({"source":[0, 1, 2, 3],
+                     "destination":[1, 2, 3, 2]})
+
+    cugraph_g.from_cudf_edgelist(df)
+
+    g = dgl.from_cugraph(cugraph_g)
+
+    assert g.device.type == 'cuda'
+    assert g.number_of_nodes() == cugraph_g.number_of_nodes()
+    assert g.number_of_edges() == cugraph_g.number_of_edges()
+
+    # assert reverse edges are not present
+    assert g.has_edges_between(0, 1)
+    assert not g.has_edges_between(1, 0)
+    assert g.has_edges_between(1, 2)
+    assert not g.has_edges_between(2, 1)
+    assert g.has_edges_between(2, 3)
+
+    # undirected graph conversion test
+    cugraph_g = cugraph.Graph(directed=False)
+    df = cudf.DataFrame({"source":[0, 1, 2, 3],
+                     "destination":[1, 2, 3, 2]})
+
+    cugraph_g.from_cudf_edgelist(df)
+
+    g = dgl.from_cugraph(cugraph_g)
+
+    assert g.device.type == 'cuda'
+    assert g.number_of_nodes() == cugraph_g.number_of_nodes()
+    # assert reverse edges are present
+    assert g.has_edges_between(0, 1)
+    assert g.has_edges_between(1, 0)
+    assert g.has_edges_between(1, 2)
+    assert g.has_edges_between(2, 1)
+    assert g.has_edges_between(2, 3)