prov/efa: Implement FI_CONTEXT2 in EFA Direct

Store the completion flags and peer address in FI_CONTEXT2 and
retrieve later when writing cq.

Signed-off-by: Jessie Yang <[email protected]>

Author: jiaxiyan

prov/efa/src/efa.h

@@ -107,6 +107,16 @@ struct efa_fabric {
 #endif
 };
 
+struct efa_context {
+	uint64_t completion_flags;
+	fi_addr_t addr;
+};
+
+#if defined(static_assert) && defined(__x86_64__)
+static_assert(sizeof(struct efa_context) <= sizeof(struct fi_context2),
+	      "efa_context must not be larger than fi_context2");
+#endif
+
 static inline
 int efa_str_to_ep_addr(const char *node, const char *service, struct efa_ep_addr *addr)
 {

prov/efa/src/efa_cq.c

@@ -36,7 +36,10 @@ static void efa_cq_construct_cq_entry(struct ibv_cq_ex *ibv_cqx,
 				      struct fi_cq_tagged_entry *entry)
 {
 	entry->op_context = (void *)ibv_cqx->wr_id;
-	entry->flags = efa_cq_opcode_to_fi_flags(ibv_wc_read_opcode(ibv_cqx));
+	if (!ibv_cq_ex->wr_id)
+		entry->flags = efa_cq_opcode_to_fi_flags(ibv_wc_read_opcode(ibv_cqx));
+	else
+		entry->flags = ((struct efa_context *) ibv_cqx->wr_id)->completion_flags;
 	entry->len = ibv_wc_read_byte_len(ibv_cqx);
 	entry->buf = NULL;
 	entry->data = 0;
@@ -81,8 +84,7 @@ static void efa_cq_handle_error(struct efa_base_ep *base_ep,
 	err_entry.prov_errno = prov_errno;
 
 	if (is_tx)
-		// TODO: get correct peer addr for TX operation
-		addr = FI_ADDR_NOTAVAIL;
+		addr = ((struct efa_context *)ibv_cqx->wr_id)->addr;
 	else
 		addr = efa_av_reverse_lookup(base_ep->av,
 					     ibv_wc_read_slid(ibv_cq_ex),

prov/efa/src/efa_msg.c

@@ -101,7 +101,11 @@ static inline ssize_t efa_post_recv(struct efa_base_ep *base_ep, const struct fi
 	wr = &base_ep->efa_recv_wr_vec[wr_index].wr;
 	wr->num_sge = msg->iov_count;
 	wr->sg_list = base_ep->efa_recv_wr_vec[wr_index].sge;
-	wr->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? msg->context : NULL);
+
+	struct efa_context *efa_context = (struct efa_context *) msg->context;
+	efa_context->completion_flags = FI_RECV | FI_MSG;
+	efa_context->addr = msg->addr;
+	wr->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? efa_context : NULL);
 
 	for (i = 0; i < msg->iov_count; i++) {
 		addr = (uintptr_t)msg->msg_iov[i].iov_base;
@@ -224,7 +228,10 @@ static inline ssize_t efa_post_send(struct efa_base_ep *base_ep, const struct fi
 		base_ep->is_wr_started = true;
 	}
 
-	qp->ibv_qp_ex->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? msg->context : NULL);
+	struct efa_context *efa_context = (struct efa_context *) msg->context;
+	efa_context->completion_flags = FI_SEND | FI_MSG;
+	efa_context->addr = msg->addr;
+	qp->ibv_qp_ex->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? efa_context : NULL);
 
 	if (flags & FI_REMOTE_CQ_DATA) {
 		ibv_wr_send_imm(qp->ibv_qp_ex, msg->data);

prov/efa/src/efa_rma.c

@@ -90,7 +90,11 @@ static inline ssize_t efa_rma_post_read(struct efa_base_ep *base_ep,
 		ibv_wr_start(qp->ibv_qp_ex);
 		base_ep->is_wr_started = true;
 	}
-	qp->ibv_qp_ex->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? msg->context : NULL);
+
+	struct efa_context *efa_context = (struct efa_context *) msg->context;
+	efa_context->completion_flags = FI_RMA | FI_READ;
+	efa_context->addr = msg->addr;
+	qp->ibv_qp_ex->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? efa_context : NULL);
 
 	/* ep->domain->info->tx_attr->rma_iov_limit is set to 1 */
 	ibv_wr_rdma_read(qp->ibv_qp_ex, msg->rma_iov[0].key, msg->rma_iov[0].addr);
@@ -225,7 +229,14 @@ static inline ssize_t efa_rma_post_write(struct efa_base_ep *base_ep,
 		ibv_wr_start(qp->ibv_qp_ex);
 		base_ep->is_wr_started = true;
 	}
-	qp->ibv_qp_ex->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? msg->context : NULL);
+
+	struct efa_context *efa_context = (struct efa_context *) msg->context;
+	efa_context->completion_flags =
+		flags & FI_REMOTE_CQ_DATA ?
+			FI_REMOTE_CQ_DATA | FI_RMA | FI_REMOTE_WRITE :
+			FI_RMA | FI_WRITE;
+	efa_context->addr = msg->addr;
+	qp->ibv_qp_ex->wr_id = (uintptr_t) ((flags & FI_COMPLETION) ? efa_context : NULL);
 
 	if (flags & FI_REMOTE_CQ_DATA) {
 		ibv_wr_rdma_write_imm(qp->ibv_qp_ex, msg->rma_iov[0].key,