Clear resource pool when running out of memory

steffenlarsen · steffenlarsen · commit 255b422f013b · 2022-03-08T13:41:21.000+03:00
Signed-off-by: Steffen Larsen &lt;steffen.larsen@intel.com&gt;
diff --git a/sycl/include/CL/sycl/detail/resource_pool.hpp b/sycl/include/CL/sycl/detail/resource_pool.hpp
@@ -184,6 +184,11 @@ class __SYCL_EXPORT ResourcePool {
     MPlatform = Platform;
   }
 
+  /// Returns true if the resource pool is enabled and false otherwise.
+  ///
+  /// \return a boolean value specifying whether the pool is enabled.
+  bool isEnabled() { return MIsPoolingEnabled; }
+
   /// Creates a managed resource from the pool.
   ///
   /// \param Range is the range of the resulting buffer.
diff --git a/sycl/source/detail/device_image_impl.hpp b/sycl/source/detail/device_image_impl.hpp
@@ -194,14 +194,12 @@ class device_image_impl {
   RT::PiMem &get_spec_const_buffer_ref() noexcept {
     std::lock_guard<std::mutex> Lock{MSpecConstAccessMtx};
     if (nullptr == MSpecConstsBuffer && !MSpecConstsBlob.empty()) {
-      const detail::plugin &Plugin = getSyclObjImpl(MContext)->getPlugin();
       // Uses PI_MEM_FLAGS_HOST_PTR_COPY instead of PI_MEM_FLAGS_HOST_PTR_USE
       // since post-enqueue cleanup might trigger destruction of
       // device_image_impl and, as a result, destruction of MSpecConstsBlob
       // while MSpecConstsBuffer is still in use.
       // TODO consider changing the lifetime of device_image_impl instead
-      memBufferCreateHelper(Plugin,
-                            detail::getSyclObjImpl(MContext)->getHandleRef(),
+      memBufferCreateHelper(detail::getSyclObjImpl(MContext),
                             PI_MEM_FLAGS_ACCESS_RW | PI_MEM_FLAGS_HOST_PTR_COPY,
                             MSpecConstsBlob.size(), MSpecConstsBlob.data(),
                             &MSpecConstsBuffer, nullptr);
diff --git a/sycl/source/detail/mem_alloc_helper.hpp b/sycl/source/detail/mem_alloc_helper.hpp
@@ -13,7 +13,9 @@
 __SYCL_INLINE_NAMESPACE(cl) {
 namespace sycl {
 namespace detail {
-void memBufferCreateHelper(const plugin &Plugin, pi_context Ctx,
+class context_impl;
+
+void memBufferCreateHelper(std::shared_ptr<context_impl> CtxImpl,
                            pi_mem_flags Flags, size_t Size, void *HostPtr,
                            pi_mem *RetMem,
                            const pi_mem_properties *Props = nullptr);
diff --git a/sycl/source/detail/memory_manager.cpp b/sycl/source/detail/memory_manager.cpp
@@ -124,9 +124,11 @@ static void waitForEvents(const std::vector<EventImplPtr> &Events) {
   }
 }
 
-void memBufferCreateHelper(const plugin &Plugin, pi_context Ctx,
-                           pi_mem_flags Flags, size_t Size, void *HostPtr,
-                           pi_mem *RetMem, const pi_mem_properties *Props) {
+static pi_result memBufferCreateNocheckHelper(const plugin &Plugin,
+                                              pi_context Ctx,
+                                              pi_mem_flags Flags, size_t Size,
+                                              void *HostPtr, pi_mem *RetMem,
+                                              const pi_mem_properties *Props) {
 #ifdef XPTI_ENABLE_INSTRUMENTATION
   uint64_t CorrID = 0;
 #endif
@@ -147,11 +149,30 @@ void memBufferCreateHelper(const plugin &Plugin, pi_context Ctx,
                            CorrID);
     }};
 #endif
-    Plugin.call<PiApiKind::piMemBufferCreate>(Ctx, Flags, Size, HostPtr, RetMem,
-                                              Props);
+    return Plugin.call_nocheck<PiApiKind::piMemBufferCreate>(
+        Ctx, Flags, Size, HostPtr, RetMem, Props);
   }
 }
 
+void memBufferCreateHelper(std::shared_ptr<context_impl> CtxImpl,
+                           pi_mem_flags Flags, size_t Size, void *HostPtr,
+                           pi_mem *RetMem, const pi_mem_properties *Props) {
+  const detail::plugin &Plugin = CtxImpl->getPlugin();
+  RT::PiResult Err = memBufferCreateNocheckHelper(
+      Plugin, CtxImpl->getHandleRef(), Flags, Size, HostPtr, RetMem, Props);
+
+  ResourcePool &Resources = CtxImpl->getResourcePool();
+  if (Err == PI_MEM_OBJECT_ALLOCATION_FAILURE && Resources.isEnabled()) {
+    // Clear resource pool and retry allocation.
+    Resources.clear();
+    Err = memBufferCreateNocheckHelper(Plugin, CtxImpl->getHandleRef(), Flags,
+                                       Size, HostPtr, RetMem, Props);
+  }
+
+  if (Err != PI_SUCCESS)
+    Plugin.reportPiError(Err, "memBufferCreateHelper()");
+}
+
 void memReleaseHelper(const plugin &Plugin, pi_mem Mem) {
   // FIXME piMemRelease does not guarante memory release. It is only true if
   // reference counter is 1. However, SYCL runtime currently only calls
@@ -361,9 +382,8 @@ MemoryManager::allocateBufferObject(ContextImplPtr TargetContext, void *UserPtr,
     CreationFlags |= PI_MEM_FLAGS_HOST_PTR_ALLOC;
 
   RT::PiMem NewMem = nullptr;
-  const detail::plugin &Plugin = TargetContext->getPlugin();
-  memBufferCreateHelper(Plugin, TargetContext->getHandleRef(), CreationFlags,
-                        Size, UserPtr, &NewMem, nullptr);
+  memBufferCreateHelper(TargetContext, CreationFlags, Size, UserPtr, &NewMem,
+                        nullptr);
   return NewMem;
 }
 
diff --git a/sycl/source/detail/resource_pool.cpp b/sycl/source/detail/resource_pool.cpp
@@ -42,19 +42,21 @@ ResourcePool::FreeEntry ResourcePool::getOrAllocateEntry(
   assert(Size && "Size must be greater than 0");
   assert(ContextImplPtr->getPlatformImpl() == MPlatform &&
          "Context platform does not match the resource pool platform.");
-  std::lock_guard<std::mutex> Lock{MMutex};
+  {
+    std::lock_guard<std::mutex> Lock{MMutex};
 
-  // Find the free entry with the smallest suitable size.
-  auto FoundFreeEntry = MFreeEntries.upper_bound(Size - 1);
+    // Find the free entry with the smallest suitable size.
+    auto FoundFreeEntry = MFreeEntries.upper_bound(Size - 1);
 
-  // If there was a fitting free entry in the pool, remove and return it.
-  const bool IsOldEntry = FoundFreeEntry != MFreeEntries.end();
-  if (IsNewEntry)
-    *IsNewEntry = !IsOldEntry;
-  if (IsOldEntry) {
-    FreeEntry Entry = *FoundFreeEntry;
-    MFreeEntries.erase(FoundFreeEntry);
-    return Entry;
+    // If there was a fitting free entry in the pool, remove and return it.
+    const bool IsOldEntry = FoundFreeEntry != MFreeEntries.end();
+    if (IsNewEntry)
+      *IsNewEntry = !IsOldEntry;
+    if (IsOldEntry) {
+      FreeEntry Entry = *FoundFreeEntry;
+      MFreeEntries.erase(FoundFreeEntry);
+      return Entry;
+    }
   }
 
   // If there was no suitable free entry we allocate memory and return it in a
@@ -63,8 +65,8 @@ ResourcePool::FreeEntry ResourcePool::getOrAllocateEntry(
   if (DataPtr)
     MemFlags |= PI_MEM_FLAGS_HOST_PTR_COPY;
   RT::PiMem NewResMem;
-  memBufferCreateHelper(MPlatform->getPlugin(), ContextImplPtr->getHandleRef(),
-                        MemFlags, Size, DataPtr, &NewResMem, nullptr);
+  memBufferCreateHelper(ContextImplPtr, MemFlags, Size, DataPtr, &NewResMem,
+                        nullptr);
   ++MAllocCount;
   return {Size, NewResMem};
 }
diff --git a/sycl/source/detail/usm/usm_impl.cpp b/sycl/source/detail/usm/usm_impl.cpp
@@ -43,12 +43,15 @@ using alloc = cl::sycl::usm::alloc;
 namespace detail {
 namespace usm {
 
-void *alignedAllocHost(size_t Alignment, size_t Size, const context &Ctxt,
-                       alloc Kind, const detail::code_location &CL) {
+static pi_result alignedAllocHostHelper(size_t Alignment, size_t Size,
+                                        const context &Ctxt, alloc Kind,
+                                        const detail::code_location &CL,
+                                        void **OutPtr) {
   XPTI_CREATE_TRACEPOINT(CL);
-  void *RetVal = nullptr;
-  if (Size == 0)
-    return nullptr;
+  if (Size == 0) {
+    *OutPtr = nullptr;
+    return PI_SUCCESS;
+  }
   if (Ctxt.is_host()) {
     if (!Alignment) {
       // worst case default
@@ -57,10 +60,11 @@ void *alignedAllocHost(size_t Alignment, size_t Size, const context &Ctxt,
 
     aligned_allocator<char> Alloc(Alignment);
     try {
-      RetVal = Alloc.allocate(Size);
+      *OutPtr = Alloc.allocate(Size);
     } catch (const std::bad_alloc &) {
       // Conform with Specification behavior
-      RetVal = nullptr;
+      *OutPtr = nullptr;
+      return PI_MEM_OBJECT_ALLOCATION_FAILURE;
     }
   } else {
     std::shared_ptr<context_impl> CtxImpl = detail::getSyclObjImpl(Ctxt);
@@ -71,13 +75,13 @@ void *alignedAllocHost(size_t Alignment, size_t Size, const context &Ctxt,
     switch (Kind) {
     case alloc::host: {
       Error = Plugin.call_nocheck<PiApiKind::piextUSMHostAlloc>(
-          &RetVal, C, nullptr, Size, Alignment);
+          OutPtr, C, nullptr, Size, Alignment);
       break;
     }
     case alloc::device:
     case alloc::shared:
     case alloc::unknown: {
-      RetVal = nullptr;
+      *OutPtr = nullptr;
       Error = PI_INVALID_VALUE;
       break;
     }
@@ -86,21 +90,41 @@ void *alignedAllocHost(size_t Alignment, size_t Size, const context &Ctxt,
     // Error is for debugging purposes.
     // The spec wants a nullptr returned, not an exception.
     if (Error != PI_SUCCESS)
-      return nullptr;
+      *OutPtr = nullptr;
+
+    return Error;
+  }
+  return PI_SUCCESS;
+}
+
+void *alignedAllocHost(size_t Alignment, size_t Size, const context &Ctxt,
+                       alloc Kind, const detail::code_location &CL) {
+  void *RetVal;
+  pi_result Err =
+      alignedAllocHostHelper(Alignment, Size, Ctxt, Kind, CL, &RetVal);
+
+  std::shared_ptr<context_impl> CtxImpl = detail::getSyclObjImpl(Ctxt);
+  ResourcePool &Resources = CtxImpl->getResourcePool();
+  if (Err == PI_OUT_OF_RESOURCES && Resources.isEnabled()) {
+    // Clear resource pool and retry allocation.
+    Resources.clear();
+    alignedAllocHostHelper(Alignment, Size, Ctxt, Kind, CL, &RetVal);
   }
   return RetVal;
 }
 
-void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
-                   const device &Dev, alloc Kind,
-                   const detail::code_location &CL) {
+static pi_result alignedAllocHelper(size_t Alignment, size_t Size,
+                                    const context &Ctxt, const device &Dev,
+                                    alloc Kind, const detail::code_location &CL,
+                                    void **OutPtr) {
   XPTI_CREATE_TRACEPOINT(CL);
-  void *RetVal = nullptr;
-  if (Size == 0)
-    return nullptr;
+  if (Size == 0) {
+    *OutPtr = nullptr;
+    return PI_SUCCESS;
+  }
   if (Ctxt.is_host()) {
     if (Kind == alloc::unknown) {
-      RetVal = nullptr;
+      *OutPtr = nullptr;
     } else {
       if (!Alignment) {
         // worst case default
@@ -109,10 +133,11 @@ void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
 
       aligned_allocator<char> Alloc(Alignment);
       try {
-        RetVal = Alloc.allocate(Size);
+        *OutPtr = Alloc.allocate(Size);
       } catch (const std::bad_alloc &) {
         // Conform with Specification behavior
-        RetVal = nullptr;
+        *OutPtr = nullptr;
+        return PI_MEM_OBJECT_ALLOCATION_FAILURE;
       }
     }
   } else {
@@ -126,18 +151,18 @@ void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
     case alloc::device: {
       Id = detail::getSyclObjImpl(Dev)->getHandleRef();
       Error = Plugin.call_nocheck<PiApiKind::piextUSMDeviceAlloc>(
-          &RetVal, C, Id, nullptr, Size, Alignment);
+          OutPtr, C, Id, nullptr, Size, Alignment);
       break;
     }
     case alloc::shared: {
       Id = detail::getSyclObjImpl(Dev)->getHandleRef();
       Error = Plugin.call_nocheck<PiApiKind::piextUSMSharedAlloc>(
-          &RetVal, C, Id, nullptr, Size, Alignment);
+          OutPtr, C, Id, nullptr, Size, Alignment);
       break;
     }
     case alloc::host:
     case alloc::unknown: {
-      RetVal = nullptr;
+      *OutPtr = nullptr;
       Error = PI_INVALID_VALUE;
       break;
     }
@@ -146,7 +171,27 @@ void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
     // Error is for debugging purposes.
     // The spec wants a nullptr returned, not an exception.
     if (Error != PI_SUCCESS)
-      return nullptr;
+      *OutPtr = nullptr;
+
+    return Error;
+  }
+
+  return PI_SUCCESS;
+}
+
+void *alignedAlloc(size_t Alignment, size_t Size, const context &Ctxt,
+                   const device &Dev, alloc Kind,
+                   const detail::code_location &CL) {
+  void *RetVal;
+  pi_result Err =
+      alignedAllocHelper(Alignment, Size, Ctxt, Dev, Kind, CL, &RetVal);
+
+  std::shared_ptr<context_impl> CtxImpl = detail::getSyclObjImpl(Ctxt);
+  ResourcePool &Resources = CtxImpl->getResourcePool();
+  if (Err == PI_OUT_OF_RESOURCES && Resources.isEnabled()) {
+    // Clear resource pool and retry allocation.
+    Resources.clear();
+    alignedAllocHelper(Alignment, Size, Ctxt, Dev, Kind, CL, &RetVal);
   }
   return RetVal;
 }
