Fix sycl.cpp and use std::array to replace map for bf16 device images

Signed-off-by: jinge90 <[email protected]>

Author: jinge90

sycl/source/detail/program_manager/program_manager.cpp

@@ -1945,7 +1945,9 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
     // library images will not be erased unless program manager is destroyed.
     {
       if (IsBfloat16DeviceLib) {
-        if (m_Bfloat16DeviceLibImages.count(Bfloat16DeviceLibVersion) > 0)
+        assert((Bfloat16DeviceLibVersion < 2) &&
+               "Invalid Bfloat16 Device Library Index.");
+        if (m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion].get())
           continue;
         size_t ImgSize =
             static_cast<size_t>(RawImg->BinaryEnd - RawImg->BinaryStart);
@@ -1960,8 +1962,8 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
           m_ExportedSymbolImages.insert(
               {ESProp->Name, DynBfloat16DeviceLibImg.get()});
         }
-        m_Bfloat16DeviceLibImages.insert(
-            {Bfloat16DeviceLibVersion, std::move(DynBfloat16DeviceLibImg)});
+        m_Bfloat16DeviceLibImages[Bfloat16DeviceLibVersion] =
+            std::move(DynBfloat16DeviceLibImg);
         continue;
       }
     }
@@ -2137,7 +2139,7 @@ void ProgramManager::removeImages(sycl_device_binaries DeviceBinary) {
       continue;
     const sycl_offload_entry EntriesB = RawImg->EntriesBegin;
     const sycl_offload_entry EntriesE = RawImg->EntriesEnd;
-    if ((EntriesB == EntriesE))
+    if (EntriesB == EntriesE)
       continue;
 
     RTDeviceBinaryImage *Img = DevImgIt->second.get();

sycl/source/detail/program_manager/program_manager.hpp

@@ -24,6 +24,7 @@
 #include <sycl/device.hpp>
 #include <sycl/kernel_bundle.hpp>
 
+#include <array>
 #include <cstdint>
 #include <map>
 #include <memory>
@@ -502,12 +503,11 @@ class ProgramManager {
       std::map<std::vector<unsigned char>, ur_kernel_handle_t>;
   std::unordered_map<std::string, MaterializedEntries> m_MaterializedKernels;
 
-  // Holds bfloat16 device library images, the key is 0 for fallback version
-  // and 1 for native version. These bfloat16 device library images are
-  // provided by compiler long time ago, we expect no further update, so
-  // keeping 1 copy should be OK.
-  std::unordered_map<uint32_t, DynRTDeviceBinaryImageUPtr>
-      m_Bfloat16DeviceLibImages;
+  // Holds bfloat16 device library images, the 1st element is for fallback
+  // version and 2nd is for native version. These bfloat16 device library
+  // images are provided by compiler long time ago, we expect no further
+  // update, so keeping 1 copy should be OK.
+  std::array<DynRTDeviceBinaryImageUPtr, 2> m_Bfloat16DeviceLibImages;
 
   friend class ::ProgramManagerTest;
 };

sycl/test-e2e/KernelCompiler/sycl.cpp

@@ -535,7 +535,7 @@ int main() {
   }
   // Run test_device_libraries twice to verify bfloat16 device library.
   return test_build_and_run(q) || test_device_code_split(q) ||
-         test_device_libraries(q) || test_device_libraries(q) ||
+         test_device_libraries(q) || test_esimd(q) ||
          test_device_libraries(q) || test_unsupported_options(q) ||
          test_error(q) || test_no_visible_ids(q) || test_warning(q);
 #else