[SYCL] Complete transition to Managed<ur_program_handle_t> RAII model (#19557)

Started in #19536.

Non-NFC because fixes a few resource leaks as can be seen in the updated
test.

Author: aelovikov-intel

sycl/source/backend.cpp

@@ -196,7 +196,7 @@ make_kernel_bundle(ur_native_handle_t NativeHandle,
   adapter_impl &Adapter = getAdapter(Backend);
   context_impl &ContextImpl = *getSyclObjImpl(TargetContext);
 
-  ur_program_handle_t UrProgram = nullptr;
+  Managed<ur_program_handle_t> UrProgram{Adapter};
   ur_program_native_properties_t Properties{};
   Properties.stype = UR_STRUCTURE_TYPE_PROGRAM_NATIVE_PROPERTIES;
   Properties.isNativeHandleOwned = !KeepOwnership;
@@ -258,18 +258,19 @@ make_kernel_bundle(ur_native_handle_t NativeHandle,
             "Program and kernel_bundle state mismatch " +
                 detail::codeToString(UR_RESULT_ERROR_INVALID_VALUE));
       if (State == bundle_state::executable) {
-        ur_program_handle_t UrLinkedProgram = nullptr;
+        Managed<ur_program_handle_t> UrLinkedProgram{Adapter};
+        ur_program_handle_t ProgramsToLink[] = {UrProgram};
         auto Res = Adapter.call_nocheck<UrApiKind::urProgramLinkExp>(
-            ContextImpl.getHandleRef(), 1u, &Dev, 1u, &UrProgram, nullptr,
+            ContextImpl.getHandleRef(), 1u, &Dev, 1u, ProgramsToLink, nullptr,
             &UrLinkedProgram);
         if (Res == UR_RESULT_ERROR_UNSUPPORTED_FEATURE) {
           Res = Adapter.call_nocheck<UrApiKind::urProgramLink>(
-              ContextImpl.getHandleRef(), 1u, &UrProgram, nullptr,
+              ContextImpl.getHandleRef(), 1u, ProgramsToLink, nullptr,
               &UrLinkedProgram);
         }
         Adapter.checkUrResult<errc::build>(Res);
         if (UrLinkedProgram != nullptr) {
-          UrProgram = UrLinkedProgram;
+          UrProgram = std::move(UrLinkedProgram);
         }
       }
       break;
@@ -301,9 +302,9 @@ make_kernel_bundle(ur_native_handle_t NativeHandle,
   // do the same to user images, since they may contain references to undefined
   // symbols (e.g. when kernel_bundle is supposed to be joined with another).
   auto KernelIDs = std::make_shared<std::vector<kernel_id>>();
-  auto DevImgImpl =
-      device_image_impl::create(nullptr, TargetContext, Devices, State,
-                                KernelIDs, UrProgram, ImageOriginInterop);
+  auto DevImgImpl = device_image_impl::create(
+      nullptr, TargetContext, Devices, State, KernelIDs, std::move(UrProgram),
+      ImageOriginInterop);
   device_image_plain DevImg{DevImgImpl};
 
   return kernel_bundle_impl::create(TargetContext, Devices, DevImg);

sycl/source/detail/adapter_impl.hpp

@@ -244,6 +244,10 @@ template <typename URResource> class Managed {
     if constexpr (std::is_same_v<URResource, ur_program_handle_t>)
       return UrApiKind::urProgramRelease;
   }();
+  static constexpr auto Retain = []() constexpr {
+    if constexpr (std::is_same_v<URResource, ur_program_handle_t>)
+      return UrApiKind::urProgramRetain;
+  }();
 
 public:
   Managed() = default;
@@ -258,6 +262,7 @@ template <typename URResource> class Managed {
   Managed &operator=(Managed &&Other) {
     if (R)
       Adapter->call<Release>(R);
+
     R = Other.R;
     Other.R = nullptr;
     Adapter = Other.Adapter;
@@ -285,6 +290,18 @@ template <typename URResource> class Managed {
     Adapter->call<Release>(R);
   }
 
+  Managed retain() {
+    assert(R && "Cannot retain unintialized resource!");
+    Adapter->call<Retain>(R);
+    return Managed{R, *Adapter};
+  }
+
+  bool operator==(const Managed &Other) const {
+    assert((!Adapter || !Other.Adapter || Adapter == Other.Adapter) &&
+           "Objects must belong to the same adapter!");
+    return R == Other.R;
+  }
+
 private:
   URResource R = nullptr;
   adapter_impl *Adapter = nullptr;

sycl/source/detail/device_image_impl.hpp

@@ -257,11 +257,11 @@ class device_image_impl
   device_image_impl(const RTDeviceBinaryImage *BinImage, context Context,
                     devices_range Devices, bundle_state State,
                     std::shared_ptr<std::vector<kernel_id>> KernelIDs,
-                    ur_program_handle_t Program, uint8_t Origins, private_tag)
+                    Managed<ur_program_handle_t> &&Program, uint8_t Origins,
+                    private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(Devices.to<std::vector<device_impl *>>()), MState(State),
-        MProgram(Program, getSyclObjImpl(MContext)->getAdapter()),
-        MKernelIDs(std::move(KernelIDs)),
+        MProgram(std::move(Program)), MKernelIDs(std::move(KernelIDs)),
         MSpecConstsDefValBlob(getSpecConstsDefValBlob()), MOrigins(Origins) {
     updateSpecConstSymMap();
     if (BinImage && (MOrigins & ImageOriginSYCLBIN)) {
@@ -287,40 +287,23 @@ class device_image_impl
       const RTDeviceBinaryImage *BinImage, const context &Context,
       devices_range Devices, bundle_state State,
       std::shared_ptr<std::vector<kernel_id>> KernelIDs,
-      ur_program_handle_t Program, const SpecConstMapT &SpecConstMap,
+      Managed<ur_program_handle_t> &&Program, const SpecConstMapT &SpecConstMap,
       const std::vector<unsigned char> &SpecConstsBlob, uint8_t Origins,
       std::optional<KernelCompilerBinaryInfo> &&RTCInfo,
       KernelNameSetT &&KernelNames,
       KernelNameToArgMaskMap &&EliminatedKernelArgMasks,
       std::unique_ptr<DynRTDeviceBinaryImage> &&MergedImageStorage, private_tag)
       : MBinImage(BinImage), MContext(std::move(Context)),
         MDevices(Devices.to<std::vector<device_impl *>>()), MState(State),
-        MProgram(Program, getSyclObjImpl(MContext)->getAdapter()),
-        MKernelIDs(std::move(KernelIDs)), MKernelNames{std::move(KernelNames)},
+        MProgram(std::move(Program)), MKernelIDs(std::move(KernelIDs)),
+        MKernelNames{std::move(KernelNames)},
         MEliminatedKernelArgMasks{std::move(EliminatedKernelArgMasks)},
         MSpecConstsBlob(SpecConstsBlob),
         MSpecConstsDefValBlob(getSpecConstsDefValBlob()),
         MSpecConstSymMap(SpecConstMap), MOrigins(Origins),
         MRTCBinInfo(std::move(RTCInfo)),
         MMergedImageStorage(std::move(MergedImageStorage)) {}
 
-  device_image_impl(const RTDeviceBinaryImage *BinImage, const context &Context,
-                    devices_range Devices, bundle_state State,
-                    ur_program_handle_t Program, syclex::source_language Lang,
-                    KernelNameSetT &&KernelNames,
-                    KernelNameToArgMaskMap &&EliminatedKernelArgMasks,
-                    private_tag)
-      : MBinImage(BinImage), MContext(std::move(Context)),
-        MDevices(Devices.to<std::vector<device_impl *>>()), MState(State),
-        MProgram(Program, getSyclObjImpl(MContext)->getAdapter()),
-        MKernelNames{std::move(KernelNames)},
-        MEliminatedKernelArgMasks{std::move(EliminatedKernelArgMasks)},
-        MSpecConstsDefValBlob(getSpecConstsDefValBlob()),
-        MOrigins(ImageOriginKernelCompiler),
-        MRTCBinInfo(KernelCompilerBinaryInfo{Lang}) {
-    updateSpecConstSymMap();
-  }
-
   device_image_impl(
       const RTDeviceBinaryImage *BinImage, const context &Context,
       devices_range Devices, bundle_state State,
@@ -366,14 +349,13 @@ class device_image_impl
   }
 
   device_image_impl(const context &Context, devices_range Devices,
-                    bundle_state State, ur_program_handle_t Program,
+                    bundle_state State, Managed<ur_program_handle_t> &&Program,
                     syclex::source_language Lang, KernelNameSetT &&KernelNames,
                     private_tag)
       : MBinImage(static_cast<const RTDeviceBinaryImage *>(nullptr)),
         MContext(std::move(Context)),
         MDevices(Devices.to<std::vector<device_impl *>>()), MState(State),
-        MProgram(Program, getSyclObjImpl(MContext)->getAdapter()),
-        MKernelNames{std::move(KernelNames)},
+        MProgram(std::move(Program)), MKernelNames{std::move(KernelNames)},
         MSpecConstsDefValBlob(getSpecConstsDefValBlob()),
         MOrigins(ImageOriginKernelCompiler),
         MRTCBinInfo(KernelCompilerBinaryInfo{Lang}) {}
@@ -771,14 +753,14 @@ class device_image_impl
 
     auto DeviceVec = Devices.to<std::vector<ur_device_handle_t>>();
 
-    ur_program_handle_t UrProgram = nullptr;
+    Managed<ur_program_handle_t> UrProgram;
     // SourceStrPtr will be null when source is Spir-V bytes.
     const std::string *SourceStrPtr = std::get_if<std::string>(&MBinImage);
-    bool FetchedFromCache = false;
     if (PersistentDeviceCodeCache::isEnabled() && SourceStrPtr) {
-      FetchedFromCache = extKernelCompilerFetchFromCache(
-          Devices, BuildOptions, *SourceStrPtr, UrProgram);
+      UrProgram =
+          extKernelCompilerFetchFromCache(Devices, BuildOptions, *SourceStrPtr);
     }
+    bool FetchedFromCache = (UrProgram != nullptr);
 
     adapter_impl &Adapter = ContextImpl.getAdapter();
 
@@ -813,7 +795,7 @@ class device_image_impl
     }
     return std::vector<std::shared_ptr<device_image_impl>>{
         device_image_impl::create(MContext, Devices, bundle_state::executable,
-                                  UrProgram, MRTCBinInfo->MLanguage,
+                                  std::move(UrProgram), MRTCBinInfo->MLanguage,
                                   std::move(KernelNameSet))};
   }
 
@@ -907,10 +889,10 @@ class device_image_impl
     return SS.str();
   }
 
-  bool extKernelCompilerFetchFromCache(
+  Managed<ur_program_handle_t> extKernelCompilerFetchFromCache(
       devices_range Devices,
       const std::vector<sycl::detail::string_view> &BuildOptions,
-      const std::string &SourceStr, ur_program_handle_t &UrProgram) const {
+      const std::string &SourceStr) const {
     sycl::detail::context_impl &ContextImpl = *getSyclObjImpl(MContext);
     adapter_impl &Adapter = ContextImpl.getAdapter();
 
@@ -924,7 +906,7 @@ class device_image_impl
         PersistentDeviceCodeCache::getCompiledKernelFromDisc(Devices, UserArgs,
                                                              SourceStr);
     if (BinProgs.empty()) {
-      return false;
+      return {};
     }
     for (auto &BinProg : BinProgs) {
       Binaries.push_back((uint8_t *)(BinProg.data()));
@@ -937,11 +919,12 @@ class device_image_impl
     Properties.count = 0;
     Properties.pMetadatas = nullptr;
 
+    Managed<ur_program_handle_t> UrProgram{Adapter};
     Adapter.call<UrApiKind::urProgramCreateWithBinary>(
         ContextImpl.getHandleRef(), DeviceHandles.size(), DeviceHandles.data(),
         Lengths.data(), Binaries.data(), &Properties, &UrProgram);
 
-    return true;
+    return UrProgram;
   }
 
   // Get the specialization constant default value blob.
@@ -1226,7 +1209,7 @@ class device_image_impl
     return Result;
   }
 
-  ur_program_handle_t
+  Managed<ur_program_handle_t>
   createProgramFromSource(devices_range Devices,
                           const std::vector<sycl::detail::string_view> &Options,
                           std::string *LogPtr) const {
@@ -1266,11 +1249,10 @@ class device_image_impl
           "languages at this time");
     }();
 
-    ur_program_handle_t UrProgram = nullptr;
+    Managed<ur_program_handle_t> UrProgram{Adapter};
     Adapter.call<UrApiKind::urProgramCreateWithIL>(ContextImpl.getHandleRef(),
                                                    spirv.data(), spirv.size(),
                                                    nullptr, &UrProgram);
-    // program created by urProgramCreateWithIL is implicitly retained.
     if (UrProgram == nullptr)
       throw sycl::exception(
           sycl::make_error_code(errc::invalid),

sycl/source/detail/kernel_bundle_impl.hpp

@@ -597,7 +597,8 @@ class kernel_bundle_impl
     for (const detail::RTDeviceBinaryImage *Image : BestImages)
       MDeviceImages.emplace_back(device_image_impl::create(
           Image, Context, Devs, ProgramManager::getBinImageState(Image),
-          /*KernelIDs=*/nullptr, /*URProgram=*/nullptr, ImageOriginSYCLBIN));
+          /*KernelIDs=*/nullptr, Managed<ur_program_handle_t>{},
+          ImageOriginSYCLBIN));
     ProgramManager::getInstance().bringSYCLDeviceImagesToState(MDeviceImages,
                                                                State);
     fillUniqueDeviceImages();

sycl/source/detail/kernel_name_based_cache_t.hpp

@@ -36,9 +36,10 @@ struct FastKernelCacheVal {
 
   FastKernelCacheVal(ur_kernel_handle_t KernelHandle, std::mutex *Mutex,
                      const KernelArgMask *KernelArgMask,
-                     ur_program_handle_t ProgramHandle, adapter_impl &Adapter)
+                     Managed<ur_program_handle_t> &&ProgramHandle,
+                     adapter_impl &Adapter)
       : MKernelHandle(KernelHandle), MMutex(Mutex),
-        MKernelArgMask(KernelArgMask), MProgramHandle(ProgramHandle, Adapter),
+        MKernelArgMask(KernelArgMask), MProgramHandle(std::move(ProgramHandle)),
         MAdapter(Adapter) {}
 
   ~FastKernelCacheVal() {

sycl/source/detail/kernel_program_cache.hpp

@@ -112,12 +112,11 @@ class KernelProgramCache {
   };
 
   struct ProgramBuildResult : public BuildResult<Managed<ur_program_handle_t>> {
-    ProgramBuildResult(adapter_impl &Adapter) {
-      Val = Managed<ur_program_handle_t>{Adapter};
-    }
-    ProgramBuildResult(adapter_impl &Adapter, BuildState InitialState) {
-      Val = Managed<ur_program_handle_t>{Adapter};
+    ProgramBuildResult() = default;
+    ProgramBuildResult(BuildState InitialState,
+                       Managed<ur_program_handle_t> &&Prog) {
       this->State.store(InitialState);
+      this->Val = std::move(Prog);
     }
 #ifdef _MSC_VER
 #pragma warning(push)
@@ -407,7 +406,7 @@ class KernelProgramCache {
     ProgramCache &ProgCache = LockedCache.get();
     auto [It, DidInsert] = ProgCache.Cache.try_emplace(CacheKey, nullptr);
     if (DidInsert) {
-      It->second = std::make_shared<ProgramBuildResult>(getAdapter());
+      It->second = std::make_shared<ProgramBuildResult>();
       // Save reference between the common key and the full key.
       CommonProgramKeyT CommonKey =
           std::make_pair(CacheKey.first.second, CacheKey.second);
@@ -424,14 +423,13 @@ class KernelProgramCache {
   //
   // Returns whether or not an insertion took place.
   bool insertBuiltProgram(const ProgramCacheKeyT &CacheKey,
-                          ur_program_handle_t Program) {
+                          Managed<ur_program_handle_t> &Program) {
     auto LockedCache = acquireCachedPrograms();
     ProgramCache &ProgCache = LockedCache.get();
     auto [It, DidInsert] = ProgCache.Cache.try_emplace(CacheKey, nullptr);
     if (DidInsert) {
-      It->second = std::make_shared<ProgramBuildResult>(getAdapter(),
-                                                        BuildState::BS_Done);
-      It->second->Val = Managed<ur_program_handle_t>{Program, getAdapter()};
+      It->second = std::make_shared<ProgramBuildResult>(BuildState::BS_Done,
+                                                        Program.retain());
       // Save reference between the common key and the full key.
       CommonProgramKeyT CommonKey =
           std::make_pair(CacheKey.first.second, CacheKey.second);
@@ -643,8 +641,7 @@ class KernelProgramCache {
   // If it is the first time the program is fetched, add it to the eviction
   // list.
   void registerProgramFetch(const ProgramCacheKeyT &CacheKey,
-                            const ur_program_handle_t &Program,
-                            const bool IsBuilt) {
+                            ur_program_handle_t Program, const bool IsBuilt) {
 
     size_t ProgramCacheEvictionThreshold =
         SYCLConfig<SYCL_IN_MEM_CACHE_EVICTION_THRESHOLD>::getProgramCacheSize();
@@ -799,9 +796,10 @@ class KernelProgramCache {
 
       // only the building thread will run this
       try {
-        // Remove `adapter_impl` from `ProgramBuildResult`'s ctors once `Build`
-        // returns `Managed<ur_platform_handle_t`:
-        *(&BuildResult->Val) = Build();
+        static_assert(
+            std::is_same_v<decltype(Build()), decltype(BuildResult->Val)>,
+            "Are we casting from Managed<URResource> to plain URResource?");
+        BuildResult->Val = Build();
 
         if constexpr (!std::is_same_v<EvictFT, void *>)
           EvictFunc(BuildResult->Val, /*IsBuilt=*/true);