[Runtime] Scheduler and Executor (#140)

* executor api

* draft executor interface

* WIP

* revert changes to avoid conflict with api change

* core scheduling logic

* WIP: build graph adj

* incidence matrix for in edges

* support incidence matrix for partial recv nodes

* improve

* build adjmat in scheduler

* graph store

* get degree bucketing schedule

* connect to c++ degree bucketing

* conceptual executor creation code

* executor comments

* fix

* more executor comments

* WIP: full send_and_recv schedule

* most schedulers

* simplify scheduler

* executors

* runtime

* builtin function base class

* adj indices and shape

* completely refactor scheduler

* rename and move bundled out to function.py

* use_edge_feature in msg func

* rewrite scheduler

* node edge executor

* connect with graph api

* handle zero degree

* misc

* fix test cases

* fix a good many bugs...

* remove old scheduler

* push and pull

* fix send recv

* c++ lint

* fix batched send recv

* hot fix for mxnet

* typo

* write back executor

* apply node edge

* clean up, doc string

* fix as requested

* refactor

* fix

* WIP

* WIP

* ir draft

* more on ir

* WIP: spmv schedule

* WIP

* recv schedule

* refactor

* WIP

* snr degree bucketing

* snr scheduler

* move prog to graph.py; rename

* unittest for send/recv

* remove some legacy codes

* WIP: update_all

* pass test_basics

* passed all current utests

* more utests; fix mx utest

* WIP: fixing zero deg initial value

* some tests

* fix 0deg problem

* fix mx

* fix mx

* some notes

* fix as requested

python/dgl/utils.py

 """Utility module."""
-from __future__ import absolute_import
+from __future__ import absolute_import, division
 
-from collections import Mapping
+from collections import Mapping, Iterable
 from functools import wraps
 import numpy as np
 
+from .base import DGLError
 from . import backend as F
 from . import ndarray as nd
 
@@ -14,18 +15,39 @@ class Index(object):
         self._initialize_data(data)
 
     def _initialize_data(self, data):
-        self._list_data = None  # a numpy type data
+        self._list_data = None   # a numpy type data or a slice
         self._user_tensor_data = dict()  # dictionary of user tensors
         self._dgl_tensor_data = None  # a dgl ndarray
         self._dispatch(data)
 
+    def __iter__(self):
+        return iter(self.tolist())
+
+    def __len__(self):
+        if self._list_data is not None and isinstance(self._list_data, slice):
+            slc = self._list_data
+            if slc.step is None:
+                return slc.stop - slc.start
+            else:
+                return (slc.stop - slc.start) // slc.step
+        elif self._list_data is not None:
+            return len(self._list_data)
+        elif len(self._user_tensor_data) > 0:
+            data = next(iter(self._user_tensor_data.values()))
+            return len(data)
+        else:
+            return len(self._dgl_tensor_data)
+
+    def __getitem__(self, i):
+        return self.tolist()[i]
+
     def _dispatch(self, data):
         """Store data based on its type."""
         if F.is_tensor(data):
             if not (F.dtype(data) == F.int64):
-                raise ValueError('Index data must be an int64 vector, but got: %s' % str(data))
+                raise DGLError('Index data must be an int64 vector, but got: %s' % str(data))
             if len(F.shape(data)) > 1:
-                raise ValueError('Index data must be 1D int64 vector, but got: %s' % str(data))
+                raise DGLError('Index data must be 1D int64 vector, but got: %s' % str(data))
             if len(F.shape(data)) == 0:
                 # a tensor of one int
                 self._dispatch(int(data))
@@ -33,19 +55,23 @@ class Index(object):
                 self._user_tensor_data[F.context(data)] = data
         elif isinstance(data, nd.NDArray):
             if not (data.dtype == 'int64' and len(data.shape) == 1):
-                raise ValueError('Index data must be 1D int64 vector, but got: %s' % str(data))
+                raise DGLError('Index data must be 1D int64 vector, but got: %s' % str(data))
             self._dgl_tensor_data = data
+        elif isinstance(data, slice):
+            # save it in the _list_data temporarily; materialize it if `tolist` is called
+            self._list_data = data
         else:
             try:
                 self._list_data = np.array([int(data)]).astype(np.int64)
             except:
                 try:
-                    data = np.array(data).astype('int64')
+                    data = np.array(data).astype(np.int64)
                     if data.ndim != 1:
-                        raise ValueError('Index data must be 1D int64 vector, but got: %s' % str(data))
+                        raise DGLError('Index data must be 1D int64 vector,'
+                                       ' but got: %s' % str(data))
                     self._list_data = data
                 except:
-                    raise ValueError('Error index data: %s' % str(data))
+                    raise DGLError('Error index data: %s' % str(data))
             self._user_tensor_data[F.cpu()] = F.zerocopy_from_numpy(self._list_data)
 
     def tolist(self):
@@ -56,6 +82,10 @@ class Index(object):
             else:
                 data = self.tousertensor()
                 self._list_data = F.zerocopy_to_numpy(data)
+        elif isinstance(self._list_data, slice):
+            # convert it to numpy array
+            slc = self._list_data
+            self._list_data = np.arange(slc.start, slc.stop, slc.step).astype(np.int64)
         return self._list_data
 
     def tousertensor(self, ctx=None):
@@ -63,9 +93,13 @@ class Index(object):
         if ctx is None:
             ctx = F.cpu()
         if len(self._user_tensor_data) == 0:
-            # zero copy from dgl tensor
-            dl = self._dgl_tensor_data.to_dlpack()
-            self._user_tensor_data[F.cpu()] = F.zerocopy_from_dlpack(dl)
+            if self._dgl_tensor_data is not None:
+                # zero copy from dgl tensor
+                dl = self._dgl_tensor_data.to_dlpack()
+                self._user_tensor_data[F.cpu()] = F.zerocopy_from_dlpack(dl)
+            else:
+                # zero copy from numpy array
+                self._user_tensor_data[F.cpu()] = F.zerocopy_from_numpy(self.tolist())
         if ctx not in self._user_tensor_data:
             # copy from cpu to another device
             data = next(iter(self._user_tensor_data.values()))
@@ -81,20 +115,9 @@ class Index(object):
             self._dgl_tensor_data = nd.from_dlpack(dl)
         return self._dgl_tensor_data
 
-    def __iter__(self):
-        return iter(self.tolist())
-
-    def __len__(self):
-        if self._list_data is not None:
-            return len(self._list_data)
-        elif len(self._user_tensor_data) > 0:
-            data = next(iter(self._user_tensor_data.values()))
-            return len(data)
-        else:
-            return len(self._dgl_tensor_data)
-
-    def __getitem__(self, i):
-        return self.tolist()[i]
+    def is_slice(self, start, stop, step=None):
+        return (isinstance(self._list_data, slice)
+                and self._list_data == slice(start, stop, step))
 
     def __getstate__(self):
         return self.tousertensor()
@@ -105,66 +128,6 @@ class Index(object):
 def toindex(x):
     return x if isinstance(x, Index) else Index(x)
 
-def node_iter(n):
-    """Return an iterator that loops over the given nodes.
-
-    Parameters
-    ----------
-    n : iterable
-        The node ids.
-    """
-    return iter(n)
-
-def edge_iter(u, v):
-    """Return an iterator that loops over the given edges.
-
-    Parameters
-    ----------
-    u : iterable
-        The src ids.
-    v : iterable
-        The dst ids.
-    """
-    if len(u) == len(v):
-        # many-many
-        for uu, vv in zip(u, v):
-            yield uu, vv
-    elif len(v) == 1:
-        # many-one
-        for uu in u:
-            yield uu, v[0]
-    elif len(u) == 1:
-        # one-many
-        for vv in v:
-            yield u[0], vv
-    else:
-        raise ValueError('Error edges:', u, v)
-
-def edge_broadcasting(u, v):
-    """Convert one-many and many-one edges to many-many.
-
-    Parameters
-    ----------
-    u : Index
-        The src id(s)
-    v : Index
-        The dst id(s)
-
-    Returns
-    -------
-    uu : Index
-        The src id(s) after broadcasting
-    vv : Index
-        The dst id(s) after broadcasting
-    """
-    if len(u) != len(v) and len(u) == 1:
-        u = toindex(F.full_1d(len(v), u[0]))
-    elif len(u) != len(v) and len(v) == 1:
-        v = toindex(F.full_1d(len(u), v[0]))
-    else:
-        assert len(u) == len(v)
-    return u, v
-
 class LazyDict(Mapping):
     """A readonly dictionary that does not materialize the storage."""
     def __init__(self, fn, keys):
@@ -185,18 +148,20 @@ class LazyDict(Mapping):
     def __len__(self):
         return len(self._keys)
 
+    def keys(self):
+        return self._keys
+
 class HybridDict(Mapping):
     """A readonly dictonary that merges several dict-like (python dict, LazyDict).
        If there are duplicate keys, early keys have priority over latter ones
     """
     def __init__(self, *dict_like_list):
         self._dict_like_list = dict_like_list
-        self._keys = None
+        self._keys = set()
+        for d in dict_like_list:
+            self._keys.update(d.keys())
 
     def keys(self):
-        if self._keys is None:
-            self._keys = sum([set(d.keys()) for d in self._dict_like_list], set())
-            self._keys = list(self._keys)
         return self._keys
 
     def __getitem__(self, key):
@@ -236,6 +201,19 @@ class ReadOnlyDict(Mapping):
 def build_relabel_map(x):
     """Relabel the input ids to continuous ids that starts from zero.
 
+    Ids are assigned new ids according to their ascending order.
+
+    Examples
+    --------
+    >>> x = [1, 5, 3, 6]
+    >>> n2o, o2n = build_relabel_map(x)
+    >>> n2o
+    [1, 3, 5, 6]
+    >>> o2n
+    [n/a, 0, n/a, 2, n/a, 3, 4]
+
+    "n/a" will be filled with 0
+
     Parameters
     ----------
     x : Index
@@ -349,17 +327,31 @@ def reorder(dict_like, index):
         new_dict[key] = F.gather_row(val, idx_ctx)
     return new_dict
 
-def parse_edges_tuple(edges):
-    """Parse the given edges and return the tuple.
+def build_coo_sparse_matrix(dat, row, col, dense_shape):
+    """Build coo sparse matrix
 
     Parameters
     ----------
-    edges : edges
-        Edges can be a pair of endpoint nodes (u, v), or a
-        tensor of edge ids. The default value is all the edges.
+    dat: Tensor
+        Data.
+    row: Tensor
+        Row index.
+    col: Tensor
+        Column index.
+    dense_shape: list or tuple of two integer
+        Dense shape of the sparse matrix
 
     Returns
     -------
-    A tuple of (u, v, eid)
+    SparseTensor
+        The sparse matrix.
     """
-    pass
+    nnz = len(row)
+    row = F.unsqueeze(row, 0)
+    col = F.unsqueeze(col, 0)
+    idx = F.cat([row, col], dim=0)
+    return F.sparse_matrix(dat, ('coo', idx), dense_shape)
+
+def is_iterable(obj):
+    """Return true if the object is an iterable."""
+    return isinstance(obj, Iterable)

src/scheduler/scheduler.cc

@@ -10,14 +10,18 @@
 namespace dgl {
 namespace sched {
 
-std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
-    const auto n_msgs = vids->shape[0];
+std::vector<IdArray> DegreeBucketing(const IdArray& msg_ids, const IdArray& vids,
+        const IdArray& recv_ids) {
+    auto n_msgs = msg_ids->shape[0];
+
     const int64_t* vid_data = static_cast<int64_t*>(vids->data);
+    const int64_t* msg_id_data = static_cast<int64_t*>(msg_ids->data);
+    const int64_t* recv_id_data = static_cast<int64_t*>(recv_ids->data);
 
-    // inedge: dst->msgs
+    // in edge: dst->msgs
     std::unordered_map<int64_t, std::vector<int64_t>> in_edges;
-    for (int64_t mid = 0; mid < n_msgs; ++mid) {
-        in_edges[vid_data[mid]].push_back(mid);
+    for (int64_t i = 0; i < n_msgs; ++i) {
+        in_edges[vid_data[i]].push_back(msg_id_data[i]);
     }
 
     // bkt: deg->dsts
@@ -26,14 +30,29 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
         bkt[it.second.size()].push_back(it.first);
     }
 
-    // initialize output
+    std::unordered_set<int64_t> zero_deg_nodes;
+    for (int64_t i = 0; i < recv_ids->shape[0]; ++i) {
+        if (in_edges.find(recv_id_data[i]) == in_edges.end()) {
+            zero_deg_nodes.insert(recv_id_data[i]);
+        }
+    }
+    auto n_zero_deg = zero_deg_nodes.size();
+
+    // calc output size
     int64_t n_deg = bkt.size();
     int64_t n_dst = in_edges.size();
+    int64_t n_mid_sec = bkt.size();  // zero deg won't affect message size
+    if (n_zero_deg > 0) {
+        n_deg += 1;
+        n_dst += n_zero_deg;
+    }
+
+    // initialize output
     IdArray degs = IdArray::Empty({n_deg}, vids->dtype, vids->ctx);
     IdArray nids = IdArray::Empty({n_dst}, vids->dtype, vids->ctx);
     IdArray nid_section = IdArray::Empty({n_deg}, vids->dtype, vids->ctx);
     IdArray mids = IdArray::Empty({n_msgs}, vids->dtype, vids->ctx);
-    IdArray mid_section = IdArray::Empty({n_deg}, vids->dtype, vids->ctx);
+    IdArray mid_section = IdArray::Empty({n_mid_sec}, vids->dtype, vids->ctx);
     int64_t* deg_ptr = static_cast<int64_t*>(degs->data);
     int64_t* nid_ptr = static_cast<int64_t*>(nids->data);
     int64_t* nsec_ptr = static_cast<int64_t*>(nid_section->data);
@@ -43,10 +62,10 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
     // fill in bucketing ordering
     for (const auto& it : bkt) {  // for each bucket
         const int64_t deg = it.first;
-        const int64_t n_dst = it.second.size();
+        const int64_t bucket_size = it.second.size();
         *deg_ptr++ = deg;
-        *nsec_ptr++ = n_dst;
-        *msec_ptr++ = deg * n_dst;
+        *nsec_ptr++ = bucket_size;
+        *msec_ptr++ = deg * bucket_size;
         for (const auto dst : it.second) {  // for each dst in this bucket
             *nid_ptr++ = dst;
             for (const auto mid : in_edges[dst]) {  // for each in edge of dst
@@ -55,6 +74,14 @@ std::vector<IdArray> DegreeBucketing(const IdArray& vids) {
         }
     }
 
+    if (n_zero_deg > 0) {
+        *deg_ptr = 0;
+        *nsec_ptr = n_zero_deg;
+        for (const auto dst : zero_deg_nodes) {
+            *nid_ptr++ = dst;
+        }
+    }
+
     std::vector<IdArray> ret;
     ret.push_back(std::move(degs));
     ret.push_back(std::move(nids));

src/scheduler/scheduler_apis.cc

@@ -13,18 +13,48 @@ using tvm::runtime::NDArray;
 
 namespace dgl {
 
-TVM_REGISTER_GLOBAL("scheduler._CAPI_DGLDegreeBucketing")
+TVM_REGISTER_GLOBAL("runtime.degree_bucketing._CAPI_DGLDegreeBucketing")
+.set_body([] (TVMArgs args, TVMRetValue* rv) {
+    const IdArray msg_ids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
+    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
+    const IdArray nids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[2]));
+
+    *rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(msg_ids, vids, nids));
+  });
+
+TVM_REGISTER_GLOBAL("runtime.degree_bucketing._CAPI_DGLDegreeBucketingForEdges")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[0]));
-    *rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(vids));
+    // XXX: better way to do arange?
+    int64_t n_msgs = vids->shape[0];
+    IdArray msg_ids = IdArray::Empty({n_msgs}, vids->dtype, vids->ctx);
+    int64_t* mid_data = static_cast<int64_t*>(msg_ids->data);
+    for (int64_t i = 0; i < n_msgs; ++i) {
+        mid_data[i] = i;
+    }
+    *rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(msg_ids, vids, vids));
   });
 
-TVM_REGISTER_GLOBAL("scheduler._CAPI_DGLDegreeBucketingFromGraph")
+TVM_REGISTER_GLOBAL("runtime.degree_bucketing._CAPI_DGLDegreeBucketingForRecvNodes")
 .set_body([] (TVMArgs args, TVMRetValue* rv) {
     GraphHandle ghandle = args[0];
     const Graph* gptr = static_cast<Graph*>(ghandle);
-    const auto& edges = gptr->Edges(false);
-    *rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(edges.dst));
+    const IdArray vids = IdArray::FromDLPack(CreateTmpDLManagedTensor(args[1]));
+    const auto& edges = gptr->InEdges(vids);
+    *rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(edges.id, edges.dst, vids));
   });
 
+TVM_REGISTER_GLOBAL("runtime.degree_bucketing._CAPI_DGLDegreeBucketingForFullGraph")
+.set_body([] (TVMArgs args, TVMRetValue* rv) {
+    GraphHandle ghandle = args[0];
+    const Graph* gptr = static_cast<Graph*>(ghandle);
+    const auto& edges = gptr->Edges(false);
+    int64_t n_vertices = gptr->NumVertices();
+    IdArray nids = IdArray::Empty({n_vertices}, edges.dst->dtype, edges.dst->ctx);
+    int64_t* nid_data = static_cast<int64_t*>(nids->data);
+    for (int64_t i = 0; i < n_vertices; ++i) {
+        nid_data[i] = i;
+    }
+    *rv = ConvertNDArrayVectorToPackedFunc(sched::DegreeBucketing(edges.id, edges.dst, nids));
+  });
 }  // namespace dgl

tests/mxnet/test_basics.py

@@ -250,12 +250,15 @@ def check_reduce_0deg(readonly):
         return {'m' : edges.src['h']}
     def _reduce(nodes):
         return {'h' : nodes.data['h'] + nodes.mailbox['m'].sum(1)}
+    def _init2(shape, dtype, ctx, ids):
+        return 2 + mx.nd.zeros(shape, dtype=dtype, ctx=ctx)
+    g.set_n_initializer(_init2, 'h')
     old_repr = mx.nd.random.normal(shape=(5, 5))
     g.set_n_repr({'h': old_repr})
     g.update_all(_message, _reduce)
     new_repr = g.ndata['h']
 
-    assert np.allclose(new_repr[1:].asnumpy(), old_repr[1:].asnumpy())
+    assert np.allclose(new_repr[1:].asnumpy(), 2+mx.nd.zeros((4, 5)).asnumpy())
     assert np.allclose(new_repr[0].asnumpy(), old_repr.sum(0).asnumpy())
 
 def test_reduce_0deg():
@@ -279,11 +282,15 @@ def check_pull_0deg(readonly):
         return {'m' : edges.src['h']}
     def _reduce(nodes):
         return {'h' : nodes.mailbox['m'].sum(1)}
-
     old_repr = mx.nd.random.normal(shape=(2, 5))
     g.set_n_repr({'h' : old_repr})
     g.pull(0, _message, _reduce)
     new_repr = g.ndata['h']
+    # TODO(minjie): this is not the intended behavior. Pull node#0
+    #   should reset node#0 to the initial value. The bug is because
+    #   current pull is implemented using send_and_recv. Since there
+    #   is no edge to node#0 so the send_and_recv is skipped. Fix this
+    #   behavior when optimizing the pull scheduler.
     assert np.allclose(new_repr[0].asnumpy(), old_repr[0].asnumpy())
     assert np.allclose(new_repr[1].asnumpy(), old_repr[1].asnumpy())
     g.pull(1, _message, _reduce)

tests/mxnet/test_graph_index.py

@@ -14,8 +14,8 @@ def generate_rand_graph(n):
     return g, ig
 
 def check_graph_equal(g1, g2):
-    adj1 = g1.adjacency_matrix(ctx=mx.cpu()) != 0
-    adj2 = g2.adjacency_matrix(ctx=mx.cpu()) != 0
+    adj1 = g1.adjacency_matrix(transpose=False, ctx=mx.cpu()) != 0
+    adj2 = g2.adjacency_matrix(transpose=False, ctx=mx.cpu()) != 0
     assert mx.nd.sum(adj1 - adj2).asnumpy() == 0
 
 def test_graph_gen():

tests/mxnet/test_specialization.py

@@ -230,18 +230,10 @@ def test_update_all_multi_fn():
     g.set_n_repr({'v1' : mx.nd.zeros(shape=(g.number_of_nodes(),)),
         'v2' : mx.nd.zeros(shape=(g.number_of_nodes(),))})
     fld = 'f2'
-    # update all, mix of builtin and UDF
-    g.update_all([fn.copy_src(src=fld, out='m1'), message_func],
-                 [fn.sum(msg='m1', out='v1'), reduce_func],
-                 None)
-    v1 = g.ndata['v1']
-    v2 = g.ndata['v2']
-    assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
 
     # run builtin with single message and reduce
     g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'), None)
     v1 = g.ndata['v1']
-    assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
 
     # 1 message, 2 reduces
     g.update_all(fn.copy_src(src=fld, out='m'), [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')], None)
@@ -284,20 +276,10 @@ def test_send_and_recv_multi_fn():
         'v3' : mx.nd.zeros(shape=(g.number_of_nodes(), D))})
     fld = 'f2'
 
-    # send and recv, mix of builtin and UDF
-    g.send_and_recv((u, v),
-                    [fn.copy_src(src=fld, out='m1'), message_func],
-                    [fn.sum(msg='m1', out='v1'), reduce_func],
-                    None)
-    v1 = g.ndata['v1']
-    v2 = g.ndata['v2']
-    assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
-
     # run builtin with single message and reduce
     g.send_and_recv((u, v), fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'),
                     None)
     v1 = g.ndata['v1']
-    assert np.allclose(v1.asnumpy(), v2.asnumpy(), rtol=1e-05, atol=1e-05)
 
     # 1 message, 2 reduces
     g.send_and_recv((u, v),

tests/pytorch/test_basics.py

 import torch as th
 from torch.autograd import Variable
 import numpy as np
+import dgl
 from dgl.graph import DGLGraph
 import utils as U
 
@@ -40,8 +41,8 @@ def generate_graph(grad=False):
     ecol = Variable(th.randn(17, D), requires_grad=grad)
     g.ndata['h'] = ncol
     g.edata['w'] = ecol
-    g.set_n_initializer(lambda shape, dtype, ctx : th.zeros(shape, dtype=dtype, device=ctx))
-    g.set_e_initializer(lambda shape, dtype, ctx : th.zeros(shape, dtype=dtype, device=ctx))
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.set_e_initializer(dgl.init.zero_initializer)
     return g
 
 def test_batch_setter_getter():
@@ -169,6 +170,18 @@ def test_batch_recv():
     assert(reduce_msg_shapes == {(1, 3, D), (3, 1, D)})
     reduce_msg_shapes.clear()
 
+def test_apply_nodes():
+    def _upd(nodes):
+        return {'h' : nodes.data['h'] * 2}
+    g = generate_graph()
+    g.register_apply_node_func(_upd)
+    old = g.ndata['h']
+    g.apply_nodes()
+    assert U.allclose(old * 2, g.ndata['h'])
+    u = th.tensor([0, 3, 4, 6])
+    g.apply_nodes(lambda nodes : {'h' : nodes.data['h'] * 0.}, u)
+    assert U.allclose(g.ndata['h'][u], th.zeros((4, D)))
+
 def test_apply_edges():
     def _upd(edges):
         return {'w' : edges.data['w'] * 2}
@@ -199,7 +212,7 @@ def test_update_routines():
     try:
         g.send_and_recv([u, v])
         assert False
-    except ValueError:
+    except:
         pass
 
     # pull
@@ -233,12 +246,17 @@ def test_reduce_0deg():
         return {'m' : edges.src['h']}
     def _reduce(nodes):
         return {'h' : nodes.data['h'] + nodes.mailbox['m'].sum(1)}
+    def _init2(shape, dtype, ctx, ids):
+        return 2 + th.zeros(shape, dtype=dtype, device=ctx)
+    g.set_n_initializer(_init2, 'h')
     old_repr = th.randn(5, 5)
     g.ndata['h'] = old_repr
     g.update_all(_message, _reduce)
     new_repr = g.ndata['h']
-
-    assert U.allclose(new_repr[1:], old_repr[1:])
+    # the first row of the new_repr should be the sum of all the node
+    # features; while the 0-deg nodes should be initialized by the
+    # initializer.
+    assert U.allclose(new_repr[1:], 2+th.zeros((4,5)))
     assert U.allclose(new_repr[0], old_repr.sum(0))
 
 def test_pull_0deg():
@@ -398,6 +416,7 @@ if __name__ == '__main__':
     test_batch_setter_autograd()
     test_batch_send()
     test_batch_recv()
+    test_apply_nodes()
     test_apply_edges()
     test_update_routines()
     test_reduce_0deg()

tests/pytorch/test_batched_graph.py

-import networkx as nx
 import dgl
 import torch as th
-import numpy as np
 import utils as U
 
 def tree1():
@@ -45,10 +43,6 @@ def tree2():
 def test_batch_unbatch():
     t1 = tree1()
     t2 = tree2()
-    n1 = t1.ndata['h']
-    n2 = t2.ndata['h']
-    e1 = t1.edata['h']
-    e2 = t2.edata['h']
 
     bg = dgl.batch([t1, t2])
     assert bg.number_of_nodes() == 10
@@ -82,7 +76,7 @@ def test_batch_unbatch1():
     assert U.allclose(t2.ndata['h'], s3.ndata['h'])
     assert U.allclose(t2.edata['h'], s3.edata['h'])
 
-def test_batch_sendrecv():
+def test_batch_send_then_recv():
     t1 = tree1()
     t2 = tree2()
 
@@ -93,12 +87,27 @@ def test_batch_sendrecv():
     v = [1, 1, 4 + 5, 4 + 5]
 
     bg.send((u, v))
-    bg.recv(v)
+    bg.recv([1, 9]) # assuming recv takes in unique nodes
 
     t1, t2 = dgl.unbatch(bg)
     assert t1.ndata['h'][1] == 7
     assert t2.ndata['h'][4] == 2
 
+def test_batch_send_and_recv():
+    t1 = tree1()
+    t2 = tree2()
+
+    bg = dgl.batch([t1, t2])
+    bg.register_message_func(lambda edges: {'m' : edges.src['h']})
+    bg.register_reduce_func(lambda nodes: {'h' : th.sum(nodes.mailbox['m'], 1)})
+    u = [3, 4, 2 + 5, 0 + 5]
+    v = [1, 1, 4 + 5, 4 + 5]
+
+    bg.send_and_recv((u, v))
+
+    t1, t2 = dgl.unbatch(bg)
+    assert t1.ndata['h'][1] == 7
+    assert t2.ndata['h'][4] == 2
 
 def test_batch_propagate():
     t1 = tree1()
@@ -147,11 +156,9 @@ def test_batch_no_edge():
     g1 = dgl.DGLGraph()
     g1.add_nodes(6)
     g1.add_edges([4, 4, 2, 2, 0], [5, 3, 3, 1, 1])
-    e1 = th.randn(5, 10)
     g2 = dgl.DGLGraph()
     g2.add_nodes(6)
     g2.add_edges([0, 1, 2, 5, 4, 5], [1 ,2 ,3, 4, 3, 0])
-    e2 = th.randn(6, 10)
     g3 = dgl.DGLGraph()
     g3.add_nodes(1)  # no edges
     g = dgl.batch([g1, g3, g2]) # should not throw an error
@@ -160,6 +167,7 @@ if __name__ == '__main__':
     test_batch_unbatch()
     test_batch_unbatch1()
     test_batched_edge_ordering()
-    test_batch_sendrecv()
+    test_batch_send_then_recv()
+    test_batch_send_and_recv()
     test_batch_propagate()
     test_batch_no_edge()

tests/pytorch/test_graph.py

@@ -39,24 +39,54 @@ def test_adjmat_speed():
     t0 = time.time()
     g.adjacency_matrix()
     dur2 = time.time() - t0
-    assert dur2 < dur1 / 5
+    print('first time {}, second time {}'.format(dur1, dur2))
+    assert dur2 < dur1
+
+def test_incmat():
+    g = dgl.DGLGraph()
+    g.add_nodes(4)
+    g.add_edge(0, 1) # 0
+    g.add_edge(0, 2) # 1
+    g.add_edge(0, 3) # 2
+    g.add_edge(2, 3) # 3
+    g.add_edge(1, 1) # 4
+    assert U.allclose(
+            g.incidence_matrix('in').to_dense(),
+            th.tensor([[0., 0., 0., 0., 0.],
+                       [1., 0., 0., 0., 1.],
+                       [0., 1., 0., 0., 0.],
+                       [0., 0., 1., 1., 0.]]))
+    assert U.allclose(
+            g.incidence_matrix('out').to_dense(),
+            th.tensor([[1., 1., 1., 0., 0.],
+                       [0., 0., 0., 0., 1.],
+                       [0., 0., 0., 1., 0.],
+                       [0., 0., 0., 0., 0.]]))
+    assert U.allclose(
+            g.incidence_matrix('both').to_dense(),
+            th.tensor([[-1., -1., -1., 0., 0.],
+                       [1., 0., 0., 0., 0.],
+                       [0., 1., 0., -1., 0.],
+                       [0., 0., 1., 1., 0.]]))
 
 def test_incmat_speed():
     n = 1000
-    p = 10 * math.log(n) / n
+    p = 2 * math.log(n) / n
     a = sp.random(n, n, p, data_rvs=lambda n: np.ones(n))
     g = dgl.DGLGraph(a)
     # the first call should contruct the adj
     t0 = time.time()
-    g.incidence_matrix()
+    g.incidence_matrix("in")
     dur1 = time.time() - t0
     # the second call should be cached and should be very fast
     t0 = time.time()
-    g.incidence_matrix()
+    g.incidence_matrix("in")
     dur2 = time.time() - t0
+    print('first time {}, second time {}'.format(dur1, dur2))
     assert dur2 < dur1
 
 if __name__ == '__main__':
     test_graph_creation()
     test_adjmat_speed()
+    test_incmat()
     test_incmat_speed()

tests/pytorch/test_specialization.py

 import torch as th
-import numpy as np
 import dgl
 import dgl.function as fn
 import utils as U
@@ -20,7 +19,7 @@ def generate_graph():
     g.set_e_repr({'e1': weights, 'e2': th.unsqueeze(weights, 1)})
     return g
 
-def test_update_all():
+def test_v2v_update_all():
     def _test(fld):
         def message_func(edges):
             return {'m' : edges.src[fld]}
@@ -64,7 +63,7 @@ def test_update_all():
     # test 2d node features
     _test('f2')
 
-def test_send_and_recv():
+def test_v2v_snr():
     u = th.tensor([0, 0, 0, 3, 4, 9])
     v = th.tensor([1, 2, 3, 9, 9, 0])
     def _test(fld):
@@ -111,7 +110,7 @@ def test_send_and_recv():
     # test 2d node features
     _test('f2')
 
-def test_update_all_multi_fn():
+def test_v2v_update_all_multi_fn():
     def message_func(edges):
         return {'m2': edges.src['f2']}
 
@@ -119,26 +118,17 @@ def test_update_all_multi_fn():
         return {'m2': edges.src['f2'] * edges.data['e2']}
 
     def reduce_func(nodes):
-        return {'v2': th.sum(nodes.mailbox['m2'], 1)}
+        return {'v1': th.sum(nodes.mailbox['m2'], 1)}
 
     g = generate_graph()
     g.set_n_repr({'v1' : th.zeros((10,)), 'v2' : th.zeros((10,))})
     fld = 'f2'
-    # update all, mix of builtin and UDF
-    g.update_all([fn.copy_src(src=fld, out='m1'), message_func],
-                 [fn.sum(msg='m1', out='v1'), reduce_func],
-                 None)
-    v1 = g.ndata['v1']
-    v2 = g.ndata['v2']
-    assert U.allclose(v1, v2)
 
-    # run builtin with single message and reduce
-    g.update_all(fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'), None)
+    g.update_all(message_func, reduce_func)
     v1 = g.ndata['v1']
-    assert U.allclose(v1, v2)
 
     # 1 message, 2 reduces
-    g.update_all(fn.copy_src(src=fld, out='m'), [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')], None)
+    g.update_all(fn.copy_src(src=fld, out='m'), [fn.sum(msg='m', out='v2'), fn.sum(msg='m', out='v3')])
     v2 = g.ndata['v2']
     v3 = g.ndata['v3']
     assert U.allclose(v1, v2)
@@ -159,7 +149,7 @@ def test_update_all_multi_fn():
     v2 = g.ndata['v2']
     assert U.allclose(v1, v2)
 
-def test_send_and_recv_multi_fn():
+def test_v2v_snr_multi_fn():
     u = th.tensor([0, 0, 0, 3, 4, 9])
     v = th.tensor([1, 2, 3, 9, 9, 0])
 
@@ -170,27 +160,15 @@ def test_send_and_recv_multi_fn():
         return {'m2': edges.src['f2'] * edges.data['e2']}
 
     def reduce_func(nodes):
-        return {'v2' : th.sum(nodes.mailbox['m2'], 1)}
+        return {'v1' : th.sum(nodes.mailbox['m2'], 1)}
 
     g = generate_graph()
     g.set_n_repr({'v1' : th.zeros((10, D)), 'v2' : th.zeros((10, D)),
         'v3' : th.zeros((10, D))})
     fld = 'f2'
 
-    # send and recv, mix of builtin and UDF
-    g.send_and_recv((u, v),
-                    [fn.copy_src(src=fld, out='m1'), message_func],
-                    [fn.sum(msg='m1', out='v1'), reduce_func],
-                    None)
+    g.send_and_recv((u, v), message_func, reduce_func)
     v1 = g.ndata['v1']
-    v2 = g.ndata['v2']
-    assert U.allclose(v1, v2)
-
-    # run builtin with single message and reduce
-    g.send_and_recv((u, v), fn.copy_src(src=fld, out='m'), fn.sum(msg='m', out='v1'),
-                    None)
-    v1 = g.ndata['v1']
-    assert U.allclose(v1, v2)
 
     # 1 message, 2 reduces
     g.send_and_recv((u, v),
@@ -219,8 +197,198 @@ def test_send_and_recv_multi_fn():
     v2 = g.ndata['v2']
     assert U.allclose(v1, v2)
 
+def test_e2v_update_all_multi_fn():
+    def _test(fld):
+        def message_func(edges):
+            return {'m1' : edges.src[fld] + edges.dst[fld],
+                    'm2' : edges.src[fld] * edges.dst[fld]}
+
+        def reduce_func(nodes):
+            return {fld : th.sum(nodes.mailbox['m1'] + nodes.mailbox['m2'], 1)}
+
+        def apply_func(nodes):
+            return {fld : 2 * nodes.data[fld]}
+
+        def apply_func_2(nodes):
+            return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
+
+        g = generate_graph()
+        # update all
+        v1 = g.get_n_repr()[fld]
+        # no specialization
+        g.update_all(message_func, reduce_func, apply_func)
+        v2 = g.get_n_repr()[fld]
+
+        # user break reduce func into 2 builtin
+        g.set_n_repr({fld : v1})
+        g.update_all(message_func,
+                     [fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
+                     apply_func_2)
+        v3 = g.get_n_repr()[fld]
+
+        assert th.allclose(v2, v3)
+
+    # test 1d node features
+    _test('f1')
+    # test 2d node features
+    _test('f2')
+
+def test_e2v_snr_multi_fn():
+    u = th.tensor([0, 0, 0, 3, 4, 9])
+    v = th.tensor([1, 2, 3, 9, 9, 0])
+    def _test(fld):
+        def message_func(edges):
+            return {'m1' : edges.src[fld] + edges.dst[fld],
+                    'm2' : edges.src[fld] * edges.dst[fld]}
+
+        def reduce_func(nodes):
+            return {fld : th.sum(nodes.mailbox['m1'] + nodes.mailbox['m2'], 1)}
+
+        def apply_func(nodes):
+            return {fld : 2 * nodes.data[fld]}
+
+        def apply_func_2(nodes):
+            return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
+
+        g = generate_graph()
+        # send_and_recv
+        v1 = g.get_n_repr()[fld]
+        # no specialization
+        g.send_and_recv((u, v), message_func, reduce_func, apply_func)
+        v2 = g.get_n_repr()[fld]
+
+        # user break reduce func into 2 builtin
+        g.set_n_repr({fld : v1})
+        g.send_and_recv((u, v), message_func,
+                        [fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
+                        apply_func_2)
+        v3 = g.get_n_repr()[fld]
+
+        assert th.allclose(v2, v3)
+
+    # test 1d node features
+    _test('f1')
+    # test 2d node features
+    _test('f2')
+
+def test_e2v_recv_multi_fn():
+    u = th.tensor([0, 0, 0, 3, 4, 9])
+    v = th.tensor([1, 2, 3, 9, 9, 0])
+    def _test(fld):
+        def message_func(edges):
+            return {'m1' : edges.src[fld] + edges.dst[fld],
+                    'm2' : edges.src[fld] * edges.dst[fld]}
+
+        def reduce_func(nodes):
+            return {fld : th.sum(nodes.mailbox['m1'] + nodes.mailbox['m2'], 1)}
+
+        def apply_func(nodes):
+            return {fld : 2 * nodes.data[fld]}
+
+        def apply_func_2(nodes):
+            return {fld : 2 * nodes.data['r1'] + 2 * nodes.data['r2']}
+
+        g = generate_graph()
+        # recv
+        v1 = g.get_n_repr()[fld]
+        # no specialization
+        g.send((u, v), message_func)
+        g.recv([0,1,2,3,9], reduce_func, apply_func)
+        v2 = g.get_n_repr()[fld]
+
+        # user break reduce func into 2 builtin
+        g.set_n_repr({fld : v1})
+        g.send((u, v), message_func)
+        g.recv([0,1,2,3,9],
+               [fn.sum(msg='m1', out='r1'), fn.sum(msg='m2', out='r2')],
+               apply_func_2)
+        v3 = g.get_n_repr()[fld]
+
+        assert th.allclose(v2, v3)
+
+    # test 1d node features
+    _test('f1')
+    # test 2d node features
+    _test('f2')
+
+def test_multi_fn_fallback():
+    # create a graph with zero in degree nodes
+    g = dgl.DGLGraph()
+    g.add_nodes(10)
+    for i in range(1, 9):
+        g.add_edge(0, i)
+        g.add_edge(i, 9)
+    g.ndata['h'] = th.randn(10, D)
+    g.edata['w1'] = th.randn(16,)
+    g.edata['w2'] = th.randn(16, D)
+    def _mfunc_hxw1(edges):
+        return {'m1' : edges.src['h'] * th.unsqueeze(edges.data['w1'], 1)}
+    def _mfunc_hxw2(edges):
+        return {'m2' : edges.src['h'] * edges.data['w2']}
+    def _rfunc_m1(nodes):
+        return {'o1' : th.sum(nodes.mailbox['m1'], 1)}
+    def _rfunc_m2(nodes):
+        return {'o2' : th.sum(nodes.mailbox['m2'], 1)}
+    def _rfunc_m1max(nodes):
+        return {'o3' : th.max(nodes.mailbox['m1'], 1)[0]}
+    def _afunc(nodes):
+        ret = {}
+        for k, v in nodes.data.items():
+            if k.startswith('o'):
+                ret[k] = 2 * v
+        return ret
+    # compute ground truth
+    g.update_all(_mfunc_hxw1, _rfunc_m1, _afunc)
+    o1 = g.ndata.pop('o1')
+    g.update_all(_mfunc_hxw2, _rfunc_m2, _afunc)
+    o2 = g.ndata.pop('o2')
+    g.update_all(_mfunc_hxw1, _rfunc_m1max, _afunc)
+    o3 = g.ndata.pop('o3')
+    # v2v spmv
+    g.update_all(fn.src_mul_edge(src='h', edge='w1', out='m1'),
+                 fn.sum(msg='m1', out='o1'),
+                 _afunc)
+    assert U.allclose(o1, g.ndata.pop('o1'))
+    # v2v fallback to e2v
+    g.update_all(fn.src_mul_edge(src='h', edge='w2', out='m2'),
+                 fn.sum(msg='m2', out='o2'),
+                 _afunc)
+    assert U.allclose(o2, g.ndata.pop('o2'))
+    # v2v fallback to degree bucketing
+    g.update_all(fn.src_mul_edge(src='h', edge='w1', out='m1'),
+                 fn.max(msg='m1', out='o3'),
+                 _afunc)
+    assert U.allclose(o3, g.ndata.pop('o3'))
+    # multi builtins, both v2v spmv
+    g.update_all([fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w1', out='m2')],
+                 [fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
+                 _afunc)
+    assert U.allclose(o1, g.ndata.pop('o1'))
+    assert U.allclose(o1, g.ndata.pop('o2'))
+    # multi builtins, one v2v spmv, one fallback to e2v
+    g.update_all([fn.src_mul_edge(src='h', edge='w1', out='m1'), fn.src_mul_edge(src='h', edge='w2', out='m2')],
+                 [fn.sum(msg='m1', out='o1'), fn.sum(msg='m2', out='o2')],
+                 _afunc)
+    assert U.allclose(o1, g.ndata.pop('o1'))
+    assert U.allclose(o2, g.ndata.pop('o2'))
+    # multi builtins, one v2v spmv, one fallback to e2v, one fallback to degree-bucketing
+    g.update_all([fn.src_mul_edge(src='h', edge='w1', out='m1'),
+                  fn.src_mul_edge(src='h', edge='w2', out='m2'),
+                  fn.src_mul_edge(src='h', edge='w1', out='m3')],
+                 [fn.sum(msg='m1', out='o1'),
+                  fn.sum(msg='m2', out='o2'),
+                  fn.max(msg='m3', out='o3')],
+                 _afunc)
+    assert U.allclose(o1, g.ndata.pop('o1'))
+    assert U.allclose(o2, g.ndata.pop('o2'))
+    assert U.allclose(o3, g.ndata.pop('o3'))
+
 if __name__ == '__main__':
-    test_update_all()
-    test_send_and_recv()
-    test_update_all_multi_fn()
-    test_send_and_recv_multi_fn()
+    test_v2v_update_all()
+    test_v2v_snr()
+    test_v2v_update_all_multi_fn()
+    test_v2v_snr_multi_fn()
+    test_e2v_update_all_multi_fn()
+    test_e2v_snr_multi_fn()
+    test_e2v_recv_multi_fn()
+    test_multi_fn_fallback()

toy.py

-###############################################################################
-# A toy example
-# -------------
-#
-# Let’s begin with the simplest graph possible with two nodes, and set
-# the node representations:
-
-import torch as th
-import dgl
-
-g = dgl.DGLGraph()
-g.add_nodes(2)
-g.add_edge(1, 0)
-
-x = th.tensor([[0.0, 0.0], [1.0, 2.0]])
-g.nodes[:].data['x'] = x
-
-###############################################################################
-# A syntax sugar for accessing feature data of all nodes
-print(g.ndata['x'])
-
-###############################################################################
-# What we want to do is simply to copy representation from node#1 to
-# node#0, but with a message passing interface. We do this like what we
-# will do over a pair of sockets, with a send and a recv interface. The
-# two user defined function (UDF) specifies the actions: deposit the
-# value into an internal key-value store with the key msg, and retrive
-# it. Note that there may be multiple incoming edges to a node, and the
-# receiving end aggregates them.
-
-def send_source(edges):  # type is dgl.EdgeBatch
-    return {'msg': edges.src['x']}
-
-def simple_reduce(nodes):  # type is dgl.NodeBatch
-    msgs = nodes.mailbox['msg']
-    return {'x' : th.sum(msgs, dim=1)}
-
-g.send((1, 0), message_func=send_source)
-g.recv(0, reduce_func=simple_reduce)
-print(g.ndata)
-
-###############################################################################
-# Some times the computation may involve representations on the edges.
-# Let’s say we want to “amplify” the message:
-
-w = th.tensor([2.0])
-g.edata['w'] = w
-
-def send_source_with_edge_weight(edges):
-    return {'msg': edges.src['x'] * edges.data['w']}
-
-g.send((1, 0), message_func=send_source_with_edge_weight)
-g.recv(0, reduce_func=simple_reduce)
-print(g.ndata)
-
-###############################################################################
-# Or we may need to involve the desination’s representation, and here
-# is one version:
-
-def simple_reduce_addup(nodes):
-    msgs = nodes.mailbox['msg']
-    return {'x' : nodes.data['x'] + th.sum(msgs, dim=1)}
-
-g.send((1, 0), message_func=send_source_with_edge_weight)
-g.recv(0, reduce_func=simple_reduce_addup)
-print(g.ndata)
-
-del g.ndata['x']
-del g.edata['w']