UCP/DEVICE: Make memh and local_addr optional for counter elements (#10945)

Author: michal-shalev

src/tools/perf/cuda/ucp_cuda_kernel.cu

@@ -109,28 +109,41 @@ public:
     const ucp_perf_cuda_params &get_params() const { return m_params; }
 
 private:
+    static bool has_counter(const ucx_perf_context_t &perf)
+    {
+        return (perf.params.command != UCX_PERF_CMD_PUT_SINGLE);
+    }
+
     void init_mem_list(const ucx_perf_context_t &perf)
     {
-        /* +1 for the counter */
-        size_t count  = perf.params.msg_size_cnt + 1;
-        size_t offset = 0;
+        size_t data_count = perf.params.msg_size_cnt;
+        size_t count      = data_count + (has_counter(perf) ? 1 : 0);
+        size_t offset     = 0;
         ucp_device_mem_list_elem_t elems[count];
 
-        for (size_t i = 0; i < count; ++i) {
-            elems[i].field_mask = UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH |
-                                  UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY |
-                                  UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR |
-                                  UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR |
-                                  UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH;
-            elems[i].memh       = perf.ucp.send_memh;
-            elems[i].rkey       = perf.ucp.rkey;
-            elems[i].local_addr = UCS_PTR_BYTE_OFFSET(perf.send_buffer, offset);
+        for (size_t i = 0; i < data_count; ++i) {
+            elems[i].field_mask  = UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH |
+                                   UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY |
+                                   UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR |
+                                   UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR |
+                                   UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH;
+            elems[i].memh        = perf.ucp.send_memh;
+            elems[i].rkey        = perf.ucp.rkey;
+            elems[i].local_addr  = UCS_PTR_BYTE_OFFSET(perf.send_buffer, offset);
             elems[i].remote_addr = perf.ucp.remote_addr + offset;
-            elems[i].length      = (i == count - 1) ? ONESIDED_SIGNAL_SIZE :
-                                                           perf.params.msg_size_list[i];
+            elems[i].length      = perf.params.msg_size_list[i];
             offset              += elems[i].length;
         }
 
+        if (has_counter(perf)) {
+            elems[data_count].field_mask  = UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY |
+                                            UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR |
+                                            UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH;
+            elems[data_count].rkey        = perf.ucp.rkey;
+            elems[data_count].remote_addr = perf.ucp.remote_addr + offset;
+            elems[data_count].length      = ONESIDED_SIGNAL_SIZE;
+        }
+
         ucp_device_mem_list_params_t params;
         params.field_mask   = UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENTS |
                               UCP_DEVICE_MEM_LIST_PARAMS_FIELD_ELEMENT_SIZE |
@@ -148,20 +161,22 @@ private:
 
     void init_elements(const ucx_perf_context_t &perf)
     {
-        /* +1 for the counter */
-        size_t count  = perf.params.msg_size_cnt + 1;
-        size_t offset = 0;
+        size_t data_count = perf.params.msg_size_cnt;
+        size_t count      = data_count + (has_counter(perf) ? 1 : 0);
 
         std::vector<unsigned> indices(count);
         std::vector<size_t> local_offsets(count, 0);
         std::vector<size_t> remote_offsets(count, 0);
         std::vector<size_t> lengths(count);
 
-        for (unsigned i = 0; i < count; ++i) {
+        for (unsigned i = 0; i < data_count; ++i) {
             indices[i] = i;
-            lengths[i] = (i == count - 1) ? ONESIDED_SIGNAL_SIZE :
-                                            perf.params.msg_size_list[i];
-            offset    += lengths[i];
+            lengths[i] = perf.params.msg_size_list[i];
+        }
+
+        if (has_counter(perf)) {
+            indices[data_count] = data_count;
+            lengths[data_count] = ONESIDED_SIGNAL_SIZE;
         }
 
         device_clone(&m_params.indices, indices.data(), count);

src/ucp/api/device/ucp_device_impl.h

@@ -232,9 +232,11 @@ UCS_F_DEVICE ucs_status_t ucp_device_counter_inc(
  * This operation can be polled on the receiver to detect completion of all the
  * operations of the batch, started during the same routine call.
  *
- * The last entry in the descriptor list contains
- * the remote memory registration descriptors to be used for the increment
- * operation.
+ * All the elements except the last one are data elements that must contain all
+ * @ref ucp_device_mem_list_elem_fields and @ref ucp_device_mem_list_elem_t.
+ *
+ * The last entry in the descriptor list contains the remote memory
+ * registration descriptors to be used for the increment operation.
  *
  * The routine returns a request that can be progressed and checked for
  * completion with @ref ucp_device_progress_req.

src/ucp/api/device/ucp_host.h

@@ -31,13 +31,16 @@ BEGIN_C_DECLS
  * The enumeration allows specifying which fields in @ref
  * ucp_device_mem_list_elem are present.
  *
+ * @note Counter elements can omit the @a UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH
+ *       and @a UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR fields.
+ *
  * It is used to enable backward compatibility support.
  */
 enum ucp_device_mem_list_elem_field {
     UCP_DEVICE_MEM_LIST_ELEM_FIELD_MEMH        = UCS_BIT(0), /**< Source memory handle */
-    UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY        = UCS_BIT(1), /**< Unpacked remote memory key */
+    UCP_DEVICE_MEM_LIST_ELEM_FIELD_RKEY        = UCS_BIT(1), /**< Unpacked remote memory key (always required) */
     UCP_DEVICE_MEM_LIST_ELEM_FIELD_LOCAL_ADDR  = UCS_BIT(2), /**< Local address */
-    UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR = UCS_BIT(3),  /**< Remote address */
+    UCP_DEVICE_MEM_LIST_ELEM_FIELD_REMOTE_ADDR = UCS_BIT(3), /**< Remote address */
     UCP_DEVICE_MEM_LIST_ELEM_FIELD_LENGTH      = UCS_BIT(4)  /**< Length of the local buffer in bytes */
 };
 
@@ -48,6 +51,8 @@ enum ucp_device_mem_list_elem_field {
  *
  * This describes a pair of local and remote memory for which a memory operation
  * can later be performed multiple times, possibly with varying memory offsets.
+ *
+ * @note Counter elements can omit the @a memh and @a local_addr fields.
  */
 typedef struct ucp_device_mem_list_elem {
     /**
@@ -80,6 +85,7 @@ typedef struct ucp_device_mem_list_elem {
 
     /**
      * Unpacked memory key for the remote memory endpoint.
+     * Always required.
      */
     ucp_rkey_h rkey;
 } ucp_device_mem_list_elem_t;