Repository: dmlc/dgl

Issue:
is_unibipartite check in DGLHeterograph constructor can take longer than message passing kernels
## 🐛 Bug

In DGL Heterograph's constructor, where in the case `ntypes` is a tuple it assumed the graph is unibipartite, and a sanity check inspected the metagraph.
https://github.com/dmlc/dgl/blob/master/python/dgl/heterograph.py#L100

This is accomplished by converting the `src` and `dst` to python `set`s and checking if they have overlap, which is quite slow (100us homogeneous DGLBlocks).

This makes sampling slow, as `to_block()` calls this constructor. It also impacts message passing, as a new DGLHeterograph is constructed per relation type per layer (via `__getitem__(self, etype)`).

### Possible fixes
1. I think the best method is to keep track of unibipartite-ness on the metagraph, and preserve the flag in `to_block()` and `__getitem__(self, etype)`, utilizing the same metagraph, avoiding calling `is_unibipartite` more than just once on the original graph.
1. The simplest way to address this is to change the check to an assertion so users who know their code is correct, can call `python -O` to disable it. I would only expect power users to know to use the `-O` flag, so we would need to emphasize this in the documentation.
1. Alternatively we could still do the check, but in C++. This will be faster, but would still repeat calculation only on the CPU for every mini-batch, so I do not favor this approach.


Discussion:

Resolving patch:
diff --git a/python/dgl/heterograph.py b/python/dgl/heterograph.py
--- a/python/dgl/heterograph.py
+++ b/python/dgl/heterograph.py
@@ -97,7 +97,7 @@ def _init(self, gidx, ntypes, etypes, node_frames, edge_frames):
                 errmsg = 'Invalid input. Expect a pair (srctypes, dsttypes) but got {}'.format(
                     ntypes)
                 raise TypeError(errmsg)
-            if not is_unibipartite(self._graph.metagraph):
+            if not self._graph.is_metagraph_unibipartite():
                 raise ValueError('Invalid input. The metagraph must be a uni-directional'
                                  ' bipartite graph.')
             self._ntypes = ntypes[0] + ntypes[1]
@@ -6204,23 +6204,6 @@ def make_canonical_etypes(etypes, ntypes, metagraph):
     rst = [(ntypes[sid], etypes[eid], ntypes[did]) for sid, did, eid in zip(src, dst, eid)]
     return rst
 
-def is_unibipartite(graph):
-    """Internal function that returns whether the given graph is a uni-directional
-    bipartite graph.
-
-    Parameters
-    ----------
-    graph : GraphIndex
-        Input graph
-
-    Returns
-    -------
-    bool
-        True if the graph is a uni-bipartite.
-    """
-    src, dst, _ = graph.edges()
-    return set(src.tonumpy()).isdisjoint(set(dst.tonumpy()))
-
 def find_src_dst_ntypes(ntypes, metagraph):
     """Internal function to split ntypes into SRC and DST categories.
 
diff --git a/python/dgl/heterograph_index.py b/python/dgl/heterograph_index.py
--- a/python/dgl/heterograph_index.py
+++ b/python/dgl/heterograph_index.py
@@ -72,6 +72,10 @@ def metagraph(self):
         """
         return _CAPI_DGLHeteroGetMetaGraph(self)
 
+    def is_metagraph_unibipartite(self):
+        """Return whether or not the graph is unibiparite."""
+        return _CAPI_DGLHeteroIsMetaGraphUniBipartite(self)
+
     def number_of_ntypes(self):
         """Return number of node types."""
         return self.metagraph.number_of_nodes()


Test patch:
diff --git a/tests/compute/test_heterograph.py b/tests/compute/test_heterograph.py
--- a/tests/compute/test_heterograph.py
+++ b/tests/compute/test_heterograph.py
@@ -1884,6 +1884,35 @@ def test_ismultigraph(idtype):
         {'A': 6, 'C': 6}, idtype=idtype, device=F.ctx())
     assert g.is_multigraph == True
 
+
+@parametrize_idtype
+def test_graph_index_is_unibipartite(idtype):
+    g1 = dgl.heterograph({('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5])},
+                         idtype=idtype, device=F.ctx())
+    assert g1._graph.is_metagraph_unibipartite()
+
+    # more complicated bipartite
+    g2 = dgl.heterograph({
+        ('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5]),
+        ('A', 'AC', 'C'): ([1, 0], [0, 0])
+    }, idtype=idtype, device=F.ctx())
+    assert g2._graph.is_metagraph_unibipartite()
+
+    g3 = dgl.heterograph({
+        ('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5]),
+        ('A', 'AC', 'C'): ([1, 0], [0, 0]),
+        ('A', 'AA', 'A'): ([0, 1], [0, 1])
+    }, idtype=idtype, device=F.ctx())
+    assert not g3._graph.is_metagraph_unibipartite()
+
+    g4 = dgl.heterograph({
+        ('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5]),
+        ('C', 'CA', 'A'): ([1, 0], [0, 0])
+    }, idtype=idtype, device=F.ctx())
+
+    assert not g4._graph.is_metagraph_unibipartite()
+
+
 @parametrize_idtype
 def test_bipartite(idtype):
     g1 = dgl.heterograph({('A', 'AB', 'B'): ([0, 0, 1], [1, 2, 5])},