[SYCL][NATIVECPU] Reference counting native_cpu mem handle kernel arguments (#17558)

This PR fixes mem handle arguments for asynchronous kernel launches.
This makes at least the following e2e tests pass on native cpu:
```
SpecConstants/2020/hierarchy.cpp
SpecConstants/2020/marray_vec.cpp
``` 
The native_cpu kernel handle now reference-counts the mem handle
arguments to prevent their potential release by the SYCL runtime before
the enqueued kernel started using them (if the thread takes longer to
start). These arguments are then freed when the kernel handle is
released.

The native_cpu mem handle class now reuses the same reference counting
as the other native_cpu handles instead of having its own reference
counting members, which were removed.

Author: uwedolinsky

unified-runtime/source/adapters/native_cpu/common.hpp

@@ -53,11 +53,6 @@ namespace ur {
 } // namespace ur
 } // namespace detail
 
-// Base class to store common data
-struct _ur_object {
-  ur_shared_mutex Mutex;
-};
-
 // Todo: replace this with a common helper once it is available
 struct RefCounted {
   std::atomic_uint32_t _refCount;
@@ -67,6 +62,11 @@ struct RefCounted {
   uint32_t getReferenceCount() const { return _refCount; }
 };
 
+// Base class to store common data
+struct _ur_object : RefCounted {
+  ur_shared_mutex Mutex;
+};
+
 template <typename T> inline void decrementOrDelete(T *refC) {
   if (refC->decrementReferenceCount() == 0)
     delete refC;

unified-runtime/source/adapters/native_cpu/kernel.cpp

@@ -13,7 +13,6 @@
 
 #include "common.hpp"
 #include "kernel.hpp"
-#include "memory.hpp"
 #include "program.hpp"
 
 UR_APIEXPORT ur_result_t UR_APICALL
@@ -271,6 +270,7 @@ urKernelSetArgMemObj(ur_kernel_handle_t hKernel, uint32_t argIndex,
     return UR_RESULT_SUCCESS;
   }
 
+  hKernel->addArgReference(hArgValue);
   hKernel->addPtrArg(hArgValue->_mem, argIndex);
   return UR_RESULT_SUCCESS;
 }

unified-runtime/source/adapters/native_cpu/kernel.hpp

@@ -9,6 +9,7 @@
 #pragma once
 
 #include "common.hpp"
+#include "memory.hpp"
 #include "nativecpu_state.hpp"
 #include "program.hpp"
 #include <cstring>
@@ -35,9 +36,14 @@ struct ur_kernel_handle_t_ : RefCounted {
   ur_kernel_handle_t_(const ur_kernel_handle_t_ &other)
       : Args(other.Args), hProgram(other.hProgram), _name(other._name),
         _subhandler(other._subhandler), _localArgInfo(other._localArgInfo),
-        ReqdWGSize(other.ReqdWGSize) {}
+        ReqdWGSize(other.ReqdWGSize) {
+    takeArgReferences(other);
+  }
 
-  ~ur_kernel_handle_t_() { free(_localMemPool); }
+  ~ur_kernel_handle_t_() {
+    removeArgReferences();
+    free(_localMemPool);
+  }
 
   ur_kernel_handle_t_(ur_program_handle_t hProgram, const char *name,
                       nativecpu_task_t subhandler,
@@ -186,10 +192,26 @@ struct ur_kernel_handle_t_ : RefCounted {
 
   void addPtrArg(void *Ptr, size_t Index) { Args.addPtrArg(Index, Ptr); }
 
+  void addArgReference(ur_mem_handle_t Arg) {
+    Arg->incrementReferenceCount();
+    ReferencedArgs.push_back(Arg);
+  }
+
+private:
+  void removeArgReferences() {
+    for (auto arg : ReferencedArgs)
+      decrementOrDelete(arg);
+  }
+  void takeArgReferences(const ur_kernel_handle_t_ &other) {
+    for (auto arg : other.ReferencedArgs)
+      addArgReference(arg);
+  }
+
 private:
   char *_localMemPool = nullptr;
   size_t _localMemPoolSize = 0;
   std::optional<native_cpu::WGSize_t> ReqdWGSize = std::nullopt;
   std::optional<native_cpu::WGSize_t> MaxWGSize = std::nullopt;
   std::optional<uint64_t> MaxLinearWGSize = std::nullopt;
+  std::vector<ur_mem_handle_t> ReferencedArgs;
 };

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

@@ -70,13 +70,8 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemRetain(ur_mem_handle_t hMem) {
 
 UR_APIEXPORT ur_result_t UR_APICALL urMemRelease(ur_mem_handle_t hMem) {
   UR_ASSERT(hMem, UR_RESULT_ERROR_INVALID_NULL_HANDLE);
+  decrementOrDelete(hMem);
 
-  hMem->decrementRefCount();
-  if (hMem->_refCount > 0) {
-    return UR_RESULT_SUCCESS;
-  }
-
-  delete hMem;
   return UR_RESULT_SUCCESS;
 }
 

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

@@ -38,14 +38,11 @@ struct ur_mem_handle_t_ : _ur_object {
     }
   }
 
-  void decrementRefCount() noexcept { _refCount--; }
-
   // Method to get type of the derived object (image or buffer)
   bool isImage() const { return this->IsImage; }
 
   char *_mem;
   bool _ownsMem;
-  std::atomic_uint32_t _refCount = {1};
 
 private:
   const bool IsImage;