UCP/RMA/FLUSH: Implement per-EP weak and strong fences

openucx · Feb 9, 2025 · f438a28 · f438a28
1 parent 1197f81
commit f438a28
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 0 deletions.
diff --git a/src/ucp/rma/flush.c b/src/ucp/rma/flush.c
@@ -748,3 +748,34 @@ UCS_PROFILE_FUNC(ucs_status_t, ucp_worker_fence, (worker), ucp_worker_h worker)
     UCP_WORKER_THREAD_CS_EXIT_CONDITIONAL(worker);
     return status;
 }
+
+ucs_status_t ucp_ep_fence_weak(ucp_ep_h ep)
+{
+    ucs_status_t status;
+    ucp_lane_index_t lane;
+
+    ucs_for_each_bit(lane, ep->ext->flush_state.unflushed_lanes) {
+        status = uct_ep_fence(ucp_ep_get_lane(ep, lane), 0);
+        if (status != UCS_OK) {
+            return status;
+        }
+    }
+
+    return UCS_OK;
+}
+
+ucs_status_t ucp_ep_fence_strong(ucp_ep_h ep)
+{
+    ucs_status_t status;
+    void *request;
+
+    request = ucp_ep_flush_internal(ep, 0, &ucp_request_null_param, NULL,
+                                    ucp_ep_flushed_callback, "ep_fence_strong");
+    status = ucp_flush_wait(ep->worker, request);
+    if (status != UCS_OK) {
+        return status;
+    }
+
+    ucp_ep_mark_flushed(ep);
+    return UCS_OK;
+}
diff --git a/src/ucp/rma/rma.h b/src/ucp/rma/rma.h
@@ -109,4 +109,8 @@ void ucp_ep_flush_remote_completed(ucp_request_t *req);
 
 void ucp_rma_sw_send_cmpl(ucp_ep_h ep);
 
+ucs_status_t ucp_ep_fence_weak(ucp_ep_h ep);
+
+ucs_status_t ucp_ep_fence_strong(ucp_ep_h ep);
+
 #endif