[SYCL] Change NativePrograms.insert to [] access (#14873)

`map.insert` doesn't insert values if the set already contains them.
This can happen when UR/PI happens to reuse the same program pointer
that it used for a previous program.

--

This was causing some tests in the PI 2 UR conversion to randomly fail,
including at least #14765 .

Fixes #14819.

sycl/source/detail/program_manager/program_manager.cpp

@@ -229,7 +229,7 @@ ProgramManager::createURProgram(const RTDeviceBinaryImage &Img,
   {
     std::lock_guard<std::mutex> Lock(MNativeProgramsMutex);
     // associate the UR program with the image it was created for
-    NativePrograms.insert({Res, &Img});
+    NativePrograms[Res] = &Img;
   }
 
   Ctx->addDeviceGlobalInitializer(Res, {Device}, &Img);
@@ -840,9 +840,9 @@ ur_program_handle_t ProgramManager::getBuiltURProgram(
 
     {
       std::lock_guard<std::mutex> Lock(MNativeProgramsMutex);
-      NativePrograms.insert({BuiltProgram.get(), &Img});
+      NativePrograms[BuiltProgram.get()] = &Img;
       for (RTDeviceBinaryImage *LinkedImg : DeviceImagesToLink) {
-        NativePrograms.insert({BuiltProgram.get(), LinkedImg});
+        NativePrograms[BuiltProgram.get()] = LinkedImg;
       }
     }
 
@@ -2500,7 +2500,7 @@ device_image_plain ProgramManager::build(const device_image_plain &DeviceImage,
 
     {
       std::lock_guard<std::mutex> Lock(MNativeProgramsMutex);
-      NativePrograms.insert({BuiltProgram.get(), &Img});
+      NativePrograms[BuiltProgram.get()] = &Img;
     }
 
     ContextImpl->addDeviceGlobalInitializer(BuiltProgram.get(), Devs, &Img);

sycl/unittests/program_manager/arg_mask/EliminatedArgMask.cpp

@@ -20,8 +20,10 @@
 
 class EAMTestKernel;
 class EAMTestKernel2;
+class EAMTestKernel3;
 constexpr const char EAMTestKernelName[] = "EAMTestKernel";
 constexpr const char EAMTestKernel2Name[] = "EAMTestKernel2";
+constexpr const char EAMTestKernel3Name[] = "EAMTestKernel3";
 constexpr unsigned EAMTestKernelNumArgs = 4;
 
 namespace sycl {
@@ -39,6 +41,12 @@ struct KernelInfo<EAMTestKernel2> : public unittest::MockKernelInfoBase {
   static constexpr const char *getName() { return EAMTestKernel2Name; }
 };
 
+template <>
+struct KernelInfo<EAMTestKernel3> : public unittest::MockKernelInfoBase {
+  static constexpr unsigned getNumParams() { return EAMTestKernelNumArgs; }
+  static constexpr const char *getName() { return EAMTestKernel3Name; }
+};
+
 } // namespace detail
 } // namespace _V1
 } // namespace sycl
@@ -90,10 +98,39 @@ static sycl::unittest::UrImage generateEAMTestKernel2Image() {
   return Img;
 }
 
+static sycl::unittest::UrImage generateEAMTestKernel3Image() {
+  using namespace sycl::unittest;
+
+  // Eliminated arguments are 2nd and 4th.
+  std::vector<unsigned char> KernelEAM{0b00001010};
+  UrProperty EAMKernelPOI = makeKernelParamOptInfo(
+      EAMTestKernel3Name, EAMTestKernelNumArgs, KernelEAM);
+  UrArray<UrProperty> ImgKPOI{std::move(EAMKernelPOI)};
+
+  UrPropertySet PropSet;
+  PropSet.insert(__SYCL_PROPERTY_SET_KERNEL_PARAM_OPT_INFO, std::move(ImgKPOI));
+
+  std::vector<unsigned char> Bin{0, 1, 2, 3, 4, 5}; // Random data
+
+  UrArray<UrOffloadEntry> Entries = makeEmptyKernels({EAMTestKernel3Name});
+
+  UrImage Img{SYCL_DEVICE_BINARY_TYPE_SPIRV,       // Format
+              __SYCL_DEVICE_BINARY_TARGET_SPIRV64, // DeviceTargetSpec
+              "",                                  // Compile options
+              "",                                  // Link options
+              std::move(Bin),
+              std::move(Entries),
+              std::move(PropSet)};
+
+  return Img;
+}
+
 static sycl::unittest::UrImage EAMImg = generateEAMTestKernelImage();
 static sycl::unittest::UrImage EAM2Img = generateEAMTestKernel2Image();
+static sycl::unittest::UrImage EAM3Img = generateEAMTestKernel3Image();
 static sycl::unittest::UrImageArray<1> EAMImgArray{&EAMImg};
 static sycl::unittest::UrImageArray<1> EAM2ImgArray{&EAM2Img};
+static sycl::unittest::UrImageArray<1> EAM3ImgArray{&EAM3Img};
 
 // ur_program_handle_t address is used as a key for ProgramManager::NativePrograms
 // storage. redefinedProgramLinkCommon makes ur_program_handle_t address equal to 0x1.
@@ -106,6 +143,17 @@ 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:
@@ -203,3 +251,58 @@ TEST(EliminatedArgMask, KernelBundleWith2Kernels) {
 
   EXPECT_EQ(*EliminatedArgMask, ExpElimArgMask);
 }
+
+// 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.
+// The ProgramContext stores a map with `ur_program_handle_t`s, which are never
+// cleared. This test ensures that newer `ur_program_handle_t`s with the same
+// values override older ones.
+TEST(EliminatedArgMask, ReuseOfHandleValues) {
+  sycl::detail::ProgramManager &PM =
+      sycl::detail::ProgramManager::getInstance();
+
+  ur_program_handle_t ProgBefore = nullptr;
+  ur_program_handle_t ProgAfter = nullptr;
+  {
+    auto Name = sycl::detail::KernelInfo<EAMTestKernel>::getName();
+    sycl::unittest::UrMock<> Mock;
+    sycl::platform Plt = sycl::platform();
+    mock::getCallbacks().set_replace_callback("urProgramCreateWithIL",
+                                              &setFixedProgramPtr);
+    mock::getCallbacks().set_replace_callback("urProgramRelease",
+                                              &releaseFixedProgramPtr);
+
+    const sycl::device Dev = Plt.get_devices()[0];
+    sycl::queue Queue{Dev};
+    auto Ctx = Queue.get_context();
+    ProgBefore = PM.getBuiltURProgram(sycl::detail::getSyclObjImpl(Ctx),
+                                      sycl::detail::getSyclObjImpl(Dev), Name);
+    auto Mask = PM.getEliminatedKernelArgMask(ProgBefore, Name);
+    EXPECT_NE(Mask, nullptr);
+    EXPECT_EQ(Mask->at(0), 1);
+  }
+
+  {
+    auto Name = sycl::detail::KernelInfo<EAMTestKernel3>::getName();
+    sycl::unittest::UrMock<> Mock;
+    sycl::platform Plt = sycl::platform();
+    mock::getCallbacks().set_replace_callback("urProgramCreateWithIL",
+                                              &setFixedProgramPtr);
+    mock::getCallbacks().set_replace_callback("urProgramRelease",
+                                              &releaseFixedProgramPtr);
+
+    const sycl::device Dev = Plt.get_devices()[0];
+    sycl::queue Queue{Dev};
+    auto Ctx = Queue.get_context();
+    ProgAfter = PM.getBuiltURProgram(sycl::detail::getSyclObjImpl(Ctx),
+                                     sycl::detail::getSyclObjImpl(Dev), Name);
+    auto Mask = PM.getEliminatedKernelArgMask(ProgAfter, Name);
+    EXPECT_NE(Mask, nullptr);
+    EXPECT_EQ(Mask->at(0), 0);
+  }
+
+  // Verify that the test is behaving correctly and that the pointer is being
+  // reused
+  EXPECT_EQ(ProgBefore, ProgAfter);
+}