[llvm] 1a22f1b - [ORC] Add withResourceKeyDo method to ResourceTracker.

[Lang Hames via llvm-commits]

Author: Lang Hames
Date: 2022-12-19T14:56:08-08:00
New Revision: 1a22f1b64679f6b6d83ba8d4e395908eee65773d

URL: https://github.com/llvm/llvm-project/commit/1a22f1b64679f6b6d83ba8d4e395908eee65773d
DIFF: https://github.com/llvm/llvm-project/commit/1a22f1b64679f6b6d83ba8d4e395908eee65773d.diff

LOG: [ORC] Add withResourceKeyDo method to ResourceTracker.

This method behaves the same as MaterializationResponsibility::withResourceKeyDo
(which now forwards to the new method): It locks the session while providing
access to the ResourceKey associated with the tracker.

Adding this method to ResourceTracker allows resources to be allocated and
tracked for a given MaterializationUnit prior to that MaterializationUnit being
materialized. E.g. Platforms can now track and remove initializers and other
symbols.

Added: 
    

Modified: 
    llvm/include/llvm/ExecutionEngine/Orc/Core.h

Removed: 
    


################################################################################
diff  --git a/llvm/include/llvm/ExecutionEngine/Orc/Core.h b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
index cb3de04b8469b..b6d1bf08988e2 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/Core.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/Core.h
@@ -70,6 +70,10 @@ class ResourceTracker : public ThreadSafeRefCountedBase<ResourceTracker> {
                                          ~static_cast<uintptr_t>(1));
   }
 
+  /// Runs the given callback under the session lock, passing in the associated
+  /// ResourceKey. This is the safe way to associate resources with trackers.
+  template <typename Func> Error withResourceKeyDo(Func &&F);
+
   /// Remove all resources associated with this key.
   Error remove();
 
@@ -530,8 +534,11 @@ class MaterializationResponsibility {
   ///        emitted or notified of an error.
   ~MaterializationResponsibility();
 
-  /// Returns the ResourceTracker for this instance.
-  template <typename Func> Error withResourceKeyDo(Func &&F) const;
+  /// Runs the given callback under the session lock, passing in the associated
+  /// ResourceKey. This is the safe way to associate resources with trackers.
+  template <typename Func> Error withResourceKeyDo(Func &&F) const {
+    return RT->withResourceKeyDo(std::forward<Func>(F));
+  }
 
   /// Returns the target JITDylib that these symbols are being materialized
   ///        into.
@@ -1770,21 +1777,20 @@ class ExecutionSession {
       JITDispatchHandlers;
 };
 
+template <typename Func> Error ResourceTracker::withResourceKeyDo(Func &&F) {
+  return getJITDylib().getExecutionSession().runSessionLocked([&]() -> Error {
+    if (isDefunct())
+      return make_error<ResourceTrackerDefunct>(this);
+    F(getKeyUnsafe());
+    return Error::success();
+  });
+}
+
 inline ExecutionSession &
 MaterializationResponsibility::getExecutionSession() const {
   return JD.getExecutionSession();
 }
 
-template <typename Func>
-Error MaterializationResponsibility::withResourceKeyDo(Func &&F) const {
-  return JD.getExecutionSession().runSessionLocked([&]() -> Error {
-    if (RT->isDefunct())
-      return make_error<ResourceTrackerDefunct>(RT);
-    F(RT->getKeyUnsafe());
-    return Error::success();
-  });
-}
-
 template <typename GeneratorT>
 GeneratorT &JITDylib::addGenerator(std::unique_ptr<GeneratorT> DefGenerator) {
   auto &G = *DefGenerator;