draft

Signed-off-by: Tikhomirova, Kseniya <[email protected]>

Author: KseniyaTikhomirova

sycl/source/detail/program_manager/program_manager.cpp

@@ -750,7 +750,8 @@ setSpecializationConstants(const std::shared_ptr<device_image_impl> &InputImpl,
   }
 }
 
-static inline void CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) {
+static inline void
+CheckAndDecompressImage([[maybe_unused]] RTDeviceBinaryImage *Img) {
 #ifndef SYCL_RT_ZSTD_NOT_AVAIABLE
   if (auto CompImg = dynamic_cast<CompressedRTDeviceBinaryImage *>(Img))
     if (CompImg->IsCompressed())
@@ -894,8 +895,10 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
 
     {
       std::lock_guard<std::mutex> Lock(MNativeProgramsMutex);
-      // NativePrograms map does not intend to keep reference to program handle, so keys in the map can be invalid (reference count went to zero and the underlying program disposed of).
-      // Protecting from incorrect values by removal of map entries with same handle (obviously invalid entries).
+      // NativePrograms map does not intend to keep reference to program handle,
+      // so keys in the map can be invalid (reference count went to zero and the
+      // underlying program disposed of). Protecting from incorrect values by
+      // removal of map entries with same handle (obviously invalid entries).
       std::ignore = NativePrograms.erase(BuiltProgram.get());
       NativePrograms.insert({BuiltProgram.get(), &Img});
       for (RTDeviceBinaryImage *LinkedImg : DeviceImagesToLink) {
@@ -2661,7 +2664,8 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
 
     // Device is not used when creating program from SPIRV, so passing only one
     // device is OK.
-    // getOrCreateURProgram adds NativePrg to NativePrograms storage (no extra ref).
+    // getOrCreateURProgram adds NativePrg to NativePrograms storage (no extra
+    // ref).
     auto [NativePrg, DeviceCodeWasInCache] = getOrCreateURProgram(
         Img, {&Img}, Context, Devs, CompileOpts + LinkOpts, SpecConsts);
 
@@ -2695,7 +2699,7 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
               URDevices, DeviceLibReqMask, ExtraProgramsToLink);
 
     emitBuiltProgramInfo(BuiltProgram.get(), ContextImpl);
- 
+
     {
       std::lock_guard<std::mutex> Lock(MNativeProgramsMutex);
       NativePrograms.insert({BuiltProgram.get(), &Img});

sycl/unittests/program_manager/arg_mask/EliminatedArgMask.cpp

@@ -108,17 +108,6 @@ inline ur_result_t redefinedProgramCreateEAM(void *pParams) {
   return UR_RESULT_SUCCESS;
 }
 
-mock::dummy_handle_t_ FixedHandle;
-inline ur_result_t setFixedProgramPtr(void *pParams) {
-  auto params = *static_cast<ur_program_create_with_il_params_t *>(pParams);
-  **params.pphProgram = reinterpret_cast<ur_program_handle_t>(&FixedHandle);
-  return UR_RESULT_SUCCESS;
-}
-inline ur_result_t releaseFixedProgramPtr(void *pParams) {
-  // Do nothing
-  return UR_RESULT_SUCCESS;
-}
-
 class MockHandler : public sycl::handler {
 
 public:
@@ -217,6 +206,41 @@ TEST(EliminatedArgMask, KernelBundleWith2Kernels) {
   EXPECT_EQ(*EliminatedArgMask, ExpElimArgMask);
 }
 
+std::vector<std::unique_ptr<mock::dummy_handle_t_>> UsedProgramHandles;
+std::vector<std::unique_ptr<mock::dummy_handle_t_>> ProgramHandlesToReuse;
+
+inline ur_result_t setFixedProgramPtr(void *pParams) {
+  auto params = *static_cast<ur_program_create_with_il_params_t *>(pParams);
+  if (ProgramHandlesToReuse.size())
+  {
+    auto it = ProgramHandlesToReuse.begin()+1;
+    std::move(ProgramHandlesToReuse.begin(), it, std::back_inserter(UsedProgramHandles));
+    ProgramHandlesToReuse.erase(ProgramHandlesToReuse.begin(), it);
+  }
+  else
+    UsedProgramHandles.push_back(std::make_unique<mock::dummy_handle_t_>(mock::createDummyHandle<ur_program_handle_t>(sizeof(unsigned))));
+  **params.pphProgram = *reinterpret_cast<ur_program_handle_t*>(UsedProgramHandles.back().get());
+  std::cout << "**params.pphProgram = "  << **params.pphProgram << std::endl;
+  return UR_RESULT_SUCCESS;
+}
+inline ur_result_t releaseFixedProgramPtr(void *pParams) {
+  ur_program_handle_t& params = *static_cast<ur_program_handle_t *>(pParams);
+  {
+    auto it = std::find_if(UsedProgramHandles.begin(), UsedProgramHandles.end(), [&params](const std::unique_ptr<mock::dummy_handle_t_>& item){ return *reinterpret_cast<ur_program_handle_t*>(item.get()) == params; });
+    if (it == UsedProgramHandles.end())
+      return UR_RESULT_SUCCESS;
+    std::cout << "releaseFixedProgramPtr = "  << params << std::endl;
+    std::move(it, it + 1, std::back_inserter(ProgramHandlesToReuse));
+    UsedProgramHandles.erase(it, it +1);
+  }
+  return UR_RESULT_SUCCESS;
+}
+
+inline ur_result_t customProgramRetain(void *pParams) {
+ // do nothing
+  return UR_RESULT_SUCCESS;
+}
+
 // It's possible for the same handle to be reused for multiple distinct programs
 // This can happen if a program is released (freeing underlying memory) and then
 // a new program happens to get given that same memory for its handle.
@@ -237,6 +261,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) {
                                               &setFixedProgramPtr);
     mock::getCallbacks().set_replace_callback("urProgramRelease",
                                               &releaseFixedProgramPtr);
+    mock::getCallbacks().set_replace_callback("urProgramRetain",
+                                              &customProgramRetain);
 
     const sycl::device Dev = Plt.get_devices()[0];
     sycl::queue Queue{Dev};
@@ -246,6 +272,7 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) {
     auto Mask = PM.getEliminatedKernelArgMask(ProgBefore, Name);
     EXPECT_NE(Mask, nullptr);
     EXPECT_EQ(Mask->at(0), 1);
+    EXPECT_EQ(UsedProgramHandles.size(), 1u);
   }
 
   {
@@ -256,6 +283,8 @@ TEST(EliminatedArgMask, ReuseOfHandleValues) {
                                               &setFixedProgramPtr);
     mock::getCallbacks().set_replace_callback("urProgramRelease",
                                               &releaseFixedProgramPtr);
+    mock::getCallbacks().set_replace_callback("urProgramRetain",
+                                            &customProgramRetain);
 
     const sycl::device Dev = Plt.get_devices()[0];
     sycl::queue Queue{Dev};