[SYCL][CUDA] Map/unmap pinned host memory (#2098)

Changes map on pinned host memory to return a pointer to the pinned memory rather than copy to new memory. Likewise, unmap of pinned memory does no write operation.

Signed-off-by: Steffen Larsen <[email protected]>

Author: steffenlarsen

sycl/plugins/cuda/pi_cuda.cpp

@@ -3975,6 +3975,8 @@ pi_result cuda_piEnqueueMemImageFill(pi_queue command_queue, pi_mem image,
 
 /// Implements mapping on the host using a BufferRead operation.
 /// Mapped pointers are stored in the pi_mem object.
+/// If the buffer uses pinned host memory a pointer to that memory is returned
+/// and no read operation is done.
 /// \TODO Untie types from OpenCL
 ///
 pi_result cuda_piEnqueueMemBufferMap(pi_queue command_queue, pi_mem buffer,
@@ -3984,11 +3986,14 @@ pi_result cuda_piEnqueueMemBufferMap(pi_queue command_queue, pi_mem buffer,
                                      pi_uint32 num_events_in_wait_list,
                                      const pi_event *event_wait_list,
                                      pi_event *event, void **ret_map) {
-
   assert(ret_map != nullptr);
   assert(command_queue != nullptr);
+  assert(buffer != nullptr);
+  assert(buffer->mem_type_ == _pi_mem::mem_type::buffer);
 
   pi_result ret_err = PI_INVALID_OPERATION;
+  const bool is_pinned = buffer->mem_.buffer_mem_.allocMode_ ==
+                         _pi_mem::mem_::buffer_mem_::alloc_mode::alloc_host_ptr;
 
   // Currently no support for overlapping regions
   if (buffer->mem_.buffer_mem_.get_map_ptr() != nullptr) {
@@ -4002,15 +4007,21 @@ pi_result cuda_piEnqueueMemBufferMap(pi_queue command_queue, pi_mem buffer,
     ret_err = PI_SUCCESS;
   }
 
-  if ((map_flags & CL_MAP_READ) || (map_flags & CL_MAP_WRITE)) {
+  if (!is_pinned && ((map_flags & CL_MAP_READ) || (map_flags & CL_MAP_WRITE))) {
+    // Pinned host memory is already on host so it doesn't need to be read.
     ret_err = cuda_piEnqueueMemBufferRead(
         command_queue, buffer, blocking_map, offset, size, hostPtr,
         num_events_in_wait_list, event_wait_list, event);
   } else {
+    ScopedContext active(command_queue->get_context());
+
+    if (is_pinned) {
+      ret_err = cuda_piEnqueueEventsWait(command_queue, num_events_in_wait_list,
+                                         event_wait_list, nullptr);
+    }
+
     if (event) {
       try {
-        ScopedContext active(command_queue->get_context());
-
         *event = _pi_event::make_native(PI_COMMAND_TYPE_MEM_BUFFER_MAP,
                                         command_queue);
         (*event)->start();
@@ -4026,6 +4037,7 @@ pi_result cuda_piEnqueueMemBufferMap(pi_queue command_queue, pi_mem buffer,
 
 /// Implements the unmap from the host, using a BufferWrite operation.
 /// Requires the mapped pointer to be already registered in the given memobj.
+/// If memobj uses pinned host memory, this will not do a write.
 ///
 pi_result cuda_piEnqueueMemUnmap(pi_queue command_queue, pi_mem memobj,
                                  void *mapped_ptr,
@@ -4037,22 +4049,33 @@ pi_result cuda_piEnqueueMemUnmap(pi_queue command_queue, pi_mem memobj,
   assert(command_queue != nullptr);
   assert(mapped_ptr != nullptr);
   assert(memobj != nullptr);
+  assert(memobj->mem_type_ == _pi_mem::mem_type::buffer);
   assert(memobj->mem_.buffer_mem_.get_map_ptr() != nullptr);
   assert(memobj->mem_.buffer_mem_.get_map_ptr() == mapped_ptr);
 
-  if ((memobj->mem_.buffer_mem_.get_map_flags() & CL_MAP_WRITE) ||
-      (memobj->mem_.buffer_mem_.get_map_flags() &
-       CL_MAP_WRITE_INVALIDATE_REGION)) {
+  const bool is_pinned = memobj->mem_.buffer_mem_.allocMode_ ==
+                         _pi_mem::mem_::buffer_mem_::alloc_mode::alloc_host_ptr;
+
+  if (!is_pinned &&
+      ((memobj->mem_.buffer_mem_.get_map_flags() & CL_MAP_WRITE) ||
+       (memobj->mem_.buffer_mem_.get_map_flags() &
+        CL_MAP_WRITE_INVALIDATE_REGION))) {
+    // Pinned host memory is only on host so it doesn't need to be written to.
     ret_err = cuda_piEnqueueMemBufferWrite(
         command_queue, memobj, true,
         memobj->mem_.buffer_mem_.get_map_offset(mapped_ptr),
         memobj->mem_.buffer_mem_.get_size(), mapped_ptr,
         num_events_in_wait_list, event_wait_list, event);
   } else {
+    ScopedContext active(command_queue->get_context());
+
+    if (is_pinned) {
+      ret_err = cuda_piEnqueueEventsWait(command_queue, num_events_in_wait_list,
+                                         event_wait_list, nullptr);
+    }
+
     if (event) {
       try {
-        ScopedContext active(command_queue->get_context());
-
         *event = _pi_event::make_native(PI_COMMAND_TYPE_MEM_BUFFER_UNMAP,
                                         command_queue);
         (*event)->start();

sycl/unittests/pi/cuda/test_mem_obj.cpp

@@ -113,3 +113,84 @@ TEST_F(CudaTestMemObj, piMemBufferCreateNoActiveContext) {
   ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piMemRelease>(memObj)),
             PI_SUCCESS);
 }
+
+TEST_F(CudaTestMemObj, piMemBufferPinnedMappedRead) {
+  const size_t memSize = sizeof(int);
+  const int value = 20;
+
+  pi_queue queue;
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piQueueCreate>(
+                context_, device_, 0, &queue)),
+            PI_SUCCESS);
+  ASSERT_NE(queue, nullptr);
+  ASSERT_EQ(queue->get_context(), context_);
+
+  pi_mem memObj;
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piMemBufferCreate>(
+                context_, PI_MEM_FLAGS_ACCESS_RW | PI_MEM_FLAGS_HOST_PTR_ALLOC,
+                memSize, nullptr, &memObj)),
+            PI_SUCCESS);
+
+  ASSERT_EQ(
+      (plugin.call_nocheck<detail::PiApiKind::piEnqueueMemBufferWrite>(
+          queue, memObj, true, 0, sizeof(int), &value, 0, nullptr, nullptr)),
+      PI_SUCCESS);
+
+  int *host_ptr = nullptr;
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piEnqueueMemBufferMap>(
+                queue, memObj, true, CL_MAP_READ, 0, sizeof(int), 0, nullptr,
+                nullptr, (void **)&host_ptr)),
+            PI_SUCCESS);
+
+  ASSERT_EQ(*host_ptr, value);
+
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piEnqueueMemUnmap>(
+                queue, memObj, host_ptr, 0, nullptr, nullptr)),
+            PI_SUCCESS);
+
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piMemRelease>(memObj)),
+            PI_SUCCESS);
+  plugin.call<detail::PiApiKind::piQueueRelease>(queue);
+}
+
+TEST_F(CudaTestMemObj, piMemBufferPinnedMappedWrite) {
+  const size_t memSize = sizeof(int);
+  const int value = 30;
+
+  pi_queue queue;
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piQueueCreate>(
+                context_, device_, 0, &queue)),
+            PI_SUCCESS);
+  ASSERT_NE(queue, nullptr);
+  ASSERT_EQ(queue->get_context(), context_);
+
+  pi_mem memObj;
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piMemBufferCreate>(
+                context_, PI_MEM_FLAGS_ACCESS_RW | PI_MEM_FLAGS_HOST_PTR_ALLOC,
+                memSize, nullptr, &memObj)),
+            PI_SUCCESS);
+
+  int *host_ptr = nullptr;
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piEnqueueMemBufferMap>(
+                queue, memObj, true, CL_MAP_WRITE, 0, sizeof(int), 0, nullptr,
+                nullptr, (void **)&host_ptr)),
+            PI_SUCCESS);
+
+  *host_ptr = value;
+
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piEnqueueMemUnmap>(
+                queue, memObj, host_ptr, 0, nullptr, nullptr)),
+            PI_SUCCESS);
+
+  int read_value = 0;
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piEnqueueMemBufferRead>(
+                queue, memObj, true, 0, sizeof(int), &read_value, 0, nullptr,
+                nullptr)),
+            PI_SUCCESS);
+
+  ASSERT_EQ(read_value, value);
+
+  ASSERT_EQ((plugin.call_nocheck<detail::PiApiKind::piMemRelease>(memObj)),
+            PI_SUCCESS);
+  plugin.call<detail::PiApiKind::piQueueRelease>(queue);
+}