feat: enable topology discovery and memory transfer with AMD GPUs with verbs;ofi_rxm provider on 8xMI300X machine

.codespellrc

@@ -1,2 +1,2 @@
 [codespell]
-ignore-words-list = crate, trait, dyn, async
+ignore-words-list = crate, trait, dyn, async, hsa

.gitignore

@@ -20,4 +20,7 @@ cufile.log
 .cache/
 
 # Asio
-subprojects/asio-*
+subprojects/asio-*
+
+build
+subprojects/.wraplock

examples/python/expanded_two_peers.py

@@ -11,10 +11,22 @@
 import numpy as np
 import torch
 
-from nixl._api import nixl_agent, nixl_agent_config
-from nixl.logging import get_logger
+try:
+    try:
+        from nixl._api import nixl_agent, nixl_agent_config
+        from nixl.logging import get_logger
+    except ImportError:
+        from rixl._api import nixl_agent, nixl_agent_config
+        from rixl.logging import get_logger
+
+    logger = get_logger(__name__)
+    NIXL_AVAILABLE = True
+except ImportError:
+    import logging
 
-logger = get_logger(__name__)
+    logger = logging.getLogger(__name__)
+    logger.error("NIXL API missing install NIXL.")
+    NIXL_AVAILABLE = False
 
 
 def parse_args():
@@ -33,6 +45,10 @@ def parse_args():
 
 
 if __name__ == "__main__":
+    if not NIXL_AVAILABLE:
+        logger.warning("Skipping example - NIXL bindings not available")
+        sys.exit(0)
+
     args = parse_args()
 
     # initiator use default port

src/plugins/libfabric/libfabric_backend.cpp

@@ -307,11 +307,12 @@ nixlLibfabricEngine::nixlLibfabricEngine(const nixlBackendInitParams *init_param
 
     NIXL_INFO << "System runtime: "
               << (runtime_ == FI_HMEM_CUDA       ? "CUDA" :
+                      runtime_ == FI_HMEM_ROCR   ? "ROCr" :
                       runtime_ == FI_HMEM_NEURON ? "NEURON" :
                                                    "SYSTEM");
 
 #ifdef HAVE_CUDA
-    if (runtime_ == FI_HMEM_CUDA) {
+    if (runtime_ == FI_HMEM_CUDA || runtime_ == FI_HMEM_ROCR) {
         // Initialize CUDA context management
         vramInitCtx();
         // CUDA address workaround
@@ -673,7 +674,7 @@ nixlLibfabricEngine::getSupportedMems() const {
     nixl_mem_list_t mems;
     mems.push_back(DRAM_SEG);
 #ifdef HAVE_CUDA
-    if (runtime_ == FI_HMEM_CUDA) {
+    if (runtime_ == FI_HMEM_CUDA || runtime_ == FI_HMEM_ROCR) {
         NIXL_DEBUG << "CUDA runtime detected, adding VRAM support";
         mems.push_back(VRAM_SEG);
     } else
@@ -708,7 +709,7 @@ nixlLibfabricEngine::registerMem(const nixlBlobDesc &mem,
     // Use system runtime type to determine device-specific operations
     if (nixl_mem == VRAM_SEG) {
 #ifdef HAVE_CUDA
-        if (runtime_ == FI_HMEM_CUDA) {
+        if (runtime_ == FI_HMEM_CUDA || runtime_ == FI_HMEM_ROCR) {
             // CUDA-specific address query
             // For multi-GPU support, skip CUDA address workaround
             if (cuda_addr_wa_) {
@@ -766,7 +767,7 @@ nixlLibfabricEngine::registerMem(const nixlBlobDesc &mem,
 
 #ifdef HAVE_CUDA
     // Set CUDA context before libfabric operations for VRAM
-    if (nixl_mem == VRAM_SEG && runtime_ == FI_HMEM_CUDA) {
+    if (nixl_mem == VRAM_SEG && (runtime_ == FI_HMEM_CUDA || runtime_ == FI_HMEM_ROCR)) {
         vramApplyCtx();
     }
 #endif

src/plugins/libfabric/libfabric_backend.h

@@ -38,9 +38,17 @@
 #include "libfabric/libfabric_common.h"
 
 #ifdef HAVE_CUDA
+#ifdef __HIP_PLATFORM_AMD__
+#include <hip/hip_runtime.h>
+#include <hip/hip_runtime_api.h>
+// Define CUDA types as HIP equivalents for AMD
+typedef hipCtx_t CUcontext;
+typedef hipDevice_t CUdevice;
+#else
 #include <cuda.h>
 #include <cuda_runtime.h>
 #endif
+#endif
 
 // Forward declarations
 class nixlLibfabricEngine;

src/plugins/libfabric/meson.build

@@ -16,6 +16,14 @@
 
 # LibFabric plugin configuration
 
+# Hipify sources that include cuda_runtime.h for ROCm/HIP builds
+hipify_cmd = find_program('hipify-perl')
+hipify_src = generator(hipify_cmd,
+                       output : ['@BASENAME@_hip.cpp'],
+                       arguments : ['@INPUT@', '@EXTRA_ARGS@', '@OUTPUT@'])
+
+libfabric_backend_hip_src = hipify_src.process('libfabric_backend.cpp', extra_args: '-o')
+
 # Enable libfabric utils layer
 libfabric_plugin_deps = [
     nixl_infra,
@@ -38,7 +46,7 @@ endif
 if 'LIBFABRIC' in static_plugins
     libfabric_backend_lib = static_library(
         'LIBFABRIC',
-        'libfabric_backend.cpp',
+        libfabric_backend_hip_src,
         'libfabric_backend.h',
         'libfabric_plugin.cpp',
         dependencies: libfabric_plugin_deps,
@@ -50,7 +58,7 @@ if 'LIBFABRIC' in static_plugins
 else
     libfabric_backend_lib = shared_library(
         'LIBFABRIC',
-        'libfabric_backend.cpp',
+        libfabric_backend_hip_src,
         'libfabric_backend.h',
         'libfabric_plugin.cpp',
         dependencies: libfabric_plugin_deps,

src/utils/libfabric/libfabric_common.cpp

@@ -102,6 +102,8 @@ getAvailableNetworkDevices() {
         return {"cxi", provider_device_map["cxi"]};
     } else if (provider_device_map.find("efa") != provider_device_map.end()) {
         return {"efa", provider_device_map["efa"]};
+    } else if (provider_device_map.find("verbs;ofi_rxm") != provider_device_map.end()) {
+        return {"verbs;ofi_rxm", provider_device_map["verbs;ofi_rxm"]};
     } else if (provider_device_map.find("tcp") != provider_device_map.end()) {
         return {"tcp", {provider_device_map["tcp"][0]}};
     } else if (provider_device_map.find("sockets") != provider_device_map.end()) {

src/utils/libfabric/libfabric_rail.cpp

@@ -1329,6 +1329,13 @@ nixlLibfabricRail::registerMemory(void *buffer,
                 mr_attr.device.cuda = device_id;
                 NIXL_DEBUG << "CUDA memory registration - iface: FI_HMEM_CUDA, device.cuda: "
                            << device_id;
+            } else if (iface == FI_HMEM_ROCR) {
+                // AMD ROCr memory registration
+                // ROCr uses HSA agent handles for device identification
+                // The device_id corresponds to the GPU index (0-based)
+                mr_attr.device.rocr = device_id;
+                NIXL_DEBUG << "ROCr memory registration - iface: FI_HMEM_ROCR, device.rocr: "
+                           << device_id;
             } else if (iface == FI_HMEM_NEURON) {
                 /*
                  * Store a sentinel; libfabric requires this to be initialized.

src/utils/libfabric/libfabric_rail_manager.cpp

@@ -206,6 +206,9 @@ nixlLibfabricRailManager::nixlLibfabricRailManager(size_t striping_threshold)
         runtime_ = FI_HMEM_CUDA;
         NIXL_INFO << "System runtime: CUDA for " << topology->getNumNvidiaAccel()
                   << " NVIDIA GPU(s)";
+    } else if (topology->getNumAmdAccel() > 0) {
+        runtime_ = FI_HMEM_ROCR;
+        NIXL_INFO << "System runtime: ROCr for " << topology->getNumAmdAccel() << " AMD GPU(s)";
     } else if (topology->getNumAwsAccel() > 0) {
         runtime_ = FI_HMEM_NEURON;
         NIXL_INFO << "System runtime: NEURON for " << topology->getNumAwsAccel()

src/utils/libfabric/libfabric_rail_manager.h

@@ -28,9 +28,13 @@
 #include "libfabric_rail.h"
 
 #ifdef HAVE_CUDA
+#ifdef __HIP_PLATFORM_AMD__
+#include <hip/hip_runtime.h>
+#else
 #include <cuda.h>
 #include <cuda_runtime.h>
 #endif
+#endif
 
 // Forward declarations
 class nixlLibfabricTopology;