UCP/DEVICE: add method to write to local counter. (#10895)

Author: ofirfarjun7

src/ucp/api/device/ucp_device_impl.h

@@ -364,6 +364,28 @@ UCS_F_DEVICE uint64_t ucp_device_counter_read(const void *counter_ptr)
 }
 
 
+/**
+ * @ingroup UCP_DEVICE
+ * @brief Write value to the counter memory area.
+ *
+ * This function can be used to set counter to a specific value.
+ *
+ * The counter memory area must be initialized with the host function
+ * @ref ucp_device_counter_init.
+ *
+ * @tparam      level       Level of cooperation of the transfer.
+ * @param [in]  counter_ptr Counter memory area.
+ * @param [in]  value       Value to write.
+ *
+ */
+template<ucs_device_level_t level = UCS_DEVICE_LEVEL_THREAD>
+UCS_F_DEVICE void ucp_device_counter_write(void *counter_ptr, uint64_t value)
+{
+    return ucs_device_atomic64_write(
+            reinterpret_cast<uint64_t*>(counter_ptr), value);
+}
+
+
 /**
  * @ingroup UCP_DEVICE
  * @brief Progress a device request containing a batch of operations.

src/ucp/core/ucp_device.c

@@ -271,7 +271,8 @@ static ucs_status_t ucp_device_mem_list_create_handle(
     }
 
     if (i == 0) {
-        ucs_error("failed to select lane");
+        ucs_error("failed to select lane for local device %s",
+                  ucs_topo_sys_device_get_name(local_sys_dev));
         return UCS_ERR_NO_RESOURCE;
     }
 

src/ucs/sys/device_code.h

@@ -47,6 +47,22 @@ UCS_F_DEVICE uint64_t ucs_device_atomic64_read(const uint64_t *ptr)
 }
 
 
+/*
+ * Write a 64-bit value to counter global memory address.
+ */
+UCS_F_DEVICE void ucs_device_atomic64_write(uint64_t *ptr, uint64_t value)
+{
+#ifdef __NVCC__
+    asm volatile("st.release.sys.u64 [%0], %1;"
+             :
+             : "l"(ptr), "l"(value)
+             : "memory");
+#else
+    *ptr = value;
+#endif
+}
+
+
 /* Helper macro to print a message from a device function including the
  * thread and block indices */
 #define ucs_device_printf(_title, _fmt, ...) \

test/gtest/ucp/cuda/test_kernels.cu

@@ -12,15 +12,15 @@
 
 namespace ucx_cuda {
 
-static __global__ void memcmp_kernel(const void* s1, const void* s2,
-                                     int* result, size_t size)
+static __global__ void
+memcmp_kernel(const void *s1, const void *s2, int *result, size_t size)
 {
     unsigned int idx = blockIdx.x * blockDim.x + threadIdx.x;
 
     *result = 0;
     for (size_t i = idx; i < size; i += blockDim.x * gridDim.x) {
-        if (reinterpret_cast<const uint8_t*>(s1)[i]
-            != reinterpret_cast<const uint8_t*>(s2)[i]) {
+        if (reinterpret_cast<const uint8_t*>(s1)[i] !=
+            reinterpret_cast<const uint8_t*>(s2)[i]) {
             *result = 1;
             break;
         }
@@ -116,6 +116,20 @@ ucp_counter_inc_kernel(const kernel_params params, ucs_status_t *status)
     *status = ucp_device_wait_req(&req);
 }
 
+static __global__ void
+ucp_counter_write_kernel(const kernel_params params, ucs_status_t *status)
+{
+    ucp_device_counter_write(params.counter.address, params.counter.value);
+    *status = UCS_OK;
+}
+
+static __global__ void
+ucp_counter_read_kernel(const kernel_params params, ucs_status_t *status)
+{
+    uint64_t value = ucp_device_counter_read(params.counter.address);
+    *status = (value == params.counter.value ? UCS_OK : UCS_ERR_IO_ERROR);
+}
+
 /**
  * @brief Compares two blocks of device memory.
  *
@@ -183,4 +197,18 @@ ucs_status_t launch_ucp_counter_inc(const kernel_params &params)
     return status.sync_read();
 }
 
+ucs_status_t launch_ucp_counter_write(const kernel_params &params)
+{
+    device_status_result_ptr status;
+    ucp_counter_write_kernel<<<1, 1>>>(params, status.device_ptr());
+    return status.sync_read();
+}
+
+ucs_status_t launch_ucp_counter_read(const kernel_params &params)
+{
+    device_status_result_ptr status;
+    ucp_counter_read_kernel<<<1, 1>>>(params, status.device_ptr());
+    return status.sync_read();
+}
+
 } // namespace ucx_cuda

test/gtest/ucp/cuda/test_kernels.h

@@ -22,6 +22,8 @@ struct kernel_params {
         } single;
         struct {
             unsigned mem_list_index;
+            void* address;
+            uint64_t value;
             uint64_t inc_value;
             uint64_t remote_address;
         } counter;
@@ -56,6 +58,10 @@ ucs_status_t launch_ucp_put_multi_partial(const kernel_params &params);
 
 ucs_status_t launch_ucp_counter_inc(const kernel_params &params);
 
+ucs_status_t launch_ucp_counter_write(const kernel_params &params);
+
+ucs_status_t launch_ucp_counter_read(const kernel_params &params);
+
 }; // namespace ucx_cuda
 
 #endif

test/gtest/ucp/test_ucp_device.cc

@@ -350,6 +350,32 @@ UCS_TEST_P(test_ucp_device, counter)
     EXPECT_EQ(1, list.dst_counter_read(mem_list_index));
 }
 
+UCS_TEST_P(test_ucp_device, local_counter)
+{
+    const size_t size = counter_size();
+    mem_list list(sender(), receiver(), size, 1);
+
+    static constexpr unsigned mem_list_index = 0;
+    list.dst_counter_init(mem_list_index);
+
+    // Perform the write
+    ucx_cuda::kernel_params params;
+    params.counter.address = list.src_ptr(mem_list_index);
+    params.counter.value   = 1;
+    ucs_status_t status    = ucx_cuda::launch_ucp_counter_write(params);
+    ASSERT_EQ(UCS_OK, status);
+
+    // Check counter value
+    EXPECT_EQ(true, mem_buffer::compare(&params.counter.value,
+                                        list.src_ptr(mem_list_index),
+                                        sizeof(params.counter.value),
+                                        UCS_MEMORY_TYPE_CUDA));
+
+    // Check counter value using device API
+    status = ucx_cuda::launch_ucp_counter_read(params);
+    ASSERT_EQ(UCS_OK, status);
+}
+
 UCS_TEST_P(test_ucp_device, create_fail)
 {
     ucp_device_mem_list_handle_h handle = nullptr;