| 36 | } |
| 37 | |
| 38 | Maybe<void> Send(const void* in, size_t elem_cnt, DataType dtype, int64_t dst, |
| 39 | DeviceType device_type, ep::Stream* stream) { |
| 40 | if (GlobalProcessCtx::Rank() == dst) { |
| 41 | auto** src_data_ptr = ThreadLocalSrcDataPtr(); |
| 42 | CHECK_OR_RETURN(*src_data_ptr == nullptr); |
| 43 | *src_data_ptr = in; |
| 44 | } else { |
| 45 | std::unique_ptr<ccl::Send> send = |
| 46 | ccl::NewCollectiveCommunication<ccl::Send>(device_type, dtype); |
| 47 | send->Launch(stream, in, elem_cnt, dst); |
| 48 | } |
| 49 | return Maybe<void>::Ok(); |
| 50 | } |
| 51 | |
| 52 | Maybe<void> Recv(void* out, size_t elem_cnt, DataType dtype, int64_t src, DeviceType device_type, |
| 53 | ep::Stream* stream) { |