Unverified Commit 9ae117d3 authored by Vibhu Jawa's avatar Vibhu Jawa Committed by GitHub
Browse files

[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: default avatarMinjie 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: default avatarMinjie Wang <minjie.wang@nyu.edu>
Co-authored-by: default avatarMinjie Wang <wmjlyjemaine@gmail.com>
parent 0d1b2c5a
......@@ -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
......
......@@ -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
......
......@@ -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
############################################################
......
......@@ -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)
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment