[SYCL] Implement sycl_ext_oneapi_memcpy2d extension (#7370)

This commit adds an implementation of the sycl_ext_oneapi_memcpy2d
extension. This includes the following:

* Three new PI API functions; piextUSMEnqueueFill2D,
piextUSMEnqueueMemset2D, and piextUSMEnqueueMemcpy2D.
* Three new PI context queries;
PI_EXT_ONEAPI_CONTEXT_INFO_USM_MEMCPY2D_SUPPORT,
PI_EXT_ONEAPI_CONTEXT_INFO_USM_FILL2D_SUPPORT, and
PI_EXT_ONEAPI_CONTEXT_INFO_USM_MEMSET2D_SUPPORT. Each of these return a
boolean value signifying if the PI context supports each of the new API
functions.
* New handler member functions; ext_oneapi_memcpy2d, ext_oneapi_copy2d,
ext_oneapi_memset2d, and ext_oneapi_fill2d. These will create new
commands which will use the new PI commands if they are supported. If
the new PI API functions are not supported, these member functions will
launch auxiliary kernels to do the required work.
* New queue shortcuts for each of the new handler members.

Signed-off-by: Larsen, Steffen <[email protected]>
Co-authored-by: Petr Vesely <[email protected]>

Author: steffenlarsen

sycl/include/sycl/detail/cg.hpp

@@ -69,6 +69,9 @@ class CG {
     CodeplayInteropTask = 13,
     CodeplayHostTask = 14,
     AdviseUSM = 15,
+    Copy2DUSM = 16,
+    Fill2DUSM = 17,
+    Memset2DUSM = 18,
   };
 
   CG(CGTYPE Type, std::vector<std::vector<char>> ArgsStorage,
@@ -394,6 +397,95 @@ class CGBarrier : public CG {
         MEventsWaitWithBarrier(std::move(EventsWaitWithBarrier)) {}
 };
 
+/// "Copy 2D USM" command group class.
+class CGCopy2DUSM : public CG {
+  void *MSrc;
+  void *MDst;
+  size_t MSrcPitch;
+  size_t MDstPitch;
+  size_t MWidth;
+  size_t MHeight;
+
+public:
+  CGCopy2DUSM(void *Src, void *Dst, size_t SrcPitch, size_t DstPitch,
+              size_t Width, size_t Height,
+              std::vector<std::vector<char>> ArgsStorage,
+              std::vector<detail::AccessorImplPtr> AccStorage,
+              std::vector<std::shared_ptr<const void>> SharedPtrStorage,
+              std::vector<AccessorImplHost *> Requirements,
+              std::vector<detail::EventImplPtr> Events,
+              detail::code_location loc = {})
+      : CG(Copy2DUSM, std::move(ArgsStorage), std::move(AccStorage),
+           std::move(SharedPtrStorage), std::move(Requirements),
+           std::move(Events), std::move(loc)),
+        MSrc(Src), MDst(Dst), MSrcPitch(SrcPitch), MDstPitch(DstPitch),
+        MWidth(Width), MHeight(Height) {}
+
+  void *getSrc() const { return MSrc; }
+  void *getDst() const { return MDst; }
+  size_t getSrcPitch() const { return MSrcPitch; }
+  size_t getDstPitch() const { return MDstPitch; }
+  size_t getWidth() const { return MWidth; }
+  size_t getHeight() const { return MHeight; }
+};
+
+/// "Fill 2D USM" command group class.
+class CGFill2DUSM : public CG {
+  std::vector<char> MPattern;
+  void *MDst;
+  size_t MPitch;
+  size_t MWidth;
+  size_t MHeight;
+
+public:
+  CGFill2DUSM(std::vector<char> Pattern, void *DstPtr, size_t Pitch,
+              size_t Width, size_t Height,
+              std::vector<std::vector<char>> ArgsStorage,
+              std::vector<detail::AccessorImplPtr> AccStorage,
+              std::vector<std::shared_ptr<const void>> SharedPtrStorage,
+              std::vector<AccessorImplHost *> Requirements,
+              std::vector<detail::EventImplPtr> Events,
+              detail::code_location loc = {})
+      : CG(Fill2DUSM, std::move(ArgsStorage), std::move(AccStorage),
+           std::move(SharedPtrStorage), std::move(Requirements),
+           std::move(Events), std::move(loc)),
+        MPattern(std::move(Pattern)), MDst(DstPtr), MPitch(Pitch),
+        MWidth(Width), MHeight(Height) {}
+  void *getDst() const { return MDst; }
+  size_t getPitch() const { return MPitch; }
+  size_t getWidth() const { return MWidth; }
+  size_t getHeight() const { return MHeight; }
+  const std::vector<char> &getPattern() const { return MPattern; }
+};
+
+/// "Memset 2D USM" command group class.
+class CGMemset2DUSM : public CG {
+  char MValue;
+  void *MDst;
+  size_t MPitch;
+  size_t MWidth;
+  size_t MHeight;
+
+public:
+  CGMemset2DUSM(char Value, void *DstPtr, size_t Pitch, size_t Width,
+                size_t Height, std::vector<std::vector<char>> ArgsStorage,
+                std::vector<detail::AccessorImplPtr> AccStorage,
+                std::vector<std::shared_ptr<const void>> SharedPtrStorage,
+                std::vector<AccessorImplHost *> Requirements,
+                std::vector<detail::EventImplPtr> Events,
+                detail::code_location loc = {})
+      : CG(Memset2DUSM, std::move(ArgsStorage), std::move(AccStorage),
+           std::move(SharedPtrStorage), std::move(Requirements),
+           std::move(Events), std::move(loc)),
+        MValue(Value), MDst(DstPtr), MPitch(Pitch), MWidth(Width),
+        MHeight(Height) {}
+  void *getDst() const { return MDst; }
+  size_t getPitch() const { return MPitch; }
+  size_t getWidth() const { return MWidth; }
+  size_t getHeight() const { return MHeight; }
+  char getValue() const { return MValue; }
+};
+
 } // namespace detail
 } // __SYCL_INLINE_VER_NAMESPACE(_V1)
 } // namespace sycl

sycl/include/sycl/detail/pi.def

@@ -141,5 +141,8 @@ _PI_API(piPluginGetLastError)
 
 _PI_API(piTearDown)
 
+_PI_API(piextUSMEnqueueFill2D)
+_PI_API(piextUSMEnqueueMemset2D)
+_PI_API(piextUSMEnqueueMemcpy2D)
 
 #undef _PI_API

sycl/include/sycl/detail/pi.h

@@ -68,9 +68,15 @@
 // 12.20 Added piextQueueCreate API to be used instead of piQueueCreate, also
 // added PI_EXT_INTEL_DEVICE_INFO_MAX_COMPUTE_QUEUE_INDICES for piDeviceGetInfo.
 // Both are needed to support sycl_ext_intel_queue_index extension.
+// 12.21 Added new piextUSMEnqueueFill2D, piextUSMEnqueueMemset2D, and
+// piextUSMEnqueueMemcpy2D functions. Added new
+// PI_EXT_ONEAPI_CONTEXT_INFO_USM_FILL2D_SUPPORT,
+// PI_EXT_ONEAPI_CONTEXT_INFO_USM_MEMSET2D_SUPPORT, and
+// PI_EXT_ONEAPI_CONTEXT_INFO_USM_MEMCPY2D_SUPPORT context info query
+// descriptors.
 
 #define _PI_H_VERSION_MAJOR 12
-#define _PI_H_VERSION_MINOR 20
+#define _PI_H_VERSION_MINOR 21
 
 #define _PI_STRING_HELPER(a) #a
 #define _PI_CONCAT(a, b) _PI_STRING_HELPER(a.b)
@@ -335,7 +341,11 @@ typedef enum {
   PI_CONTEXT_INFO_REFERENCE_COUNT = 0x1080,
   // Atomics capabilities extensions
   PI_CONTEXT_INFO_ATOMIC_MEMORY_ORDER_CAPABILITIES = 0x10010,
-  PI_CONTEXT_INFO_ATOMIC_MEMORY_SCOPE_CAPABILITIES = 0x10011
+  PI_CONTEXT_INFO_ATOMIC_MEMORY_SCOPE_CAPABILITIES = 0x10011,
+  // Native 2D USM memory operation support
+  PI_EXT_ONEAPI_CONTEXT_INFO_USM_FILL2D_SUPPORT = 0x30000,
+  PI_EXT_ONEAPI_CONTEXT_INFO_USM_MEMSET2D_SUPPORT = 0x30001,
+  PI_EXT_ONEAPI_CONTEXT_INFO_USM_MEMCPY2D_SUPPORT = 0x30002
 } _pi_context_info;
 
 typedef enum {
@@ -1809,6 +1819,62 @@ __SYCL_EXPORT pi_result piextUSMGetMemAllocInfo(
     pi_context context, const void *ptr, pi_mem_alloc_info param_name,
     size_t param_value_size, void *param_value, size_t *param_value_size_ret);
 
+/// USM 2D fill API
+///
+/// \param queue is the queue to submit to
+/// \param ptr is the ptr to fill
+/// \param pitch is the total width of the destination memory including padding
+/// \param pattern is a pointer with the bytes of the pattern to set
+/// \param pattern_size is the size in bytes of the pattern
+/// \param width is width in bytes of each row to fill
+/// \param height is height the columns to fill
+/// \param num_events_in_waitlist is the number of events to wait on
+/// \param events_waitlist is an array of events to wait on
+/// \param event is the event that represents this operation
+__SYCL_EXPORT pi_result piextUSMEnqueueFill2D(pi_queue queue, void *ptr,
+                                              size_t pitch, size_t pattern_size,
+                                              const void *pattern, size_t width,
+                                              size_t height,
+                                              pi_uint32 num_events_in_waitlist,
+                                              const pi_event *events_waitlist,
+                                              pi_event *event);
+
+/// USM 2D Memset API
+///
+/// \param queue is the queue to submit to
+/// \param ptr is the ptr to fill
+/// \param pitch is the total width of the destination memory including padding
+/// \param value the value to fill into the region in \param ptr
+/// \param width is width in bytes of each row to fill
+/// \param height is height the columns to fill
+/// \param num_events_in_waitlist is the number of events to wait on
+/// \param events_waitlist is an array of events to wait on
+/// \param event is the event that represents this operation
+__SYCL_EXPORT pi_result piextUSMEnqueueMemset2D(
+    pi_queue queue, void *ptr, size_t pitch, int value, size_t width,
+    size_t height, pi_uint32 num_events_in_waitlist,
+    const pi_event *events_waitlist, pi_event *event);
+
+/// USM 2D Memcpy API
+///
+/// \param queue is the queue to submit to
+/// \param blocking is whether this operation should block the host
+/// \param dst_ptr is the location the data will be copied
+/// \param dst_pitch is the total width of the destination memory including
+/// padding
+/// \param src_ptr is the data to be copied
+/// \param src_pitch is the total width of the source memory including padding
+/// \param width is width in bytes of each row to be copied
+/// \param height is height the columns to be copied
+/// \param num_events_in_waitlist is the number of events to wait on
+/// \param events_waitlist is an array of events to wait on
+/// \param event is the event that represents this operation
+__SYCL_EXPORT pi_result piextUSMEnqueueMemcpy2D(
+    pi_queue queue, pi_bool blocking, void *dst_ptr, size_t dst_pitch,
+    const void *src_ptr, size_t src_pitch, size_t width, size_t height,
+    pi_uint32 num_events_in_waitlist, const pi_event *events_waitlist,
+    pi_event *event);
+
 /// API to get Plugin internal data, opaque to SYCL RT. Some devices whose
 /// device code is compiled by the host compiler (e.g. CPU emulators) may use it
 /// to access some device code functionality implemented in/behind the plugin.

sycl/include/sycl/detail/pi.hpp

@@ -112,6 +112,7 @@ using PiDeviceType = ::pi_device_type;
 using PiDeviceInfo = ::pi_device_info;
 using PiDeviceBinaryType = ::pi_device_binary_type;
 using PiContext = ::pi_context;
+using PiContextInfo = ::pi_context_info;
 using PiProgram = ::pi_program;
 using PiKernel = ::pi_kernel;
 using PiQueue = ::pi_queue;

sycl/include/sycl/feature_test.hpp.in

@@ -71,6 +71,7 @@ __SYCL_INLINE_VER_NAMESPACE(_V1) {
 #define SYCL_EXT_ONEAPI_QUEUE_EMPTY 1
 #define SYCL_EXT_ONEAPI_USER_DEFINED_REDUCTIONS 1
 #define SYCL_EXT_ONEAPI_WEAK_OBJECT 1
+#define SYCL_EXT_ONEAPI_MEMCPY2D 1
 #cmakedefine01 SYCL_BUILD_PI_CUDA
 #if SYCL_BUILD_PI_CUDA
 #define SYCL_EXT_ONEAPI_BACKEND_CUDA 1