Fix dependency resolution

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/source/detail/device_image_impl.hpp

@@ -746,19 +746,28 @@ class device_image_impl {
           MRTCBinInfo->MIncludePairs, BuildOptions, LogPtr);
 
       auto &PM = detail::ProgramManager::getInstance();
-      std::vector<std::shared_ptr<device_image_impl>> Result;
-      Result.reserve(Binaries->NumDeviceBinaries);
+
+      // Add all binaries and keep the images for processing.
+      std::vector<std::pair<RTDeviceBinaryImage *,
+                            std::shared_ptr<std::vector<kernel_id>>>>
+          NewImages;
+      NewImages.reserve(Binaries->NumDeviceBinaries);
       for (int I = 0; I < Binaries->NumDeviceBinaries; I++) {
         sycl_device_binary Binary = &(Binaries->DeviceBinaries[I]);
-
         RTDeviceBinaryImage *NewImage = nullptr;
         auto KernelIDs = std::make_shared<std::vector<kernel_id>>();
         PM.addImage(Binary, &NewImage, KernelIDs.get());
+        if (NewImage)
+          NewImages.push_back(
+              std::make_pair(std::move(NewImage), std::move(KernelIDs)));
+      }
 
-        // If the image is empty, we can skip it.
-        if (!NewImage)
-          continue;
-
+      // Now bring all images into the proper state. Note that we do this in a
+      // separate pass over NewImages to make sure dependency images have been
+      // registered beforehand.
+      std::vector<std::shared_ptr<device_image_impl>> Result;
+      Result.reserve(NewImages.size());
+      for (auto &[NewImage, KernelIDs] : NewImages) {
         std::set<std::string> KernelNames;
         std::unordered_map<std::string, std::string> MangledKernelNames;
         std::unordered_set<std::string> DeviceGlobalIDSet;
@@ -843,7 +852,26 @@ class device_image_impl {
             std::move(KernelNames), std::move(MangledKernelNames),
             std::string{Prefix}, std::move(DGRegs));
 
-        DevImgPlainWithDeps ImgWithDeps{DevImgImpl};
+        // Resolve dependencies.
+        // TODO: Consider making a collectDeviceImageDeps variant that takes a
+        //       set reference and inserts into that instead.
+        std::set<RTDeviceBinaryImage *> ImgDeps;
+        for (const device &Device : Devices) {
+          std::set<RTDeviceBinaryImage *> DevImgDeps =
+              PM.collectDeviceImageDeps(*NewImage, Device);
+          ImgDeps.insert(DevImgDeps.begin(), DevImgDeps.end());
+        }
+
+        // Pack main image and dependencies together.
+        std::vector<device_image_plain> NewImageAndDeps;
+        NewImageAndDeps.reserve(1 + ImgDeps.size());
+        NewImageAndDeps.push_back(std::move(
+            createSyclObjFromImpl<device_image_plain>(std::move(DevImgImpl))));
+        for (RTDeviceBinaryImage *ImgDep : ImgDeps)
+          NewImageAndDeps.push_back(PM.createDependencyImage(
+              MContext, Devices, ImgDep, bundle_state::input));
+
+        DevImgPlainWithDeps ImgWithDeps(std::move(NewImageAndDeps));
         PM.bringSYCLDeviceImageToState(ImgWithDeps, bundle_state::executable);
         Result.push_back(getSyclObjImpl(ImgWithDeps.getMain()));
       }

sycl/source/detail/program_manager/program_manager.cpp

@@ -2698,32 +2698,36 @@ ProgramManager::getSYCLDeviceImagesWithCompatibleState(
     Images.reserve(Deps.size() + 1);
     Images.push_back(
         createSyclObjFromImpl<device_image_plain>(std::move(MainImpl)));
-    for (RTDeviceBinaryImage *Dep : Deps) {
-      std::shared_ptr<std::vector<sycl::kernel_id>> DepKernelIDs;
-      {
-        std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
-        // 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) &&
-             "State mismatch between main image and its dependency");
-      DeviceImageImplPtr DepImpl = std::make_shared<detail::device_image_impl>(
-          Dep, Ctx, Devs, ImgInfoPair.second.State, DepKernelIDs,
-          /*PIProgram=*/nullptr);
-
+    for (RTDeviceBinaryImage *Dep : Deps)
       Images.push_back(
-          createSyclObjFromImpl<device_image_plain>(std::move(DepImpl)));
-    }
+          createDependencyImage(Ctx, Devs, Dep, ImgInfoPair.second.State));
     SYCLDeviceImages.push_back(std::move(Images));
   }
 
   return SYCLDeviceImages;
 }
 
+device_image_plain ProgramManager::createDependencyImage(
+    const context &Ctx, const std::vector<device> &Devs,
+    RTDeviceBinaryImage *DepImage, bundle_state DepState) {
+  std::shared_ptr<std::vector<sycl::kernel_id>> DepKernelIDs;
+  {
+    std::lock_guard<std::mutex> KernelIDsGuard(m_KernelIDsMutex);
+    // For device library images, they are not in m_BinImg2KernelIDs since
+    // no kernel is included.
+    auto DepIt = m_BinImg2KernelIDs.find(DepImage);
+    if (DepIt != m_BinImg2KernelIDs.end())
+      DepKernelIDs = DepIt->second;
+  }
+
+  assert(DepState == getBinImageState(DepImage) &&
+         "State mismatch between main image and its dependency");
+  DeviceImageImplPtr DepImpl = std::make_shared<detail::device_image_impl>(
+      DepImage, Ctx, Devs, DepState, DepKernelIDs, /*PIProgram=*/nullptr);
+
+  return createSyclObjFromImpl<device_image_plain>(std::move(DepImpl));
+}
+
 void ProgramManager::bringSYCLDeviceImageToState(
     DevImgPlainWithDeps &DeviceImage, bundle_state TargetState) {
   device_image_plain &MainImg = DeviceImage.getMain();

sycl/source/detail/program_manager/program_manager.hpp

@@ -297,6 +297,12 @@ class ProgramManager {
       const context &Ctx, const std::vector<device> &Devs,
       bundle_state TargetState, const std::vector<kernel_id> &KernelIDs = {});
 
+  // Creates a new dependency image for a given dependency binary image.
+  device_image_plain createDependencyImage(const context &Ctx,
+                                           const std::vector<device> &Devs,
+                                           RTDeviceBinaryImage *DepImage,
+                                           bundle_state DepState);
+
   // Bring image to the required state. Does it inplace
   void bringSYCLDeviceImageToState(DevImgPlainWithDeps &DeviceImage,
                                    bundle_state TargetState);
@@ -363,6 +369,12 @@ class ProgramManager {
   std::set<RTDeviceBinaryImage *>
   getRawDeviceImages(const std::vector<kernel_id> &KernelIDs);
 
+  std::set<RTDeviceBinaryImage *>
+  collectDeviceImageDeps(const RTDeviceBinaryImage &Img, const device &Dev);
+  std::set<RTDeviceBinaryImage *>
+  collectDeviceImageDepsForImportedSymbols(const RTDeviceBinaryImage &Img,
+                                           const device &Dev);
+
 private:
   ProgramManager(ProgramManager const &) = delete;
   ProgramManager &operator=(ProgramManager const &) = delete;
@@ -386,11 +398,6 @@ class ProgramManager {
   /// Add info on kernels using local arg into cache
   void cacheKernelImplicitLocalArg(RTDeviceBinaryImage &Img);
 
-  std::set<RTDeviceBinaryImage *>
-  collectDeviceImageDeps(const RTDeviceBinaryImage &Img, const device &Dev);
-  std::set<RTDeviceBinaryImage *>
-  collectDeviceImageDepsForImportedSymbols(const RTDeviceBinaryImage &Img,
-                                           const device &Dev);
   std::set<RTDeviceBinaryImage *>
   collectDependentDeviceImagesForVirtualFunctions(
       const RTDeviceBinaryImage &Img, const device &Dev);