UCT/ROCM: initial commit for device initiated ipc put

Introduces device initiated IPC put operation for the ROCm component.

- Add device-side rocm_ipc.h with HIP __device__ put/atomic_add/get_ptr,
  including a strided int4/int2 vectorized copy at thread/warp/block levels.
- Add MD mem_elem_pack/rkey_ptr that map the remote IPC handle via a shared,
  lazily-initialized process-wide handle cache; iface/ep now expose a device ep.

Author: mshanthagit

config/hip.am

@@ -0,0 +1,39 @@
+#
+# Copyright (c) Advanced Micro Devices, Inc. 2026. ALL RIGHTS RESERVED.
+# See file LICENSE for terms.
+#
+
+SUFFIXES = .hip
+
+HIPCC ?= hipcc
+
+HIPCC_CMD = $(HIPCC) $(DEFS) -DUCT_DEVICE_CODE_HIP -fPIE -I$(top_srcdir)/src -I$(top_builddir)/src $(BASE_CXXFLAGS) $(CXXFLAGS) $(HIP_CPPFLAGS) $(HIP_CXXFLAGS) $(HIPCC_EXTRA_FLAGS) -Wno-c++20-extensions -c $< -MT $@ -MF $(DEPDIR)/hip/$@.d -MMD -o $@
+HIPCC_LT_CMD = $(LIBTOOL) --tag=CXX --mode=compile $(HIPCC_CMD)
+
+define hipcc-build
+	@$(MKDIR_P) $(shell dirname $@) $(shell dirname $(DEPDIR)/hip/$@)
+	@$(if $(filter false,$(AM_V_P)),echo "  HIPCC    $@")
+	@$(if $(filter .o,$(suffix $@)),$(HIPCC_CMD),$(HIPCC_LT_CMD)) $($(1)) $(if $(filter false,$(AM_V_P)), >/dev/null)
+endef
+
+define hipcc-source
+	EXTRA_DIST += $(2)
+$(1): $(2)
+	$$(call hipcc-build,$(3))
+endef
+
+# Default rules when no target-specific compile flags are required
+.hip.o:
+	$(call hipcc-build)
+
+.hip.lo:
+	$(call hipcc-build)
+
+HIP_DEP_FILES := $(shell find $(DEPDIR)/hip/ -type f -name *.d 2>/dev/null)
+-include $(HIP_DEP_FILES)
+
+clean-local:
+	-rm -rf $(DEPDIR)/hip
+
+distclean-local:
+	-rm -rf $(DEPDIR)/hip

src/ucp/core/ucp_device.c

@@ -1,5 +1,6 @@
 /**
  * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025-2026. ALL RIGHTS RESERVED.
+ * Copyright (C) Advanced Micro Devices, Inc. 2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -106,8 +107,8 @@ static uct_allocated_memory_t ucp_device_mem_handle_hash_remove(void *handle)
 }
 
 static ucs_status_t
-ucp_device_detect_local_sys_dev(ucp_context_h context,
-                                ucs_memory_type_t mem_type,
+ucp_device_detect_local_sys_dev(const ucp_context_h context,
+                                const ucs_memory_type_t mem_type,
                                 ucs_sys_device_t *local_sys_dev)
 {
     ucs_memory_info_t mem_info;
@@ -310,7 +311,7 @@ static ucs_status_t ucp_device_local_mem_list_create_handle(
 
 static ucs_status_t ucp_device_local_mem_list_params_check(
         const ucp_device_mem_list_params_t *params, ucs_memory_type_t mem_type,
-        ucs_sys_device_t *local_sys_dev)
+        ucs_sys_device_t local_sys_dev)
 {
     const ucp_device_mem_list_elem_t *local_elements = UCS_PARAM_VALUE(
             UCP_DEVICE_MEM_LIST_PARAMS_FIELD, params, elements, ELEMENTS, NULL);
@@ -325,7 +326,6 @@ static ucs_status_t ucp_device_local_mem_list_params_check(
     const ucp_device_mem_list_elem_t *element;
     ucp_mem_h memh;
     size_t i;
-    ucs_status_t status;
 
     if ((local_elements == NULL) || (element_size == 0) ||
         (num_elements == 0) || (worker == NULL)) {
@@ -335,12 +335,6 @@ static ucs_status_t ucp_device_local_mem_list_params_check(
         return UCS_ERR_INVALID_PARAM;
     }
 
-    status = ucp_device_detect_local_sys_dev(worker->context, mem_type,
-                                             local_sys_dev);
-    if (status != UCS_OK) {
-        return status;
-    }
-
     for (i = 0; i < num_elements; i++) {
         element = (const ucp_device_mem_list_elem_t*)
                 UCS_PTR_BYTE_OFFSET(local_elements, i * element_size);
@@ -351,26 +345,59 @@ static ucs_status_t ucp_device_local_mem_list_params_check(
             return UCS_ERR_INVALID_PARAM;
         }
 
-        if (memh->sys_dev != *local_sys_dev) {
+        if (memh->sys_dev != local_sys_dev) {
             ucs_error("mismatched local sys_dev for element=%zu", i);
             return UCS_ERR_UNSUPPORTED;
         }
     }
 
+    return UCS_OK;
+}
+
+static ucs_status_t
+ucp_device_detect_export_mem_type(const ucp_context_h context,
+                                  ucs_memory_type_t *export_mem_type_p,
+                                  ucs_sys_device_t *sys_dev_p)
+{
+    ucs_status_t status = UCS_ERR_UNSUPPORTED;
+    ucs_memory_type_t mem_type;
+
+    ucs_for_each_bit(mem_type, UCP_DEVICE_MEM_TYPES) {
+        status = ucp_device_detect_local_sys_dev(context, mem_type, sys_dev_p);
+        if (status == UCS_OK) {
+            *export_mem_type_p = mem_type;
+            break;
+        }
+    }
+
     return status;
 }
 
 ucs_status_t
 ucp_device_local_mem_list_create(const ucp_device_mem_list_params_t *params,
                                  ucp_device_local_mem_list_h *mem_list_h)
 {
-    const ucs_memory_type_t export_mem_type = UCS_MEMORY_TYPE_CUDA;
-    ucs_status_t status;
-    uct_allocated_memory_t mem;
+    const ucp_worker_h worker = UCS_PARAM_VALUE(UCP_DEVICE_MEM_LIST_PARAMS_FIELD,
+                                               params, worker, WORKER, NULL);
+    ucs_memory_type_t export_mem_type;
     ucs_sys_device_t local_sys_dev;
+    uct_allocated_memory_t mem;
+    ucs_status_t status;
+
+    if (worker == NULL) {
+        ucs_error("missing worker in local mem list params");
+        return UCS_ERR_INVALID_PARAM;
+    }
+
+    status = ucp_device_detect_export_mem_type(worker->context,
+                                               &export_mem_type,
+                                               &local_sys_dev);
+    if (status != UCS_OK) {
+        return status;
+    }
 
     status = ucp_device_local_mem_list_params_check(params, export_mem_type,
-                                                    &local_sys_dev);
+                                                    local_sys_dev);
     if (status != UCS_OK) {
         ucs_error("failed to check local mem list params: %s",
                   ucs_status_string(status));
@@ -534,32 +561,20 @@ ucp_device_remote_mem_list_fill(const ucp_device_mem_list_elem_t *ucp_element,
 }
 
 static ucs_status_t ucp_device_remote_mem_list_create_handle(
-        const ucp_device_mem_list_params_t *params, ucs_memory_type_t mem_type,
-        uct_allocated_memory_t *mem)
+        const ucp_device_mem_list_params_t *params, const ucp_ep_h ep,
+        const ucs_memory_type_t mem_type, uct_allocated_memory_t *mem,
+        const ucs_sys_device_t local_sys_dev)
 {
-    const ucp_ep_h ep = ucp_device_remote_mem_list_get_first_ep(params);
     size_t uct_elem_size;
     size_t handle_size         = 0;
     ucp_tl_bitmap_t tl_bitmap[UCP_DEVICE_TL_TYPE_LAST] = {};
     const ucp_device_mem_list_elem_t *ucp_element;
     ucp_device_remote_mem_list_t *handle;
     uct_device_remote_mem_elem_t *uct_element;
-    ucs_sys_device_t local_sys_dev;
     size_t i, num_lanes;
     ucs_status_t status;
     int tl_type;
 
-    if (ep == NULL) {
-        ucs_error("no ep found in remote mem list");
-        return UCS_ERR_INVALID_PARAM;
-    }
-
-    status = ucp_device_detect_local_sys_dev(ep->worker->context, mem_type,
-                                             &local_sys_dev);
-    if (status != UCS_OK) {
-        return status;
-    }
-
     ucp_device_get_tl_bitmap(ep->worker, tl_bitmap, local_sys_dev);
 
     /* handle->num_lanes is the least common multiple of both lane types, so:
@@ -582,7 +597,7 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle(
     uct_elem_size = sizeof(uct_device_remote_mem_elem_t) +
                     (sizeof(uct_device_remote_tl_elem_t) * num_lanes);
     handle_size   = sizeof(*handle) + (params->num_elements * uct_elem_size);
-    handle      = ucs_calloc(1, handle_size, "ucp_device_remote_mem_list_t");
+    handle        = ucs_calloc(1, handle_size, "ucp_device_remote_mem_list_t");
     if (handle == NULL) {
         ucs_error("failed to allocate ucp_device_remote_mem_list_t");
         return UCS_ERR_NO_MEMORY;
@@ -601,7 +616,7 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle(
 
             if (tl_type == UCP_DEVICE_TL_TYPE_LAST) {
                 ucs_error("lane not found for element %zd", i);
-                status =  UCS_ERR_INVALID_PARAM;
+                status = UCS_ERR_INVALID_PARAM;
                 goto out;
             }
 
@@ -616,10 +631,10 @@ static ucs_status_t ucp_device_remote_mem_list_create_handle(
         ucp_element = UCS_PTR_BYTE_OFFSET(ucp_element, params->element_size);
     }
 
-    handle->version = UCP_DEVICE_MEM_LIST_VERSION_V1;
-    handle->length  = params->num_elements;
+    handle->version   = UCP_DEVICE_MEM_LIST_VERSION_V1;
+    handle->length    = params->num_elements;
     handle->num_lanes = num_lanes;
-    status          = ucp_device_mem_list_export_handle(
+    status            = ucp_device_mem_list_export_handle(
             ep->worker, handle, handle_size, mem_type, local_sys_dev, mem,
             "ucp_device_remote_mem_list_handle_t");
 
@@ -683,17 +698,31 @@ ucs_status_t
 ucp_device_remote_mem_list_create(const ucp_device_mem_list_params_t *params,
                                   ucp_device_remote_mem_list_h *mem_list_h)
 {
-    const ucs_memory_type_t export_mem_type = UCS_MEMORY_TYPE_CUDA;
-    ucs_status_t status;
+    const ucp_ep_h ep = ucp_device_remote_mem_list_get_first_ep(params);
+    ucs_memory_type_t export_mem_type;
+    ucs_sys_device_t sys_dev;
     uct_allocated_memory_t mem;
+    ucs_status_t status;
+
+    if (ep == NULL) {
+        ucs_error("no ep found in remote mem list params");
+        return UCS_ERR_INVALID_PARAM;
+    }
+
+    status = ucp_device_detect_export_mem_type(ep->worker->context,
+                                               &export_mem_type, &sys_dev);
+    if (status != UCS_OK) {
+        return status;
+    }
 
     status = ucp_device_remote_mem_list_params_check(params);
     if (status != UCS_OK) {
         return status;
     }
 
-    status = ucp_device_remote_mem_list_create_handle(params, export_mem_type,
-                                                      &mem);
+    status = ucp_device_remote_mem_list_create_handle(params, ep,
+                                                      export_mem_type,
+                                                      &mem, sys_dev);
     if (status != UCS_OK) {
         /*
          * Do not log error for UCS_ERR_NOT_CONNECTED because it is expected

src/ucp/core/ucp_mm.h

@@ -281,5 +281,7 @@ static UCS_F_ALWAYS_INLINE size_t ucp_memh_length(const ucp_mem_h memh)
                                     UCS_BIT(UCS_MEMORY_TYPE_ROCM) | \
                                     UCS_BIT(UCS_MEMORY_TYPE_ZE_DEVICE) | \
                                     UCS_BIT(UCS_MEMORY_TYPE_ZE_MANAGED)))
+#define UCP_DEVICE_MEM_TYPES \
+    (UCS_BIT(UCS_MEMORY_TYPE_CUDA) | UCS_BIT(UCS_MEMORY_TYPE_ROCM))
 
 #endif

src/ucp/wireup/select.c

@@ -1,6 +1,7 @@
 /**
  * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2001-2026. ALL RIGHTS RESERVED.
  * Copyright (C) Los Alamos National Security, LLC. 2019 ALL RIGHTS RESERVED.
+ * Copyright (C) Advanced Micro Devices, Inc. 2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -2620,7 +2621,8 @@ ucp_wireup_add_device_lanes(const ucp_wireup_select_params_t *select_params,
     ucp_wireup_select_bw_info_t bw_info = {};
     const uint64_t mem_type_bitmaps[]   = {UCS_BIT(UCS_MEMORY_TYPE_CUDA),
                                            UCS_BIT(UCS_MEMORY_TYPE_CUDA) |
-                                                   UCS_BIT(UCS_MEMORY_TYPE_HOST)};
+                                                   UCS_BIT(UCS_MEMORY_TYPE_HOST),
+                                           UCS_BIT(UCS_MEMORY_TYPE_ROCM)};
     int found_lane                      = 0;
     size_t i;
     ucp_tl_bitmap_t mem_type_tl_bitmap;

src/ucs/sys/device_code.h

@@ -1,5 +1,6 @@
 /**
  * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025. ALL RIGHTS RESERVED.
+ * Copyright (C) Advanced Micro Devices, Inc. 2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -13,11 +14,11 @@
 /*
  * Declare GPU specific functions
  */
-#ifdef __NVCC__
+#if defined(__NVCC__) || defined(__HIPCC__)
 #define UCS_F_DEVICE __device__ __forceinline__ static
 #else
 #define UCS_F_DEVICE static inline
-#endif /* __NVCC__ */
+#endif /* __NVCC__ || __HIPCC__ */
 
 
 #ifndef UCP_DEVICE_ENABLE_PARAMS_CHECK
@@ -32,6 +33,9 @@
 /* Number of threads in a warp */
 #define UCS_DEVICE_NUM_THREADS_IN_WARP 32
 
+/* Check if _n is aligned to a power-of-2 boundary _p */
+#define UCS_DEVICE_IS_ALIGNED_POW2(_n, _p) (!((_n) & ((_p) - 1)))
+
 
 /* nvcc does not provide __builtin_ia32_prefetch used by GCC's x86 intrinsic headers.
    Redirect to the generic __builtin_prefetch so those headers compile. */

src/uct/api/device/uct_device_impl.h

@@ -1,5 +1,6 @@
 /**
  * Copyright (c) NVIDIA CORPORATION & AFFILIATES, 2025-2026. ALL RIGHTS RESERVED.
+ * Copyright (C) Advanced Micro Devices, Inc. 2026. ALL RIGHTS RESERVED.
  *
  * See file LICENSE for terms.
  */
@@ -20,6 +21,14 @@
 #define UCT_CUDA_IPC_SUPPORTED 0
 #endif
 
+#if (defined(HAVE_ROCM) || \
+    (!defined(HAVE_CONFIG_H) && __has_include(<uct/rocm/ipc/rocm_ipc.h>)))
+#include <uct/rocm/ipc/rocm_ipc.h>
+#define UCT_ROCM_IPC_SUPPORTED 1
+#else
+#define UCT_ROCM_IPC_SUPPORTED 0
+#endif
+
 #if (defined(HAVE_GDA) || \
     (!defined(HAVE_CONFIG_H) && \
      __has_include(<uct/ib/mlx5/gdaki/gdaki.cuh>) && \
@@ -30,8 +39,10 @@
 #define UCT_RC_MLX5_GDA_SUPPORTED 0
 #endif
 
-#if __has_include(<uct/ib/mlx5/gdaki/d2p.cuh>) && \
-    __has_include(<infiniband/mlx5dv.h>)
+#if (defined(HAVE_GDA) || \
+    (!defined(HAVE_CONFIG_H) && \
+     __has_include(<uct/ib/mlx5/gdaki/d2p.cuh>) && \
+     __has_include(<infiniband/mlx5dv.h>)))
 #include <uct/ib/mlx5/gdaki/d2p.cuh>
 #define UCT_D2P_SUPPORTED 1
 #else
@@ -45,6 +56,9 @@ union uct_device_completion {
 #if UCT_CUDA_IPC_SUPPORTED
     uct_cuda_ipc_completion_t cuda_ipc;
 #endif
+#if UCT_ROCM_IPC_SUPPORTED
+    uct_rocm_ipc_completion_t rocm_ipc;
+#endif
 };
 
 
@@ -105,6 +119,12 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_put(
                                         channel_id, flags, comp);
     }
 #endif
+#if UCT_ROCM_IPC_SUPPORTED
+    if (device_ep->uct_tl_id == UCT_DEVICE_TL_ROCM_IPC) {
+        return uct_rocm_ipc_ep_put<level>(device_ep, mem_elem, address,
+                                          remote_address, length, flags, comp);
+    }
+#endif
 
     return UCS_ERR_UNSUPPORTED;
 }
@@ -163,6 +183,12 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_atomic_add(
                                                flags, comp);
     }
 #endif
+#if UCT_ROCM_IPC_SUPPORTED
+    if (device_ep->uct_tl_id == UCT_DEVICE_TL_ROCM_IPC) {
+        return uct_rocm_ipc_ep_atomic_add<level>(device_ep, mem_elem, inc_value,
+                                                 remote_address, flags, comp);
+    }
+#endif
 
     return UCS_ERR_UNSUPPORTED;
 }
@@ -191,6 +217,11 @@ UCS_F_DEVICE ucs_status_t uct_device_ep_get_ptr(
         return uct_cuda_ipc_ep_get_ptr(device_ep, mem_elem, address, addr_p);
     }
 #endif
+#if UCT_ROCM_IPC_SUPPORTED
+    if (device_ep->uct_tl_id == UCT_DEVICE_TL_ROCM_IPC) {
+        return uct_rocm_ipc_ep_get_ptr(device_ep, mem_elem, address, addr_p);
+    }
+#endif
 
     return UCS_ERR_UNSUPPORTED;
 }