UCP/RNDV: Skip PUT fence for ordered transports

src/ucp/core/ucp_context.c

@@ -483,8 +483,8 @@ static ucs_config_field_t ucp_context_config_table[] = {
    "When using rendezvous put protocol, force using a flush operation to ensure\n"
    "remote data delivery before sending ATP message.\n"
    "If flush mode is not forced, and the underlying transport supports both active\n"
-   "messages and put operations, the protocol will do {put,fence,ATP} on the same\n"
-   "lane without waiting for remote completion.",
+   "messages and put operations, the protocol sends ATP on the same lanes. A fence\n"
+   "is used unless the transport guarantees that ATP is ordered after prior puts.",
    ucs_offsetof(ucp_context_config_t, rndv_put_force_flush), UCS_CONFIG_TYPE_BOOL},
 
   {"PROTO_EMULATION_ENABLE", "y",

src/ucp/rndv/rndv_put.c

@@ -277,10 +277,10 @@ ucp_proto_rndv_put_common_probe(const ucp_proto_init_params_t *init_params,
     const uct_iface_attr_t *iface_attr;
     ucp_lane_index_t lane_idx, lane;
     ucp_proto_rndv_put_priv_t rpriv;
-    int send_atp, use_fence;
+    int send_atp, use_atp_lanes;
     ucp_proto_perf_t *perf;
     ucs_status_t status;
-    unsigned atp_map;
+    ucp_lane_map_t atp_map, ordered_atp_map;
 
     if (!ucp_proto_rndv_op_check(init_params, UCP_OP_ID_RNDV_SEND,
                                  support_ppln) ||
@@ -298,7 +298,8 @@ ucp_proto_rndv_put_common_probe(const ucp_proto_init_params_t *init_params,
     send_atp = !ucp_proto_rndv_init_params_is_ppln_frag(init_params);
 
     /* Check which lanes support sending ATP */
-    atp_map = 0;
+    atp_map         = 0;
+    ordered_atp_map = 0;
     for (lane_idx = 0; lane_idx < rpriv.bulk.mpriv.num_lanes; ++lane_idx) {
         lane       = rpriv.bulk.mpriv.lanes[lane_idx].super.lane;
         iface_attr = ucp_proto_common_get_iface_attr(init_params, lane);
@@ -307,13 +308,16 @@ ucp_proto_rndv_put_common_probe(const ucp_proto_init_params_t *init_params,
             ((iface_attr->cap.flags & UCT_IFACE_FLAG_AM_BCOPY) &&
              (iface_attr->cap.am.max_bcopy >= atp_size))) {
             atp_map |= UCS_BIT(lane);
+            if (iface_attr->cap.flags & UCT_IFACE_FLAG_PUT_AM_ORDER) {
+                ordered_atp_map |= UCS_BIT(lane);
+            }
         }
     }
 
-    /* Use fence only if all lanes support sending ATP and flush is not forced
+    /* Use data lanes if all lanes support sending ATP and flush is not forced
      */
-    use_fence = send_atp && !context->config.ext.rndv_put_force_flush &&
-                (rpriv.bulk.mpriv.lane_map == atp_map);
+    use_atp_lanes = send_atp && !context->config.ext.rndv_put_force_flush &&
+                    (rpriv.bulk.mpriv.lane_map == atp_map);
 
     /* All lanes can send ATP - invalidate am_lane, to use mpriv->lanes.
      * Otherwise, would need to flush all lanes and send ATP on:
@@ -326,12 +330,15 @@ ucp_proto_rndv_put_common_probe(const ucp_proto_init_params_t *init_params,
      *   size is not an issue anymore.
      * - Control lane if none of the lanes support sending ATP
      */
-    if (use_fence) {
-        /* Send fence followed by ATP on all lanes */
+    if (use_atp_lanes) {
+        /* Send ATP on all lanes, using a fence if ordering is not guaranteed */
         rpriv.bulk.super.lane = UCP_NULL_LANE;
         rpriv.put_comp_cb     = comp_cb;
         rpriv.atp_comp_cb     = NULL;
-        rpriv.stage_after_put = UCP_PROTO_RNDV_PUT_STAGE_FENCED_ATP;
+        rpriv.stage_after_put =
+                (rpriv.bulk.mpriv.lane_map == ordered_atp_map) ?
+                UCP_PROTO_RNDV_PUT_STAGE_ATP :
+                UCP_PROTO_RNDV_PUT_STAGE_FENCED_ATP;
         rpriv.flush_map       = 0;
         rpriv.atp_map         = rpriv.bulk.mpriv.lane_map;
     } else {
@@ -385,8 +392,11 @@ ucp_proto_rndv_put_common_query(const ucp_proto_query_params_t *params,
         return UCP_PROTO_RNDV_PUT_DESC;
     } else if (rpriv->flush_map != 0) {
         return "flushed " UCP_PROTO_RNDV_PUT_DESC;
-    } else {
+    } else if (rpriv->stage_after_put ==
+               UCP_PROTO_RNDV_PUT_STAGE_FENCED_ATP) {
         return "fenced " UCP_PROTO_RNDV_PUT_DESC;
+    } else {
+        return UCP_PROTO_RNDV_PUT_DESC;
     }
 }
 

src/uct/api/uct.h

@@ -429,6 +429,10 @@ typedef enum uct_atomic_op {
                                                        channel with remote peer is broken, even if there
                                                        are no outstanding send operations */
 
+        /* Operation ordering */
+#define UCT_IFACE_FLAG_PUT_AM_ORDER    UCS_BIT(47) /**< Active messages are delivered only after
+                                                       prior PUT operations on the same endpoint */
+
         /* Tag matching operations */
 #define UCT_IFACE_FLAG_TAG_EAGER_SHORT UCS_BIT(50) /**< Hardware tag matching short eager support */
 #define UCT_IFACE_FLAG_TAG_EAGER_BCOPY UCS_BIT(51) /**< Hardware tag matching bcopy eager support */

src/uct/base/uct_iface_vfs.c

@@ -25,6 +25,7 @@ static const uct_vfs_flag_info_t uct_iface_vfs_cap_infos[] = {
     {UCT_IFACE_FLAG_PUT_SHORT, "put_short"},
     {UCT_IFACE_FLAG_PUT_BCOPY, "put_bcopy"},
     {UCT_IFACE_FLAG_PUT_ZCOPY, "put_zcopy"},
+    {UCT_IFACE_FLAG_PUT_AM_ORDER, "put_am_order"},
     {UCT_IFACE_FLAG_GET_SHORT, "get_short"},
     {UCT_IFACE_FLAG_GET_BCOPY, "get_bcopy"},
     {UCT_IFACE_FLAG_GET_ZCOPY, "get_zcopy"},

src/uct/ib/rc/base/rc_iface.c

@@ -186,6 +186,7 @@ ucs_status_t uct_rc_iface_query(uct_rc_iface_t *iface,
                                   UCT_IFACE_FLAG_PUT_ZCOPY       |
                                   UCT_IFACE_FLAG_GET_BCOPY       |
                                   UCT_IFACE_FLAG_GET_ZCOPY       |
+                                  UCT_IFACE_FLAG_PUT_AM_ORDER    |
                                   UCT_IFACE_FLAG_PENDING         |
                                   UCT_IFACE_FLAG_CONNECT_TO_EP   |
                                   UCT_IFACE_FLAG_CB_SYNC         |

test/gtest/ucp/test_ucp_proto_mock.cc

@@ -765,7 +765,7 @@ UCS_TEST_P(test_ucp_proto_mock_rcx, memtype_copy_enable,
 
     check_ep_config(sender(), {
         {0,       0, "rendezvous no data fetch", ""},
-        {1,      64, "rendezvous zero-copy fenced write to remote",
+        {1,      64, "rendezvous zero-copy write to remote",
                      "rc_mlx5/mock_0:1"},
         {21992, INF, "rendezvous zero-copy read from remote",
                      "rc_mlx5/mock_0:1"},
@@ -846,7 +846,7 @@ UCS_TEST_P(test_ucp_proto_mock_rcx, rndv_4_paths,
         {3814,    283699, "rendezvous zero-copy read from remote",
          "12% on rc_mlx5/mock_1:1/path0, 14% on rc_mlx5/mock_0:1/path0, "
          "14% on rc_mlx5/mock_0:1/path1, 12% on rc_mlx5/mock_1:1/path1, 14%"},
-        {283700,  INF,    "rendezvous zero-copy fenced write to remote",
+        {283700,  INF,    "rendezvous zero-copy write to remote",
          "12% on rc_mlx5/mock_1:1/path0, 14% on rc_mlx5/mock_0:1/path0, "
          "14% on rc_mlx5/mock_0:1/path1, 12% on rc_mlx5/mock_1:1/path1, 14%"},
     }, key);
@@ -869,6 +869,49 @@ UCS_TEST_P(test_ucp_proto_mock_rcx, rma_put_2_lanes,
 
 UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_proto_mock_rcx, rcx, "rc_x")
 
+class test_ucp_proto_mock_rcx_unordered : public test_ucp_proto_mock_rcx {
+public:
+    virtual void init() override
+    {
+        /* Device with higher BW and latency */
+        add_mock_iface("mock_0:1", [](uct_iface_attr_t &iface_attr) {
+            iface_attr.cap.am.max_short  = 2000;
+            iface_attr.cap.put.max_short = 2048;
+            iface_attr.bandwidth.shared  = 28e9;
+            iface_attr.latency.c         = 600e-9;
+            iface_attr.latency.m         = 1e-9;
+            iface_attr.cap.get.max_zcopy = 16384;
+        });
+        /* Device with smaller BW but lower latency and unordered PUT-to-AM */
+        add_mock_iface("mock_1:1", [](uct_iface_attr_t &iface_attr) {
+            iface_attr.cap.am.max_short  = 208;
+            iface_attr.cap.put.max_short = 2048;
+            iface_attr.bandwidth.shared  = 24e9;
+            iface_attr.latency.c         = 500e-9;
+            iface_attr.latency.m         = 1e-9;
+            iface_attr.cap.flags &= ~UCT_IFACE_FLAG_PUT_AM_ORDER;
+        });
+        test_ucp_proto_mock::init();
+    }
+};
+
+UCS_TEST_P(test_ucp_proto_mock_rcx_unordered, rndv_put_mixed_ordering,
+           "IB_NUM_PATHS?=1", "MAX_RNDV_LANES=2", "RNDV_THRESH=0",
+           "RNDV_SCHEME=put_zcopy")
+{
+    ucp_proto_select_key_t key = any_key();
+    key.param.op_id_flags      = UCP_OP_ID_AM_SEND;
+    key.param.op_attr          = 0;
+
+    check_ep_config(sender(), {
+        {1, INF, "rendezvous zero-copy fenced write to remote",
+         "47% on rc_mlx5/mock_1:1 and 53% on rc_mlx5/mock_0:1"},
+    }, key);
+}
+
+UCP_INSTANTIATE_TEST_CASE_TLS(test_ucp_proto_mock_rcx_unordered,
+                              rcx_unordered, "rc_x")
+
 class test_ucp_proto_mock_rcx2 : public test_ucp_proto_mock {
 public:
     test_ucp_proto_mock_rcx2()
@@ -956,7 +999,7 @@ UCS_TEST_P(test_ucp_proto_mock_rcx3, single_lane_no_zcopy,
     check_ep_config(sender(), {
         {1,    94,    "rendezvous fragmented copy-in copy-out", "rc_mlx5/mock_0:1"},
         {95,   53753, "rendezvous zero-copy read from remote",  "rc_mlx5/mock_1:1"},
-        {53754, INF,  "rendezvous zero-copy fenced write to remote",
+        {53754, INF,  "rendezvous zero-copy write to remote",
          "54% on rc_mlx5/mock_0:1 and 46% on rc_mlx5/mock_1:1"},
     }, key);
 }
@@ -1269,10 +1312,10 @@ UCS_TEST_P(test_ucp_proto_mock_gpu, cuda_managed_ppln_host_frag,
         {0, 0,    "short",   "rc_mlx5/mock"},
         {1, 8246, "copy-in", "rc_mlx5/mock"},
         {8247, 512 * UCS_KBYTE,
-         "rendezvous cuda_copy, fenced write to remote, frag host, cuda_copy, frag host",
+         "rendezvous cuda_copy, write to remote, frag host, cuda_copy, frag host",
          "rc_mlx5/mock"},
         {(512 * UCS_KBYTE) + 1, INF,
-         "rendezvous pipeline cuda_copy, fenced write to remote, frag host, cuda_copy, frag host",
+         "rendezvous pipeline cuda_copy, write to remote, frag host, cuda_copy, frag host",
          "rc_mlx5/mock"},
         }, key);
 }

test/gtest/ucp/test_ucp_request.cc

@@ -737,8 +737,10 @@ class test_proto_reset_atp : public test_proto_reset {
 
             ucp_request_t *req = (ucp_request_t*)ureq - 1;
             return (req->flags & UCP_REQUEST_FLAG_PROTO_SEND) &&
-                   (req->send.proto_stage ==
-                    UCP_PROTO_RNDV_PUT_STAGE_FENCED_ATP);
+                   ((req->send.proto_stage ==
+                     UCP_PROTO_RNDV_PUT_STAGE_ATP) ||
+                    (req->send.proto_stage ==
+                     UCP_PROTO_RNDV_PUT_STAGE_FENCED_ATP));
         });
 
 
@@ -775,5 +777,6 @@ UCS_TEST_P(test_proto_reset_atp, rndv_put, "RNDV_THRESH=0",
 }
 
 UCP_INSTANTIATE_TEST_CASE_TLS(test_proto_reset_atp, ib, "ib")
+UCP_INSTANTIATE_TEST_CASE_TLS(test_proto_reset_atp, srd, "srd")
 
 #endif