[SYCL] Filter implicit kernel bundle images

The runtime currently uses implicit kernel bundles to allow for
operations such as setting specialization constants. Since the handler
does not know the kernel to launch before it is being launched, the
implicit kernel bundle created is not limited to specific kernels.
As an effect of this, the kernel bundle may try to build more device
images than it needs.
These changes make the runtime filter the device images of implicit
kernel bundles as soon as the relevant kernel is known.

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

Author: steffenlarsen

sycl/include/CL/sycl/handler.hpp

@@ -730,7 +730,8 @@ class __SYCL_EXPORT handler {
     // If the kernel lambda is callable with a kernel_handler argument, manifest
     // the associated kernel handler.
     if (IsCallableWithKernelHandler) {
-      getOrInsertHandlerKernelBundle(/*Insert=*/true);
+      getOrInsertFilteredHandlerKernelBundle(/*Insert=*/true,
+                                             {get_kernel_id<KernelName>()});
     }
   }
 
@@ -1272,6 +1273,9 @@ class __SYCL_EXPORT handler {
   std::shared_ptr<detail::kernel_bundle_impl>
   getOrInsertHandlerKernelBundle(bool Insert) const;
 
+  std::shared_ptr<detail::kernel_bundle_impl>
+  getOrInsertFilteredHandlerKernelBundle(bool Insert, kernel_id KernelId) const;
+
   void setHandlerKernelBundle(
       const std::shared_ptr<detail::kernel_bundle_impl> &NewKernelBundleImpPtr);
 

sycl/source/detail/kernel_bundle_impl.hpp

@@ -496,6 +496,16 @@ class kernel_bundle_impl {
 
   bool isInterop() const { return MIsInterop; }
 
+  void filterImages(const kernel_id &KernelId) {
+    std::vector<device_image_plain> FilteredDeviceImages;
+    std::copy_if(MDeviceImages.begin(), MDeviceImages.end(),
+                 std::back_inserter(FilteredDeviceImages),
+                 [&KernelId](const device_image_plain &Img) {
+                   return Img.has_kernel(KernelId);
+                 });
+    MDeviceImages = FilteredDeviceImages;
+  }
+
 private:
   context MContext;
   std::vector<device> MDevices;

sycl/source/handler.cpp

@@ -128,6 +128,67 @@ handler::getOrInsertHandlerKernelBundle(bool Insert) const {
   return KernelBundleImpPtr;
 }
 
+// Returns a shared_ptr to kernel_bundle stored in the extended members vector.
+//
+// If there is no kernel_bundle created:
+// returns newly created kernel_bundle with KernelId if Insert is true
+// returns shared_ptr(nullptr) if Insert is false
+//
+// If there already existed a kernel_bundle, the underlying device images are
+// filtered such that only the ones containing KernelId remain in the
+// kernel_bundle.
+std::shared_ptr<detail::kernel_bundle_impl>
+handler::getOrInsertFilteredHandlerKernelBundle(bool Insert,
+                                                kernel_id KernelId) const {
+
+  std::lock_guard<std::mutex> Lock(
+      detail::GlobalHandler::instance().getHandlerExtendedMembersMutex());
+
+  assert(!MSharedPtrStorage.empty());
+
+  std::shared_ptr<std::vector<detail::ExtendedMemberT>> ExtendedMembersVec =
+      detail::convertToExtendedMembers(MSharedPtrStorage[0]);
+  // Look for the kernel bundle in extended members
+  std::shared_ptr<detail::kernel_bundle_impl> KernelBundleImpPtr;
+  for (const detail::ExtendedMemberT &EMember : *ExtendedMembersVec)
+    if (detail::ExtendedMembersType::HANDLER_KERNEL_BUNDLE == EMember.MType) {
+      KernelBundleImpPtr =
+          std::static_pointer_cast<detail::kernel_bundle_impl>(EMember.MData);
+      break;
+    }
+
+  // No kernel bundle yet, create one
+  if (!KernelBundleImpPtr) {
+    if (Insert) {
+      KernelBundleImpPtr =
+          detail::getSyclObjImpl(get_kernel_bundle<bundle_state::input>(
+              MQueue->get_context(), {KernelId}));
+      if (KernelBundleImpPtr->empty()) {
+        KernelBundleImpPtr =
+            detail::getSyclObjImpl(get_kernel_bundle<bundle_state::executable>(
+                MQueue->get_context(), {KernelId}));
+      }
+
+      detail::ExtendedMemberT EMember = {
+          detail::ExtendedMembersType::HANDLER_KERNEL_BUNDLE,
+          KernelBundleImpPtr};
+      ExtendedMembersVec->push_back(EMember);
+    }
+    return KernelBundleImpPtr;
+  }
+
+  auto HandlerImplMember = (*ExtendedMembersVec)[0];
+  assert(detail::ExtendedMembersType::HANDLER_IMPL == HandlerImplMember.MType);
+  auto HandlerImpl =
+      std::static_pointer_cast<detail::handler_impl>(HandlerImplMember.MData);
+
+  // Kernel bundles set explicitly by the user must not be filtered
+  if (!HandlerImpl->isStateExplicitKernelBundle())
+    KernelBundleImpPtr->filterImages(KernelId);
+
+  return KernelBundleImpPtr;
+}
+
 // Sets kernel bundle to the provided one. Either replaces existing one or
 // create a new entry in the extended members vector.
 void handler::setHandlerKernelBundle(

sycl/test/abi/sycl_symbols_linux.dump

@@ -4320,6 +4320,7 @@ _ZNK2cl4sycl7context9getNativeEv
 _ZNK2cl4sycl7handler14getHandlerImplEv
 _ZNK2cl4sycl7handler27isStateExplicitKernelBundleEv
 _ZNK2cl4sycl7handler30getOrInsertHandlerKernelBundleEb
+_ZNK2cl4sycl7handler38getOrInsertFilteredHandlerKernelBundleEbNS0_9kernel_idE
 _ZNK2cl4sycl7program10get_kernelENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE
 _ZNK2cl4sycl7program10get_kernelENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEEb
 _ZNK2cl4sycl7program10has_kernelENSt7__cxx1112basic_stringIcSt11char_traitsIcESaIcEEE