[UR] Fix buffer allocation logic

unified-runtime/source/adapters/cuda/memory.cpp

@@ -19,28 +19,17 @@
 
 /// Creates a UR Memory object using a CUDA memory allocation.
 /// Can trigger a manual copy depending on the mode.
-/// \TODO Implement USE_HOST_PTR using cuHostRegister
-/// https://github.com/intel/llvm/issues/9789
-///
 UR_APIEXPORT ur_result_t UR_APICALL urMemBufferCreate(
     ur_context_handle_t hContext, ur_mem_flags_t flags, size_t size,
     const ur_buffer_properties_t *pProperties, ur_mem_handle_t *phBuffer) {
-  // Currently, USE_HOST_PTR is not implemented using host register
-  // since this triggers a weird segfault after program ends.
-  // Setting this constant to true enables testing that behavior.
-  const bool EnableUseHostPtr = false;
-  const bool PerformInitialCopy =
-      (flags & UR_MEM_FLAG_ALLOC_COPY_HOST_POINTER) ||
-      ((flags & UR_MEM_FLAG_USE_HOST_POINTER) && !EnableUseHostPtr);
+  const bool PerformInitialCopy = (flags & UR_MEM_FLAG_ALLOC_COPY_HOST_POINTER);
   ur_mem_handle_t MemObj = nullptr;
 
   try {
     auto HostPtr = pProperties ? pProperties->pHost : nullptr;
     BufferMem::AllocMode AllocMode = BufferMem::AllocMode::Classic;
 
-    if ((flags & UR_MEM_FLAG_USE_HOST_POINTER) && EnableUseHostPtr) {
-      UR_CHECK_ERROR(
-          cuMemHostRegister(HostPtr, size, CU_MEMHOSTREGISTER_DEVICEMAP));
+    if (flags & UR_MEM_FLAG_USE_HOST_POINTER) {
       AllocMode = BufferMem::AllocMode::UseHostPtr;
     } else if (flags & UR_MEM_FLAG_ALLOC_HOST_POINTER) {
       HostPtr = umfPoolMalloc(hContext->MemoryPoolHost, size);
@@ -396,8 +385,12 @@ ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t Mem,
       // Host allocation has already been made
       UR_CHECK_ERROR(cuMemHostGetDevicePointer(&DevPtr, Buffer.HostPtr, 0));
     } else if (Buffer.MemAllocMode == BufferMem::AllocMode::UseHostPtr) {
-      UR_CHECK_ERROR(cuMemHostRegister(Buffer.HostPtr, Buffer.Size,
-                                       CU_MEMHOSTALLOC_DEVICEMAP));
+      auto const result = cuMemHostRegister(Buffer.HostPtr, Buffer.Size,
+                                            CU_MEMHOSTREGISTER_DEVICEMAP);
+      if (result == CUDA_SUCCESS)
+        Buffer.HostPtrRegisteredByUR = true;
+      else if (result != CUDA_ERROR_HOST_MEMORY_ALREADY_REGISTERED)
+        UR_CHECK_ERROR(result);
       UR_CHECK_ERROR(cuMemHostGetDevicePointer(&DevPtr, Buffer.HostPtr, 0));
     } else {
       *(void **)&DevPtr = umfPoolMalloc(hDevice->MemoryPoolDevice, Buffer.Size);

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

@@ -90,14 +90,23 @@ struct BufferMem {
   std::unordered_map<void *, BufferMap> PtrToBufferMap;
 
   AllocMode MemAllocMode;
+  // True if UR performed the cuMemHostRegister call and owns the registration.
+  bool HostPtrRegisteredByUR = false;
 
   BufferMem(ur_context_handle_t Context, ur_mem_handle_t OuterMemStruct,
             AllocMode Mode, void *HostPtr, size_t Size)
       : Ptrs(Context->getDevices().size(), native_type{0}),
         OuterMemStruct{OuterMemStruct}, HostPtr{HostPtr}, Size{Size},
         MemAllocMode{Mode} {};
 
-  BufferMem(const BufferMem &Buffer) = default;
+  BufferMem(const BufferMem &Buffer)
+      : Ptrs{Buffer.Ptrs}, Parent{Buffer.Parent},
+        OuterMemStruct{Buffer.OuterMemStruct}, HostPtr{Buffer.HostPtr},
+        Size{Buffer.Size}, PtrToBufferMap{Buffer.PtrToBufferMap},
+        MemAllocMode{Buffer.MemAllocMode} {
+    // HostPtrRegisteredByUR is intentionally not copied: ownership of the
+    // cuMemHostRegister call belongs to the original BufferMem only.
+  }
 
   native_type getPtrWithOffset(const ur_device_handle_t Device, size_t Offset);
 
@@ -165,7 +174,8 @@ struct BufferMem {
       }
       break;
     case AllocMode::UseHostPtr:
-      UR_CHECK_ERROR(cuMemHostUnregister(HostPtr));
+      if (HostPtrRegisteredByUR)
+        UR_CHECK_ERROR(cuMemHostUnregister(HostPtr));
       break;
     case AllocMode::AllocHostPtr:
       UMF_CHECK_ERROR(umfFree((void *)HostPtr));

unified-runtime/source/adapters/hip/memory.cpp

@@ -87,25 +87,17 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemRelease(ur_mem_handle_t hMem) {
 
 /// Creates a UR Memory object using a HIP memory allocation.
 /// Can trigger a manual copy depending on the mode.
-/// \TODO Implement USE_HOST_PTR using hipHostRegister
-/// https://github.com/intel/llvm/issues/9789
 UR_APIEXPORT ur_result_t UR_APICALL urMemBufferCreate(
     ur_context_handle_t hContext, ur_mem_flags_t flags, size_t size,
     const ur_buffer_properties_t *pProperties, ur_mem_handle_t *phBuffer) {
-  // Currently, USE_HOST_PTR is not implemented using host register
-  // since this triggers a weird segfault after program ends.
-  // Setting this constant to true enables testing that behavior.
-  const bool EnableUseHostPtr = false;
-  const bool PerformInitialCopy =
-      (flags & UR_MEM_FLAG_ALLOC_COPY_HOST_POINTER) ||
-      ((flags & UR_MEM_FLAG_USE_HOST_POINTER) && !EnableUseHostPtr);
+  const bool PerformInitialCopy = (flags & UR_MEM_FLAG_ALLOC_COPY_HOST_POINTER);
   ur_result_t Result = UR_RESULT_SUCCESS;
   ur_mem_handle_t RetMemObj = nullptr;
 
   try {
     auto HostPtr = pProperties ? pProperties->pHost : nullptr;
     BufferMem::AllocMode AllocMode = BufferMem::AllocMode::Classic;
-    if ((flags & UR_MEM_FLAG_USE_HOST_POINTER) && EnableUseHostPtr) {
+    if (flags & UR_MEM_FLAG_USE_HOST_POINTER) {
       AllocMode = BufferMem::AllocMode::UseHostPtr;
     } else if (flags & UR_MEM_FLAG_ALLOC_HOST_POINTER) {
       UR_CHECK_ERROR(hipHostMalloc(&HostPtr, size));
@@ -453,8 +445,12 @@ ur_result_t allocateMemObjOnDeviceIfNeeded(ur_mem_handle_t Mem,
       // Host allocation has already been made
       UR_CHECK_ERROR(hipHostGetDevicePointer(&DevPtr, Buffer.HostPtr, 0));
     } else if (Buffer.MemAllocMode == BufferMem::AllocMode::UseHostPtr) {
-      UR_CHECK_ERROR(
-          hipHostRegister(Buffer.HostPtr, Buffer.Size, hipHostRegisterMapped));
+      auto const result =
+          hipHostRegister(Buffer.HostPtr, Buffer.Size, hipHostRegisterMapped);
+      if (result == hipSuccess)
+        Buffer.HostPtrRegisteredByUR = true;
+      else if (result != hipErrorHostMemoryAlreadyRegistered)
+        UR_CHECK_ERROR(result);
       UR_CHECK_ERROR(hipHostGetDevicePointer(&DevPtr, Buffer.HostPtr, 0));
     } else {
       UR_CHECK_ERROR(hipMalloc(&DevPtr, Buffer.Size));

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

@@ -77,6 +77,8 @@ struct BufferMem {
   std::unordered_map<void *, BufferMap> PtrToBufferMap;
 
   AllocMode MemAllocMode;
+  // True if UR performed the hipHostRegister call and owns the registration.
+  bool HostPtrRegisteredByUR = false;
 
 private:
   // Vector of HIP pointers
@@ -89,6 +91,15 @@ struct BufferMem {
         PtrToBufferMap{}, MemAllocMode{Mode},
         Ptrs(Context->Devices.size(), native_type{0}) {};
 
+  BufferMem(const BufferMem &Buffer)
+      : Parent{Buffer.Parent}, OuterMemStruct{Buffer.OuterMemStruct},
+        HostPtr{Buffer.HostPtr}, Size{Buffer.Size},
+        PtrToBufferMap{Buffer.PtrToBufferMap},
+        MemAllocMode{Buffer.MemAllocMode}, Ptrs{Buffer.Ptrs} {
+    // HostPtrRegisteredByUR intentionally not copied: ownership of the
+    // hipHostRegister call belongs to the original BufferMem only.
+  }
+
   // This will allocate memory on device if there isn't already an active
   // allocation on the device
   native_type getPtr(const ur_device_handle_t Device) {
@@ -161,7 +172,8 @@ struct BufferMem {
       }
       break;
     case AllocMode::UseHostPtr:
-      UR_CHECK_ERROR(hipHostUnregister(HostPtr));
+      if (HostPtrRegisteredByUR)
+        UR_CHECK_ERROR(hipHostUnregister(HostPtr));
       break;
     case AllocMode::AllocHostPtr:
       UR_CHECK_ERROR(hipHostFree(HostPtr));

unified-runtime/test/conformance/memory/urMemBufferCreate.cpp

@@ -122,8 +122,8 @@ TEST_P(urMemBufferCreateTest, CopyHostPointer) {
 TEST_P(urMemBufferCreateTest, UseHostPointer) {
   // These all copy memory instead of mapping it
   // https://github.com/intel/llvm/issues/18836
-  UUR_KNOWN_FAILURE_ON(uur::LevelZero{}, uur::LevelZeroV2{}, uur::HIP{},
-                       uur::CUDA{}, uur::OpenCL{"Intel(R) UHD Graphics 770"},
+  UUR_KNOWN_FAILURE_ON(uur::LevelZero{}, uur::LevelZeroV2{},
+                       uur::OpenCL{"Intel(R) UHD Graphics 770"},
                        uur::Offload{});
 
   std::vector<unsigned char> dataWrite{};