[NN] Add fast path for GateGCNConv when it has only one edge type (#2994)

* fix gatedgcn

* fix lint

python/dgl/nn/pytorch/conv/gatedgraphconv.py

@@ -61,6 +61,7 @@ class GatedGraphConv(nn.Module):
             [ 0.6393,  0.3447,  0.3893,  0.4279,  0.3342,  0.3809,  0.0406,  0.5030,
             0.1342,  0.0425]], grad_fn=<AddBackward0>)
     """
+
     def __init__(self,
                  in_feats,
                  out_feats,
@@ -110,7 +111,7 @@ class GatedGraphConv(nn.Module):
         """
         self._allow_zero_in_degree = set_value
 
-    def forward(self, graph, feat, etypes):
+    def forward(self, graph, feat, etypes=None):
         """
 
         Description
@@ -125,9 +126,10 @@ class GatedGraphConv(nn.Module):
             The input feature of shape :math:`(N, D_{in})` where :math:`N`
             is the number of nodes of the graph and :math:`D_{in}` is the
             input feature size.
-        etypes : torch.LongTensor
+        etypes : torch.LongTensor, or None
             The edge type tensor of shape :math:`(E,)` where :math:`E` is
-            the number of edges of the graph.
+            the number of edges of the graph. When there's only one edge type,
+            this argument can be skipped
 
         Returns
         -------
@@ -139,21 +141,33 @@ class GatedGraphConv(nn.Module):
             assert graph.is_homogeneous, \
                 "not a homogeneous graph; convert it with to_homogeneous " \
                 "and pass in the edge type as argument"
-            assert etypes.min() >= 0 and etypes.max() < self._n_etypes, \
-                "edge type indices out of range [0, {})".format(self._n_etypes)
-            zero_pad = feat.new_zeros((feat.shape[0], self._out_feats - feat.shape[1]))
+            if self._n_etypes != 1:
+                assert etypes.min() >= 0 and etypes.max() < self._n_etypes, \
+                    "edge type indices out of range [0, {})".format(
+                        self._n_etypes)
+
+            zero_pad = feat.new_zeros(
+                (feat.shape[0], self._out_feats - feat.shape[1]))
             feat = th.cat([feat, zero_pad], -1)
 
             for _ in range(self._n_steps):
-                graph.ndata['h'] = feat
-                for i in range(self._n_etypes):
-                    eids = th.nonzero(etypes == i, as_tuple=False).view(-1).type(graph.idtype)
-                    if len(eids) > 0:
-                        graph.apply_edges(
-                            lambda edges: {'W_e*h': self.linears[i](edges.src['h'])},
-                            eids
-                        )
-                graph.update_all(fn.copy_e('W_e*h', 'm'), fn.sum('m', 'a'))
-                a = graph.ndata.pop('a') # (N, D)
+                if self._n_etypes == 1 and etypes is None:
+                    # Fast path when graph has only one edge type
+                    graph.ndata['h'] = self.linears[0](feat)
+                    graph.update_all(fn.copy_u('h', 'm'), fn.sum('m', 'a'))
+                    a = graph.ndata.pop('a')  # (N, D)
+                else:
+                    graph.ndata['h'] = feat
+                    for i in range(self._n_etypes):
+                        eids = th.nonzero(
+                            etypes == i, as_tuple=False).view(-1).type(graph.idtype)
+                        if len(eids) > 0:
+                            graph.apply_edges(
+                                lambda edges: {
+                                    'W_e*h': self.linears[i](edges.src['h'])},
+                                eids
+                            )
+                    graph.update_all(fn.copy_e('W_e*h', 'm'), fn.sum('m', 'a'))
+                    a = graph.ndata.pop('a')  # (N, D)
                 feat = self.gru(a, feat)
             return feat

tests/pytorch/test_nn.py

@@ -725,6 +725,23 @@ def test_gated_graph_conv(g, idtype):
     # current we only do shape check
     assert h.shape[-1] == 10
 
+@parametrize_dtype
+@pytest.mark.parametrize('g', get_cases(['homo'], exclude=['zero-degree']))
+def test_gated_graph_conv_one_etype(g, idtype):
+    ctx = F.ctx()
+    g = g.astype(idtype).to(ctx)
+    ggconv = nn.GatedGraphConv(5, 10, 5, 1)
+    etypes = th.zeros(g.number_of_edges())
+    feat = F.randn((g.number_of_nodes(), 5))
+    ggconv = ggconv.to(ctx)
+    etypes = etypes.to(ctx)
+
+    h = ggconv(g, feat, etypes)
+    h2 = ggconv(g, feat)
+    # current we only do shape check
+    assert F.allclose(h, h2)
+    assert h.shape[-1] == 10
+
 @parametrize_dtype
 @pytest.mark.parametrize('g', get_cases(['homo', 'block-bipartite'], exclude=['zero-degree']))
 def test_nn_conv(g, idtype):
@@ -1113,6 +1130,7 @@ if __name__ == '__main__':
     test_gin_conv()
     test_agnn_conv()
     test_gated_graph_conv()
+    test_gated_graph_conv_one_etype()
     test_nn_conv()
     test_gmm_conv()
     test_dotgat_conv()