UCP/RMA: Make put/rndv and get/rndv a fallback of zcopy

src/ucp/core/ucp_context.c

@@ -400,10 +400,6 @@ static ucs_config_field_t ucp_context_config_table[] = {
    "even if invalidation workflow isn't supported",
    ucs_offsetof(ucp_context_config_t, rndv_errh_ppln_enable), UCS_CONFIG_TYPE_BOOL},
 
-  {"RMA_PPLN_ENABLE", "n",
-   "Force-enable the RMA rendezvous put/get protocols.",
-   ucs_offsetof(ucp_context_config_t, rma_ppln_enable), UCS_CONFIG_TYPE_BOOL},
-
   {"FLUSH_WORKER_EPS", "y",
    "Enable flushing the worker by flushing its endpoints. Allows completing\n"
    "the flush operation in a bounded time even if there are new requests on\n"

src/ucp/core/ucp_context.h

@@ -130,8 +130,6 @@ typedef struct ucp_context_config {
     int                                    rndv_shm_cuda_staging_force;
     /** Enable error handling for rndv pipeline protocol */
     int                                    rndv_errh_ppln_enable;
-    /** Force-enable the RMA rendezvous put/get protocols */
-    int                                    rma_ppln_enable;
     /** Threshold for using tag matching offload capabilities. Smaller buffers
      *  will not be posted to the transport. */
     size_t                                 tm_thresh;

src/ucp/proto/AGENTS.md

@@ -62,6 +62,12 @@ should plug in here rather than going through the legacy direct paths.
 - When two protos tie, the earlier-listed (lower index in
   `ucp_protocols[]`, i.e. earlier in the `UCP_PROTO_FOR_EACH` list) wins;
   rely on this for deterministic test output.
+- `proto_class`/`fallback_class` (`ucp_proto_t`) override the cost model: a
+  proto is dropped from any range where a proto listing its class as a
+  fallback is available, even if cheaper. Keep the relation acyclic; to select
+  a fallback proto in a test, exclude the preferred proto with `UCX_PROTOS`. If
+  configuration disables every proto in a range, the fallback rule still
+  applies after they are re-enabled.
 - Stage callbacks must drain all pending `UCS_INPROGRESS` on completion
   before transitioning. The framework will call `progress[stage]` again
   on `UCS_ERR_NO_RESOURCE`.

src/ucp/proto/proto.h

@@ -85,6 +85,19 @@ enum {
 };
 
 
+/**
+ * Protocol classes, used to prioritize protocols which implement the same
+ * operation by a different data path. A class must not be a fallback, directly
+ * or indirectly, of a protocol which belongs to it. Classes are matched within
+ * a single selection key, which has one operation, so protocols of different
+ * operations can share the same class.
+ */
+enum {
+    /* RMA by a rendezvous protocol rather than a direct data path */
+    UCP_PROTO_CLASS_RMA_RNDV = UCS_BIT(0)
+};
+
+
 /**
  * Parameters for protocol initialization function
  */
@@ -202,6 +215,15 @@ struct ucp_proto {
     const char               *desc; /* Protocol description */
     unsigned                 flags; /* Protocol flags for special handling */
 
+    /* Bitmap of UCP_PROTO_CLASS_xxx classes this protocol belongs to */
+    unsigned                 proto_class;
+
+    /* Bitmap of UCP_PROTO_CLASS_xxx classes which are a fallback for this
+     * protocol. Protocols of these classes are not selected on message sizes
+     * where this protocol is available.
+     */
+    unsigned                 fallback_class;
+
     /* Bitmap of UCS_BIT(UCP_DATATYPE_xxx) classes this protocol supports.
      * Probe is skipped for any other dt_class. Must be non-zero.
      */

src/ucp/proto/proto_select.c

@@ -44,6 +44,69 @@ static const void *ucp_proto_select_init_priv_buf(
     return &ucs_array_elem(&proto_init->priv_buf, proto->priv_offset);
 }
 
+/* Print the protocols which are disabled on the current message size range */
+static void ucp_proto_select_trace_disabled(
+        const ucp_proto_select_init_protocols_t *proto_init,
+        const ucs_dynamic_bitmap_t *disabled_proto_mask)
+{
+    UCS_STRING_BUFFER_ONSTACK(strb, UCP_PROTO_CONFIG_STR_MAX);
+    const ucp_proto_init_elem_t *proto;
+    unsigned proto_idx;
+
+    if (!ucs_log_is_enabled(UCS_LOG_LEVEL_TRACE) ||
+        ucs_dynamic_bitmap_is_zero(disabled_proto_mask)) {
+        return;
+    }
+
+    UCS_DYNAMIC_BITMAP_FOR_EACH_BIT(proto_idx, disabled_proto_mask) {
+        proto = &ucs_array_elem(&proto_init->protocols, proto_idx);
+        ucs_string_buffer_appendf(&strb, "%s,",
+                                  ucp_proto_id_field(proto->proto_id, name));
+    }
+
+    ucs_string_buffer_rtrim(&strb, ",");
+    ucs_trace("disabled: %s", ucs_string_buffer_cstr(&strb));
+}
+
+/*
+ * Disables the protocols which are a fallback for a protocol available on the
+ * current message size range.
+ */
+static void ucp_proto_select_disable_fallback(
+        const ucp_proto_select_init_protocols_t *proto_init,
+        const ucs_dynamic_bitmap_t *proto_mask,
+        ucs_dynamic_bitmap_t *disabled_proto_mask)
+{
+    const ucp_proto_init_elem_t *proto;
+    unsigned overridden_classes, proto_class, fallback_class;
+    unsigned proto_idx;
+
+    overridden_classes = 0;
+    UCS_DYNAMIC_BITMAP_FOR_EACH_BIT(proto_idx, proto_mask) {
+        proto          = &ucs_array_elem(&proto_init->protocols, proto_idx);
+        proto_class    = ucp_proto_id_field(proto->proto_id, proto_class);
+        fallback_class = ucp_proto_id_field(proto->proto_id, fallback_class);
+        ucs_assertv((proto_class & fallback_class) == 0,
+                    "%s: proto_class 0x%x overlaps fallback_class 0x%x",
+                    ucp_proto_id_field(proto->proto_id, name), proto_class,
+                    fallback_class);
+
+        if (ucs_dynamic_bitmap_get(disabled_proto_mask, proto_idx)) {
+            continue;
+        }
+
+        overridden_classes |= fallback_class;
+    }
+
+    UCS_DYNAMIC_BITMAP_FOR_EACH_BIT(proto_idx, proto_mask) {
+        proto       = &ucs_array_elem(&proto_init->protocols, proto_idx);
+        proto_class = ucp_proto_id_field(proto->proto_id, proto_class);
+        if (proto_class & overridden_classes) {
+            ucs_dynamic_bitmap_set(disabled_proto_mask, proto_idx);
+        }
+    }
+}
+
 /*
  * Fills 'proto_mask' and 'perf_list' with candidate protocols for the next
  * range, and sets *max_length_p to the end of that range.
@@ -140,19 +203,24 @@ static ucs_status_t ucp_proto_thresholds_next_range(
                   proto->cfg_priority, max_prio_proto_name, max_cfg_priority);
     }
 
-    /* Remove disabled protocols. 'disabled_proto_mask' must be contained in
-     * 'valid_proto_mask'. */
+    /* If all protocols were disabled, we couldn't have any configured protocol
+     * (because that protocol would be enabled). In this case we allow using
+     * disabled protocols as well.
+     */
     if (ucs_dynamic_bitmap_is_equal(proto_mask, &disabled_proto_mask)) {
-        /* If all protocols were disabled, we couldn't have any configured
-         * protocol (because that protocol would be enabled). In this case we
-         * allow using disabled protocols as well.
-         */
         ucs_assert(max_cfg_priority == 0);
-    } else {
-        ucs_dynamic_bitmap_not_inplace(&disabled_proto_mask,
-                                       ucs_dynamic_bitmap_num_bits(proto_mask));
-        ucs_dynamic_bitmap_and_inplace(proto_mask, &disabled_proto_mask);
+        ucs_dynamic_bitmap_reset_all(&disabled_proto_mask);
     }
+
+    ucp_proto_select_disable_fallback(proto_init, proto_mask,
+                                      &disabled_proto_mask);
+    ucp_proto_select_trace_disabled(proto_init, &disabled_proto_mask);
+
+    /* Remove disabled protocols. 'disabled_proto_mask' is contained in
+     * 'proto_mask', and the fallback rule never disables all protocols. */
+    ucs_dynamic_bitmap_not_inplace(&disabled_proto_mask,
+                                   ucs_dynamic_bitmap_num_bits(proto_mask));
+    ucs_dynamic_bitmap_and_inplace(proto_mask, &disabled_proto_mask);
     ucs_assert(!ucs_dynamic_bitmap_is_zero(proto_mask));
 
     /* Add data to perf_list */

src/ucp/rma/get_offload.c

@@ -250,13 +250,14 @@ static ucs_status_t ucp_proto_get_offload_zcopy_reset(ucp_request_t *req)
 }
 
 ucp_proto_t ucp_get_offload_zcopy_proto = {
-    .name     = "get/zcopy",
-    .desc     = UCP_PROTO_ZCOPY_DESC,
-    .flags    = 0,
-    .dt_mask  = UCP_DT_MASK_CONTIG_IOV,
-    .probe    = ucp_proto_get_offload_zcopy_probe,
-    .query    = ucp_proto_multi_query,
-    .progress = {ucp_proto_get_offload_zcopy_progress},
-    .abort    = ucp_proto_request_zcopy_abort,
-    .reset    = ucp_proto_get_offload_zcopy_reset
+    .name           = "get/zcopy",
+    .desc           = UCP_PROTO_ZCOPY_DESC,
+    .flags          = 0,
+    .fallback_class = UCP_PROTO_CLASS_RMA_RNDV,
+    .dt_mask        = UCP_DT_MASK_CONTIG_IOV,
+    .probe          = ucp_proto_get_offload_zcopy_probe,
+    .query          = ucp_proto_multi_query,
+    .progress       = {ucp_proto_get_offload_zcopy_progress},
+    .abort          = ucp_proto_request_zcopy_abort,
+    .reset          = ucp_proto_get_offload_zcopy_reset
 };

src/ucp/rma/put_offload.c

@@ -346,15 +346,16 @@ ucp_proto_put_offload_zcopy_probe(const ucp_proto_init_params_t *init_params)
 }
 
 ucp_proto_t ucp_put_offload_zcopy_proto = {
-    .name     = "put/offload/zcopy",
-    .desc     = UCP_PROTO_ZCOPY_DESC,
-    .flags    = 0,
-    .dt_mask  = UCP_DT_MASK_CONTIG_IOV,
-    .probe    = ucp_proto_put_offload_zcopy_probe,
-    .query    = ucp_proto_multi_query,
-    .progress = {ucp_proto_put_offload_zcopy_progress},
-    .abort    = ucp_proto_request_zcopy_abort,
-    .reset    = ucp_proto_offload_zcopy_reset
+    .name           = "put/offload/zcopy",
+    .desc           = UCP_PROTO_ZCOPY_DESC,
+    .flags          = 0,
+    .fallback_class = UCP_PROTO_CLASS_RMA_RNDV,
+    .dt_mask        = UCP_DT_MASK_CONTIG_IOV,
+    .probe          = ucp_proto_put_offload_zcopy_probe,
+    .query          = ucp_proto_multi_query,
+    .progress       = {ucp_proto_put_offload_zcopy_progress},
+    .abort          = ucp_proto_request_zcopy_abort,
+    .reset          = ucp_proto_offload_zcopy_reset
 };
 
 static void

src/ucp/rma/rma_rndv.c

@@ -19,29 +19,18 @@
 #include <ucp/proto/proto_init.h>
 #include <ucp/proto/proto_single.inl>
 #include <ucp/rndv/proto_rndv.inl>
-#include <ucs/arch/cpu.h>
 
 
 #define UCP_PROTO_RMA_RNDV_RTS_NAME             "RMA_RTS"
 #define UCP_PROTO_RMA_RNDV_MIN_DST_VERSION      22
 #define UCP_PROTO_RMA_RNDV_ZERO_GET_PENALTY     1e-3
-#define UCP_PROTO_RMA_RNDV_PUT_FALLBACK_PENALTY 1e-3
-#define UCP_PROTO_RMA_RNDV_GET_FALLBACK_PENALTY 1e-3
 
 
 static int
 ucp_proto_rma_rndv_probe_check(const ucp_proto_init_params_t *init_params,
                                ucp_operation_id_t op_id)
 {
     const ucp_proto_select_param_t *sel_param = init_params->select_param;
-    const ucp_context_h context               = init_params->worker->context;
-
-    /* TODO: We prefer to use direct zcopy when possible, remove this check when
-     * prioritization of protocols is implemented. */
-    if (!context->config.ext.rma_ppln_enable &&
-        (ucs_arch_get_cpu_model() != UCS_CPU_MODEL_NVIDIA_VERA)) {
-        return 0;
-    }
 
     if (!ucp_proto_init_check_op(init_params, UCS_BIT(op_id)) ||
         ucp_proto_rndv_init_params_is_ppln_frag(init_params) ||
@@ -171,10 +160,7 @@ ucp_proto_put_rndv_probe(const ucp_proto_init_params_t *init_params)
     ucp_proto_rndv_ctrl_init_params_t params = {
         .super.super         = *init_params,
         .super.latency       = 0,
-        /* Prefer direct PUT zcopy when it is available; keep PUT/RNDV as a
-         * fallback for cases where the peer can only pull the data. */
-        .super.overhead      = context->config.ext.proto_overhead_rndv_rts +
-                               UCP_PROTO_RMA_RNDV_PUT_FALLBACK_PENALTY,
+        .super.overhead      = context->config.ext.proto_overhead_rndv_rts,
         .super.cfg_thresh    = context->config.ext.zcopy_thresh,
         .super.cfg_priority  = 5,
         .super.min_length    = 0,
@@ -260,14 +246,10 @@ ucp_proto_get_rndv_zero_length_variant(const ucp_proto_init_elem_t *proto)
 static double ucp_proto_get_rndv_variant_overhead(ucp_context_h context,
                                                   ucp_proto_init_elem_t *proto)
 {
-    /* Prefer direct GET zcopy when it is available; keep GET/RNDV as a
-     * fallback for cases where the peer can only send the data. */
-    double overhead = context->config.ext.proto_overhead_rndv_rtr +
-                      UCP_PROTO_RMA_RNDV_GET_FALLBACK_PENALTY;
+    double overhead = context->config.ext.proto_overhead_rndv_rtr;
 
+    /* Avoid RNDV for a zero-length GET when no zcopy protocol replaces it */
     if (ucp_proto_get_rndv_zero_length_variant(proto)) {
-        /* Keep zero-only RNDV receive variants available, but make direct GET
-         * protocols preferable for zero-length RMA GET. */
         overhead += UCP_PROTO_RMA_RNDV_ZERO_GET_PENALTY;
     }
 
@@ -635,25 +617,27 @@ ucs_status_t ucp_rma_rndv_process_rts(ucp_worker_h worker,
 }
 
 ucp_proto_t ucp_put_rndv_proto = {
-    .name     = "put/rndv",
-    .desc     = UCP_PROTO_RNDV_DESC,
-    .flags    = 0,
-    .dt_mask  = UCS_BIT(UCP_DATATYPE_CONTIG),
-    .probe    = ucp_proto_put_rndv_probe,
-    .query    = ucp_proto_rma_rndv_query,
-    .progress = {ucp_proto_put_rndv_progress},
-    .abort    = ucp_proto_rndv_rts_abort,
-    .reset    = ucp_proto_rndv_rts_reset
+    .name        = "put/rndv",
+    .desc        = UCP_PROTO_RNDV_DESC,
+    .flags       = 0,
+    .proto_class = UCP_PROTO_CLASS_RMA_RNDV,
+    .dt_mask     = UCS_BIT(UCP_DATATYPE_CONTIG),
+    .probe       = ucp_proto_put_rndv_probe,
+    .query       = ucp_proto_rma_rndv_query,
+    .progress    = {ucp_proto_put_rndv_progress},
+    .abort       = ucp_proto_rndv_rts_abort,
+    .reset       = ucp_proto_rndv_rts_reset
 };
 
 ucp_proto_t ucp_get_rndv_proto = {
-    .name     = "get/rndv",
-    .desc     = UCP_PROTO_RNDV_DESC,
-    .flags    = 0,
-    .dt_mask  = UCS_BIT(UCP_DATATYPE_CONTIG),
-    .probe    = ucp_proto_get_rndv_probe,
-    .query    = ucp_proto_rma_rndv_query,
-    .progress = {ucp_proto_get_rndv_progress},
-    .abort    = ucp_proto_get_rndv_abort,
-    .reset    = ucp_proto_get_rndv_reset
+    .name        = "get/rndv",
+    .desc        = UCP_PROTO_RNDV_DESC,
+    .flags       = 0,
+    .proto_class = UCP_PROTO_CLASS_RMA_RNDV,
+    .dt_mask     = UCS_BIT(UCP_DATATYPE_CONTIG),
+    .probe       = ucp_proto_get_rndv_probe,
+    .query       = ucp_proto_rma_rndv_query,
+    .progress    = {ucp_proto_get_rndv_progress},
+    .abort       = ucp_proto_get_rndv_abort,
+    .reset       = ucp_proto_get_rndv_reset
 };

src/ucs/datastruct/bitmap.h

@@ -105,13 +105,14 @@ size_t ucs_bitmap_bits_fns(const ucs_bitmap_word_t *bits, size_t num_words,
                            size_t start_index, size_t bit_count);
 
 
-/* Helper function to set all bitmap bits to a given value, avoiding a call to
- * memset() if the value is known to be 0, to workaround a compiler warning.
+/* Helper function to set all bitmap bits to a given value. An empty bitmap is
+ * skipped, since 'bits' may be NULL in that case, and memset() does not accept
+ * a NULL pointer even with a zero length.
  */
 static UCS_F_ALWAYS_INLINE void
 ucs_bitmap_bits_memset(ucs_bitmap_word_t *bits, int value, size_t num_words)
 {
-    if (__builtin_constant_p(num_words) && (num_words == 0)) {
+    if (num_words == 0) {
         return;
     }
 

test/gtest/ucp/test_ucp_proto.cc

@@ -999,9 +999,11 @@ UCS_TEST_P(test_ucp_proto_cuda_async_non_reg, cuda_async_registrable_filter)
     ucp_datatype_iter_mem_dereg(&dt_iter, UCP_DT_MASK_ALL);
 }
 
+/* Remove the GET zcopy protocol, which replaces GET/RNDV on registrable
+ * memory, so that GET/RNDV is always selected */
 UCS_TEST_P(test_ucp_proto_cuda_async_non_reg,
            cuda_async_rndv_get_zcopy_proto_filter, "RNDV_THRESH=0",
-           "RNDV_SCHEME=get_zcopy", "RMA_PPLN_ENABLE=y")
+           "RNDV_SCHEME=get_zcopy", "PROTOS=^get/zcopy")
 {
     /* Keep the real CUDA allocation small, but inspect a large protocol range
      * where RMA GET/RNDV is selected. */