[SYCL] Keep multiple copies for bf16 device library image

sycl/source/detail/program_manager/program_manager.cpp

@@ -1837,84 +1837,27 @@ ProgramManager::kernelImplicitLocalArgPos(const std::string &KernelName) const {
   return {};
 }
 
-static bool shouldSkipEmptyImage(sycl_device_binary RawImg, bool IsRTC) {
-  // For bfloat16 device library image, we should keep it. However, in some
-  // scenario, __sycl_register_lib can be called multiple times and the same
-  // bfloat16 device library image may be handled multiple times which is not
-  // needed. 2 static bool variables are created to record whether native or
-  // fallback bfloat16 device library image has been handled, if yes, we just
-  // need to skip it.
-  // We cannot prevent redundant loads of device library images if they are part
-  // of a runtime-compiled device binary, as these will be freed when the
-  // corresponding kernel bundle is destroyed. Hence, normal kernels cannot rely
-  // on the presence of RTC device library images.
+static bool shouldSkipEmptyImage(sycl_device_binary RawImg) {
+  // For bfloat16 device library image, we should keep it although it doesn't
+  // include any kernel.
   sycl_device_binary_property_set ImgPS;
-  static bool IsNativeBF16DeviceLibHandled = false;
-  static bool IsFallbackBF16DeviceLibHandled = false;
   for (ImgPS = RawImg->PropertySetsBegin; ImgPS != RawImg->PropertySetsEnd;
        ++ImgPS) {
     if (ImgPS->Name &&
-        !strcmp(__SYCL_PROPERTY_SET_DEVICELIB_METADATA, ImgPS->Name)) {
-      sycl_device_binary_property ImgP;
-      for (ImgP = ImgPS->PropertiesBegin; ImgP != ImgPS->PropertiesEnd;
-           ++ImgP) {
-        if (ImgP->Name && !strcmp("bfloat16", ImgP->Name) &&
-            (ImgP->Type == SYCL_PROPERTY_TYPE_UINT32))
-          break;
-      }
-      if (ImgP == ImgPS->PropertiesEnd)
-        return true;
-
-      // A valid bfloat16 device library image is found here.
-      // If it originated from RTC, we cannot skip it, but do not mark it as
-      // being present.
-      if (IsRTC)
-        return false;
-
-      // Otherwise, we need to check whether it has been handled already.
-      uint32_t BF16NativeVal = DeviceBinaryProperty(ImgP).asUint32();
-      if (((BF16NativeVal == 0) && IsFallbackBF16DeviceLibHandled) ||
-          ((BF16NativeVal == 1) && IsNativeBF16DeviceLibHandled))
-        return true;
-
-      if (BF16NativeVal == 0)
-        IsFallbackBF16DeviceLibHandled = true;
-      else
-        IsNativeBF16DeviceLibHandled = true;
-
+        !strcmp(__SYCL_PROPERTY_SET_DEVICELIB_METADATA, ImgPS->Name))
       return false;
-    }
   }
-  return true;
-}
 
-static bool isCompiledAtRuntime(sycl_device_binaries DeviceBinary) {
-  // Check whether the first device binary contains a legacy format offload
-  // entry with a `$` in its name.
-  if (DeviceBinary->NumDeviceBinaries > 0) {
-    sycl_device_binary Binary = DeviceBinary->DeviceBinaries;
-    if (Binary->EntriesBegin != Binary->EntriesEnd) {
-      sycl_offload_entry Entry = Binary->EntriesBegin;
-      if (!Entry->IsNewOffloadEntryType() &&
-          std::string_view{Entry->name}.find('$') != std::string_view::npos) {
-        return true;
-      }
-    }
-  }
-  return false;
+  return true;
 }
 
 void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
   const bool DumpImages = std::getenv("SYCL_DUMP_IMAGES") && !m_UseSpvFile;
-  const bool IsRTC = isCompiledAtRuntime(DeviceBinary);
   for (int I = 0; I < DeviceBinary->NumDeviceBinaries; I++) {
     sycl_device_binary RawImg = &(DeviceBinary->DeviceBinaries[I]);
     const sycl_offload_entry EntriesB = RawImg->EntriesBegin;
     const sycl_offload_entry EntriesE = RawImg->EntriesEnd;
-    // If the image does not contain kernels, skip it unless it is one of the
-    // bfloat16 device libraries, and it wasn't loaded before or resulted from
-    // runtime compilation.
-    if ((EntriesB == EntriesE) && shouldSkipEmptyImage(RawImg, IsRTC))
+    if ((EntriesB == EntriesE) && shouldSkipEmptyImage(RawImg))
       continue;
 
     std::unique_ptr<RTDeviceBinaryImage> Img;
@@ -1946,6 +1889,25 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
     // Fill maps for kernel bundles
     std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
 
+    // For bfloat16 device library image, it doesn't include any kernel, device
+    // global, virtual function, so just skip adding it to any related maps. We
+    // only need to 1) add exported symbols to m_ExportedSymbolImages, and 2)
+    // add the device image to m_Bfloat16DeviceLibImages.
+    {
+      auto Bfloat16DeviceLibProp = Img->getDeviceLibMetadata();
+      if (Bfloat16DeviceLibProp.isAvailable()) {
+        uint32_t LibVersion = DeviceBinaryProperty(*(Bfloat16DeviceLibProp.begin())).asUint32();
+        if (m_Bfloat16DeviceLibImages.count(LibVersion) > 0)
+          continue;
+        for (const sycl_device_binary_property &ESProp :
+             Img->getExportedSymbols()) {
+          m_ExportedSymbolImages.insert({ESProp->Name, Img.get()});
+        }
+        m_Bfloat16DeviceLibImages.insert({LibVersion, std::move(Img)});
+        continue;
+      }
+    }
+
     // Register all exported symbols
     for (const sycl_device_binary_property &ESProp :
          Img->getExportedSymbols()) {
@@ -2110,19 +2072,14 @@ void ProgramManager::addImages(sycl_device_binaries DeviceBinary) {
 }
 
 void ProgramManager::removeImages(sycl_device_binaries DeviceBinary) {
-  bool IsRTC = isCompiledAtRuntime(DeviceBinary);
   for (int I = 0; I < DeviceBinary->NumDeviceBinaries; I++) {
     sycl_device_binary RawImg = &(DeviceBinary->DeviceBinaries[I]);
     auto DevImgIt = m_DeviceImages.find(RawImg);
     if (DevImgIt == m_DeviceImages.end())
       continue;
     const sycl_offload_entry EntriesB = RawImg->EntriesBegin;
     const sycl_offload_entry EntriesE = RawImg->EntriesEnd;
-    // Skip clean up if there are no offload entries, unless `DeviceBinary`
-    // resulted from runtime compilation: Then, this is one of the `bfloat16`
-    // device libraries, so we want to make sure that the image and its exported
-    // symbols are removed from the program manager's maps.
-    if (EntriesB == EntriesE && !IsRTC)
+    if ((EntriesB == EntriesE))
       continue;
 
     RTDeviceBinaryImage *Img = DevImgIt->second.get();
@@ -2650,7 +2607,11 @@ ProgramManager::getSYCLDeviceImagesWithCompatibleState(
       std::shared_ptr<std::vector<sycl::kernel_id>> DepKernelIDs;
       {
         std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
-        DepKernelIDs = m_BinImg2KernelIDs[Dep];
+        // For device library images, they are not in m_BinImg2KernelIDs since
+        // no kernel is included.
+        auto DepIt = m_BinImg2KernelIDs.find(Dep);
+        if (DepIt != m_BinImg2KernelIDs.end())
+          DepKernelIDs = DepIt->second;
       }
 
       assert(ImgInfoPair.second.State == getBinImageState(Dep) &&
@@ -2863,6 +2824,11 @@ static void mergeImageData(const std::vector<device_image_plain> &Imgs,
   for (const device_image_plain &Img : Imgs) {
     const std::shared_ptr<device_image_impl> &DeviceImageImpl =
         getSyclObjImpl(Img);
+    auto BinImgRef = DeviceImageImpl->get_bin_image_ref();
+    // For bfloat16 deice library image, no kernels, spec const are included,
+    // so we just skip merging data.
+    if (BinImgRef && BinImgRef->getDeviceLibMetadata().isAvailable())
+      continue;
     // Duplicates are not expected here, otherwise urProgramLink should fail
     KernelIDs.insert(KernelIDs.end(),
                      DeviceImageImpl->get_kernel_ids_ptr()->begin(),

sycl/source/detail/program_manager/program_manager.hpp

@@ -498,6 +498,13 @@ 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, RTDeviceBinaryImageUPtr>
+      m_Bfloat16DeviceLibImages;
+
   friend class ::ProgramManagerTest;
 };
 } // namespace detail