Check kernel info ownership with mixed image origins

Author: sergey-semenov

sycl/source/detail/kernel_impl.cpp

@@ -59,11 +59,11 @@ kernel_impl::kernel_impl(Managed<ur_kernel_handle_t> &&Kernel,
       MKernelBundleImpl(KernelBundleImpl.shared_from_this()),
       MIsInterop(MDeviceImageImpl->getOriginMask() & ImageOriginInterop),
       MKernelArgMaskPtr{ArgMask}, MCacheMutex{CacheMutex},
-      MOwnsDeviceKernelInfo(MDeviceImageImpl->getOriginMask() &
-                            ~ImageOriginSYCLOffline),
+      MOwnsDeviceKernelInfo(checkOwnsDeviceKernelInfo()),
       MDeviceKernelInfo(MOwnsDeviceKernelInfo
                             ? createCompileTimeKernelInfo(getName())
                             : createCompileTimeKernelInfo()) {
+
   // Enable USM indirect access for interop and non-sycl-jit source kernels.
   // sycl-jit kernels will enable this if needed through the regular kernel
   // path.
@@ -123,6 +123,17 @@ std::string_view kernel_impl::getName() const {
   return MName;
 }
 
+bool kernel_impl::checkOwnsDeviceKernelInfo() {
+  // If the image originates from something other than standard offline
+  // compilation, this kernel needs to own its info structure.
+  // We could also have a mixed origin image, in which case the device kernel
+  // info might reside in program manager.
+  return MDeviceImageImpl->getOriginMask() != ImageOriginSYCLOffline &&
+         (!(MDeviceImageImpl->getOriginMask() & ImageOriginSYCLOffline) ||
+          !ProgramManager::getInstance().tryGetDeviceKernelInfo(
+              static_cast<KernelNameStrT>(getName())));
+}
+
 bool kernel_impl::isBuiltInKernel(device_impl &Device) const {
   auto BuiltInKernels = Device.get_info<info::device::built_in_kernel_ids>();
   if (BuiltInKernels.empty())

sycl/source/detail/kernel_impl.hpp

@@ -239,6 +239,7 @@ class kernel_impl {
   std::mutex *getCacheMutex() const { return MCacheMutex; }
   std::string_view getName() const;
 
+  bool checkOwnsDeviceKernelInfo();
   DeviceKernelInfo &getDeviceKernelInfo() {
     return MOwnsDeviceKernelInfo
                ? MDeviceKernelInfo

sycl/source/detail/program_manager/program_manager.cpp

@@ -1833,6 +1833,13 @@ ProgramManager::getDeviceKernelInfo(KernelNameStrRefT KernelName) {
   return It->second;
 }
 
+DeviceKernelInfo *
+ProgramManager::tryGetDeviceKernelInfo(KernelNameStrRefT KernelName) {
+  std::lock_guard<std::mutex> Guard(m_DeviceKernelInfoMapMutex);
+  auto It = m_DeviceKernelInfoMap.find(KernelName);
+  return It != m_DeviceKernelInfoMap.end() ? &It->second : nullptr;
+}
+
 static bool isBfloat16DeviceLibImage(sycl_device_binary RawImg,
                                      uint32_t *LibVersion = nullptr) {
   sycl_device_binary_property_set ImgPS;

sycl/source/detail/program_manager/program_manager.hpp

@@ -372,6 +372,7 @@ class ProgramManager {
 
   DeviceKernelInfo &getDeviceKernelInfo(const CompileTimeKernelInfoTy &Info);
   DeviceKernelInfo &getDeviceKernelInfo(KernelNameStrRefT KernelName);
+  DeviceKernelInfo *tryGetDeviceKernelInfo(KernelNameStrRefT KernelName);
 
   std::set<const RTDeviceBinaryImage *>
   getRawDeviceImages(const std::vector<kernel_id> &KernelIDs);