Don't assume ctx contains all devs in platform

A context might only contain a subset of the devices in a platform. Each
ur_mem_handle_t_ contains some vectors of length ctx.NumDevices.
Device->getIndex() will return the index of a device in a platform.
Therefore, when we are using the ur_mem_handle_t vectors we must make
sure the ctxIdx doesn't exceed ctx.NumDevices.

Author: Hugh Delaney

source/adapters/hip/memory.cpp

@@ -459,7 +459,8 @@ ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t Mem,
 
   if (Mem->isBuffer()) {
     auto &Buffer = std::get<BufferMem>(Mem->Mem);
-    hipDeviceptr_t &DevPtr = Buffer.Ptrs[hDevice->getIndex()];
+    hipDeviceptr_t &DevPtr =
+        Buffer.Ptrs[hDevice->getIndex() % Buffer.Ptrs.size()];
 
     // Allocation has already been made
     if (DevPtr != BufferMem::native_type{0}) {
@@ -482,12 +483,12 @@ ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t Mem,
     try {
       auto &Image = std::get<SurfaceMem>(Mem->Mem);
       // Allocation has already been made
-      if (Image.Arrays[hDevice->getIndex()]) {
+      if (Image.Arrays[hDevice->getIndex() % Image.Arrays.size()]) {
         return UR_RESULT_SUCCESS;
       }
       UR_CHECK_ERROR(hipArray3DCreate(
           reinterpret_cast<hipCUarray *>(&ImageArray), &Image.ArrayDesc));
-      Image.Arrays[hDevice->getIndex()] = ImageArray;
+      Image.Arrays[hDevice->getIndex() % Image.Arrays.size()] = ImageArray;
       // HIP_RESOURCE_DESC is a union of different structs, shown here
       // We need to fill it as described here to use it for a surface or texture
       // HIP_RESOURCE_DESC::resType must be HIP_RESOURCE_TYPE_ARRAY and
@@ -499,7 +500,7 @@ ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t Mem,
       ImageResDesc.resType = hipResourceTypeArray;
 
       UR_CHECK_ERROR(hipCreateSurfaceObject(&Surface, &ImageResDesc));
-      Image.SurfObjs[hDevice->getIndex()] = Surface;
+      Image.SurfObjs[hDevice->getIndex() % Image.SurfObjs.size()] = Surface;
     } catch (ur_result_t Err) {
       if (ImageArray) {
         UR_CHECK_ERROR(hipFreeArray(ImageArray));
@@ -610,9 +611,10 @@ ur_result_t migrateMemoryToDeviceIfNeeded(ur_mem_handle_t Mem,
   UR_ASSERT(hDevice, UR_RESULT_ERROR_INVALID_NULL_HANDLE);
   // Device allocation has already been initialized with most up to date
   // data in buffer
-  if (Mem->HaveMigratedToDeviceSinceLastWrite[hDevice->getIndex()]) {
+  if (Mem->HaveMigratedToDeviceSinceLastWrite
+          [hDevice->getIndex() %
+           Mem->HaveMigratedToDeviceSinceLastWrite.size()])
     return UR_RESULT_SUCCESS;
-  }
 
   ScopedContext Active(hDevice);
   if (Mem->isBuffer()) {
@@ -621,6 +623,8 @@ ur_result_t migrateMemoryToDeviceIfNeeded(ur_mem_handle_t Mem,
     UR_CHECK_ERROR(migrateImageToDevice(Mem, hDevice));
   }
 
-  Mem->HaveMigratedToDeviceSinceLastWrite[hDevice->getIndex()] = true;
+  Mem->HaveMigratedToDeviceSinceLastWrite
+      [hDevice->getIndex() % Mem->HaveMigratedToDeviceSinceLastWrite.size()] =
+      true;
   return UR_RESULT_SUCCESS;
 }

source/adapters/hip/memory.hpp

@@ -102,7 +102,8 @@ struct BufferMem {
       throw Err;
     }
     return reinterpret_cast<native_type>(
-        reinterpret_cast<uint8_t *>(Ptrs[Device->getIndex()]) + Offset);
+        reinterpret_cast<uint8_t *>(Ptrs[Device->getIndex() % Ptrs.size()]) +
+        Offset);
   }
 
   // This will allocate memory on device if there isn't already an active
@@ -266,7 +267,7 @@ struct SurfaceMem {
         Err != UR_RESULT_SUCCESS) {
       throw Err;
     }
-    return Arrays[Device->getIndex()];
+    return Arrays[Device->getIndex() % Arrays.size()];
   }
 
   // Will allocate a new surface on device if not already allocated
@@ -276,7 +277,7 @@ struct SurfaceMem {
         Err != UR_RESULT_SUCCESS) {
       throw Err;
     }
-    return SurfObjs[Device->getIndex()];
+    return SurfObjs[Device->getIndex() % SurfObjs.size()];
   }
 
   ur_mem_type_t getImageType() const noexcept { return ImageDesc.type; }
@@ -510,8 +511,9 @@ struct ur_mem_handle_t_ {
     urEventRetain(NewEvent);
     LastEventWritingToMemObj = NewEvent;
     for (const auto &Device : Context->getDevices()) {
-      HaveMigratedToDeviceSinceLastWrite[Device->getIndex()] =
-          Device == NewEvent->getQueue()->getDevice();
+      HaveMigratedToDeviceSinceLastWrite
+          [Device->getIndex() % HaveMigratedToDeviceSinceLastWrite.size()] =
+              Device == NewEvent->getQueue()->getDevice();
     }
   }
 };