add device image map in ProgramManager

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

Author: jinge90

sycl/source/detail/compiler.hpp

@@ -46,6 +46,8 @@
   "SYCL/specialization constants default values"
 /// PropertySetRegistry::SYCL_DEVICELIB_REQ_MASK defined in PropertySetIO.h
 #define __SYCL_PROPERTY_SET_DEVICELIB_REQ_MASK "SYCL/devicelib req mask"
+/// PropertySetRegistry::SYCL_DEVICELIB_METADATA defined in PropertySetIO.h
+#define __SYCL_PROPERTY_SET_DEVICELIB_METADATA "SYCL/devicelib metadata"
 /// PropertySetRegistry::SYCL_KERNEL_PARAM_OPT_INFO defined in PropertySetIO.h
 #define __SYCL_PROPERTY_SET_KERNEL_PARAM_OPT_INFO "SYCL/kernel param opt"
 /// PropertySetRegistry::SYCL_KERNEL_PROGRAM_METADATA defined in PropertySetIO.h

sycl/source/detail/device_binary_image.cpp

@@ -180,6 +180,7 @@ void RTDeviceBinaryImage::init(sycl_device_binary Bin) {
   SpecConstDefaultValuesMap.init(
       Bin, __SYCL_PROPERTY_SET_SPEC_CONST_DEFAULT_VALUES_MAP);
   DeviceLibReqMask.init(Bin, __SYCL_PROPERTY_SET_DEVICELIB_REQ_MASK);
+  DeviceLibMetaData.init(Bin, __SYCL_PROPERTY_SET_DEVICELIB_METADATA);
   KernelParamOptInfo.init(Bin, __SYCL_PROPERTY_SET_KERNEL_PARAM_OPT_INFO);
   AssertUsed.init(Bin, __SYCL_PROPERTY_SET_SYCL_ASSERT_USED);
   ProgramMetadata.init(Bin, __SYCL_PROPERTY_SET_PROGRAM_METADATA);

sycl/source/detail/device_binary_image.hpp

@@ -214,6 +214,9 @@ class RTDeviceBinaryImage {
     return SpecConstDefaultValuesMap;
   }
   const PropertyRange &getDeviceLibReqMask() const { return DeviceLibReqMask; }
+  const PropertyRange &getDeviceLibMetaData() const {
+    return DeviceLibMetaData;
+  }
   const PropertyRange &getKernelParamOptInfo() const {
     return KernelParamOptInfo;
   }
@@ -246,6 +249,7 @@ class RTDeviceBinaryImage {
   RTDeviceBinaryImage::PropertyRange SpecConstIDMap;
   RTDeviceBinaryImage::PropertyRange SpecConstDefaultValuesMap;
   RTDeviceBinaryImage::PropertyRange DeviceLibReqMask;
+  RTDeviceBinaryImage::PropertyRange DeviceLibMetaData;
   RTDeviceBinaryImage::PropertyRange KernelParamOptInfo;
   RTDeviceBinaryImage::PropertyRange AssertUsed;
   RTDeviceBinaryImage::PropertyRange ProgramMetadata;

sycl/source/detail/program_manager/program_manager.cpp

@@ -752,7 +752,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())
@@ -1772,9 +1773,25 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
     sycl_device_binary RawImg = &(DeviceBinary->DeviceBinaries[I]);
     const sycl_offload_entry EntriesB = RawImg->EntriesBegin;
     const sycl_offload_entry EntriesE = RawImg->EntriesEnd;
-    // Treat the image as empty one
-    if (EntriesB == EntriesE)
+    // Treat the image as empty one or devicelib images
+    if (EntriesB == EntriesE) {
+      std::unique_ptr<RTDeviceBinaryImage> Img =
+          std::make_unique<RTDeviceBinaryImage>(RawImg);
+      const RTDeviceBinaryImage::PropertyRange &DeviceLibMetaProp =
+          Img->getDeviceLibMetaData();
+      if (DeviceLibMetaProp.isAvailable()) {
+        std::lock_guard<std::mutex> DeviceLibImagesGuard(
+            m_DeviceLibImagesMutex);
+        unsigned DeviceLibMetaData =
+            DeviceBinaryProperty(*(DeviceLibMetaProp.begin())).asUint32();
+        // Add device library image to device image map only when current
+        // key has not been inserted.
+        if (m_DeviceLibImages.find(DeviceLibMetaData) ==
+            m_DeviceLibImages.end())
+          m_DeviceLibImages[DeviceLibMetaData] = std::move(Img);
+      }
       continue;
+    }
 
     std::unique_ptr<RTDeviceBinaryImage> Img;
     if (isDeviceImageCompressed(RawImg))

sycl/source/detail/program_manager/program_manager.hpp

@@ -343,6 +343,10 @@ class ProgramManager {
   /// identified by its name.
   using RTDeviceBinaryImageUPtr = std::unique_ptr<RTDeviceBinaryImage>;
 
+  std::mutex m_DeviceLibImagesMutex;
+
+  std::unordered_map<unsigned, RTDeviceBinaryImageUPtr> m_DeviceLibImages;
+
   /// Maps names of kernels to their unique kernel IDs.
   /// TODO: Use std::unordered_set with transparent hash and equality functions
   ///       when C++20 is enabled for the runtime library.