[SYCL][Bindless] Add image_mem_handle to image_mem_handle devices copies. (#12449)

cppchedy · przemektmalon · web-flow · commit 1e2e6baaf860 · 2024-04-25T11:16:54.000+02:00
Add support for device to device copy for `image_mem_handle` to
`image_mem_handle`.

---------

Co-authored-by: Przemek Malon &lt;przemek.malon@codeplay.com&gt;
diff --git a/sycl/doc/extensions/experimental/sycl_ext_oneapi_bindless_images.asciidoc b/sycl/doc/extensions/experimental/sycl_ext_oneapi_bindless_images.asciidoc
@@ -810,6 +810,12 @@ public:
     size_t DeviceRowPitch,
     sycl::range<3> HostExtent,
     sycl::range<3> CopyExtent);
+
+  // Simple device to device copy
+  void ext_oneapi_copy(
+      ext::oneapi::experimental::image_mem_handle Src,
+      ext::oneapi::experimental::image_mem_handle Dest,
+      const ext::oneapi::experimental::image_descriptor &ImageDesc);
 };
 
 class queue {
@@ -935,14 +941,32 @@ public:
       size_t DeviceRowPitch,
       sycl::range<3> HostExtent,
       sycl::range<3> CopyExtent);
+  
+  // Simple device to device copy
+  event ext_oneapi_copy(
+      ext::oneapi::experimental::image_mem_handle Src,
+      ext::oneapi::experimental::image_mem_handle Dest,
+      const ext::oneapi::experimental::image_descriptor &ImageDesc);
+  event ext_oneapi_copy(
+      ext::oneapi::experimental::image_mem_handle Src,
+      ext::oneapi::experimental::image_mem_handle Dest,
+      const ext::oneapi::experimental::image_descriptor &ImageDesc,
+      event DepEvent);
+  event ext_oneapi_copy(
+      ext::oneapi::experimental::image_mem_handle Src,
+      ext::oneapi::experimental::image_mem_handle Dest,
+      const ext::oneapi::experimental::image_descriptor &ImageDesc,
+      const std::vector<event> &DepEvents);
 };
 }
 ```
 
 To enable the copying of images an `ext_oneapi_copy` function is proposed as a 
 method of the queue and handler. It can be used to copy image memory, whether 
 allocated through USM or using an `image_mem_handle`, from host to 
-device, or device to host. For the `ext_oneapi_copy` variants that do not take 
+device, or device to host. Device to device copies are currently supported only 
+through `image_mem_handle` allocations. 
+For the `ext_oneapi_copy` variants that do not take 
 offsets and extents, the image descriptor passed to the `ext_oneapi_copy` API 
 is used to determine the pixel size, dimensions, and extent in memory of the 
 image to copy. If performing sub-region copy, the size of the memory region is 
@@ -2588,6 +2612,8 @@ These features still need to be handled:
                    wording around what types are allowed to be read or written.
                  - Allow `read_image` and `read_mipmap` to return a 
                    user-defined type.
+|5.1|2024-01-17| - Added overload for `ext_oneapi_copy` enabling device to device
+                   copies using `image_mem_handle`.
 |5.1|2023-12-06| - Added unique addressing modes per dimension to the 
                    `bindless_image_sampler`
 |5.2|2024-02-14| - Image read and write functions now accept 3-component 
diff --git a/sycl/include/sycl/handler.hpp b/sycl/include/sycl/handler.hpp
@@ -3225,6 +3225,18 @@ class __SYCL_EXPORT handler {
       const ext::oneapi::experimental::image_descriptor &DeviceImgDesc,
       size_t DeviceRowPitch);
 
+  /// Copies data from device to device memory, where \p Src and \p Dest
+  /// are opaque image memory handles.
+  /// An exception is thrown if either \p Src or \p Dest is incomplete
+  ///
+  /// \param Src is an opaque image memory handle to the source memory.
+  /// \param Dest is an opaque image memory handle to the destination memory.
+  /// \param ImageDesc is the source image descriptor
+  void
+  ext_oneapi_copy(ext::oneapi::experimental::image_mem_handle Src,
+                  ext::oneapi::experimental::image_mem_handle Dest,
+                  const ext::oneapi::experimental::image_descriptor &ImageDesc);
+
   /// Copies data from one memory region to another, where \p Src and \p Dest
   /// are USM pointers. Allows for a sub-region copy, where \p SrcOffset ,
   /// \p DestOffset , and \p Extent are used to determine the sub-region.
diff --git a/sycl/include/sycl/queue.hpp b/sycl/include/sycl/queue.hpp
@@ -1838,6 +1838,74 @@ class __SYCL_EXPORT queue : public detail::OwnerLessBase<queue> {
         CodeLoc);
   }
 
+  /// Copies data from device to device memory, where \p Src and \p Dest
+  /// are opaque image memory handles.
+  /// An exception is thrown if either \p Src or \p Dest is incomplete
+  ///
+  /// \param Src is an opaque image memory handle to the source memory.
+  /// \param Dest is an opaque image memory handle to the destination memory.
+  /// \param ImageDesc is the source image descriptor
+  /// \param DepEvent is an events that specifies the kernel dependency.
+  /// \return an event representing the copy operation.
+  event ext_oneapi_copy(
+      ext::oneapi::experimental::image_mem_handle Src,
+      ext::oneapi::experimental::image_mem_handle Dest,
+      const ext::oneapi::experimental::image_descriptor &ImageDesc,
+      event DepEvent,
+      const detail::code_location &CodeLoc = detail::code_location::current()) {
+    detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
+    return submit(
+        [&](handler &CGH) {
+          CGH.depends_on(DepEvent);
+          CGH.ext_oneapi_copy(Src, Dest, ImageDesc);
+        },
+        CodeLoc);
+  }
+
+  /// Copies data from device to device memory, where \p Src and \p Dest
+  /// are opaque image memory handles.
+  /// An exception is thrown if either \p Src or \p Dest is incomplete
+  ///
+  /// \param Src is an opaque image memory handle to the source memory.
+  /// \param Dest is an opaque image memory handle to the destination memory.
+  /// \param ImageDesc is the source image descriptor
+  /// \param DepEvents is a vector of events that specifies the kernel
+  /// dependencies.
+  /// \return an event representing the copy operation.
+  event ext_oneapi_copy(
+      ext::oneapi::experimental::image_mem_handle Src,
+      ext::oneapi::experimental::image_mem_handle Dest,
+      const ext::oneapi::experimental::image_descriptor &ImageDesc,
+      const std::vector<event> &DepEvents,
+      const detail::code_location &CodeLoc = detail::code_location::current()) {
+    detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
+    return submit(
+        [&](handler &CGH) {
+          CGH.depends_on(DepEvents);
+          CGH.ext_oneapi_copy(Src, Dest, ImageDesc);
+        },
+        CodeLoc);
+  }
+
+  /// Copies data from device to device memory, where \p Src and \p Dest
+  /// are opaque image memory handles.
+  /// An exception is thrown if either \p Src or \p Dest is incomplete
+  ///
+  /// \param Src is an opaque image memory handle to the source memory.
+  /// \param Dest is an opaque image memory handle to the destination memory.
+  /// \param ImageDesc is the source image descriptor
+  /// \return an event representing the copy operation.
+  event ext_oneapi_copy(
+      ext::oneapi::experimental::image_mem_handle Src,
+      ext::oneapi::experimental::image_mem_handle Dest,
+      const ext::oneapi::experimental::image_descriptor &ImageDesc,
+      const detail::code_location &CodeLoc = detail::code_location::current()) {
+    detail::tls_code_loc_t TlsCodeLocCapture(CodeLoc);
+    return submit(
+        [&](handler &CGH) { CGH.ext_oneapi_copy(Src, Dest, ImageDesc); },
+        CodeLoc);
+  }
+
   /// Copies data from one memory region to another, where \p Src and \p Dest
   /// are USM pointers. Allows for a sub-region copy, where \p SrcOffset ,
   /// \p DestOffset , and \p Extent are used to determine the sub-region.
diff --git a/sycl/plugins/unified_runtime/CMakeLists.txt b/sycl/plugins/unified_runtime/CMakeLists.txt
@@ -114,8 +114,15 @@ if(SYCL_PI_UR_USE_FETCH_CONTENT)
   )
 
   fetch_adapter_source(cuda
-    ${UNIFIED_RUNTIME_REPO}
-    ${UNIFIED_RUNTIME_TAG}
+    "https://github.com/oneapi-src/unified-runtime.git"
+    # commit 7fcfe3ad8882fee23d83fa0fc4c4c944262a9ea3
+    # Merge: b37fa2c4 f9fb1167
+    # Author: Kenneth Benzie (Benie) <k.benzie@codeplay.com>
+    # Date:   Wed Apr 24 10:38:00 2024 +0100
+    #     Merge pull request #1265 from cppchedy/chedy/device-to-device-copy
+    #
+    #     [Bindless][Exp] Add support for device to device copies between CuArrays
+    7fcfe3ad8882fee23d83fa0fc4c4c944262a9ea3
   )
 
   fetch_adapter_source(hip
diff --git a/sycl/source/detail/memory_manager.cpp b/sycl/source/detail/memory_manager.cpp
@@ -1776,7 +1776,9 @@ void MemoryManager::copy_image_bindless(
   assert((Flags == (sycl::detail::pi::PiImageCopyFlags)
                        ext::oneapi::experimental::image_copy_flags::HtoD ||
           Flags == (sycl::detail::pi::PiImageCopyFlags)
-                       ext::oneapi::experimental::image_copy_flags::DtoH) &&
+                       ext::oneapi::experimental::image_copy_flags::DtoH ||
+          Flags == (sycl::detail::pi::PiImageCopyFlags)
+                       ext::oneapi::experimental::image_copy_flags::DtoD) &&
          "Invalid flags passed to copy_image_bindless.");
   if (!Dst || !Src)
     throw sycl::exception(
diff --git a/sycl/source/handler.cpp b/sycl/source/handler.cpp
@@ -1167,6 +1167,57 @@ void handler::ext_oneapi_copy(
   setType(detail::CG::CopyImage);
 }
 
+void handler::ext_oneapi_copy(
+    ext::oneapi::experimental::image_mem_handle Src,
+    ext::oneapi::experimental::image_mem_handle Dest,
+    const ext::oneapi::experimental::image_descriptor &ImageDesc) {
+  throwIfGraphAssociated<
+      ext::oneapi::experimental::detail::UnsupportedGraphFeatures::
+          sycl_ext_oneapi_bindless_images>();
+  ImageDesc.verify();
+
+  MSrcPtr = Src.raw_handle;
+  MDstPtr = Dest.raw_handle;
+
+  sycl::detail::pi::PiMemImageDesc PiDesc = {};
+  PiDesc.image_width = ImageDesc.width;
+  PiDesc.image_height = ImageDesc.height;
+  PiDesc.image_depth = ImageDesc.depth;
+  PiDesc.image_array_size = ImageDesc.array_size;
+  if (ImageDesc.array_size > 1) {
+    // Image Array.
+    PiDesc.image_type = ImageDesc.height > 0 ? PI_MEM_TYPE_IMAGE2D_ARRAY
+                                             : PI_MEM_TYPE_IMAGE1D_ARRAY;
+
+    // Cubemap.
+    PiDesc.image_type =
+        ImageDesc.type == sycl::ext::oneapi::experimental::image_type::cubemap
+            ? PI_MEM_TYPE_IMAGE_CUBEMAP
+            : PiDesc.image_type;
+  } else {
+    PiDesc.image_type = ImageDesc.depth > 0
+                            ? PI_MEM_TYPE_IMAGE3D
+                            : (ImageDesc.height > 0 ? PI_MEM_TYPE_IMAGE2D
+                                                    : PI_MEM_TYPE_IMAGE1D);
+  }
+
+  sycl::detail::pi::PiMemImageFormat PiFormat;
+  PiFormat.image_channel_data_type =
+      sycl::_V1::detail::convertChannelType(ImageDesc.channel_type);
+  PiFormat.image_channel_order =
+      sycl::_V1::detail::convertChannelOrder(ImageDesc.channel_order);
+
+  MImpl->MSrcOffset = {0, 0, 0};
+  MImpl->MDestOffset = {0, 0, 0};
+  MImpl->MCopyExtent = {ImageDesc.width, ImageDesc.height, ImageDesc.depth};
+  MImpl->MHostExtent = {ImageDesc.width, ImageDesc.height, ImageDesc.depth};
+  MImpl->MImageDesc = PiDesc;
+  MImpl->MImageFormat = PiFormat;
+  MImpl->MImageCopyFlags =
+      sycl::detail::pi::PiImageCopyFlags::PI_IMAGE_COPY_DEVICE_TO_DEVICE;
+  setType(detail::CG::CopyImage);
+}
+
 void handler::ext_oneapi_copy(
     ext::oneapi::experimental::image_mem_handle Src, sycl::range<3> SrcOffset,
     const ext::oneapi::experimental::image_descriptor &SrcImgDesc, void *Dest,
diff --git a/sycl/test-e2e/bindless_images/device_to_device_copy.cpp b/sycl/test-e2e/bindless_images/device_to_device_copy.cpp
@@ -0,0 +1,119 @@
+// REQUIRES: linux
+// REQUIRES: cuda
+
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+#include <iostream>
+#include <sycl/sycl.hpp>
+
+// Uncomment to print additional test information
+// #define VERBOSE_PRINT
+
+namespace syclexp = sycl::ext::oneapi::experimental;
+
+void copy_image_mem_handle_to_image_mem_handle(
+    syclexp::image_descriptor &desc, const std::vector<float> &testData,
+    sycl::device dev, sycl::queue q, std::vector<float> &out) {
+  syclexp::image_mem imgMemSrc(desc, dev, q.get_context());
+  syclexp::image_mem imgMemDst(desc, dev, q.get_context());
+
+  q.ext_oneapi_copy((void *)testData.data(), imgMemSrc.get_handle(), desc);
+  q.wait_and_throw();
+
+  q.ext_oneapi_copy(imgMemSrc.get_handle(), imgMemDst.get_handle(), desc);
+  q.wait_and_throw();
+
+  q.ext_oneapi_copy(imgMemDst.get_handle(), (void *)out.data(), desc);
+  q.wait_and_throw();
+}
+
+bool check_test(const std::vector<float> &out,
+                const std::vector<float> &expected) {
+  assert(out.size() == expected.size());
+  bool validated = true;
+  for (int i = 0; i < out.size(); i++) {
+    bool mismatch = false;
+    if (out[i] != expected[i]) {
+      mismatch = true;
+      validated = false;
+    }
+
+    if (mismatch) {
+#ifdef VERBOSE_PRINT
+      std::cout << "Result mismatch! Expected: " << expected[i]
+                << ", Actual: " << out[i] << std::endl;
+#else
+      break;
+#endif
+    }
+  }
+  return validated;
+}
+
+template <sycl::image_channel_order channelOrder,
+          sycl::image_channel_type channelType, int dim,
+          syclexp::image_type type = syclexp::image_type::standard>
+bool run_copy_test_with(sycl::device &dev, sycl::queue &q,
+                        sycl::range<dim> dims) {
+  std::vector<float> dataSequence(dims.size());
+  std::vector<float> out(dims.size());
+
+  std::vector<float> expected(dims.size());
+
+  std::iota(dataSequence.begin(), dataSequence.end(), 0);
+  std::iota(expected.begin(), expected.end(), 0);
+
+  syclexp::image_descriptor desc;
+
+  if constexpr (type == syclexp::image_type::standard) {
+    desc = syclexp::image_descriptor(dims, channelOrder, channelType);
+  } else {
+    desc = syclexp::image_descriptor(
+        {dims[0], dim > 2 ? dims[1] : 0}, channelOrder, channelType,
+        syclexp::image_type::array, 1, dim > 2 ? dims[2] : dims[1]);
+  }
+
+  copy_image_mem_handle_to_image_mem_handle(desc, dataSequence, dev, q, out);
+
+  return check_test(out, expected);
+}
+
+int main() {
+
+  sycl::device dev;
+  sycl::queue q(dev);
+  auto ctxt = q.get_context();
+
+  // Standard images copies
+  bool validated = run_copy_test_with<sycl::image_channel_order::r,
+                                      sycl::image_channel_type::fp32, 2>(
+      dev, q, {2048 * 8, 2048 * 8});
+
+  validated &= run_copy_test_with<sycl::image_channel_order::r,
+                                  sycl::image_channel_type::fp32, 1>(
+      dev, q, {512 * 256});
+
+  validated &= run_copy_test_with<sycl::image_channel_order::r,
+                                  sycl::image_channel_type::fp32, 3>(
+      dev, q, {2048, 2048, 64});
+
+  // Layered images copies
+  validated &=
+      run_copy_test_with<sycl::image_channel_order::r,
+                         sycl::image_channel_type::fp32, 2,
+                         syclexp::image_type::array>(dev, q, {956, 38});
+  validated &=
+      run_copy_test_with<sycl::image_channel_order::r,
+                         sycl::image_channel_type::fp32, 3,
+                         syclexp::image_type::array>(dev, q, {2048, 2048, 64});
+
+  if (!validated) {
+    std::cout << "Tests failed";
+    return 1;
+  }
+
+  std::cout << "Tests passed";
+
+  return 0;
+}
diff --git a/sycl/test/abi/sycl_symbols_linux.dump b/sycl/test/abi/sycl_symbols_linux.dump
@@ -3577,6 +3577,7 @@ _ZN4sycl3_V17handler13getKernelNameEv
 _ZN4sycl3_V17handler14setNDRangeUsedEb
 _ZN4sycl3_V17handler15ext_oneapi_copyENS0_3ext6oneapi12experimental16image_mem_handleENS0_5rangeILi3EEERKNS4_16image_descriptorEPvS7_S7_S7_
 _ZN4sycl3_V17handler15ext_oneapi_copyENS0_3ext6oneapi12experimental16image_mem_handleEPvRKNS4_16image_descriptorE
+_ZN4sycl3_V17handler15ext_oneapi_copyENS0_3ext6oneapi12experimental16image_mem_handleES5_RKNS4_16image_descriptorE
 _ZN4sycl3_V17handler15ext_oneapi_copyEPvNS0_3ext6oneapi12experimental16image_mem_handleERKNS5_16image_descriptorE
 _ZN4sycl3_V17handler15ext_oneapi_copyEPvNS0_5rangeILi3EEES2_S4_RKNS0_3ext6oneapi12experimental16image_descriptorEmS4_S4_
 _ZN4sycl3_V17handler15ext_oneapi_copyEPvNS0_5rangeILi3EEES4_NS0_3ext6oneapi12experimental16image_mem_handleES4_RKNS7_16image_descriptorES4_
diff --git a/sycl/test/abi/sycl_symbols_windows.dump b/sycl/test/abi/sycl_symbols_windows.dump
@@ -4067,6 +4067,7 @@
 ?ext_oneapi_copy@handler@_V1@sycl@@QEAAXPEAXUimage_mem_handle@experimental@oneapi@ext@23@AEBUimage_descriptor@56723@@Z
 ?ext_oneapi_copy@handler@_V1@sycl@@QEAAXPEAXV?$range@$02@23@01AEBUimage_descriptor@experimental@oneapi@ext@23@_K11@Z
 ?ext_oneapi_copy@handler@_V1@sycl@@QEAAXPEAXV?$range@$02@23@1Uimage_mem_handle@experimental@oneapi@ext@23@1AEBUimage_descriptor@67823@1@Z
+?ext_oneapi_copy@handler@_V1@sycl@@QEAAXUimage_mem_handle@experimental@oneapi@ext@23@0AEBUimage_descriptor@56723@@Z
 ?ext_oneapi_copy@handler@_V1@sycl@@QEAAXUimage_mem_handle@experimental@oneapi@ext@23@PEAXAEBUimage_descriptor@56723@@Z
 ?ext_oneapi_copy@handler@_V1@sycl@@QEAAXUimage_mem_handle@experimental@oneapi@ext@23@V?$range@$02@23@AEBUimage_descriptor@56723@PEAX111@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@PEAX0AEBUimage_descriptor@experimental@oneapi@ext@23@_KAEBUcode_location@detail@23@@Z
@@ -4081,6 +4082,9 @@
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@PEAXV?$range@$02@23@1Uimage_mem_handle@experimental@oneapi@ext@23@1AEBUimage_descriptor@78923@1AEBUcode_location@detail@23@@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@PEAXV?$range@$02@23@1Uimage_mem_handle@experimental@oneapi@ext@23@1AEBUimage_descriptor@78923@1AEBV?$vector@Vevent@_V1@sycl@@V?$allocator@Vevent@_V1@sycl@@@std@@@std@@AEBUcode_location@detail@23@@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@PEAXV?$range@$02@23@1Uimage_mem_handle@experimental@oneapi@ext@23@1AEBUimage_descriptor@78923@1V423@AEBUcode_location@detail@23@@Z
+?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@0AEBUimage_descriptor@67823@AEBUcode_location@detail@23@@Z
+?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@0AEBUimage_descriptor@67823@AEBV?$vector@Vevent@_V1@sycl@@V?$allocator@Vevent@_V1@sycl@@@std@@@std@@AEBUcode_location@detail@23@@Z
+?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@0AEBUimage_descriptor@67823@V423@AEBUcode_location@detail@23@@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@PEAXAEBUimage_descriptor@67823@AEBUcode_location@detail@23@@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@PEAXAEBUimage_descriptor@67823@AEBV?$vector@Vevent@_V1@sycl@@V?$allocator@Vevent@_V1@sycl@@@std@@@std@@AEBUcode_location@detail@23@@Z
 ?ext_oneapi_copy@queue@_V1@sycl@@QEAA?AVevent@23@Uimage_mem_handle@experimental@oneapi@ext@23@PEAXAEBUimage_descriptor@67823@V423@AEBUcode_location@detail@23@@Z
