fix(rdma): send periodic control messages to sync sender/receiver

We hit a bug where the sender sends a long run of eager messages and
ends up outpacing the receiver by more than the width of the message
buffer, causing an error.

In this patch, receiver will send a control message to sender at least
every `msgbuff_size - max_requests` messages, and sender will pause if
it hasn't received a control message within this duration.

Signed-off-by: Eric Raut <eraut@amazon.com>

Author: rauteric

include/nccl_ofi_rdma.h

@@ -535,6 +535,7 @@ typedef struct nccl_net_ofi_rdma_send_comm {
 	/* Counters for total sent and received control messages */
 	uint64_t n_ctrl_received;
 	uint64_t n_ctrl_expected;
+	uint16_t last_ctrl_received;
 
 	bool comm_active;
 
@@ -615,6 +616,8 @@ typedef struct nccl_net_ofi_rdma_recv_comm {
 	uint64_t n_ctrl_sent;
 	uint64_t n_ctrl_delivered;
 
+	uint16_t last_ctrl_sent;
+
 	/* Number of rails */
 	int num_rails;
 

src/nccl_ofi_rdma.c

@@ -32,6 +32,9 @@
 /* Message buffer size -- maximum span of simultaneous inflight messages */
 #define NCCL_OFI_RDMA_MSGBUFF_SIZE 256
 
+static_assert(NCCL_OFI_RDMA_MSGBUFF_SIZE > NCCL_OFI_MAX_REQUESTS,
+	"Message buffer size must be larger than max_requests");
+
 /* Maximum number of comms open simultaneously. Eventually this will be
    runtime-expandable */
 #define NCCL_OFI_RDMA_MAX_COMMS    (1 << NCCL_OFI_RDMA_COMM_ID_BITS)
@@ -91,6 +94,27 @@
 #define GET_RDMA_WRITE_IMM_DATA(comm_id, seq, nseg) \
 	((seq) | ((comm_id) << NCCL_OFI_RDMA_SEQ_BITS) | ((nseg) << (NCCL_OFI_RDMA_SEQ_BITS + NCCL_OFI_RDMA_COMM_ID_BITS)))
 
+static inline uint16_t msg_seq_num_distance(uint16_t right, uint16_t left)
+{
+	assert((right <= MSG_SEQ_NUM_MASK) && (left <= MSG_SEQ_NUM_MASK));
+	return (right - left) & MSG_SEQ_NUM_MASK;
+}
+
+static inline bool receiver_out_of_sync(uint16_t msg_seq_num, uint16_t recv_last_seq_num)
+{
+	/**
+	 * If receiver has started message n, then the last incomplete is no
+	 * less than n - max_requests + 1, and so the receiver's msgbuff has
+	 * space for at least message (last_incomplete + msgbuff_size - 1)
+	 *
+	 * Therefore, the condition for sender to wait for receiver is:
+	 *
+	 * distance(msg_seq_num, last_ctrl_received) == msgbuff_size - max_requests
+	 */
+	return msg_seq_num_distance(msg_seq_num, recv_last_seq_num) ==
+		(NCCL_OFI_RDMA_MSGBUFF_SIZE - NCCL_OFI_MAX_REQUESTS);
+}
+
 /** Global variables **/
 
 /* Maximum size of an eager message (see OFI_NCCL_EAGER_MAX_SIZE) */
@@ -1035,7 +1059,7 @@ static inline int handle_eager_recv(nccl_net_ofi_rdma_recv_comm_t *r_comm,
 	}
 
 	if (OFI_UNLIKELY(stat != NCCL_OFI_MSGBUFF_INPROGRESS)) {
-		NCCL_OFI_WARN("Unexpected message status (%d) (ctrl recv)", (int)stat);
+		NCCL_OFI_WARN("Unexpected message status (%d) for msg %hu", (int)stat, msg_seq_num);
 		return -EINVAL;
 	}
 
@@ -1213,6 +1237,11 @@ static inline int handle_bounce_recv(nccl_net_ofi_rdma_device_t *device, int rai
 
 		nccl_net_ofi_mutex_lock(&s_comm->ctrl_recv_lock);
 		s_comm->n_ctrl_received += 1;
+		if (msg_seq_num_distance((uint16_t)ctrl_msg->msg_seq_num,
+					 s_comm->last_ctrl_received)
+		    <= NCCL_OFI_RDMA_MSGBUFF_SIZE) {
+			s_comm->last_ctrl_received =  ctrl_msg->msg_seq_num;
+		}
 		nccl_net_ofi_mutex_unlock(&s_comm->ctrl_recv_lock);
 
 		break;
@@ -3286,6 +3315,8 @@ static int recv(nccl_net_ofi_recv_comm_t *recv_comm, int n, void **buffers,
 
 	NCCL_OFI_TRACE_RECV(dev_id, r_comm->local_comm_id, sizes[0], req, base_req);
 
+	bool send_ctrl = false;
+
 	if (eager) {
 		if (recv_data->eager_copy_req == NULL) {
 			/* If we don't need to do eager copy, this recv is already complete */
@@ -3302,7 +3333,20 @@ static int recv(nccl_net_ofi_recv_comm_t *recv_comm, int n, void **buffers,
 				goto error;
 			}
 		}
+
+		/* Send a ctrl message if we haven't sent one recently enough
+		   (see corresponding check in send) */
+		if (receiver_out_of_sync(msg_seq_num, r_comm->last_ctrl_sent)) {
+			send_ctrl = true;
+		} else {
+			send_ctrl = false;
+		}
 	} else {
+		send_ctrl = true;
+	}
+
+	if (send_ctrl)
+	{
 		nccl_net_ofi_rdma_req_t *send_ctrl_req =
 			allocate_send_ctrl_req(r_comm, device, dev_id, msg_seq_num,
 					       buffers[0], sizes[0], mr_handles[0]);
@@ -3322,6 +3366,8 @@ static int recv(nccl_net_ofi_recv_comm_t *recv_comm, int n, void **buffers,
 			/* TODO: Remove req from message buffer */
 			goto error;
 		}
+
+		r_comm->last_ctrl_sent = msg_seq_num;
 	}
 
 	/* Return request to NCCL */
@@ -4173,6 +4219,7 @@ static nccl_net_ofi_rdma_recv_comm_t *prepare_recv_comm(nccl_net_ofi_rdma_listen
 	memset(&r_comm->cleanup_list_elem, 0, sizeof(r_comm->cleanup_list_elem));
 	r_comm->n_ctrl_sent = 0;
 	r_comm->n_ctrl_delivered = 0;
+	r_comm->last_ctrl_sent = (uint16_t)((-1) & MSG_SEQ_NUM_MASK);
 
 	/* Allocate recv communicator ID */
 	comm_id = nccl_ofi_idpool_allocate_id(device->comm_idpool);
@@ -5389,6 +5436,10 @@ static int send(nccl_net_ofi_send_comm_t *send_comm, void *data, int size, int t
 		goto error;
 	}
 
+	nccl_net_ofi_mutex_lock(&s_comm->ctrl_recv_lock);
+	uint16_t last_ctrl_received = s_comm->last_ctrl_received;
+	nccl_net_ofi_mutex_unlock(&s_comm->ctrl_recv_lock);
+
 	dev_id = s_comm->base.base.dev_id;
 
 	ep = (nccl_net_ofi_rdma_ep_t *)s_comm->base.base.ep;
@@ -5405,6 +5456,12 @@ static int send(nccl_net_ofi_send_comm_t *send_comm, void *data, int size, int t
 		goto error;
 	}
 
+	if (receiver_out_of_sync(msg_seq_num, last_ctrl_received)) {
+		/* Wait for receiver to catch up */
+		ret = ofi_process_cq(ep);
+		goto free_req;
+	}
+
 	/*
 	 * TODO: Use NCCL provided tags when using grouped receives aka
 	 * props->maxRecvs > 1.
@@ -5940,6 +5997,7 @@ static inline int create_send_comm(nccl_net_ofi_conn_handle_t *handle,
 	ret_s_comm->received_close_message = false;
 	ret_s_comm->n_ctrl_received = 0;
 	ret_s_comm->n_ctrl_expected = 0;
+	ret_s_comm->last_ctrl_received = (uint16_t)((-1) & MSG_SEQ_NUM_MASK);
 
 	/* Store communicator ID from handle in communicator */
 	if (OFI_UNLIKELY(handle->comm_id >= device->num_comm_ids)) {