Fix/hmem reg new

src/plugins/libfabric/libfabric_backend.cpp

@@ -76,6 +76,110 @@ nrtQueryAddr(const void *va, std::string *efa_bdf) {
     return -1;
 }
 
+// ROCr (AMD GPU) runtime library loading
+void *
+dlopen_libhsa() {
+    // Try libhsa-runtime64.so first (HSA Runtime)
+    static void *handle = dlopen("libhsa-runtime64.so.1", RTLD_NOW | RTLD_LAZY);
+    if (handle) {
+        return handle;
+    }
+    // Fallback to versioned library
+    handle = dlopen("libhsa-runtime64.so", RTLD_NOW | RTLD_LAZY);
+    return handle;
+}
+
+void *
+dlopen_libhip() {
+    // Try libamdhip64.so (HIP Runtime) as alternative
+    static void *handle = dlopen("libamdhip64.so.6", RTLD_NOW | RTLD_LAZY);
+    if (handle) {
+        return handle;
+    }
+    // Fallback to unversioned library
+    handle = dlopen("libamdhip64.so", RTLD_NOW | RTLD_LAZY);
+    return handle;
+}
+
+template<class Fn>
+Fn *
+_load_hsa_symbol(const char *fn_name, Fn *) {
+    void *hsa_handle = dlopen_libhsa();
+    if (hsa_handle) {
+        return reinterpret_cast<Fn *>(dlsym(hsa_handle, fn_name));
+    }
+    return nullptr;
+}
+
+template<class Fn>
+Fn *
+_load_hip_symbol(const char *fn_name, Fn *) {
+    void *hip_handle = dlopen_libhip();
+    if (hip_handle) {
+        return reinterpret_cast<Fn *>(dlsym(hip_handle, fn_name));
+    }
+    return nullptr;
+}
+
+#define LOAD_HSA_SYMBOL(sym) _load_hsa_symbol(#sym, &sym)
+#define LOAD_HIP_SYMBOL(sym) _load_hip_symbol(#sym, &sym)
+
+// HIP API function signatures for dynamic loading
+using hipGetDeviceProperties_fn = int (*)(void *prop, int device);
+using hipPointerGetAttribute_fn = int (*)(void *data, int attribute, const void *ptr);
+using hipDeviceGetPCIBusId_fn = int (*)(char *pciBusId, int len, int device);
+
+// ROCr (AMD GPU) memory address query using HIP APIs
+int
+rocrQueryAddr(const void *va, std::string *efa_bdf) {
+    // Strategy: Use HIP runtime to query PCI bus ID from memory address
+    // HIP API: hipPointerGetAttribute -> get device ID -> hipDeviceGetPCIBusId
+
+    // Load HIP symbols dynamically
+    hipPointerGetAttribute_fn hipPointerGetAttribute = nullptr;
+    hipDeviceGetPCIBusId_fn hipDeviceGetPCIBusId = nullptr;
+
+    void *hip_handle = dlopen_libhip();
+    if (hip_handle) {
+        hipPointerGetAttribute = reinterpret_cast<hipPointerGetAttribute_fn>(
+            dlsym(hip_handle, "hipPointerGetAttribute"));
+        hipDeviceGetPCIBusId = reinterpret_cast<hipDeviceGetPCIBusId_fn>(
+            dlsym(hip_handle, "hipDeviceGetPCIBusId"));
+    }
+
+    if (!hipPointerGetAttribute || !hipDeviceGetPCIBusId) {
+        // HIP library not available - this is acceptable, fall back to all rails
+        NIXL_TRACE << "HIP runtime library not available - using all rails for AMD GPU memory";
+        return -1;
+    }
+
+    // Query device ID from memory pointer
+    int device_id = -1;
+    // hipPointerAttribute_t::device = 2
+    constexpr int hipPointerAttributeDevice = 2;
+    int ret = hipPointerGetAttribute(&device_id, hipPointerAttributeDevice, va);
+    if (ret != 0 || device_id < 0) {
+        NIXL_TRACE << "Failed to query device ID from memory address " << va
+                   << " - using all rails";
+        return -1;
+    }
+
+    // Query PCI bus ID from device ID
+    char pci_bus_id[64];
+    ret = hipDeviceGetPCIBusId(pci_bus_id, sizeof(pci_bus_id), device_id);
+    if (ret != 0) {
+        NIXL_TRACE << "Failed to query PCI bus ID for AMD GPU device " << device_id
+                   << " - using all rails";
+        return -1;
+    }
+
+    // Format PCI bus ID (HIP returns format like "0000:59:00.0")
+    efa_bdf->assign(pci_bus_id);
+    NIXL_DEBUG << "ROCr query: memory " << va << " -> device " << device_id
+               << " -> PCI " << *efa_bdf;
+    return 0;
+}
+
 } // namespace
 
 #ifdef HAVE_CUDA
@@ -301,6 +405,7 @@ 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");
 
@@ -727,7 +832,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,6 +871,19 @@ nixlLibfabricEngine::registerMem(const nixlBlobDesc &mem,
             NIXL_DEBUG << "Queried PCI bus ID: " << pci_bus_id << " for GPU " << mem.devId;
         }
 #endif
+        if (runtime_ == FI_HMEM_ROCR) {
+            // AMD ROCr-specific address query
+            int ret = rocrQueryAddr((void *)mem.addr, &pci_bus_id);
+            if (ret) {
+                NIXL_TRACE << "ROCr memory address query not available - using all rails for AMD GPU "
+                           << mem.devId;
+                // Fall back to all rails - this is acceptable for initial enablement
+                // Future optimization: Implement rocrQueryAddr using HIP/HSA APIs
+            } else {
+                NIXL_DEBUG << "Queried PCI bus ID: " << pci_bus_id << " for AMD GPU "
+                           << mem.devId;
+            }
+        }
         if (runtime_ == FI_HMEM_NEURON) {
             // Neuron-specific address query
             int ret = nrtQueryAddr((void *)mem.addr, &pci_bus_id);

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_rail.cpp

@@ -1285,6 +1285,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

@@ -51,6 +51,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;