Repository: dmlc/dgl

Issue:
[Regression] dgl.batch() in 0.5.x is slower than 0.4.x
## 🐛 Bug

Not sure if this is a bug, but here is the result:

## To Reproduce

Benchmarking script:
```python
import dgl
import numpy as np
import time
import itertools

for E, B in itertools.product([1000, 10000, 100000], [10, 100, 1000]):
    gs = [dgl.DGLGraph((np.random.randint(0, 100, E), np.random.randint(0, 100, E)), readonly=True) for _ in range(B)]
    t0 = time.time()
    for _ in range(20):
        g = dgl.batch(gs)
    tt = time.time()
    print(E, B, tt - t0)
```

Result on a g3.16x instance (although the code works on CPU):

|E|B|0.5.x|0.4.x|
|:-:|:-:|:-:|:-:|
| 1000 | 10 | 0.0637 | 0.00382 |
| 1000 | 100 | 0.447 | 0.0276 |
| 1000 | 1000 | 4.44 | 0.353 |
| 10000 | 10 | 0.0843 | 0.0138 |
| 10000 | 100 | 0.790 | 0.194 |
| 10000 | 1000 | 10.0 | 2.85 |
| 100000 | 10 | 0.376 | 0.120 |
| 100000 | 100 | 4.11 | 2.38 |
| 100000 | 1000 | 41.0 | 26.2 |

## Context

Vijay and I noticed a significant slowdown on the training of GATs for graph regression on ZINC in Benchmarking GNNs.  The slowdown goes as much as 100% (depending on the machine).  With `CUDA_LAUNCH_BLOCKING=1`, I further figured out that fetching a batched graph from the dataloader increases from 13ms to 89ms.  The forward propagation also increases from 29ms to 41ms per iteration, and the backward propagation increases from 31ms to 43ms, but that's a separate issue.


Discussion:

Resolving patch:
diff --git a/python/dgl/backend/tensorflow/tensor.py b/python/dgl/backend/tensorflow/tensor.py
--- a/python/dgl/backend/tensorflow/tensor.py
+++ b/python/dgl/backend/tensorflow/tensor.py
@@ -145,7 +145,8 @@ def to_backend_ctx(dglctx):
 
 
 def astype(input, ty):
-    return tf.cast(input, dtype=ty)
+    with tf.device(input.device):
+        return tf.cast(input, dtype=ty)
 
 
 def asnumpy(input):
diff --git a/python/dgl/batch.py b/python/dgl/batch.py
--- a/python/dgl/batch.py
+++ b/python/dgl/batch.py
@@ -1,12 +1,14 @@
 """Utilities for batching/unbatching graphs."""
 from collections.abc import Mapping
-from collections import defaultdict
 
 from . import backend as F
 from .base import ALL, is_all, DGLError, dgl_warning
+from .heterograph_index import disjoint_union
+from .heterograph import DGLHeteroGraph
 from . import convert
 from . import utils
 
+
 __all__ = ['batch', 'unbatch', 'batch_hetero', 'unbatch_hetero']
 
 def batch(graphs, ndata=ALL, edata=ALL, *, node_attrs=None, edge_attrs=None):
@@ -156,61 +158,44 @@ def batch(graphs, ndata=ALL, edata=ALL, *, node_attrs=None, edge_attrs=None):
         dgl_warning('Arguments edge_attrs has been deprecated. Please use'
                     ' edata instead.')
         edata = edge_attrs
-    if not (is_all(ndata) or isinstance(ndata, list)):
+    if not (is_all(ndata) or isinstance(ndata, list) or ndata is None):
         raise DGLError('Invalid argument ndata: must be a string list but got {}.'.format(
             type(ndata)))
-    if not (is_all(edata) or isinstance(edata, list)):
+    if not (is_all(edata) or isinstance(edata, list) or edata is None):
         raise DGLError('Invalid argument edata: must be a string list but got {}.'.format(
             type(edata)))
     if any(g.is_block for g in graphs):
         raise DGLError("Batching a block is not supported.")
 
-    utils.check_all_same_device(graphs, 'graphs')
-    utils.check_all_same_idtype(graphs, 'graphs')
-    relations = graphs[0].canonical_etypes
-    ntypes = graphs[0].ntypes
-    idtype = graphs[0].idtype
-    device = graphs[0].device
-
-    # Batch graph structure for each relation graph
-    edge_dict = defaultdict(list)
-    num_nodes_dict = defaultdict(int)
-    for g in graphs:
-        for rel in relations:
-            srctype, etype, dsttype = rel
-            u, v = g.edges(order='eid', etype=rel)
-            src = u + num_nodes_dict[srctype]
-            dst = v + num_nodes_dict[dsttype]
-            edge_dict[rel].append((src, dst))
-        for ntype in ntypes:
-            num_nodes_dict[ntype] += g.number_of_nodes(ntype)
-    for rel in relations:
-        src, dst = zip(*edge_dict[rel])
-        edge_dict[rel] = (F.cat(src, 0), F.cat(dst, 0))
-    retg = convert.heterograph(edge_dict, num_nodes_dict, idtype=idtype, device=device)
+    relations = list(sorted(graphs[0].canonical_etypes))
+    ntypes = list(sorted(graphs[0].ntypes))
+    etypes = [etype for _, etype, _ in relations]
+
+    gidx = disjoint_union(graphs[0]._graph.metagraph, [g._graph for g in graphs])
+    retg = DGLHeteroGraph(gidx, ntypes, etypes)
 
     # Compute batch num nodes
     bnn = {}
-    for ntype in graphs[0].ntypes:
+    for ntype in ntypes:
         bnn[ntype] = F.cat([g.batch_num_nodes(ntype) for g in graphs], 0)
     retg.set_batch_num_nodes(bnn)
 
     # Compute batch num edges
     bne = {}
-    for etype in graphs[0].canonical_etypes:
+    for etype in relations:
         bne[etype] = F.cat([g.batch_num_edges(etype) for g in graphs], 0)
     retg.set_batch_num_edges(bne)
 
     # Batch node feature
     if ndata is not None:
-        for ntype in graphs[0].ntypes:
+        for ntype in ntypes:
             feat_dicts = [g.nodes[ntype].data for g in graphs if g.number_of_nodes(ntype) > 0]
             ret_feat = _batch_feat_dicts(feat_dicts, ndata, 'nodes["{}"].data'.format(ntype))
             retg.nodes[ntype].data.update(ret_feat)
 
     # Batch edge feature
     if edata is not None:
-        for etype in graphs[0].canonical_etypes:
+        for etype in relations:
             feat_dicts = [g.edges[etype].data for g in graphs if g.number_of_edges(etype) > 0]
             ret_feat = _batch_feat_dicts(feat_dicts, edata, 'edges[{}].data'.format(etype))
             retg.edges[etype].data.update(ret_feat)


Test patch:
diff --git a/tests/compute/test_batched_heterograph.py b/tests/compute/test_batched_heterograph.py
--- a/tests/compute/test_batched_heterograph.py
+++ b/tests/compute/test_batched_heterograph.py
@@ -325,6 +325,6 @@ def test_unbatch2(idtype):
     #test_topology('int32')
     #test_batching_batched('int32')
     #test_batched_features('int32')
-    #test_empty_relation('int32')
+    # test_empty_relation('int64')
     #test_to_device('int32')
     pass