Merge pull request #1240 from cppchedy/chedy/fix-bi-copy-image

[Bindless][exp] 1D Image copy fix
oneapi-src · Feb 5, 2024 · c0b1f13 · c0b1f13
2 parents 475ce8b + cf806c3
commit c0b1f13
Showing 1 changed file with 48 additions and 8 deletions.
diff --git a/source/adapters/cuda/image.cpp b/source/adapters/cuda/image.cpp
@@ -628,15 +628,37 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
  ScopedContext Active(hQueue->getContext());
  CUstream Stream = hQueue->getNextTransferStream();
  enqueueEventsWait(hQueue, Stream, numEventsInWaitList, phEventWaitList);
+
  // We have to use a different copy function for each image dimensionality.
 
  if (imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_HOST_TO_DEVICE) {
  if (pImageDesc->type == UR_MEM_TYPE_IMAGE1D) {
+ CUmemorytype memType;
+
+ // Check what type of memory is pDst. If cuPointerGetAttribute returns
+ // somthing different from CUDA_SUCCESS then we know that pDst memory
+ // type is a CuArray. Otherwise, it's CU_MEMORYTYPE_DEVICE.
+ bool isCudaArray =
+ cuPointerGetAttribute(&memType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
+ (CUdeviceptr)pDst) != CUDA_SUCCESS;
+
  size_t CopyExtentBytes = PixelSizeBytes * copyExtent.width;
  char *SrcWithOffset = (char *)pSrc + (srcOffset.x * PixelSizeBytes);
- UR_CHECK_ERROR(
- cuMemcpyHtoAAsync((CUarray)pDst, dstOffset.x * PixelSizeBytes,
- (void *)SrcWithOffset, CopyExtentBytes, Stream));
+
+ if (isCudaArray) {
+ UR_CHECK_ERROR(cuMemcpyHtoAAsync(
+ (CUarray)pDst, dstOffset.x * PixelSizeBytes,
+ (void *)SrcWithOffset, CopyExtentBytes, Stream));
+ } else if (memType == CU_MEMORYTYPE_DEVICE) {
+ void *DstWithOffset =
+ (void *)((char *)pDst + (PixelSizeBytes * dstOffset.x));
+ UR_CHECK_ERROR(cuMemcpyHtoDAsync((CUdeviceptr)DstWithOffset,
+ (void *)SrcWithOffset,
+ CopyExtentBytes, Stream));
+ } else {
+ // This should be unreachable.
+ return UR_RESULT_ERROR_INVALID_VALUE;
+ }
  } else if (pImageDesc->type == UR_MEM_TYPE_IMAGE2D) {
  CUDA_MEMCPY2D cpy_desc = {};
  cpy_desc.srcMemoryType = CUmemorytype_enum::CU_MEMORYTYPE_HOST;
@@ -679,13 +701,31 @@ UR_APIEXPORT ur_result_t UR_APICALL urBindlessImagesImageCopyExp(
  }
  } else if (imageCopyFlags == UR_EXP_IMAGE_COPY_FLAG_DEVICE_TO_HOST) {
  if (pImageDesc->type == UR_MEM_TYPE_IMAGE1D) {
+ CUmemorytype memType;
+ // Check what type of memory is pSrc. If cuPointerGetAttribute returns
+ // somthing different from CUDA_SUCCESS then we know that pSrc memory
+ // type is a CuArray. Otherwise, it's CU_MEMORYTYPE_DEVICE.
+ bool isCudaArray =
+ cuPointerGetAttribute(&memType, CU_POINTER_ATTRIBUTE_MEMORY_TYPE,
+ (CUdeviceptr)pSrc) != CUDA_SUCCESS;
+
  size_t CopyExtentBytes = PixelSizeBytes * copyExtent.width;
- size_t src_offset_bytes = PixelSizeBytes * srcOffset.x;
- void *dst_with_offset =
+ void *DstWithOffset =
  (void *)((char *)pDst + (PixelSizeBytes * dstOffset.x));
- UR_CHECK_ERROR(cuMemcpyAtoHAsync(dst_with_offset, (CUarray)pSrc,
- src_offset_bytes, CopyExtentBytes,
- Stream));
+
+ if (isCudaArray) {
+ UR_CHECK_ERROR(cuMemcpyAtoHAsync(DstWithOffset, (CUarray)pSrc,
+ PixelSizeBytes * srcOffset.x,
+ CopyExtentBytes, Stream));
+ } else if (memType == CU_MEMORYTYPE_DEVICE) {
+ char *SrcWithOffset = (char *)pSrc + (srcOffset.x * PixelSizeBytes);
+ UR_CHECK_ERROR(cuMemcpyDtoHAsync(DstWithOffset,
+ (CUdeviceptr)SrcWithOffset,
+ CopyExtentBytes, Stream));
+ } else {
+ // This should be unreachable.
+ return UR_RESULT_ERROR_INVALID_VALUE;
+ }
  } else if (pImageDesc->type == UR_MEM_TYPE_IMAGE2D) {
  CUDA_MEMCPY2D cpy_desc = {};
  cpy_desc.srcXInBytes = srcOffset.x;