[SYCL] Attach auxiliary resources to memory objects

sycl/include/CL/sycl/detail/buffer_impl.hpp

@@ -169,6 +169,7 @@ class __SYCL_EXPORT buffer_impl final : public SYCLMemObjT {
   ~buffer_impl() {
     try {
       BaseT::updateHostMemory();
+      BaseT::detachResources();
     } catch (...) {
     }
     destructorNotification(this);

sycl/include/CL/sycl/detail/image_impl.hpp

@@ -229,6 +229,7 @@ class __SYCL_EXPORT image_impl final : public SYCLMemObjT {
   ~image_impl() {
     try {
       BaseT::updateHostMemory();
+      BaseT::detachResources();
     } catch (...) {
     }
   }

sycl/include/CL/sycl/detail/sycl_mem_obj_t.hpp

@@ -203,6 +203,9 @@ class __SYCL_EXPORT SYCLMemObjT : public SYCLMemObjI {
   // members must be alive.
   void updateHostMemory();
 
+  // Detach additional resources associated with the memory object.
+  void detachResources() const;
+
 public:
   __SYCL_DLL_LOCAL bool useHostPtr() {
     return has_property<property::buffer::use_host_ptr>() ||

sycl/include/CL/sycl/handler.hpp

@@ -473,6 +473,7 @@ class __SYCL_EXPORT handler {
   /// They are then forwarded to command group and destroyed only after
   /// the command group finishes the work on device/host.
   /// The 'MSharedPtrStorage' suits that need.
+  /// NOTE: This is no longer in use and should be removed with next ABI break.
   ///
   /// @param ReduObj is a pointer to object that must be stored.
   void addReduction(const std::shared_ptr<const void> &ReduObj) {

sycl/include/sycl/ext/oneapi/reduction.hpp

@@ -366,6 +366,17 @@ template <typename T> struct AreAllButLastReductions<T> {
   static constexpr bool value = !std::is_base_of<reduction_impl_base, T>::value;
 };
 
+/// Helper for attaching a resource to the lifetime of the memory associated
+/// with accessor.
+__SYCL_EXPORT void attachLifetime(std::shared_ptr<queue_impl> &Context,
+                                  std::shared_ptr<const void> &Resource,
+                                  detail::AccessorBaseHost &AttachTo);
+
+/// Helper for attaching a resource to the lifetime of USM memory.
+__SYCL_EXPORT void attachLifetime(std::shared_ptr<queue_impl> &Context,
+                                  std::shared_ptr<const void> &Resource,
+                                  void *AttachTo);
+
 /// This class encapsulates the reduction variable/accessor,
 /// the reduction operator and an optional operator identity.
 template <typename T, class BinaryOperation, int Dims, bool IsUSM,
@@ -645,7 +656,7 @@ class reduction_impl : private reduction_impl_base {
 
   accessor<T, buffer_dim, access::mode::read>
   getReadAccToPreviousPartialReds(handler &CGH) const {
-    CGH.addReduction(MOutBufPtr);
+    attachResourceLifetimeToMem(CGH, MOutBufPtr);
     return {*MOutBufPtr, CGH};
   }
 
@@ -673,7 +684,7 @@ class reduction_impl : private reduction_impl_base {
   std::enable_if_t<!IsOneWG, rw_accessor_type>
   getWriteMemForPartialReds(size_t Size, handler &CGH) {
     MOutBufPtr = std::make_shared<buffer<T, buffer_dim>>(range<1>(Size));
-    CGH.addReduction(MOutBufPtr);
+    attachResourceLifetimeToMem(CGH, MOutBufPtr);
     return createHandlerWiredReadWriteAccessor(CGH, *MOutBufPtr);
   }
 
@@ -691,7 +702,7 @@ class reduction_impl : private reduction_impl_base {
 
     // Create a new output buffer and return an accessor to it.
     MOutBufPtr = std::make_shared<buffer<T, buffer_dim>>(range<1>(Size));
-    CGH.addReduction(MOutBufPtr);
+    attachResourceLifetimeToMem(CGH, MOutBufPtr);
     return createHandlerWiredReadWriteAccessor(CGH, *MOutBufPtr);
   }
 
@@ -707,19 +718,19 @@ class reduction_impl : private reduction_impl_base {
       return *MRWAcc;
 
     auto RWReduVal = std::make_shared<T>(MIdentity);
-    CGH.addReduction(RWReduVal);
+    attachResourceLifetimeToMem(CGH, RWReduVal);
     MOutBufPtr = std::make_shared<buffer<T, 1>>(RWReduVal.get(), range<1>(1));
-    CGH.addReduction(MOutBufPtr);
+    attachResourceLifetimeToMem(CGH, MOutBufPtr);
     return createHandlerWiredReadWriteAccessor(CGH, *MOutBufPtr);
   }
 
   accessor<int, 1, access::mode::read_write, access::target::device,
            access::placeholder::false_t>
   getReadWriteAccessorToInitializedGroupsCounter(handler &CGH) {
     auto CounterMem = std::make_shared<int>(0);
-    CGH.addReduction(CounterMem);
+    attachResourceLifetimeToMem(CGH, CounterMem);
     auto CounterBuf = std::make_shared<buffer<int, 1>>(CounterMem.get(), 1);
-    CGH.addReduction(CounterBuf);
+    attachResourceLifetimeToMem(CGH, CounterBuf);
     return {*CounterBuf, CGH};
   }
 
@@ -767,6 +778,23 @@ class reduction_impl : private reduction_impl_base {
     return Acc;
   }
 
+  /// Attaches the resource to the lifetime of the associated memory of the
+  /// reduction.
+  void attachResourceLifetimeToMem(handler &CGH,
+                                   std::shared_ptr<const void> Resource) const {
+#ifndef __SYCL_DEVICE_ONLY__
+    if (is_usm)
+      detail::attachLifetime(CGH.MQueue, Resource, MUSMPointer);
+    else if (MDWAcc != nullptr)
+      detail::attachLifetime(CGH.MQueue, Resource, *MDWAcc);
+    else
+      detail::attachLifetime(CGH.MQueue, Resource, *MRWAcc);
+#else
+    (void)CGH;
+    (void)Resource;
+#endif
+  }
+
   /// Identity of the BinaryOperation.
   /// The result of BinaryOperation(X, MIdentity) is equal to X for any X.
   const T MIdentity;

sycl/source/detail/context_impl.cpp

@@ -115,6 +115,25 @@ cl_context context_impl::get() const {
 bool context_impl::is_host() const { return MHostContext; }
 
 context_impl::~context_impl() {
+  // In case a user is leaking a memory object we may still have attached
+  // resources. These resources will be released with the context, but we need
+  // to do it before we release the backend context to avoid confusing errors.
+  std::unordered_map<const SYCLMemObjI *,
+                     std::vector<std::shared_ptr<const void>>>
+      MemObjLTARes;
+  {
+    std::lock_guard<std::mutex> lock(MMemObjLifetimeAttachedResourcesMutex);
+    std::swap(MemObjLTARes, MMemObjLifetimeAttachedResources);
+  }
+  MemObjLTARes.clear();
+  std::unordered_map<const void *, std::vector<std::shared_ptr<const void>>>
+      USMLTARes;
+  {
+    std::lock_guard<std::mutex> lock(MUSMLifetimeAttachedResourcesMutex);
+    std::swap(USMLTARes, MUSMLifetimeAttachedResources);
+  }
+  USMLTARes.clear();
+
   for (auto LibProg : MCachedLibPrograms) {
     assert(LibProg.second && "Null program must not be kept in the cache");
     getPlugin().call<PiApiKind::piProgramRelease>(LibProg.second);
@@ -206,6 +225,58 @@ pi_native_handle context_impl::getNative() const {
   return Handle;
 }
 
+void context_impl::attachLifetimeToMemObj(std::shared_ptr<const void> &Resource,
+                                          const SYCLMemObjI *AttachTo) {
+  std::lock_guard<std::mutex> lock(MMemObjLifetimeAttachedResourcesMutex);
+  auto AttachedResourcesIt = MMemObjLifetimeAttachedResources.find(AttachTo);
+  if (AttachedResourcesIt != MMemObjLifetimeAttachedResources.end())
+    AttachedResourcesIt->second.push_back(Resource);
+  else
+    MMemObjLifetimeAttachedResources.insert({AttachTo, {Resource}});
+}
+
+void context_impl::detachMemObjLifetimeResources(
+    const SYCLMemObjI *AttachedTo) {
+  // Swap the attached resources and let them go out of scope without the lock.
+  // This is required as they could potentially have their own attached
+  // resources they need to detach.
+  std::vector<std::shared_ptr<const void>> AttachedResources;
+  {
+    std::lock_guard<std::mutex> lock(MMemObjLifetimeAttachedResourcesMutex);
+    auto AttachedResourcesIt =
+        MMemObjLifetimeAttachedResources.find(AttachedTo);
+    if (AttachedResourcesIt == MMemObjLifetimeAttachedResources.end())
+      return;
+    std::swap(AttachedResourcesIt->second, AttachedResources);
+    MMemObjLifetimeAttachedResources.erase(AttachedResourcesIt);
+  }
+}
+
+void context_impl::attachLifetimeToUSM(std::shared_ptr<const void> &Resource,
+                                       const void *AttachTo) {
+  std::lock_guard<std::mutex> lock(MUSMLifetimeAttachedResourcesMutex);
+  auto AttachedResourcesIt = MUSMLifetimeAttachedResources.find(AttachTo);
+  if (AttachedResourcesIt != MUSMLifetimeAttachedResources.end())
+    AttachedResourcesIt->second.push_back(Resource);
+  else
+    MUSMLifetimeAttachedResources.insert({AttachTo, {Resource}});
+}
+
+void context_impl::detachUSMLifetimeResources(const void *AttachedTo) {
+  // Swap the attached resources and let them go out of scope without the lock.
+  // This is required as they could potentially have their own attached
+  // resources they need to detach.
+  std::vector<std::shared_ptr<const void>> AttachedResources;
+  {
+    std::lock_guard<std::mutex> lock(MUSMLifetimeAttachedResourcesMutex);
+    auto AttachedResourcesIt = MUSMLifetimeAttachedResources.find(AttachedTo);
+    if (AttachedResourcesIt == MUSMLifetimeAttachedResources.end())
+      return;
+    std::swap(AttachedResourcesIt->second, AttachedResources);
+    MUSMLifetimeAttachedResources.erase(AttachedResourcesIt);
+  }
+}
+
 } // namespace detail
 } // namespace sycl
 } // __SYCL_INLINE_NAMESPACE(cl)

sycl/source/detail/context_impl.hpp

@@ -10,6 +10,7 @@
 #include <CL/sycl/detail/common.hpp>
 #include <CL/sycl/detail/os_util.hpp>
 #include <CL/sycl/detail/pi.hpp>
+#include <CL/sycl/detail/sycl_mem_obj_i.hpp>
 #include <CL/sycl/exception_list.hpp>
 #include <CL/sycl/info/info_desc.hpp>
 #include <CL/sycl/property_list.hpp>
@@ -167,6 +168,30 @@ class context_impl {
   /// \return a native handle.
   pi_native_handle getNative() const;
 
+  /// Attach a resource to a memory object.
+  ///
+  /// \param Resource is the resource to attach to the memory object
+  /// \param AttachTo is the memory object to attach the resource to
+  void attachLifetimeToMemObj(std::shared_ptr<const void> &Resource,
+                              const SYCLMemObjI *AttachTo);
+
+  /// Detach all resources attached to a memory object.
+  ///
+  /// \param AttachedTo is the memory object to detach resources from
+  void detachMemObjLifetimeResources(const SYCLMemObjI *AttachedTo);
+
+  /// Attach a resource to a USM pointer.
+  ///
+  /// \param Resource is the resource to attach to the USM pointer
+  /// \param AttachTo is the USM pointer to attach the resource to
+  void attachLifetimeToUSM(std::shared_ptr<const void> &Resource,
+                           const void *AttachTo);
+
+  /// Detach all resources attached to a USM pointer.
+  ///
+  /// \param AttachedTo is the USM pointer to detach resources from
+  void detachUSMLifetimeResources(const void *AttachedTo);
+
 private:
   async_handler MAsyncHandler;
   std::vector<device> MDevices;
@@ -177,6 +202,22 @@ class context_impl {
   std::map<std::pair<DeviceLibExt, RT::PiDevice>, RT::PiProgram>
       MCachedLibPrograms;
   mutable KernelProgramCache MKernelProgramCache;
+
+  /// Matches SYCL memory objects to attached resources.
+  /// TODO: On ABI break this could be made part of SYCLMemObjT instead.
+  std::unordered_map<const SYCLMemObjI *,
+                     std::vector<std::shared_ptr<const void>>>
+      MMemObjLifetimeAttachedResources;
+
+  /// Protects m_MemObjLifetimeAttachedResources.
+  std::mutex MMemObjLifetimeAttachedResourcesMutex;
+
+  /// Matches USM pointers to attached resources.
+  std::unordered_map<const void *, std::vector<std::shared_ptr<const void>>>
+      MUSMLifetimeAttachedResources;
+
+  /// Protects m_USMLifetimeAttachedResources.
+  std::mutex MUSMLifetimeAttachedResourcesMutex;
 };
 
 } // namespace detail

sycl/source/detail/reduction.cpp

@@ -110,6 +110,19 @@ reduGetMaxWGSize(std::shared_ptr<sycl::detail::queue_impl> Queue,
   return WGSize;
 }
 
+__SYCL_EXPORT void attachLifetime(std::shared_ptr<queue_impl> &Queue,
+                                  std::shared_ptr<const void> &Resource,
+                                  detail::AccessorBaseHost &AttachTo) {
+  Queue->getContextImplPtr()->attachLifetimeToMemObj(Resource,
+                                  getSyclObjImpl(AttachTo)->MSYCLMemObj);
+}
+
+__SYCL_EXPORT void attachLifetime(std::shared_ptr<queue_impl> &Queue,
+                                  std::shared_ptr<const void> &Resource,
+                                  void *AttachTo) {
+  Queue->getContextImplPtr()->attachLifetimeToUSM(Resource, AttachTo);
+}
+
 } // namespace detail
 } // namespace oneapi
 } // namespace ext

sycl/source/detail/sycl_mem_obj_t.cpp

@@ -90,6 +90,15 @@ void SYCLMemObjT::updateHostMemory() {
         pi::cast<RT::PiMem>(MInteropMemObject));
   }
 }
+
+// TODO: With ABI break the attached resources can be held by this type. When
+// that happens this will be obsolete as the resources will automatically be
+// destroyed with the object.
+void SYCLMemObjT::detachResources() const {
+  if (MInteropContext)
+    MInteropContext->detachMemObjLifetimeResources(this);
+}
+
 const plugin &SYCLMemObjT::getPlugin() const {
   assert((MInteropContext != nullptr) &&
          "Trying to get Plugin from SYCLMemObjT with nullptr ContextImpl.");

sycl/source/detail/usm/usm_impl.cpp

@@ -164,6 +164,8 @@ void free(void *Ptr, const context &Ctxt, const detail::code_location &CL) {
     const detail::plugin &Plugin = CtxImpl->getPlugin();
     Plugin.call<PiApiKind::piextUSMFree>(C, Ptr);
   }
+  // Detach resources.
+  detail::getSyclObjImpl(Ctxt)->detachUSMLifetimeResources(Ptr);
 }
 
 // For ABI compatibility

sycl/test/abi/sycl_symbols_linux.dump

@@ -3697,6 +3697,8 @@ _ZN2cl4sycl3ext6oneapi10level_zero12make_programERKNS0_7contextEm
 _ZN2cl4sycl3ext6oneapi10level_zero13make_platformEm
 _ZN2cl4sycl3ext6oneapi15filter_selectorC1ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
 _ZN2cl4sycl3ext6oneapi15filter_selectorC2ERKNSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
+_ZN2cl4sycl3ext6oneapi6detail14attachLifetimeERSt10shared_ptrINS0_6detail10queue_implEERS4_IKvEPv
+_ZN2cl4sycl3ext6oneapi6detail14attachLifetimeERSt10shared_ptrINS0_6detail10queue_implEERS4_IKvERNS5_16AccessorBaseHostE
 _ZN2cl4sycl3ext6oneapi6detail16reduGetMaxWGSizeESt10shared_ptrINS0_6detail10queue_implEEm
 _ZN2cl4sycl3ext6oneapi6detail17reduComputeWGSizeEmmRm
 _ZN2cl4sycl3ext6oneapi6detail33reduGetMaxNumConcurrentWorkGroupsESt10shared_ptrINS0_6detail10queue_implEE
@@ -4148,6 +4150,7 @@ _ZNK2cl4sycl6detail10image_implILi3EE4sizeEv
 _ZNK2cl4sycl6detail10image_implILi3EE7getTypeEv
 _ZNK2cl4sycl6detail10image_implILi3EE9get_countEv
 _ZNK2cl4sycl6detail10image_implILi3EE9get_rangeEv
+_ZNK2cl4sycl6detail11SYCLMemObjT15detachResourcesEv
 _ZNK2cl4sycl6detail11SYCLMemObjT9getPluginEv
 _ZNK2cl4sycl6detail11SYCLMemObjT9isInteropEv
 _ZNK2cl4sycl6detail11stream_impl22get_max_statement_sizeEv

sycl/test/abi/sycl_symbols_windows.dump

@@ -1313,6 +1313,8 @@
 ?atanpi@__host_std@cl@@YA?AVhalf@half_impl@detail@sycl@2@V34562@@Z
 ?atanpi@__host_std@cl@@YAMM@Z
 ?atanpi@__host_std@cl@@YANN@Z
+?attachLifetime@detail@oneapi@ext@sycl@cl@@YAXAEAV?$shared_ptr@$$CBX@std@@AEAVAccessorBaseHost@145@@Z
+?attachLifetime@detail@oneapi@ext@sycl@cl@@YAXAEAV?$shared_ptr@$$CBX@std@@PEAX@Z
 ?barrier@handler@sycl@cl@@QEAAXAEBV?$vector@Vevent@sycl@cl@@V?$allocator@Vevent@sycl@cl@@@std@@@std@@@Z
 ?barrier@handler@sycl@cl@@QEAAXXZ
 ?begin@exception_list@sycl@cl@@QEBA?AV?$_Vector_const_iterator@V?$_Vector_val@U?$_Simple_types@Vexception_ptr@std@@@std@@@std@@@std@@XZ
@@ -1689,6 +1691,7 @@
 ?depends_on@handler@sycl@cl@@QEAAXAEBV?$vector@Vevent@sycl@cl@@V?$allocator@Vevent@sycl@cl@@@std@@@std@@@Z
 ?depends_on@handler@sycl@cl@@QEAAXVevent@23@@Z
 ?destructorNotification@buffer_impl@detail@sycl@cl@@QEAAXPEAX@Z
+?detachResources@SYCLMemObjT@detail@sycl@cl@@IEBAXXZ
 ?determineHostPtr@SYCLMemObjT@detail@sycl@cl@@IEAAXAEBV?$shared_ptr@Vcontext_impl@detail@sycl@cl@@@std@@_NAEAPEAXAEA_N@Z
 ?device_has@queue@sycl@cl@@QEBA_NW4aspect@23@@Z
 ?die@pi@detail@sycl@cl@@YAXPEBD@Z