[ORC] Replace ThreadSafeContext::getContext with withContextDo. (#146819)

lhames · web-flow · commit 0bfa0bcd79f7 · 2025-07-03T17:03:39.000+10:00
This removes ThreadSafeContext::Lock, ThreadSafeContext::getLock, and
ThreadSafeContext::getContext, and replaces them with a
ThreadSafeContext::withContextDo method (and const override).

The new method can be used to access an existing
ThreadSafeContext-wrapped LLVMContext in a safe way:

ThreadSafeContext TSCtx = ... ;
TSCtx.withContextDo([](LLVMContext *Ctx) {
  // this closure has exclusive access to Ctx.
});

The new API enforces correct locking, whereas the old APIs relied on
manual locking (which almost no in-tree code preformed, relying instead
on incidental exclusive access to the ThreadSafeContext).
diff --git a/clang/lib/Interpreter/Interpreter.cpp b/clang/lib/Interpreter/Interpreter.cpp
@@ -373,8 +373,11 @@ Interpreter::Interpreter(std::unique_ptr<CompilerInstance> Instance,
   auto LLVMCtx = std::make_unique<llvm::LLVMContext>();
   TSCtx = std::make_unique<llvm::orc::ThreadSafeContext>(std::move(LLVMCtx));
 
-  Act = std::make_unique<IncrementalAction>(*CI, *TSCtx->getContext(), ErrOut,
-                                            *this, std::move(Consumer));
+  Act = TSCtx->withContextDo([&](llvm::LLVMContext *Ctx) {
+    return std::make_unique<IncrementalAction>(*CI, *Ctx, ErrOut, *this,
+                                               std::move(Consumer));
+  });
+
   if (ErrOut)
     return;
   CI->ExecuteAction(*Act);
@@ -495,10 +498,10 @@ Interpreter::createWithCUDA(std::unique_ptr<CompilerInstance> CI,
   std::unique_ptr<Interpreter> Interp = std::move(*InterpOrErr);
 
   llvm::Error Err = llvm::Error::success();
-  llvm::LLVMContext &LLVMCtx = *Interp->TSCtx->getContext();
 
-  auto DeviceAct =
-      std::make_unique<IncrementalAction>(*DCI, LLVMCtx, Err, *Interp);
+  auto DeviceAct = Interp->TSCtx->withContextDo([&](llvm::LLVMContext *Ctx) {
+    return std::make_unique<IncrementalAction>(*DCI, *Ctx, Err, *Interp);
+  });
 
   if (Err)
     return std::move(Err);
diff --git a/llvm/examples/OrcV2Examples/LLJITWithThinLTOSummaries/LLJITWithThinLTOSummaries.cpp b/llvm/examples/OrcV2Examples/LLJITWithThinLTOSummaries/LLJITWithThinLTOSummaries.cpp
@@ -169,7 +169,9 @@ Expected<ThreadSafeModule> loadModule(StringRef Path,
 
   MemoryBufferRef BitcodeBufferRef = (**BitcodeBuffer).getMemBufferRef();
   Expected<std::unique_ptr<Module>> M =
-      parseBitcodeFile(BitcodeBufferRef, *TSCtx.getContext());
+      TSCtx.withContextDo([&](LLVMContext *Ctx) {
+        return parseBitcodeFile(BitcodeBufferRef, *Ctx);
+      });
   if (!M)
     return M.takeError();
 
diff --git a/llvm/examples/OrcV2Examples/OrcV2CBindingsBasicUsage/OrcV2CBindingsBasicUsage.c b/llvm/examples/OrcV2Examples/OrcV2CBindingsBasicUsage/OrcV2CBindingsBasicUsage.c
@@ -22,11 +22,8 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Create a new LLVM module.
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
@@ -57,6 +54,9 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   //  - Free the builder.
   LLVMDisposeBuilder(Builder);
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
+
   // Our demo module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
diff --git a/llvm/examples/OrcV2Examples/OrcV2CBindingsDumpObjects/OrcV2CBindingsDumpObjects.c b/llvm/examples/OrcV2Examples/OrcV2CBindingsDumpObjects/OrcV2CBindingsDumpObjects.c
@@ -31,8 +31,7 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  LLVMContextRef Ctx = LLVMContextCreate();
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
   LLVMTypeRef ParamTypes[] = {LLVMInt32Type(), LLVMInt32Type()};
   LLVMTypeRef SumFunctionType =
@@ -45,6 +44,8 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   LLVMValueRef SumArg1 = LLVMGetParam(SumFunction, 1);
   LLVMValueRef Result = LLVMBuildAdd(Builder, SumArg0, SumArg1, "result");
   LLVMBuildRet(Builder, Result);
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
   LLVMOrcDisposeThreadSafeContext(TSCtx);
   return TSM;
diff --git a/llvm/examples/OrcV2Examples/OrcV2CBindingsIRTransforms/OrcV2CBindingsIRTransforms.c b/llvm/examples/OrcV2Examples/OrcV2CBindingsIRTransforms/OrcV2CBindingsIRTransforms.c
@@ -32,8 +32,7 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  LLVMContextRef Ctx = LLVMContextCreate();
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
   LLVMTypeRef ParamTypes[] = {LLVMInt32Type(), LLVMInt32Type()};
   LLVMTypeRef SumFunctionType =
@@ -47,6 +46,7 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   LLVMValueRef Result = LLVMBuildAdd(Builder, SumArg0, SumArg1, "result");
   LLVMBuildRet(Builder, Result);
   LLVMDisposeBuilder(Builder);
+  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
   LLVMOrcDisposeThreadSafeContext(TSCtx);
   return TSM;
diff --git a/llvm/examples/OrcV2Examples/OrcV2CBindingsLazy/OrcV2CBindingsLazy.c b/llvm/examples/OrcV2Examples/OrcV2CBindingsLazy/OrcV2CBindingsLazy.c
@@ -67,11 +67,9 @@ const char MainMod[] =
 LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
                                 const char *Name,
                                 LLVMOrcThreadSafeModuleRef *TSM) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
 
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext for the Module.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Wrap Source in a MemoryBuffer
   LLVMMemoryBufferRef MB =
@@ -85,6 +83,10 @@ LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
     // TODO: LLVMDisposeMessage(ErrMsg);
   }
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
+
   // Our module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   *TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
diff --git a/llvm/examples/OrcV2Examples/OrcV2CBindingsMCJITLikeMemoryManager/OrcV2CBindingsMCJITLikeMemoryManager.c b/llvm/examples/OrcV2Examples/OrcV2CBindingsMCJITLikeMemoryManager/OrcV2CBindingsMCJITLikeMemoryManager.c
@@ -150,11 +150,8 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Create a new LLVM module.
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
@@ -182,6 +179,10 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   //  - Build the return instruction.
   LLVMBuildRet(Builder, Result);
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
+
   // Our demo module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
diff --git a/llvm/examples/OrcV2Examples/OrcV2CBindingsRemovableCode/OrcV2CBindingsRemovableCode.c b/llvm/examples/OrcV2Examples/OrcV2CBindingsRemovableCode/OrcV2CBindingsRemovableCode.c
@@ -22,11 +22,8 @@ int handleError(LLVMErrorRef Err) {
 }
 
 LLVMOrcThreadSafeModuleRef createDemoModule(void) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Create a new LLVM module.
   LLVMModuleRef M = LLVMModuleCreateWithNameInContext("demo", Ctx);
@@ -57,6 +54,10 @@ LLVMOrcThreadSafeModuleRef createDemoModule(void) {
   //  - Free the builder.
   LLVMDisposeBuilder(Builder);
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx =
+      LLVMOrcCreateNewThreadSafeContextFromLLVMContext(Ctx);
+
   // Our demo module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   LLVMOrcThreadSafeModuleRef TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
diff --git a/llvm/examples/OrcV2Examples/OrcV2CBindingsVeryLazy/OrcV2CBindingsVeryLazy.c b/llvm/examples/OrcV2Examples/OrcV2CBindingsVeryLazy/OrcV2CBindingsVeryLazy.c
@@ -74,11 +74,8 @@ LLVMErrorRef applyDataLayout(void *Ctx, LLVMModuleRef M) {
 LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
                                 const char *Name,
                                 LLVMOrcThreadSafeModuleRef *TSM) {
-  // Create a new ThreadSafeContext and underlying LLVMContext.
-  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
-
-  // Get a reference to the underlying LLVMContext.
-  LLVMContextRef Ctx = LLVMOrcThreadSafeContextGetContext(TSCtx);
+  // Create an LLVMContext.
+  LLVMContextRef Ctx = LLVMContextCreate();
 
   // Wrap Source in a MemoryBuffer
   LLVMMemoryBufferRef MB =
@@ -93,6 +90,9 @@ LLVMErrorRef parseExampleModule(const char *Source, size_t Len,
     return Err;
   }
 
+  // Create a new ThreadSafeContext to hold the context.
+  LLVMOrcThreadSafeContextRef TSCtx = LLVMOrcCreateNewThreadSafeContext();
+
   // Our module is now complete. Wrap it and our ThreadSafeContext in a
   // ThreadSafeModule.
   *TSM = LLVMOrcCreateNewThreadSafeModule(M, TSCtx);
diff --git a/llvm/include/llvm-c/Orc.h b/llvm/include/llvm-c/Orc.h
@@ -1062,20 +1062,32 @@ LLVMErrorRef LLVMOrcCreateStaticLibrarySearchGeneratorForPath(
     const char *FileName);
 
 /**
- * Create a ThreadSafeContext containing a new LLVMContext.
+ * Create a ThreadSafeContextRef containing a new LLVMContext.
  *
  * Ownership of the underlying ThreadSafeContext data is shared: Clients
- * can and should dispose of their ThreadSafeContext as soon as they no longer
- * need to refer to it directly. Other references (e.g. from ThreadSafeModules)
- * will keep the data alive as long as it is needed.
+ * can and should dispose of their ThreadSafeContextRef as soon as they no
+ * longer need to refer to it directly. Other references (e.g. from
+ * ThreadSafeModules) will keep the underlying data alive as long as it is
+ * needed.
  */
 LLVMOrcThreadSafeContextRef LLVMOrcCreateNewThreadSafeContext(void);
 
 /**
- * Get a reference to the wrapped LLVMContext.
+ * Create a ThreadSafeContextRef from a given LLVMContext, which must not be
+ * associated with any existing ThreadSafeContext.
+ *
+ * The underlying ThreadSafeContext will take ownership of the LLVMContext
+ * object, so clients should not free the LLVMContext passed to this
+ * function.
+ *
+ * Ownership of the underlying ThreadSafeContext data is shared: Clients
+ * can and should dispose of their ThreadSafeContextRef as soon as they no
+ * longer need to refer to it directly. Other references (e.g. from
+ * ThreadSafeModules) will keep the underlying data alive as long as it is
+ * needed.
  */
-LLVMContextRef
-LLVMOrcThreadSafeContextGetContext(LLVMOrcThreadSafeContextRef TSCtx);
+LLVMOrcThreadSafeContextRef
+LLVMOrcCreateNewThreadSafeContextFromLLVMContext(LLVMContextRef Ctx);
 
 /**
  * Dispose of a ThreadSafeContext.
diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
@@ -36,16 +36,6 @@ class ThreadSafeContext {
   };
 
 public:
-  // RAII based lock for ThreadSafeContext.
-  class [[nodiscard]] Lock {
-  public:
-    Lock(std::shared_ptr<State> S) : S(std::move(S)), L(this->S->Mutex) {}
-
-  private:
-    std::shared_ptr<State> S;
-    std::unique_lock<std::recursive_mutex> L;
-  };
-
   /// Construct a null context.
   ThreadSafeContext() = default;
 
@@ -56,17 +46,20 @@ class ThreadSafeContext {
            "Can not construct a ThreadSafeContext from a nullptr");
   }
 
-  /// Returns a pointer to the LLVMContext that was used to construct this
-  /// instance, or null if the instance was default constructed.
-  LLVMContext *getContext() { return S ? S->Ctx.get() : nullptr; }
-
-  /// Returns a pointer to the LLVMContext that was used to construct this
-  /// instance, or null if the instance was default constructed.
-  const LLVMContext *getContext() const { return S ? S->Ctx.get() : nullptr; }
+  template <typename Func> decltype(auto) withContextDo(Func &&F) {
+    if (auto TmpS = S) {
+      std::lock_guard<std::recursive_mutex> Lock(TmpS->Mutex);
+      return F(TmpS->Ctx.get());
+    } else
+      return F((LLVMContext *)nullptr);
+  }
 
-  Lock getLock() const {
-    assert(S && "Can not lock an empty ThreadSafeContext");
-    return Lock(S);
+  template <typename Func> decltype(auto) withContextDo(Func &&F) const {
+    if (auto TmpS = S) {
+      std::lock_guard<std::recursive_mutex> Lock(TmpS->Mutex);
+      return F(const_cast<const LLVMContext *>(TmpS->Ctx.get()));
+    } else
+      return F((const LLVMContext *)nullptr);
   }
 
 private:
@@ -89,10 +82,7 @@ class ThreadSafeModule {
     // *before* the context that it depends on.
     // We also need to lock the context to make sure the module tear-down
     // does not overlap any other work on the context.
-    if (M) {
-      auto L = TSCtx.getLock();
-      M = nullptr;
-    }
+    TSCtx.withContextDo([this](LLVMContext *Ctx) { M = nullptr; });
     M = std::move(Other.M);
     TSCtx = std::move(Other.TSCtx);
     return *this;
@@ -111,45 +101,39 @@ class ThreadSafeModule {
 
   ~ThreadSafeModule() {
     // We need to lock the context while we destruct the module.
-    if (M) {
-      auto L = TSCtx.getLock();
-      M = nullptr;
-    }
+    TSCtx.withContextDo([this](LLVMContext *Ctx) { M = nullptr; });
   }
 
   /// Boolean conversion: This ThreadSafeModule will evaluate to true if it
   /// wraps a non-null module.
-  explicit operator bool() const {
-    if (M) {
-      assert(TSCtx.getContext() &&
-             "Non-null module must have non-null context");
-      return true;
-    }
-    return false;
-  }
+  explicit operator bool() const { return !!M; }
 
   /// Locks the associated ThreadSafeContext and calls the given function
   /// on the contained Module.
   template <typename Func> decltype(auto) withModuleDo(Func &&F) {
-    assert(M && "Can not call on null module");
-    auto Lock = TSCtx.getLock();
-    return F(*M);
+    return TSCtx.withContextDo([&](LLVMContext *) {
+      assert(M && "Can not call on null module");
+      return F(*M);
+    });
   }
 
   /// Locks the associated ThreadSafeContext and calls the given function
   /// on the contained Module.
   template <typename Func> decltype(auto) withModuleDo(Func &&F) const {
-    assert(M && "Can not call on null module");
-    auto Lock = TSCtx.getLock();
-    return F(*M);
+    return TSCtx.withContextDo([&](const LLVMContext *) {
+      assert(M && "Can not call on null module");
+      return F(*M);
+    });
   }
 
   /// Locks the associated ThreadSafeContext and calls the given function,
   /// passing the contained std::unique_ptr<Module>. The given function should
   /// consume the Module.
   template <typename Func> decltype(auto) consumingModuleDo(Func &&F) {
-    auto Lock = TSCtx.getLock();
-    return F(std::move(M));
+    return TSCtx.withContextDo([&](LLVMContext *) {
+      assert(M && "Can not call on null module");
+      return F(std::move(M));
+    });
   }
 
   /// Get a raw pointer to the contained module without locking the context.
diff --git a/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp b/llvm/lib/ExecutionEngine/Orc/OrcV2CBindings.cpp
@@ -729,9 +729,9 @@ LLVMOrcThreadSafeContextRef LLVMOrcCreateNewThreadSafeContext(void) {
   return wrap(new ThreadSafeContext(std::make_unique<LLVMContext>()));
 }
 
-LLVMContextRef
-LLVMOrcThreadSafeContextGetContext(LLVMOrcThreadSafeContextRef TSCtx) {
-  return wrap(unwrap(TSCtx)->getContext());
+LLVMOrcThreadSafeContextRef
+LLVMOrcCreateNewThreadSafeContextFromLLVMContext(LLVMContextRef Ctx) {
+  return wrap(new ThreadSafeContext(std::unique_ptr<LLVMContext>(unwrap(Ctx))));
 }
 
 void LLVMOrcDisposeThreadSafeContext(LLVMOrcThreadSafeContextRef TSCtx) {
diff --git a/llvm/lib/ExecutionEngine/Orc/Speculation.cpp b/llvm/lib/ExecutionEngine/Orc/Speculation.cpp
@@ -60,8 +60,6 @@ void IRSpeculationLayer::emit(std::unique_ptr<MaterializationResponsibility> R,
                               ThreadSafeModule TSM) {
 
   assert(TSM && "Speculation Layer received Null Module ?");
-  assert(TSM.getContext().getContext() != nullptr &&
-         "Module with null LLVMContext?");
 
   // Instrumentation of runtime calls, lock the Module
   TSM.withModuleDo([this, &R](Module &M) {
diff --git a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
@@ -53,9 +53,11 @@ ThreadSafeModule cloneToNewContext(const ThreadSafeModule &TSM,
         "cloned module buffer");
     ThreadSafeContext NewTSCtx(std::make_unique<LLVMContext>());
 
-    auto ClonedModule = cantFail(
-        parseBitcodeFile(ClonedModuleBufferRef, *NewTSCtx.getContext()));
-    ClonedModule->setModuleIdentifier(M.getName());
+    auto ClonedModule = NewTSCtx.withContextDo([&](LLVMContext *Ctx) {
+      auto TmpM = cantFail(parseBitcodeFile(ClonedModuleBufferRef, *Ctx));
+      TmpM->setModuleIdentifier(M.getName());
+      return TmpM;
+    });
     return ThreadSafeModule(std::move(ClonedModule), std::move(NewTSCtx));
   });
 }
diff --git a/llvm/tools/lli/lli.cpp b/llvm/tools/lli/lli.cpp
diff --git a/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/RTDyldObjectLinkingLayerTest.cpp
diff --git a/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ReOptimizeLayerTest.cpp
diff --git a/llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp b/llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp
