[UR][Offload] Implement buffer map/unmap (#19054)

Implement buffer map/unmap for Offload. Pretty much just a copy of the
CUDA/HIP adapters implementation, with some workarounds for current
offload limitations.

Author: callumfare

unified-runtime/source/adapters/offload/enqueue.cpp

@@ -165,3 +165,85 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferWrite(
 
   return UR_RESULT_SUCCESS;
 }
+
+ur_result_t enqueueNoOp(ur_queue_handle_t hQueue, ur_event_handle_t *phEvent) {
+  // This path is a no-op, but we can't output a real event because
+  // Offload doesn't currently support creating arbitrary events, and we
+  // don't know the last real event in the queue. Instead we just have to
+  // wait on the whole queue and then return an empty (implicitly
+  // finished) event.
+  *phEvent = ur_event_handle_t_::createEmptyEvent();
+  return urQueueFinish(hQueue);
+}
+
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferMap(
+    ur_queue_handle_t hQueue, ur_mem_handle_t hBuffer, bool blockingMap,
+    ur_map_flags_t mapFlags, size_t offset, size_t size,
+    uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+    ur_event_handle_t *phEvent, void **ppRetMap) {
+
+  auto &BufferImpl = std::get<BufferMem>(hBuffer->Mem);
+  auto MapPtr = BufferImpl.mapToPtr(size, offset, mapFlags);
+
+  if (!MapPtr) {
+    return UR_RESULT_ERROR_INVALID_MEM_OBJECT;
+  }
+
+  const bool IsPinned =
+      BufferImpl.MemAllocMode == BufferMem::AllocMode::AllocHostPtr;
+
+  ur_result_t Result = UR_RESULT_SUCCESS;
+  if (!IsPinned &&
+      ((mapFlags & UR_MAP_FLAG_READ) || (mapFlags & UR_MAP_FLAG_WRITE))) {
+    // Pinned host memory is already on host so it doesn't need to be read.
+    Result = urEnqueueMemBufferRead(hQueue, hBuffer, blockingMap, offset, size,
+                                    MapPtr, numEventsInWaitList,
+                                    phEventWaitList, phEvent);
+  } else {
+    if (IsPinned) {
+      // TODO: Ignore the event waits list for now. When urEnqueueEventsWait is
+      // implemented we can call it on the wait list.
+    }
+
+    if (phEvent) {
+      enqueueNoOp(hQueue, phEvent);
+    }
+  }
+  *ppRetMap = MapPtr;
+
+  return Result;
+}
+
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemUnmap(
+    ur_queue_handle_t hQueue, ur_mem_handle_t hMem, void *pMappedPtr,
+    uint32_t numEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+    ur_event_handle_t *phEvent) {
+  auto &BufferImpl = std::get<BufferMem>(hMem->Mem);
+
+  auto *Map = BufferImpl.getMapDetails(pMappedPtr);
+  UR_ASSERT(Map != nullptr, UR_RESULT_ERROR_INVALID_MEM_OBJECT);
+
+  const bool IsPinned =
+      BufferImpl.MemAllocMode == BufferMem::AllocMode::AllocHostPtr;
+
+  ur_result_t Result = UR_RESULT_SUCCESS;
+  if (!IsPinned && ((Map->MapFlags & UR_MAP_FLAG_WRITE) ||
+                    (Map->MapFlags & UR_MAP_FLAG_WRITE_INVALIDATE_REGION))) {
+    // Pinned host memory is only on host so it doesn't need to be written to.
+    Result = urEnqueueMemBufferWrite(
+        hQueue, hMem, true, Map->MapOffset, Map->MapSize, pMappedPtr,
+        numEventsInWaitList, phEventWaitList, phEvent);
+  } else {
+    if (IsPinned) {
+      // TODO: Ignore the event waits list for now. When urEnqueueEventsWait is
+      // implemented we can call it on the wait list.
+    }
+
+    if (phEvent) {
+      enqueueNoOp(hQueue, phEvent);
+    }
+  }
+  BufferImpl.unmap(pMappedPtr);
+
+  return Result;
+}

unified-runtime/source/adapters/offload/event.cpp

@@ -14,7 +14,7 @@
 #include "event.hpp"
 #include "ur2offload.hpp"
 
-UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hKernel,
+UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hEvent,
                                                    ur_event_info_t propName,
                                                    size_t propSize,
                                                    void *pPropValue,
@@ -23,7 +23,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hKernel,
 
   switch (propName) {
   case UR_EVENT_INFO_REFERENCE_COUNT:
-    return ReturnValue(hKernel->RefCount.load());
+    return ReturnValue(hEvent->RefCount.load());
   default:
     return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
   }
@@ -42,9 +42,11 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventGetProfilingInfo(ur_event_handle_t,
 UR_APIEXPORT ur_result_t UR_APICALL
 urEventWait(uint32_t numEvents, const ur_event_handle_t *phEventWaitList) {
   for (uint32_t i = 0; i < numEvents; i++) {
-    auto Res = olWaitEvent(phEventWaitList[i]->OffloadEvent);
-    if (Res) {
-      return offloadResultToUR(Res);
+    if (phEventWaitList[i]->OffloadEvent) {
+      auto Res = olWaitEvent(phEventWaitList[i]->OffloadEvent);
+      if (Res) {
+        return offloadResultToUR(Res);
+      }
     }
   }
   return UR_RESULT_SUCCESS;

unified-runtime/source/adapters/offload/event.hpp

@@ -17,4 +17,13 @@
 
 struct ur_event_handle_t_ : RefCounted {
   ol_event_handle_t OffloadEvent;
+  ur_command_t Type;
+
+  static ur_event_handle_t createEmptyEvent() {
+    auto *Event = new ur_event_handle_t_();
+    // Null event represents an empty event. Waiting on it is a no-op.
+    Event->OffloadEvent = nullptr;
+
+    return Event;
+  }
 };

unified-runtime/source/adapters/offload/memory.hpp

@@ -22,6 +22,23 @@ struct BufferMem {
     AllocHostPtr,
   };
 
+  struct BufferMap {
+    size_t MapSize;
+    size_t MapOffset;
+    ur_map_flags_t MapFlags;
+    // Allocated host memory used exclusively for this map.
+    std::unique_ptr<unsigned char[]> MapMem;
+
+    BufferMap(size_t MapSize, size_t MapOffset, ur_map_flags_t MapFlags)
+        : MapSize(MapSize), MapOffset(MapOffset), MapFlags(MapFlags),
+          MapMem(nullptr) {}
+
+    BufferMap(size_t MapSize, size_t MapOffset, ur_map_flags_t MapFlags,
+              std::unique_ptr<unsigned char[]> &&MapMem)
+        : MapSize(MapSize), MapOffset(MapOffset), MapFlags(MapFlags),
+          MapMem(std::move(MapMem)) {}
+  };
+
   ur_mem_handle_t Parent;
   // Underlying device pointer
   void *Ptr;
@@ -30,6 +47,7 @@ struct BufferMem {
   size_t Size;
 
   AllocMode MemAllocMode;
+  std::unordered_map<void *, BufferMap> PtrToBufferMap;
 
   BufferMem(ur_mem_handle_t Parent, BufferMem::AllocMode Mode, void *Ptr,
             void *HostPtr, size_t Size)
@@ -38,6 +56,37 @@ struct BufferMem {
 
   void *get() const noexcept { return Ptr; }
   size_t getSize() const noexcept { return Size; }
+
+  BufferMap *getMapDetails(void *Map) {
+    auto Details = PtrToBufferMap.find(Map);
+    if (Details != PtrToBufferMap.end()) {
+      return &Details->second;
+    }
+    return nullptr;
+  }
+
+  void *mapToPtr(size_t MapSize, size_t MapOffset,
+                 ur_map_flags_t MapFlags) noexcept {
+
+    void *MapPtr = nullptr;
+    // If the buffer already has a host pointer we can just use it, otherwise
+    // create a new host allocation
+    if (HostPtr == nullptr) {
+      auto MapMem = std::make_unique<unsigned char[]>(MapSize);
+      MapPtr = MapMem.get();
+      PtrToBufferMap.insert(
+          {MapPtr, BufferMap(MapSize, MapOffset, MapFlags, std::move(MapMem))});
+    } else {
+      MapPtr = static_cast<char *>(HostPtr) + MapOffset;
+      PtrToBufferMap.insert({MapPtr, BufferMap(MapSize, MapOffset, MapFlags)});
+    }
+    return MapPtr;
+  }
+
+  void unmap(void *MapPtr) noexcept {
+    assert(MapPtr != nullptr);
+    PtrToBufferMap.erase(MapPtr);
+  }
 };
 
 struct ur_mem_handle_t_ : RefCounted {

unified-runtime/source/adapters/offload/ur_interface_loader.cpp

@@ -176,15 +176,15 @@ UR_DLLEXPORT ur_result_t UR_APICALL urGetEnqueueProcAddrTable(
   pDdiTable->pfnMemBufferCopy = nullptr;
   pDdiTable->pfnMemBufferCopyRect = nullptr;
   pDdiTable->pfnMemBufferFill = nullptr;
-  pDdiTable->pfnMemBufferMap = nullptr;
+  pDdiTable->pfnMemBufferMap = urEnqueueMemBufferMap;
   pDdiTable->pfnMemBufferRead = urEnqueueMemBufferRead;
   pDdiTable->pfnMemBufferReadRect = nullptr;
   pDdiTable->pfnMemBufferWrite = urEnqueueMemBufferWrite;
   pDdiTable->pfnMemBufferWriteRect = nullptr;
   pDdiTable->pfnMemImageCopy = nullptr;
   pDdiTable->pfnMemImageRead = nullptr;
   pDdiTable->pfnMemImageWrite = nullptr;
-  pDdiTable->pfnMemUnmap = nullptr;
+  pDdiTable->pfnMemUnmap = urEnqueueMemUnmap;
   pDdiTable->pfnUSMFill2D = urEnqueueUSMFill2D;
   pDdiTable->pfnUSMFill = nullptr;
   pDdiTable->pfnUSMAdvise = nullptr;