[CUDA][Bindless] Add support for device to device pitched copies and host to host copies

DBDuncan · DBDuncan · commit 3a31ffee22b0 · 2025-01-27T14:40:46.000Z
Extends bindless images copies to support device to device pitched copies and host to host copies
diff --git a/include/ur_api.h b/include/ur_api.h
@@ -8977,13 +8977,15 @@ typedef enum ur_exp_image_copy_flag_t {
   UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_HOST = UR_BIT(1),
   /// Device to device
   UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE = UR_BIT(2),
+  /// Host to host
+  UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST = UR_BIT(3),
   /// @cond
   UR_EXP_IMAGE_COPY_FLAG_FORCE_UINT32 = 0x7fffffff
   /// @endcond
 
 } ur_exp_image_copy_flag_t;
 /// @brief Bit Mask for validating ur_exp_image_copy_flags_t
-#define UR_EXP_IMAGE_COPY_FLAGS_MASK 0xfffffff8
+#define UR_EXP_IMAGE_COPY_FLAGS_MASK 0xfffffff0
 
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Sampler cubemap seamless filtering mode.
diff --git a/include/ur_print.hpp b/include/ur_print.hpp
@@ -10388,6 +10388,9 @@ inline std::ostream &operator<<(std::ostream &os,
   case UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE:
     os << "UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE";
     break;
+  case UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST:
+    os << "UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST";
+    break;
   default:
     os << "unknown enumerator";
     break;
@@ -10436,6 +10439,17 @@ inline ur_result_t printFlag<ur_exp_image_copy_flag_t>(std::ostream &os,
     }
     os << UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE;
   }
+
+  if ((val & UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST) ==
+      (uint32_t)UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST) {
+    val ^= (uint32_t)UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST;
+    if (!first) {
+      os << " | ";
+    } else {
+      first = false;
+    }
+    os << UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST;
+  }
   if (val != 0) {
     std::bitset<32> bits(val);
     if (!first) {
diff --git a/scripts/core/EXP-BINDLESS-IMAGES.rst b/scripts/core/EXP-BINDLESS-IMAGES.rst
@@ -108,6 +108,7 @@ Enums
     * ${X}_EXP_IMAGE_COPY_FLAG_HOST_TO_DEVICE
     * ${X}_EXP_IMAGE_COPY_FLAG_DEVICE_TO_HOST
     * ${X}_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE
+    * ${X}_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST
 
 * ${x}_exp_sampler_cubemap_filter_mode_t
     * ${X}_EXP_SAMPLER_CUBEMAP_FILTER_MODE_SEAMLESS
@@ -253,6 +254,10 @@ Changelog
 +----------+-------------------------------------------------------------+
 | 18.0     | Added BindlessImagesMapExternalLinearMemoryExp function.    |
 +----------+-------------------------------------------------------------+
+| 19.0     || Added ${X}_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST                | 
+|          || Added support for DtoD usm pitch copies                    |
+|          || Added support for HtoH copies                              |
++----------+-------------------------------------------------------------+
 
 Contributors
 --------------------------------------------------------------------------------
diff --git a/scripts/core/exp-bindless-images.yml b/scripts/core/exp-bindless-images.yml
@@ -167,6 +167,8 @@ etors:
     desc: "Device to host"
   - name: DEVICE_TO_DEVICE
     desc: "Device to device"
+  - name: HOST_TO_HOST
+    desc: "Host to host"
 --- #--------------------------------------------------------------------------
 type: enum
 extend: True
diff --git a/source/adapters/cuda/enqueue.cpp b/source/adapters/cuda/enqueue.cpp
@@ -46,34 +46,6 @@ ur_result_t enqueueEventsWait(ur_queue_handle_t CommandQueue, CUstream Stream,
   }
 }
 
-template <typename PtrT>
-void getUSMHostOrDevicePtr(PtrT USMPtr, CUmemorytype *OutMemType,
-                           CUdeviceptr *OutDevPtr, PtrT *OutHostPtr) {
-  // do not throw if cuPointerGetAttribute returns CUDA_ERROR_INVALID_VALUE
-  // checks with PI_CHECK_ERROR are not suggested
-  CUresult Ret = cuPointerGetAttribute(
-      OutMemType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, (CUdeviceptr)USMPtr);
-  // ARRAY, UNIFIED types are not supported!
-  assert(*OutMemType != CU_MEMORYTYPE_ARRAY &&
-         *OutMemType != CU_MEMORYTYPE_UNIFIED);
-
-  // pointer not known to the CUDA subsystem (possibly a system allocated ptr)
-  if (Ret == CUDA_ERROR_INVALID_VALUE) {
-    *OutMemType = CU_MEMORYTYPE_HOST;
-    *OutDevPtr = 0;
-    *OutHostPtr = USMPtr;
-
-    // todo: resets the above "non-stick" error
-  } else if (Ret == CUDA_SUCCESS) {
-    *OutDevPtr = (*OutMemType == CU_MEMORYTYPE_DEVICE)
-                     ? reinterpret_cast<CUdeviceptr>(USMPtr)
-                     : 0;
-    *OutHostPtr = (*OutMemType == CU_MEMORYTYPE_HOST) ? USMPtr : nullptr;
-  } else {
-    UR_CHECK_ERROR(Ret);
-  }
-}
-
 ur_result_t setCuMemAdvise(CUdeviceptr DevPtr, size_t Size,
                            ur_usm_advice_flags_t URAdviceFlags,
                            CUdevice Device) {
diff --git a/source/adapters/cuda/enqueue.hpp b/source/adapters/cuda/enqueue.hpp
@@ -9,6 +9,7 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include "common.hpp"
 #include <cassert>
 #include <cuda.h>
 #include <ur_api.h>
@@ -17,6 +18,34 @@ ur_result_t enqueueEventsWait(ur_queue_handle_t CommandQueue, CUstream Stream,
                               uint32_t NumEventsInWaitList,
                               const ur_event_handle_t *EventWaitList);
 
+template <typename PtrT>
+void getUSMHostOrDevicePtr(PtrT USMPtr, CUmemorytype *OutMemType,
+                           CUdeviceptr *OutDevPtr, PtrT *OutHostPtr) {
+  // do not throw if cuPointerGetAttribute returns CUDA_ERROR_INVALID_VALUE
+  // checks with PI_CHECK_ERROR are not suggested
+  CUresult Ret = cuPointerGetAttribute(
+      OutMemType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE, (CUdeviceptr)USMPtr);
+  // ARRAY, UNIFIED types are not supported!
+  assert(*OutMemType != CU_MEMORYTYPE_ARRAY &&
+         *OutMemType != CU_MEMORYTYPE_UNIFIED);
+
+  // pointer not known to the CUDA subsystem (possibly a system allocated ptr)
+  if (Ret == CUDA_ERROR_INVALID_VALUE) {
+    *OutMemType = CU_MEMORYTYPE_HOST;
+    *OutDevPtr = 0;
+    *OutHostPtr = USMPtr;
+
+    // todo: resets the above "non-stick" error
+  } else if (Ret == CUDA_SUCCESS) {
+    *OutDevPtr = (*OutMemType == CU_MEMORYTYPE_DEVICE)
+                     ? reinterpret_cast<CUdeviceptr>(USMPtr)
+                     : 0;
+    *OutHostPtr = (*OutMemType == CU_MEMORYTYPE_HOST) ? USMPtr : nullptr;
+  } else {
+    UR_CHECK_ERROR(Ret);
+  }
+}
+
 void guessLocalWorkSize(ur_device_handle_t Device, size_t *ThreadsPerBlock,
                         const size_t *GlobalWorkSize, const uint32_t WorkDim,
                         ur_kernel_handle_t Kernel);
diff --git a/source/adapters/cuda/image.cpp b/source/adapters/cuda/image.cpp
@@ -583,7 +583,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
     const ur_event_handle_t *phEventWaitList, ur_event_handle_t *phEvent) {
   UR_ASSERT((imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_HOST_TO_DEVICE ||
              imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_HOST ||
-             imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE),
+             imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE ||
+             imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST),
             UR_RESULT_ERROR_INVALID_VALUE);
   UR_ASSERT(pSrcImageFormat->channelOrder == pDstImageFormat->channelOrder,
             UR_RESULT_ERROR_INVALID_ARGUMENT);
@@ -651,6 +652,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
         cpy_desc.srcY = pCopyRegion->srcOffset.y;
         cpy_desc.dstXInBytes = pCopyRegion->dstOffset.x * PixelSizeBytes;
         cpy_desc.dstY = pCopyRegion->dstOffset.y;
+        cpy_desc.WidthInBytes = PixelSizeBytes * pCopyRegion->copyExtent.width;
+        cpy_desc.Height = pCopyRegion->copyExtent.height;
         cpy_desc.srcPitch = pSrcImageDesc->width * PixelSizeBytes;
         if (pDstImageDesc->rowPitch == 0) {
           cpy_desc.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
@@ -661,8 +664,6 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
           cpy_desc.dstDevice = (CUdeviceptr)pDst;
           cpy_desc.dstPitch = pDstImageDesc->rowPitch;
         }
-        cpy_desc.WidthInBytes = PixelSizeBytes * pCopyRegion->copyExtent.width;
-        cpy_desc.Height = pCopyRegion->copyExtent.height;
         UR_CHECK_ERROR(cuMemcpy2DAsync(&cpy_desc, Stream));
       } else if (pDstImageDesc->type == UR_MEM_TYPE_IMAGE3D) {
         CUDA_MEMCPY3D cpy_desc = {};
@@ -740,22 +741,20 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
         cpy_desc.srcY = pCopyRegion->srcOffset.y;
         cpy_desc.dstXInBytes = pCopyRegion->dstOffset.x * PixelSizeBytes;
         cpy_desc.dstY = pCopyRegion->dstOffset.y;
+        cpy_desc.WidthInBytes = PixelSizeBytes * pCopyRegion->copyExtent.width;
+        cpy_desc.Height = pCopyRegion->copyExtent.height;
+        cpy_desc.dstPitch = pDstImageDesc->width * PixelSizeBytes;
         cpy_desc.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_HOST;
         cpy_desc.dstHost = pDst;
         if (pSrcImageDesc->rowPitch == 0) {
           cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
           cpy_desc.srcArray = as_CUArray(pSrc);
         } else {
           // Pitched memory
-          cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_DEVICE;
           cpy_desc.srcPitch = pSrcImageDesc->rowPitch;
+          cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_DEVICE;
           cpy_desc.srcDevice = (CUdeviceptr)pSrc;
         }
-        cpy_desc.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_HOST;
-        cpy_desc.dstHost = pDst;
-        cpy_desc.dstPitch = pDstImageDesc->width * PixelSizeBytes;
-        cpy_desc.WidthInBytes = PixelSizeBytes * pCopyRegion->copyExtent.width;
-        cpy_desc.Height = pCopyRegion->copyExtent.height;
         UR_CHECK_ERROR(cuMemcpy2DAsync(&cpy_desc, Stream));
       } else if (pSrcImageDesc->type == UR_MEM_TYPE_IMAGE3D) {
         CUDA_MEMCPY3D cpy_desc = {};
@@ -797,7 +796,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
         UR_CHECK_ERROR(cuMemcpy3DAsync(&cpy_desc, Stream));
       }
     } else {
-      // imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE
+      // imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_DEVICE ||
+      // imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_HOST_TO_HOST
 
       // we don't support copying between different image types.
       if (pSrcImageDesc->type != pDstImageDesc->type) {
@@ -810,30 +810,67 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
       // synchronous because of the explicit call to cuStreamSynchronize at
       // the end
       if (pSrcImageDesc->type == UR_MEM_TYPE_IMAGE1D) {
+        // Check what type of memory pSrc and pDst are to set the correct
+        // attributes of cpy_desc.
+        // If cuPointerGetAttribute returns something different from
+        // CUDA_SUCCESS then we know that the memory type is a CuArray.
+        // Otherwise, it's CU_MEMORYTYPE_DEVICE.
+        CUmemorytype memType;
+        bool isSrcCudaArray =
+            cuPointerGetAttribute(&memType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
+                                  (CUdeviceptr)pSrc) != CUDA_SUCCESS;
+        bool isDstCudaArray =
+            cuPointerGetAttribute(&memType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
+                                  (CUdeviceptr)pDst) != CUDA_SUCCESS;
+
         CUDA_MEMCPY2D cpy_desc = {};
         cpy_desc.srcXInBytes = pCopyRegion->srcOffset.x * PixelSizeBytes;
         cpy_desc.srcY = 0;
         cpy_desc.dstXInBytes = pCopyRegion->dstOffset.x * PixelSizeBytes;
         cpy_desc.dstY = 0;
-        cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
-        cpy_desc.srcArray = as_CUArray(pSrc);
-        cpy_desc.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
-        cpy_desc.dstArray = (CUarray)pDst;
         cpy_desc.WidthInBytes = PixelSizeBytes * pCopyRegion->copyExtent.width;
         cpy_desc.Height = 1;
+        if (isSrcCudaArray) {
+          cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
+          cpy_desc.srcArray = as_CUArray(pSrc);
+        } else {
+          getUSMHostOrDevicePtr(pSrc, &cpy_desc.srcMemoryType,
+                                &cpy_desc.srcDevice, &cpy_desc.srcHost);
+        }
+        if (isDstCudaArray) {
+          cpy_desc.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
+          cpy_desc.dstArray = (CUarray)pDst;
+        } else {
+          getUSMHostOrDevicePtr(pDst, &cpy_desc.dstMemoryType,
+                                &cpy_desc.dstDevice, &cpy_desc.dstHost);
+        }
         UR_CHECK_ERROR(cuMemcpy2DAsync(&cpy_desc, Stream));
       } else if (pSrcImageDesc->type == UR_MEM_TYPE_IMAGE2D) {
         CUDA_MEMCPY2D cpy_desc = {};
         cpy_desc.srcXInBytes = pCopyRegion->srcOffset.x * PixelSizeBytes;
         cpy_desc.srcY = pCopyRegion->srcOffset.y;
         cpy_desc.dstXInBytes = pCopyRegion->dstOffset.x * PixelSizeBytes;
         cpy_desc.dstY = pCopyRegion->dstOffset.y;
-        cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
-        cpy_desc.srcArray = as_CUArray(pSrc);
-        cpy_desc.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
-        cpy_desc.dstArray = (CUarray)pDst;
         cpy_desc.WidthInBytes = PixelSizeBytes * pCopyRegion->copyExtent.width;
         cpy_desc.Height = pCopyRegion->copyExtent.height;
+        if (pSrcImageDesc->rowPitch == 0) {
+          cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
+          cpy_desc.srcArray = as_CUArray(pSrc);
+        } else {
+          // Pitched memory
+          cpy_desc.srcPitch = pSrcImageDesc->rowPitch;
+          getUSMHostOrDevicePtr(pSrc, &cpy_desc.srcMemoryType,
+                                &cpy_desc.srcDevice, &cpy_desc.srcHost);
+        }
+        if (pDstImageDesc->rowPitch == 0) {
+          cpy_desc.dstMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_ARRAY;
+          cpy_desc.dstArray = (CUarray)pDst;
+        } else {
+          // Pitched memory
+          cpy_desc.dstPitch = pDstImageDesc->rowPitch;
+          getUSMHostOrDevicePtr(pDst, &cpy_desc.dstMemoryType,
+                                &cpy_desc.dstDevice, &cpy_desc.dstHost);
+        }
         UR_CHECK_ERROR(cuMemcpy2DAsync(&cpy_desc, Stream));
       } else if (pSrcImageDesc->type == UR_MEM_TYPE_IMAGE3D) {
         CUDA_MEMCPY3D cpy_desc = {};
