Use separate streams for pushsend/recv kernels in UB p2p exchanges

[erhoo82]

This fix uses separate streams that overlap these pooling kernels.
This can resolve potential deadlock involved with send/recv orders and also improves performance.

[ksivaman]

LGTM

[timmoon10]

I see that _stream_send and _stream_compute[:] wait on _stream_recv and the main stream waits on _stream_compute[last_compute_stream_id]. I don't see if the main stream is ever synchronized with _stream_send and _stream_compute[:] though. Could we run into race conditions?

[erhoo82]

I think it is fine. We need to make sure the dependency of main_stream on the stream that executes the last kernel in the module, which is _stream_compute[last_compute_stream_id].

The compute kernel on _stream_compute[last_compute_stream_id] depends on the preceding kernel on _stream_recv and _stream_send executes memcopy for backdrop at last, which doesn't have any dependency.

@timmoon10 Let me know this makes sense to you.

[timmoon10]

Ok, it makes sense to me that _stream_compute is correctly synchronized. _stream_compute[i] waits on _stream_compute[(i-1)%_steam_compute.size()], so they are effectively one stream. I'm still worried about _stream_send since users of B_copy (on the main stream) might have incorrect values if _stream_send is very slow doing the cudaMemcpyAsync.

[ksivaman]

/te-ci

[erhoo82]

The memcopy is for storing activation for backprop thus this is only valid execution path in fprop kernel. Since the following layers do not have direct dependency to this data, extended execution of memcopy is fine and this is intended for potential overlap with following kernels.

[timmoon10]

I see, so we have a race condition but we don't expect it to cause problems in practice since the data is used much later during bprop. I see the benefit for inter-layer overlapping, but it feels delicate. Could we record an event on _stream_send and synchronize when B_copy is consumed? Or alternatively use a dedicated stream for _stream_send for all layers and synchronize it at the end of fprop/beginning of bprop?

[erhoo82]

Let me actually add a stream wait when using B_copy.

[timmoon10]

LGTM

[ksivaman]

/te-ci