UCP/RMA/FLUSH: PR fixes 11.0

openucx · Mar 10, 2025 · 1cfc7d8 · 1cfc7d8
1 parent c28dd5a
commit 1cfc7d8
Showing 1 changed file with 6 additions and 4 deletions.
diff --git a/src/ucp/rma/put_offload.c b/src/ucp/rma/put_offload.c
@@ -27,10 +27,12 @@ static ucs_status_t ucp_proto_put_offload_short_progress(uct_pending_req_t *self
     ucs_status_t status;
     uct_rkey_t tl_rkey;
 
-    if (!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED)) {
-        ucp_proto_single_rma_init_func(req);
-        req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED;
-    }
+    /* Use a signed boolean mask (-!!) to update unflushed lanes branchlessly.
+     * Evaluates to 0xFFFFFFFFFFFFFFFF if initialized, 0 otherwise. */
+    req->send.ep->ext->unflushed_lanes |=
+            UCS_BIT(spriv->super.lane) &
+            -!!(req->flags & UCP_REQUEST_FLAG_PROTO_INITIALIZED);
+    req->flags |= UCP_REQUEST_FLAG_PROTO_INITIALIZED;
 
     tl_rkey = ucp_rkey_get_tl_rkey(req->send.rma.rkey, spriv->super.rkey_index);
     status  = uct_ep_put_short(ucp_ep_get_fast_lane(ep, spriv->super.lane),