Store HostPtr in ur mem

We might want to write this mem into an image on another device.

Author: Hugh Delaney

source/adapters/cuda/enqueue.cpp

@@ -1098,6 +1098,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemImageWrite(
   std::ignore = slicePitch;
 
   UR_ASSERT(hImage->isImage(), UR_RESULT_ERROR_INVALID_MEM_OBJECT);
+  auto &Image = std::get<SurfaceMem>(hImage->Mem);
+  // FIXME: We are assuming that the lifetime of host ptr lives as long as the
+  // image
+  if (!Image.HostPtr)
+    Image.HostPtr = pSrc;
 
   ur_result_t Result = UR_RESULT_SUCCESS;
 
@@ -1107,8 +1112,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemImageWrite(
     Result = enqueueEventsWait(hQueue, CuStream, numEventsInWaitList,
                                phEventWaitList);
 
-    CUarray Array =
-        std::get<SurfaceMem>(hImage->Mem).getArray(hQueue->getDevice());
+    CUarray Array = Image.getArray(hQueue->getDevice());
 
     CUDA_ARRAY_DESCRIPTOR ArrayDesc;
     UR_CHECK_ERROR(cuArrayGetDescriptor(&ArrayDesc, Array));
@@ -1126,7 +1130,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemImageWrite(
       UR_CHECK_ERROR(RetImplEvent->start());
     }
 
-    ur_mem_type_t ImgType = std::get<SurfaceMem>(hImage->Mem).getType();
+    ur_mem_type_t ImgType = Image.getType();
     if (ImgType == UR_MEM_TYPE_IMAGE1D) {
       UR_CHECK_ERROR(
           cuMemcpyHtoAAsync(Array, ByteOffsetX, pSrc, BytesToCopy, CuStream));

source/adapters/cuda/memory.cpp

@@ -528,35 +528,35 @@ ur_result_t migrateImageToDevice(ur_mem_handle_t Mem,
   // dimensionality
   if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE2D) {
     memset(&CpyDesc2D, 0, sizeof(CpyDesc2D));
-    CpyDesc2D.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_HOST;
     CpyDesc2D.srcHost = Image.HostPtr;
     CpyDesc2D.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
     CpyDesc2D.dstArray = ImageArray;
     CpyDesc2D.WidthInBytes = PixelSizeBytes * Image.ImageDesc.width;
     CpyDesc2D.Height = Image.ImageDesc.height;
-    UR_CHECK_ERROR(cuMemcpy2D(&CpyDesc2D));
   } else if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE3D) {
     memset(&CpyDesc3D, 0, sizeof(CpyDesc3D));
-    CpyDesc3D.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_HOST;
     CpyDesc3D.srcHost = Image.HostPtr;
     CpyDesc3D.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
     CpyDesc3D.dstArray = ImageArray;
     CpyDesc3D.WidthInBytes = PixelSizeBytes * Image.ImageDesc.width;
     CpyDesc3D.Height = Image.ImageDesc.height;
     CpyDesc3D.Depth = Image.ImageDesc.depth;
-    UR_CHECK_ERROR(cuMemcpy3D(&CpyDesc3D));
   }
 
   if (Mem->LastEventWritingToMemObj == nullptr) {
-    if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE1D) {
-      UR_CHECK_ERROR(
-          cuMemcpyHtoA(ImageArray, 0, Image.HostPtr, ImageSizeBytes));
-    } else if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE2D) {
-      CpyDesc2D.srcHost = Image.HostPtr;
-      UR_CHECK_ERROR(cuMemcpy2D(&CpyDesc2D));
-    } else if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE3D) {
-      CpyDesc3D.srcHost = Image.HostPtr;
-      UR_CHECK_ERROR(cuMemcpy3D(&CpyDesc3D));
+    if (Image.HostPtr) {
+      if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE1D) {
+        UR_CHECK_ERROR(
+            cuMemcpyHtoA(ImageArray, 0, Image.HostPtr, ImageSizeBytes));
+      } else if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE2D) {
+        CpyDesc2D.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_HOST;
+        CpyDesc2D.srcHost = Image.HostPtr;
+        UR_CHECK_ERROR(cuMemcpy2D(&CpyDesc2D));
+      } else if (Image.ImageDesc.type == UR_MEM_TYPE_IMAGE3D) {
+        CpyDesc3D.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_HOST;
+        CpyDesc3D.srcHost = Image.HostPtr;
+        UR_CHECK_ERROR(cuMemcpy3D(&CpyDesc3D));
+      }
     }
   } else if (Mem->LastEventWritingToMemObj->getQueue()->getDevice() !=
              hDevice) {