prov/shm: fix RMA read and write bandwidth regressions

Read fix: extend smr_do_fast_rma to use sender-side CMA for RMA reads
at all sizes (not just > SMR_INJECT_SIZE). Delivery is inherently
complete when process_vm_readv returns. This avoids the receiver-side
CMA round-trip that regresses 20-25% on Graviton and AMD at 1-4096B.

Write fix: allow FI_REMOTE_CQ_DATA through smr_rma_fast so writedata
uses sender-side CMA instead of the receiver-side IOV path with its
expensive atomic return queue. The sender does process_vm_writev
directly, then posts ofi_op_write_async with cq_data. The receiver
generates the remote CQ entry on seeing SMR_REMOTE_CQ_DATA in the
write_async notification.

Both target buffers are always registered MRs with pinned pages, making
sender-side CMA safe for reads and writes.

Signed-off-by: Yin Li <yinliq@amazon.com>

Author: yinliaws

prov/shm/src/smr_progress.c

@@ -1334,6 +1334,16 @@ static void smr_progress_cmd(struct smr_ep *ep)
 			ret = smr_progress_cmd_rma(ep, cmd);
 			break;
 		case ofi_op_write_async:
+			if (cmd->hdr.smr_flags & SMR_REMOTE_CQ_DATA) {
+				smr_complete_rx(ep, NULL, ofi_op_write,
+					smr_rx_cq_flags(0, cmd->hdr.smr_flags),
+					cmd->hdr.size, NULL,
+					cmd->hdr.rx_id, 0, cmd->hdr.cq_data);
+			} else {
+				ofi_ep_peer_rx_cntr_inc(&ep->util_ep,
+							cmd->hdr.op);
+			}
+			break;
 		case ofi_op_read_async:
 			ofi_ep_peer_rx_cntr_inc(&ep->util_ep, cmd->hdr.op);
 			break;

prov/shm/src/smr_rma.c

@@ -44,17 +44,20 @@ static void smr_add_rma_cmd(struct smr_region *peer_smr,
 static void smr_format_rma_resp(struct smr_cmd *cmd, int64_t peer_id,
 				const struct fi_rma_iov *rma_iov, size_t count,
 				size_t total_len, uint32_t op,
-				uint64_t op_flags)
+				uint64_t op_flags, uint64_t data)
 {
-	smr_generic_format(cmd, 0, peer_id, op, 0, 0, op_flags);
+	uint8_t smr_flags = 0;
+	if (op_flags & FI_REMOTE_CQ_DATA)
+		smr_flags |= SMR_REMOTE_CQ_DATA;
+	smr_generic_format(cmd, 0, peer_id, op, 0, data, smr_flags);
 	cmd->hdr.size = total_len;
 }
 
 static ssize_t smr_rma_fast(struct smr_ep *ep, struct smr_region *peer_smr,
 			    const struct iovec *iov, size_t iov_count,
 			    const struct fi_rma_iov *rma_iov, size_t rma_count,
 			    void **desc, int rx_id, int tx_id, void *context,
-			    uint32_t op, uint64_t op_flags)
+			    uint32_t op, uint64_t op_flags, uint64_t data)
 {
 	struct iovec vma_iovec[SMR_IOV_LIMIT], rma_iovec[SMR_IOV_LIMIT];
 	struct ofi_xpmem_client *xpmem;
@@ -91,7 +94,7 @@ static ssize_t smr_rma_fast(struct smr_ep *ep, struct smr_region *peer_smr,
 
 	smr_format_rma_resp(&ce->cmd, rx_id, rma_iov, rma_count, total_len,
 			    (op == ofi_op_write) ? ofi_op_write_async :
-			    ofi_op_read_async, op_flags);
+			    ofi_op_read_async, op_flags, data);
 
 	smr_cmd_queue_commit(ce, pos);
 
@@ -106,15 +109,24 @@ static ssize_t smr_rma_fast(struct smr_ep *ep, struct smr_region *peer_smr,
 
 static inline bool smr_do_fast_rma(struct smr_ep *ep, uint64_t op_flags,
 				   size_t rma_count, size_t total_len,
-				   struct smr_region *peer_smr)
+				   struct smr_region *peer_smr, uint32_t op)
 {
 	struct smr_domain *domain;
 
 	domain = container_of(ep->util_ep.domain, struct smr_domain,
 			      util_domain);
 
-	return domain->fast_rma && !(op_flags &
-		    (FI_REMOTE_CQ_DATA | FI_DELIVERY_COMPLETE)) &&
+	/* For reads, sender-side CMA is safe at all sizes: delivery is
+	 * inherently complete when process_vm_readv returns (data is in
+	 * local buffer), and the target memory is always a registered MR
+	 * with pinned pages. */
+	if (op == ofi_op_read_req && total_len <= SMR_INJECT_SIZE)
+		return domain->fast_rma &&
+		       !(op_flags & FI_REMOTE_CQ_DATA) &&
+		       rma_count == 1 && smr_vma_enabled(ep, peer_smr);
+
+	return domain->fast_rma &&
+		     !(op_flags & FI_DELIVERY_COMPLETE) &&
 		     rma_count == 1 && smr_vma_enabled(ep, peer_smr) &&
 		     total_len > SMR_INJECT_SIZE;
 
@@ -153,10 +165,10 @@ static ssize_t smr_generic_rma(
 		goto unlock;
 
 	total_len = ofi_total_iov_len(iov, iov_count);
-	if (smr_do_fast_rma(ep, op_flags, rma_count, total_len, peer_smr)) {
+	if (smr_do_fast_rma(ep, op_flags, rma_count, total_len, peer_smr, op)) {
 		ret = smr_rma_fast(ep, peer_smr, iov, iov_count, rma_iov,
 				   rma_count, desc, rx_id, tx_id, context, op,
-				   op_flags);
+				   op_flags, data);
 		goto unlock;
 	}