Repository: dmlc/dgl
Pull request: [Bugfix] Fix multiple send recv
## Description
This PR fixes #75 and other small bugs (see below) and supports multiple recv after send.

Previously, each recv call clears message graph so following recv won't see any message. This PR fix the bug by removing message graph and put message in edge space (but still in a separate message frame). Now message frame has the same size as edge frame and whether there is a message on an edge is indicated by a special "boolean" index in DGLGraph (graph._msg_index).

Now users can:
- call multiple send (on the edges or not), each using a different message function (with different message fields)
- call multiple recv

## Checklist
- [x] The PR title starts with [$CATEGORY] (such as [Model], [Doc], [Feature]])
- [x] Changes are complete (i.e. I finished coding on this PR)
- [x] All changes have test coverage
- [x] Code is well-documented
- [x] To the my best knowledge, examples are either not affected by this change,
      or have been fixed to be compatible with this change
- [x] Related issue is referred in this PR

## Changes
- [x] fix multiple send (using different fields)
- [x] fix multiple recv and remove message graph
- [x] fix frame append bug when some feature is not in the new frame
- [x] fix bug in send when edges is ALL


Problem statement:
[BUG] Consecutive sends without recv
Here's the test I was running on `cpp` branch.
```python
def test_send_twice():
    g = DGLGraph()
    g.add_nodes(2)
    g.add_edge(0, 1)
    def _message_a(src, edge):
        return {'a': src['a']}
    def _message_b(src, edge):
        return {'b': src['b']}
    def _reduce(node, msgs):
        return {'a': msgs['a'].sum(1), 'b': msgs['b'].sum(1)}

    old_a = th.randn(2, 5)
    old_b = th.randn(2, 5)
    g.set_n_repr({'a': old_a, 'b': old_b})

    g.send(0, 1, _message_a)
    g.send(0, 1, _message_b)
    g.recv([1], _reduce)

    new_repr = g.get_n_repr()

    assert th.allclose(new_repr['a'][1], old_a[0])
    assert th.allclose(new_repr['b'][1], old_b[0])
```

First, if I do two consecutive sends with different message keys, DGLGraph throws a `KeyError: 'b'` regardless of whether we are sending on the same node pairs/edges.

The next issue is that we are adding edges into a "message graph" during sends.  This does not make a lot of sense on multigraphs if we allow sending between the same node pairs twice without recv, which will leave us with two edges in the message graph.  Actually, in the test case above, I stepped into the `_batch_recv` method and found that the degree is 2 in the degree bucketing procedure.

So do we want to allow (1) consecutive sends without recv, and (2) consecutive sends between the same node pairs without recv?

If we were to allow both, I guess ultimately we would want to have the message graph represented as an edge subgraph, or somehow maintain a set of "already-added-edges" in the message graph.

Thoughts?

Hints from issue discussion:
The issue is because we didn't check whether the edges have been triggered or not. I think it's a good time to think about the case thoroughly (suppose we still want to keep the send/recv flexibility here):
* If the given edges have been triggered before, we should use `update_rows` instead of `append` to save the messages.
* If the given edges are new, use `append`. If the messages are of different keys, pad zero for other edges.
* If the given edges are partly new, split the edges into two sets and treat them differently as above.

**Do we need more APIs?**
I guess the `has_edges` API is sufficient to implement the above logic.

**Efficiency?**
Not efficient at all. However, I think flexibility is more important for these cases. Note that, if user calls `send_and_recv`, the messages are transient and thus shall not go through all the message graph shenanigans. This is what @ylfdq1118 is working on right now.

**Multigraph?**
Maybe user should use edge id as the argument to `send`? In such case, the message graph will also be a multigraph. Will this cause problem?
My main argument is that since we are supporting multigraphs, we either need to represent the message graph as a subgraph of the original graph, or check whether an edge is already in the message graph before addition (by maintaining a Python set of triggered edges); we could no longer just `add_edges` into the message graphs as it is right now.

I don't think `has_edges()` by itself is appropriate for this purpose because it only asks whether an edge is connecting two nodes.
How do you plan to maintain the storage? The good thing about the current design is that the message frame is always appending.
I guess we can still keep the current design of appending to message frames; all we need is to keep track of which edge IDs we have triggered.  We call `add_edges` only for those edges that are not triggered.  This can essentially combine with what we will do on message frames:

* For the edges triggered before, call `update_rows`
* For others, `add_edges` and `append`.

I expect that using edge subgraphs as message graphs will involve lots of overhauling so I don't recommend it for now.

Also,
> * If the given edges are new, use `append`. If the messages are of different keys, pad zero for other edges.

I'm confused of this.  Why do we want to pad zeros?  And I think that this is not even necessarily correct.  What if we are reducing by max and the messages are negative?
My two cents: if we still want to have the flexibility of allowing multiple sends on different keys, it is very likely that we need separate message graphs for each key.

For example:

* The test code above will have a message graph for key `a` with edge `(0, 1)`, and another message graph for key `b` with edge `(0, 1)`.
* If we are running `g.send(0, 2, _message_b)` instead, then message graph `a` should have edge `(0, 1)` but message graph `b` instead should have `(0, 2)`.  If we are recv'ing on 0, the `msgs` argument should be a dictionary with `a` from `(0, 1)` and `b` from `(0, 2)`.
Bring @GaiYu0 into the loop. This is the issue we've discussed before. The problem is the in-message-degree will be different for different message keys (if we don't pad the missing messages). Quan's solution is to extend the API to let user tell us which message key is required. Since this is a corner case, I am afraid that this might slow down the common case due to extra check and overhead.
Another related question here is how should we support degree-padding other than degree-bucketing? I remembered this is how JTNN people write their codes. Is this correct @BarclayII ?
> Bring @GaiYu0 into the loop. This is the issue we've discussed before. The problem is the in-message-degree will be different for different message keys (if we don't pad the missing messages). Quan's solution is to extend the API to let user tell us which message key is required. Since this is a corner case, I am afraid that this might slow down the common case due to extra check and overhead.

I wasn't thinking of extending the API of any sorts.  I prototyped my solution in `cpp-multisend` branch and you can see the diff.  Duplicate edges are not handled yet there and I believe more optimization is needed.  But you can see that the user interface there is not changed whatsoever.

> Another related question here is how should we support degree-padding other than degree-bucketing? I remembered this is how JTNN people write their codes. Is this correct @BarclayII ?

For degree padding:

* I think it only works if we are using builtin reducers like `sum` or `max`.  For different reducers we also need to pad different values (0 for `sum`, -inf for `max`).
* I would imagine that we have another general scheduler at the same level as `degree_bucketing`, except that the reduce function is called only once with padded messages.
* We can figure out the maximum in-degree and pad the list of incoming messages `[(in_degree[i], ...) for i in range(N)]` to a tensor `(N, max(in_degree), ...)`.
* We can allow the users to select which strategy they with to use during reduce (`degree_bucketing` or `degree_padding`).  Maybe we can add an argument `preferred_scheduler` in `recv()` and other calls involving reducers?
Oh I see. I misunderstood your solution. Can we simplify it to forbid user to have misaligned message fields in multiple sends? I'm afraid maintaining one message graph for one message field is very costly.
I think forbidding this is fair as user can call send on field `a` first, recv it and then call send on field `b`, and we are no problem with that.

Commit message: test multi-send-recv after conversion

Patch to review:
diff --git a/python/dgl/frame.py b/python/dgl/frame.py
index 4451c8fcdcb2..0d5374a61800 100644
--- a/python/dgl/frame.py
+++ b/python/dgl/frame.py
@@ -366,6 +366,20 @@ def _append(self, other):
             # directly updating columns.
             self._columns = {key: Column.create(data) for key, data in other.items()}
         else:
+            # pad columns that no new were provided
+            for key, col in self.items():
+                if key not in other:
+                    scheme = col.scheme
+                    ctx = F.context(col.data)
+                    if self.get_initializer(key) is None:
+                        self._warn_and_set_initializer()
+                    new_data = self.get_initializer(key)(
+                            (other.num_rows,) + scheme.shape, scheme.dtype,
+                            ctx, slice(self._num_rows,
+                                       self._num_rows + other.num_rows)
+                    )
+                    other[key] = new_data
+            # append other to self
             for key, col in other.items():
                 if key not in self._columns:
                     # the column does not exist; init a new column
diff --git a/python/dgl/graph.py b/python/dgl/graph.py
index c3193d9a5fa6..36357b0d760c 100644
--- a/python/dgl/graph.py
+++ b/python/dgl/graph.py
@@ -179,7 +179,7 @@ def __init__(self,
         # graph
         self._readonly=readonly
         self._graph = create_graph_index(graph_data, multigraph, readonly)
-        # frame
+        # node and edge frame
         if node_frame is None:
             self._node_frame = FrameRef(Frame(num_rows=self.number_of_nodes()))
         else:
@@ -188,10 +188,14 @@ def __init__(self,
             self._edge_frame = FrameRef(Frame(num_rows=self.number_of_edges()))
         else:
             self._edge_frame = edge_frame
-        # msg graph & frame
-        self._msg_graph = create_graph_index(multigraph=multigraph)
-        self._msg_frame = FrameRef()
-        self.reset_messages()
+        # message indicator:
+        # if self._msg_index[eid] == 1, then edge eid has message
+        self._msg_index = utils.zero_index(size=self.number_of_edges())
+        # message frame
+        self._msg_frame = FrameRef(Frame(num_rows=self.number_of_edges()))
+        # set initializer for message frame
+        # all fields including message indicator should use zero_initializer
+        self._msg_frame.set_initializer(dgl.init.zero_initializer)
         # registered functions
         self._message_func = None
         self._reduce_func = None
@@ -243,7 +247,6 @@ def add_nodes(self, num, data=None):
                 [1., 1., 1., 1.]])
         """
         self._graph.add_nodes(num)
-        self._msg_graph.add_nodes(num)
         if data is None:
             # Initialize feature placeholders if there are features existing
             self._node_frame.add_rows(num)
@@ -303,6 +306,9 @@ def add_edge(self, u, v, data=None):
             self._edge_frame.add_rows(1)
         else:
             self._edge_frame.append(data)
+        # resize msg_index and msg_frame
+        self._msg_index = self._msg_index.add_elements(1)
+        self._msg_frame.add_rows(1)
 
     def add_edges(self, u, v, data=None):
         """Add multiple edges for list of source nodes u and destination nodes
@@ -359,6 +365,10 @@ def add_edges(self, u, v, data=None):
             self._edge_frame.add_rows(max(len(u), len(v)))
         else:
             self._edge_frame.append(data)
+        # initialize feature placeholder for messages
+        num = max(len(u), len(v))
+        self._msg_index = self._msg_index.add_elements(num)
+        self._msg_frame.add_rows(num)
 
     def clear(self):
         """Remove all nodes and edges, as well as their features, from the
@@ -382,7 +392,7 @@ def clear(self):
         self._graph.clear()
         self._node_frame.clear()
         self._edge_frame.clear()
-        self._msg_graph.clear()
+        self._msg_index = utils.zero_index(0)
         self._msg_frame.clear()
 
     def clear_cache(self):
@@ -394,12 +404,6 @@ def clear_cache(self):
         """
         self._graph.clear_cache()
 
-    def reset_messages(self):
-        """Clear all messages."""
-        self._msg_graph.clear()
-        self._msg_frame.clear()
-        self._msg_graph.add_nodes(self.number_of_nodes())
-
     def number_of_nodes(self):
         """Return the number of nodes in the graph.
 
@@ -1168,7 +1172,9 @@ def from_networkx(self, nx_graph, node_attrs=None, edge_attrs=None):
         self._graph.from_networkx(nx_graph)
         self._node_frame.add_rows(self.number_of_nodes())
         self._edge_frame.add_rows(self.number_of_edges())
-        self._msg_graph.add_nodes(self._graph.number_of_nodes())
+        self._msg_index = utils.zero_index(self.number_of_edges())
+        self._msg_frame.add_rows(self.number_of_edges())
+
         # copy attributes
         def _batcher(lst):
             if F.is_tensor(lst[0]):
@@ -1225,7 +1231,8 @@ def from_scipy_sparse_matrix(self, a):
         self._graph.from_scipy_sparse_matrix(a)
         self._node_frame.add_rows(self.number_of_nodes())
         self._edge_frame.add_rows(self.number_of_edges())
-        self._msg_graph.add_nodes(self._graph.number_of_nodes())
+        self._msg_index = utils.zero_index(self.number_of_edges())
+        self._msg_frame.add_rows(self.number_of_edges())
 
     def node_attr_schemes(self):
         """Return the node feature schemes.
@@ -1934,7 +1941,7 @@ def send(self, edges=ALL, message_func="default"):
             message_func = self._message_func
 
         if is_all(edges):
-            eid = ALL
+            eid = utils.toindex(slice(0, self.number_of_edges()))
             u, v, _ = self._graph.edges()
         elif isinstance(edges, tuple):
             u, v = edges
@@ -1946,14 +1953,15 @@ def send(self, edges=ALL, message_func="default"):
             eid = utils.toindex(edges)
             u, v, _ = self._graph.find_edges(eid)
 
+        if len(eid) == 0:
+            # no edge to be triggered
+            return
+
         with ir.prog() as prog:
             scheduler.schedule_send(graph=self, u=u, v=v, eid=eid,
                                     message_func=message_func)
             Runtime.run(prog)
 
-        # update message graph and frame
-        self._msg_graph.add_edges(u, v)
-
     def recv(self,
              v=ALL,
              reduce_func="default",
@@ -2039,10 +2047,6 @@ def recv(self,
             apply_node_func = self._apply_node_func
         assert reduce_func is not None
 
-        if self._msg_frame.num_rows == 0:
-            # no message has ever been sent
-            return
-
         if is_all(v):
             v = F.arange(0, self.number_of_nodes())
         elif isinstance(v, int):
@@ -2060,9 +2064,6 @@ def recv(self,
                                     inplace=inplace)
             Runtime.run(prog)
 
-        # FIXME(minjie): multi send bug
-        self.reset_messages()
-
     def send_and_recv(self,
                       edges,
                       message_func="default",
diff --git a/python/dgl/runtime/ir/executor.py b/python/dgl/runtime/ir/executor.py
index 87facddf29a0..e49fdcb3562f 100644
--- a/python/dgl/runtime/ir/executor.py
+++ b/python/dgl/runtime/ir/executor.py
@@ -34,6 +34,7 @@ class OpCode(object):
     WRITE_DICT_ = 24
     APPEND_ROW_ = 25
     WRITE_ROW_INPLACE_ = 26
+    CLEAR_FRAME_ = 27
 
 class Executor(object):
     @abstractmethod
@@ -645,3 +646,32 @@ def run(self):
 def APPEND_ROW_(fd1, fd2):
     reg = IR_REGISTRY[OpCode.APPEND_ROW_]
     get_current_prog().issue(reg['executor_cls'](fd1, fd2))
+
+class ClearFrame_Executor(Executor):
+    def __init__(self, fd):
+        self.fd = fd
+
+    def opcode(self):
+        return OpCode.CLEAR_FRAME_
+
+    def arg_vars(self):
+        return []
+
+    def ret_var(self):
+        return None
+
+    def run(self):
+        frame = self.fd.data
+        num_rows = frame.num_rows
+        frame.clear()
+        frame.add_rows(num_rows)
+
+IR_REGISTRY[OpCode.CLEAR_FRAME_] = {
+    'name': 'CLEAR_FRAME_',
+    'args_type': [],
+    'ret_type': None,
+    'executor_cls': ClearFrame_Executor,
+}
+def CLEAR_FRAME_(fd):
+    reg = IR_REGISTRY[OpCode.CLEAR_FRAME_]
+    get_current_prog().issue(reg['executor_cls'](fd))
diff --git a/python/dgl/runtime/scheduler.py b/python/dgl/runtime/scheduler.py
index 212fd7271913..c7a06ea86e91 100644
--- a/python/dgl/runtime/scheduler.py
+++ b/python/dgl/runtime/scheduler.py
@@ -44,15 +44,16 @@ def schedule_send(graph, u, v, eid, message_func):
     # TODO(minjie): support builtin message func
     message_func = _standardize_func_usage(message_func, 'message')
     # vars
-    nf = var.FEAT_DICT(graph._node_frame)
-    ef = var.FEAT_DICT(graph._edge_frame)
-    mf = var.FEAT_DICT(graph._msg_frame)
-    u = var.IDX(u)
-    v = var.IDX(v)
-    eid = var.IDX(eid)
-    msg = _gen_send(graph, nf, ef, u, v, eid, message_func)
-    # TODO: handle duplicate messages
-    ir.APPEND_ROW_(mf, msg)
+    var_nf = var.FEAT_DICT(graph._node_frame)
+    var_ef = var.FEAT_DICT(graph._edge_frame)
+    var_mf = var.FEAT_DICT(graph._msg_frame)
+    var_u = var.IDX(u)
+    var_v = var.IDX(v)
+    var_eid = var.IDX(eid)
+    msg = _gen_send(graph, var_nf, var_ef, var_u, var_v, var_eid, message_func)
+    ir.WRITE_ROW_(var_mf, var_eid, msg)
+    # set message indicator to 1
+    graph._msg_index = graph._msg_index.set_items(eid, 1)
 
 def schedule_recv(graph,
                   recv_nodes,
@@ -74,25 +75,37 @@ def schedule_recv(graph,
     inplace: bool
         If True, the update will be done in place
     """
-    src, dst, mid = graph._msg_graph.in_edges(recv_nodes)
-    if len(mid) == 0:
-        # All recv nodes are 0-degree nodes; downgrade to apply nodes.
+    _, _, eid = graph._graph.in_edges(recv_nodes)
+    if len(eid) > 0:
+        e_mask = graph._msg_index.get_items(eid).nonzero()
+        eid = eid.get_items(e_mask)
+    if len(eid) == 0:
+        # Downgrade to apply nodes if
+        #   1) all recv nodes are 0-degree nodes
+        #   2) no send has been called
         if apply_func is not None:
             schedule_apply_nodes(graph, recv_nodes, apply_func, inplace)
     else:
+        src, dst, _ = graph._graph.find_edges(eid)
         var_nf = var.FEAT_DICT(graph._node_frame, name='nf')
         # sort and unique the argument
         recv_nodes, _ = F.sort_1d(F.unique(recv_nodes.tousertensor()))
         recv_nodes = utils.toindex(recv_nodes)
         var_recv_nodes = var.IDX(recv_nodes, name='recv_nodes')
         # reduce
-        reduced_feat = _gen_reduce(graph, reduce_func, (src, dst, mid), recv_nodes)
+        reduced_feat = _gen_reduce(graph, reduce_func, (src, dst, eid),
+                                   recv_nodes)
         # apply
-        final_feat = _apply_with_accum(graph, var_recv_nodes, var_nf, reduced_feat, apply_func)
+        final_feat = _apply_with_accum(graph, var_recv_nodes, var_nf,
+                                       reduced_feat, apply_func)
         if inplace:
             ir.WRITE_ROW_INPLACE_(var_nf, var_recv_nodes, final_feat)
         else:
             ir.WRITE_ROW_(var_nf, var_recv_nodes, final_feat)
+        # set message indicator to 0
+        graph._msg_index = graph._msg_index.set_items(eid, 0)
+        if not graph._msg_index.has_nonzero():
+            ir.CLEAR_FRAME_(var.FEAT_DICT(graph._msg_frame, name='mf'))
 
 def schedule_snr(graph,
                  edge_tuples,
@@ -426,13 +439,14 @@ def _gen_reduce(graph, reduce_func, edge_tuples, recv_nodes):
     recv_nodes : utils.Index
     """
     call_type = "recv"
-    _, dst, mid = edge_tuples
+    _, dst, eid = edge_tuples
     rfunc = _standardize_func_usage(reduce_func, 'reduce')
     rfunc_is_list = utils.is_iterable(rfunc)
     # Create a tmp frame to hold the feature data.
     # The frame has the same size and schemes of the
     # node frame.
-    # TODO(minjie): should replace this with an IR call to make the program stateless.
+    # TODO(minjie): should replace this with an IR call to make the program
+    # stateless.
     tmpframe = FrameRef(frame_like(graph._node_frame._frame, len(recv_nodes)))
 
     # vars
@@ -444,8 +458,8 @@ def _gen_reduce(graph, reduce_func, edge_tuples, recv_nodes):
         # UDF message + builtin reducer
         # analyze e2v spmv
         spmv_rfunc, rfunc = spmv.analyze_e2v_spmv(graph, rfunc)
-        # FIXME: refactor this when fixing the multi-recv bug
-        inc = spmv.build_inc_matrix_eid(graph._msg_frame.num_rows, mid, dst, recv_nodes)
+        inc = spmv.build_inc_matrix_eid(graph._msg_frame.num_rows, eid, dst,
+                                        recv_nodes)
         spmv.gen_e2v_spmv_schedule(inc, spmv_rfunc, msg, out)
 
         if len(rfunc) == 0:
@@ -456,7 +470,7 @@ def _gen_reduce(graph, reduce_func, edge_tuples, recv_nodes):
         rfunc = BundledFunction(rfunc)
 
     # gen degree bucketing schedule for UDF recv
-    db.gen_degree_bucketing_schedule(graph, rfunc, mid, dst,
+    db.gen_degree_bucketing_schedule(graph, rfunc, eid, dst,
             recv_nodes, nf, msg, out)
     return out
 
diff --git a/python/dgl/utils.py b/python/dgl/utils.py
index defb10590204..c53de7995e88 100644
--- a/python/dgl/utils.py
+++ b/python/dgl/utils.py
@@ -75,6 +75,7 @@ def _dispatch(self, data):
                     raise DGLError('Error index data: %s' % str(data))
             self._user_tensor_data[F.cpu()] = F.zerocopy_from_numpy(self._pydata)
 
+
     def tonumpy(self):
         """Convert to a numpy ndarray."""
         if self._pydata is None:
@@ -126,9 +127,61 @@ def __getstate__(self):
     def __setstate__(self, state):
         self._initialize_data(state)
 
+    def get_items(self, index):
+        tensor = self.tousertensor()
+        if isinstance(index, Index):
+            if isinstance(index._pydata, slice):
+                # XXX: assuming tensor supports index using slices
+                #      cannot call tousertensor because for slice type it would
+                #      further call tonumpy, which overwrites _pydata and causes
+                #      future bugs...
+                index = index._pydata
+            else:
+                index = index.tousertensor()
+        return Index(tensor[index])
+
+    def set_items(self, index, value):
+        tensor = self.tousertensor()
+        if isinstance(index, Index):
+            if isinstance(index._pydata, slice):
+                # XXX: assuming tensor supports index using slices
+                #      cannot call tousertensor because for slice type it would
+                #      further call tonumpy, which overwrites _pydata and causes
+                #      future bugs...
+                index = index._pydata
+            else:
+                index = index.tousertensor()
+        if isinstance(value, Index):
+            value = value.tousertensor()
+        tensor[index] = value
+        return Index(tensor)
+
+    def add_elements(self, num):
+        if num == 0:
+            return self
+        new_items = F.zeros((num,), dtype=F.int64, ctx=F.cpu())
+        if len(self) == 0:
+            return Index(new_items)
+        else:
+            tensor = self.tousertensor()
+            tensor = F.cat((tensor, new_items), dim=0)
+            return Index(tensor)
+
+    def nonzero(self):
+        tensor = self.tousertensor()
+        mask = F.nonzero_1d(tensor != 0)
+        return Index(mask)
+
+    def has_nonzero(self):
+        tensor = self.tousertensor()
+        return F.sum(tensor, 0) > 0
+
 def toindex(x):
     return x if isinstance(x, Index) else Index(x)
 
+def zero_index(size):
+    return Index(F.zeros((size,), dtype=F.int64, ctx=F.cpu()))
+
 class LazyDict(Mapping):
     """A readonly dictionary that does not materialize the storage."""
     def __init__(self, fn, keys):
diff --git a/tests/pytorch/test_basics.py b/tests/pytorch/test_basics.py
index 9ab18712b4ce..a2451416ea8a 100644
--- a/tests/pytorch/test_basics.py
+++ b/tests/pytorch/test_basics.py
@@ -490,34 +490,6 @@ def _init2(shape, dtype, ctx, ids):
     # non-0deg check: not touched
     assert U.allclose(new[1], old[1])
 
-def _disabled_test_send_twice():
-    # TODO(minjie): please re-enable this unittest after the send code problem is fixed.
-    g = DGLGraph()
-    g.add_nodes(3)
-    g.add_edge(0, 1)
-    g.add_edge(2, 1)
-    def _message_a(edges):
-        return {'a': edges.src['a']}
-    def _message_b(edges):
-        return {'a': edges.src['a'] * 3}
-    def _reduce(nodes):
-        return {'a': nodes.mailbox['a'].max(1)[0]}
-
-    old_repr = th.randn(3, 5)
-    g.ndata['a'] = old_repr
-    g.send((0, 1), _message_a)
-    g.send((0, 1), _message_b)
-    g.recv(1, _reduce)
-    new_repr = g.ndata['a']
-    assert U.allclose(new_repr[1], old_repr[0] * 3)
-
-    g.ndata['a'] = old_repr
-    g.send((0, 1), _message_a)
-    g.send((2, 1), _message_b)
-    g.recv(1, _reduce)
-    new_repr = g.ndata['a']
-    assert U.allclose(new_repr[1], th.stack([old_repr[0], old_repr[2] * 3], 0).max(0)[0])
-
 def test_send_multigraph():
     g = DGLGraph(multigraph=True)
     g.add_nodes(3)
@@ -614,6 +586,10 @@ def test_dynamic_addition():
     g.edges[4].data['h1'] = th.randn(1, D)
     assert g.edata['h1'].shape[0] == g.edata['h2'].shape[0] == 5
 
+    # test add edge with part of the features
+    g.add_edge(2, 1, {'h1': th.randn(1, D)})
+    assert len(g.edata['h1']) == len(g.edata['h2'])
+
 
 def test_repr():
     G = dgl.DGLGraph()
diff --git a/tests/pytorch/test_multi_send_recv.py b/tests/pytorch/test_multi_send_recv.py
new file mode 100644
index 000000000000..6a52844308ac
--- /dev/null
+++ b/tests/pytorch/test_multi_send_recv.py
@@ -0,0 +1,349 @@
+import torch as th
+from torch.autograd import Variable
+import numpy as np
+import dgl
+from dgl.graph import DGLGraph
+import utils as U
+from collections import defaultdict as ddict
+import scipy.sparse as sp
+
+D = 5
+
+def message_func(edges):
+    assert len(edges.src['h'].shape) == 2
+    assert edges.src['h'].shape[1] == D
+    return {'m' : edges.src['h']}
+
+def reduce_func(nodes):
+    msgs = nodes.mailbox['m']
+    assert len(msgs.shape) == 3
+    assert msgs.shape[2] == D
+    return {'accum' : th.sum(msgs, 1)}
+
+def apply_node_func(nodes):
+    return {'h' : nodes.data['h'] + nodes.data['accum']}
+
+def generate_graph(grad=False):
+    g = DGLGraph()
+    g.add_nodes(10) # 10 nodes.
+    # create a graph where 0 is the source and 9 is the sink
+    # 16 edges
+    for i in range(1, 9):
+        g.add_edge(0, i)
+        g.add_edge(i, 9)
+    ncol = Variable(th.randn(10, D), requires_grad=grad)
+    ecol = Variable(th.randn(16, D), requires_grad=grad)
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.set_e_initializer(dgl.init.zero_initializer)
+    g.ndata['h'] = ncol
+    g.edata['w'] = ecol
+    return g
+
+def test_multi_send():
+    g = generate_graph()
+    def _fmsg(edges):
+        assert edges.src['h'].shape == (5, D)
+        return {'m' : edges.src['h']}
+    g.register_message_func(_fmsg)
+    # many-many send
+    u = th.tensor([0, 0, 0, 0, 0])
+    v = th.tensor([1, 2, 3, 4, 5])
+    g.send((u, v))
+    # duplicate send
+    u = th.tensor([0])
+    v = th.tensor([1, 2, 3, 4, 5])
+    g.send((u, v))
+    # send more
+    u = th.tensor([1, 2, 3, 4, 5])
+    v = th.tensor([9])
+    g.send((u, v))
+
+    # check if message indicator is as expected
+    expected = th.zeros((g.number_of_edges(),), dtype=th.int64)
+    eid = g.edge_ids([0, 0, 0, 0, 0, 1, 2, 3, 4, 5],
+                     [1, 2, 3, 4, 5, 9, 9, 9, 9, 9])
+    expected[eid] = 1
+    assert th.equal(g._msg_index.tousertensor(), expected)
+
+def test_multi_recv():
+    # basic recv test
+    g = generate_graph()
+    h = g.ndata['h']
+    g.register_message_func(message_func)
+    g.register_reduce_func(reduce_func)
+    g.register_apply_node_func(apply_node_func)
+    expected = th.zeros((g.number_of_edges(),), dtype=th.int64)
+    # two separate round of send and recv
+    u = [4, 5, 6]
+    v = [9]
+    g.send((u, v))
+    eid = g.edge_ids(u, v)
+    expected[eid] = 1
+    assert th.equal(g._msg_index.tousertensor(), expected)
+    g.recv(v)
+    expected[eid] = 0
+    assert th.equal(g._msg_index.tousertensor(), expected)
+
+    u = [0]
+    v = [1, 2, 3]
+    g.send((u, v))
+    eid = g.edge_ids(u, v)
+    expected[eid] = 1
+    assert th.equal(g._msg_index.tousertensor(), expected)
+    g.recv(v)
+    expected[eid] = 0
+    assert th.equal(g._msg_index.tousertensor(), expected)
+
+    h1 = g.ndata['h']
+
+    # one send, two recv
+    g.ndata['h'] = h
+    u = th.tensor([0, 0, 0, 4, 5, 6])
+    v = th.tensor([1, 2, 3, 9, 9, 9])
+    g.send((u, v))
+    eid = g.edge_ids(u, v)
+    expected[eid] = 1
+    assert th.equal(g._msg_index.tousertensor(), expected)
+    u = [4, 5, 6]
+    v = [9]
+    g.recv(v)
+    eid = g.edge_ids(u, v)
+    expected[eid] = 0
+    assert th.equal(g._msg_index.tousertensor(), expected)
+    u = [0]
+    v = [1, 2, 3]
+    g.recv(v)
+    eid = g.edge_ids(u, v)
+    expected[eid] = 0
+    assert th.equal(g._msg_index.tousertensor(), expected)
+
+    h2 = g.ndata['h']
+    assert U.allclose(h1, h2)
+
+def test_multi_recv_0deg():
+    # test recv with 0deg nodes;
+    g = DGLGraph()
+    def _message(edges):
+        return {'m' : edges.src['h']}
+    def _reduce(nodes):
+        return {'h' : nodes.data['h'] + nodes.mailbox['m'].sum(1)}
+    def _apply(nodes):
+        return {'h' : nodes.data['h'] * 2}
+    def _init2(shape, dtype, ctx, ids):
+        return 2 + th.zeros(shape, dtype=dtype, device=ctx)
+    g.register_message_func(_message)
+    g.register_reduce_func(_reduce)
+    g.register_apply_node_func(_apply)
+    g.set_n_initializer(_init2)
+    g.add_nodes(2)
+    g.add_edge(0, 1)
+    # recv both 0deg and non-0deg nodes
+    old = th.randn((2, 5))
+    g.ndata['h'] = old
+    g.send((0, 1))
+    g.recv([0, 1])
+    new = g.ndata['h']
+    # 0deg check: initialized with the func and got applied
+    assert U.allclose(new[0], th.full((5,), 4))
+    # non-0deg check
+    assert U.allclose(new[1], th.sum(old, 0) * 2)
+
+    # recv again on zero degree node
+    g.recv([0])
+    assert U.allclose(g.nodes[0].data['h'], th.full((5,), 8))
+
+    # recv again on node with no incoming message
+    g.recv([1])
+    assert U.allclose(g.nodes[1].data['h'], th.sum(old, 0) * 4)
+
+def test_send_twice_different_shape():
+    g = generate_graph()
+    def _message_1(edges):
+        return {'h': edges.src['h']}
+    def _message_2(edges):
+        return {'h': th.cat((edges.src['h'], edges.data['w']), dim=1)}
+    g.send(message_func=_message_1)
+    g.send(message_func=_message_2)
+
+def test_send_twice_different_msg():
+    g = DGLGraph()
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.add_nodes(3)
+    g.add_edge(0, 1)
+    g.add_edge(2, 1)
+    def _message_a(edges):
+        return {'a': edges.src['a']}
+    def _message_b(edges):
+        return {'a': edges.src['a'] * 3}
+    def _reduce(nodes):
+        return {'a': nodes.mailbox['a'].max(1)[0]}
+
+    old_repr = th.randn(3, 5)
+    g.ndata['a'] = old_repr
+    g.send((0, 1), _message_a)
+    g.send((0, 1), _message_b)
+    g.recv(1, _reduce)
+    new_repr = g.ndata['a']
+    assert U.allclose(new_repr[1], old_repr[0] * 3)
+
+    g.ndata['a'] = old_repr
+    g.send((0, 1), _message_a)
+    g.send((2, 1), _message_b)
+    g.recv(1, _reduce)
+    new_repr = g.ndata['a']
+    assert U.allclose(new_repr[1], th.stack([old_repr[0], old_repr[2] * 3], 0).max(0)[0])
+
+def test_send_twice_different_field():
+    g = DGLGraph()
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.add_nodes(2)
+    g.add_edge(0, 1)
+    def _message_a(edges):
+        return {'a': edges.src['a']}
+    def _message_b(edges):
+        return {'b': edges.src['b']}
+    def _reduce(nodes):
+        return {'a': nodes.mailbox['a'].sum(1), 'b': nodes.mailbox['b'].sum(1)}
+    old_a = th.randn(2, 5)
+    old_b = th.randn(2, 5)
+    g.set_n_repr({'a': old_a, 'b': old_b})
+    g.send((0, 1), _message_a)
+    g.send((0, 1), _message_b)
+    g.recv([1], _reduce)
+    new_repr = g.get_n_repr()
+    assert th.allclose(new_repr['a'][1], old_a[0])
+    assert th.allclose(new_repr['b'][1], old_b[0])
+
+def test_dynamic_addition():
+    N = 3
+    D = 1
+
+    g = DGLGraph()
+    def _init(shape, dtype, ctx, ids):
+        return th.randn(shape, dtype=dtype, device=ctx)
+    g.set_n_initializer(_init)
+    g.set_e_initializer(_init)
+
+    def _message(edges):
+        return {'m' : edges.src['h1'] + edges.dst['h2'] + edges.data['h1'] +
+                edges.data['h2']}
+    def _reduce(nodes):
+        return {'h' : nodes.mailbox['m'].sum(1)}
+    def _apply(nodes):
+        return {'h' : nodes.data['h']}
+
+    g.register_message_func(_message)
+    g.register_reduce_func(_reduce)
+    g.register_apply_node_func(_apply)
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.set_e_initializer(dgl.init.zero_initializer)
+
+    # add nodes and edges
+    g.add_nodes(N)
+    g.ndata.update({'h1': th.randn(N, D),
+                    'h2': th.randn(N, D)})
+    g.add_nodes(3)
+    g.add_edge(0, 1)
+    g.add_edge(1, 0)
+    g.edata.update({'h1': th.randn(2, D),
+                    'h2': th.randn(2, D)})
+    g.send()
+    expected = th.ones((g.number_of_edges(),), dtype=th.int64)
+    assert th.equal(g._msg_index.tousertensor(), expected)
+
+    # add more edges
+    g.add_edges([0, 2], [2, 0], {'h1': th.randn(2, D)})
+    g.send(([0, 2], [2, 0]))
+    g.recv(0)
+
+    g.add_edge(1, 2)
+    g.edges[4].data['h1'] = th.randn(1, D)
+    g.send((1, 2))
+    g.recv([1, 2])
+
+    h = g.ndata.pop('h')
+
+    # a complete round of send and recv
+    g.send()
+    g.recv()
+    assert U.allclose(h, g.ndata['h'])
+
+def test_recv_no_send():
+    g = generate_graph()
+    g.recv(1, reduce_func)
+    # test recv after clear
+    g.clear()
+    g.add_nodes(3)
+    g.add_edges([0, 1], [1, 2])
+    g.set_n_initializer(dgl.init.zero_initializer)
+    g.ndata['h'] = th.randn(3, D)
+    g.send((1, 2), message_func)
+    expected = th.zeros((2,), dtype=th.int64)
+    expected[1] = 1
+    assert th.equal(g._msg_index.tousertensor(), expected)
+    g.recv(2, reduce_func)
+    expected[1] = 0
+    assert th.equal(g._msg_index.tousertensor(), expected)
+
+def test_send_recv_after_conversion():
+    # test send and recv after converting from a graph with edges
+
+    g = generate_graph()
+
+    # nx graph
+    nxg = g.to_networkx(node_attrs=['h'])
+    g1 = DGLGraph()
+    # some random node and edges
+    g1.add_nodes(4)
+    g1.add_edges([1, 2], [2, 3])
+    g1.set_n_initializer(dgl.init.zero_initializer)
+    g1.from_networkx(nxg, node_attrs=['h'])
+
+    # sparse matrix
+    row, col= g.all_edges()
+    data = range(len(row))
+    n = g.number_of_nodes()
+    a = sp.coo_matrix((data, (row, col)), shape=(n, n))
+    g2 = DGLGraph()
+    # some random node and edges
+    g2.add_nodes(5)
+    g2.add_edges([1, 2, 4], [2, 3, 0])
+    g2.set_n_initializer(dgl.init.zero_initializer)
+    g2.from_scipy_sparse_matrix(a)
+    g2.ndata['h'] = g.ndata['h']
+
+    # on dgl graph
+    g.send(message_func=message_func)
+    g.recv([0, 1, 3, 5], reduce_func=reduce_func,
+           apply_node_func=apply_node_func)
+    g.recv([0, 2, 4, 8], reduce_func=reduce_func,
+           apply_node_func=apply_node_func)
+
+    # nx
+    g1.send(message_func=message_func)
+    g1.recv([0, 1, 3, 5], reduce_func=reduce_func,
+            apply_node_func=apply_node_func)
+    g1.recv([0, 2, 4, 8], reduce_func=reduce_func,
+            apply_node_func=apply_node_func)
+
+    # sparse matrix
+    g2.send(message_func=message_func)
+    g2.recv([0, 1, 3, 5], reduce_func=reduce_func,
+            apply_node_func=apply_node_func)
+    g2.recv([0, 2, 4, 8], reduce_func=reduce_func,
+            apply_node_func=apply_node_func)
+
+    assert U.allclose(g.ndata['h'], g1.ndata['h'])
+    assert U.allclose(g.ndata['h'], g2.ndata['h'])
+
+
+if __name__ == '__main__':
+    test_multi_send()
+    test_multi_recv()
+    test_multi_recv_0deg()
+    test_dynamic_addition()
+    test_send_twice_different_shape()
+    test_send_twice_different_msg()
+    test_send_twice_different_field()
+    test_recv_no_send()
+    test_send_recv_after_conversion()


Code review comments:
diff_hunk: @@ -188,10 +188,14 @@ def __init__(self,
             self._edge_frame = FrameRef(Frame(num_rows=self.number_of_edges()))
         else:
             self._edge_frame = edge_frame
-        # msg graph & frame
-        self._msg_graph = create_graph_index(multigraph=multigraph)
-        self._msg_frame = FrameRef()
-        self.reset_messages()
+        # message indicator:
+        # if self._msg_index[eid] == 1, then edge eid has message
+        self._msg_index = utils.zero_index(size=self.number_of_edges())
+        # message frame
+        self._msg_frame = FrameRef(Frame(num_rows=self.number_of_edges()))
+        # set initializer for message frame
+        # all fields including message indicator should use zero_initializer
path: python/dgl/graph.py
text: @user1:
The indicator is not part of the frame anymore.

diff_hunk: @@ -126,9 +127,61 @@ def __getstate__(self):
     def __setstate__(self, state):
         self._initialize_data(state)
 
+    def get_items(self, index):
path: python/dgl/utils.py
text: @user1:
We need docstring for `get_items` and `set_items`. What is the meaning of `index`? I found you have three scenarios (1) advanced index, (2) boolean mask and (3) slice. If so, please make sure that the correct backend function is called (e.g. `gather_rows` and `narrow`). Boolean mask is currently missing in the backend. You could add it.

@author:
I don't think we have case (2).

There are two possible cases: advanced index and slice. I feel like tensors from most (all) backends support slicing. And for slice type, the step may not be zero, so I cannot use narrow.

@author:
But you are right, I am not sure whether this two functions should be complete. For now, I only want them to accept index of type utils.Index. But it could also accept native python types like int, slice, list

diff_hunk: @@ -126,9 +127,61 @@ def __getstate__(self):
     def __setstate__(self, state):
         self._initialize_data(state)
 
+    def get_items(self, index):
+        tensor = self.tousertensor()
+        if isinstance(index, Index):
+            if isinstance(index._pydata, slice):
+                # XXX: assuming tensor supports index using slices
+                #      cannot call tousertensor because for slice type it would
+                #      further call tonumpy, which overwrites _pydata and causes
+                #      future bugs...
path: python/dgl/utils.py
text: @user1:
What does this mean?

@author:
Originally when I implemented get_items and set_items, I write 
```
if isintance(index, Index):
    index = index.tousertensor()
```
And I just found this is buggy, because, if index._pydata is a slice, index.tousertensor will can tonumpy, which overwrites _pydata. As a result, index._pydata becomes a numpy.ndarray... And this can potentially cause bugs

diff_hunk: @@ -74,25 +75,37 @@ def schedule_recv(graph,
     inplace: bool
         If True, the update will be done in place
     """
-    src, dst, mid = graph._msg_graph.in_edges(recv_nodes)
-    if len(mid) == 0:
-        # All recv nodes are 0-degree nodes; downgrade to apply nodes.
+    _, _, eid = graph._graph.in_edges(recv_nodes)
+    if len(eid) > 0:
+        e_mask = graph._msg_index.get_items(eid).nonzero()
+        eid = eid.get_items(e_mask)
+    if len(eid) == 0:
+        # Downgrade to apply nodes if
+        #   1) all recv nodes are 0-degree nodes
+        #   2) no send has been called
         if apply_func is not None:
             schedule_apply_nodes(graph, recv_nodes, apply_func, inplace)
     else:
+        src, dst, _ = graph._graph.find_edges(eid)
path: python/dgl/runtime/scheduler.py
text: @user1:
lingfan: Mistakenly removed your comment...

@author:
You are right about using mask on src, dst to avoid find_edge

diff_hunk: @@ -366,6 +366,20 @@ def _append(self, other):
             # directly updating columns.
             self._columns = {key: Column.create(data) for key, data in other.items()}
         else:
+            # pad columns that no new were provided
path: python/dgl/frame.py
text: @user1:
pad columns that are not provided in the other frame with initial values

diff_hunk: @@ -645,3 +646,32 @@ def run(self):
 def APPEND_ROW_(fd1, fd2):
     reg = IR_REGISTRY[OpCode.APPEND_ROW_]
     get_current_prog().issue(reg['executor_cls'](fd1, fd2))
+
+class ClearFrame_Executor(Executor):
+    def __init__(self, fd):
+        self.fd = fd
+
+    def opcode(self):
+        return OpCode.CLEAR_FRAME_
+
+    def arg_vars(self):
+        return []
path: python/dgl/runtime/ir/executor.py
text: @user1:
`return [self.fd]`.

diff_hunk: @@ -74,25 +75,37 @@ def schedule_recv(graph,
     inplace: bool
         If True, the update will be done in place
     """
-    src, dst, mid = graph._msg_graph.in_edges(recv_nodes)
-    if len(mid) == 0:
-        # All recv nodes are 0-degree nodes; downgrade to apply nodes.
+    _, _, eid = graph._graph.in_edges(recv_nodes)
+    if len(eid) > 0:
+        e_mask = graph._msg_index.get_items(eid).nonzero()
+        eid = eid.get_items(e_mask)
path: python/dgl/runtime/scheduler.py
text: @user1:
I guess this does not work for MX. MX does not support using a bool mask to index the value.

@author:
Oh, it's not a boolean mask. Instead just nonzero index. Let me rename `e_mask` to `nonzero_idx`

diff_hunk: @@ -359,6 +365,10 @@ def add_edges(self, u, v, data=None):
             self._edge_frame.add_rows(max(len(u), len(v)))
         else:
             self._edge_frame.append(data)
+        # initialize feature placeholder for messages
+        num = max(len(u), len(v))
path: python/dgl/graph.py
text: @user1:
Put this `num` before the `if` statement so it could be used in the first branch.