Enable peering among all devices on the system. (iree-org#19555)

We have to do this for models that are using multiple devices but are not using them as a single logical device. Signed-off-by: Andrew Woloszyn <[email protected]>
nod-ai · Dec 23, 2024 · 1f19761 · 1f19761
1 parent 0184eee
commit 1f19761
Showing 1 changed file with 8 additions and 5 deletions.
diff --git a/runtime/src/iree/hal/drivers/hip/hip_device.c b/runtime/src/iree/hal/drivers/hip/hip_device.c
@@ -455,13 +455,16 @@ iree_status_t iree_hal_hip_device_create(
     }
 
     if (iree_status_is_ok(status)) {
-      for (iree_host_size_t j = 0;
-           j < device_count && iree_status_is_ok(status); ++j) {
-        if (i == j) {
+      int hip_device_count = 0;
+      status = IREE_HIP_CALL_TO_STATUS(
+          symbols, hipGetDeviceCount(&hip_device_count), "hipGetDeviceCount");
+
+      for (int j = 0; j < hip_device_count && iree_status_is_ok(status); ++j) {
+        if (j == device->devices[i].hip_device) {
           continue;
         }
-        status = IREE_HIP_CALL_TO_STATUS(
-            symbols, hipDeviceEnablePeerAccess(devices[j], 0));
+        status =
+            IREE_HIP_CALL_TO_STATUS(symbols, hipDeviceEnablePeerAccess(j, 0));
       }
     }
   }