[llvm] [ORC] Support non-moving cloneToContext (PR #148688)

[Jeaye Wilkerson via llvm-commits]

https://github.com/jeaye created https://github.com/llvm/llvm-project/pull/148688

There are two usability concerns with the new `cloneToContext` from https://github.com/llvm/llvm-project/pull/146852.

1. This API assumes the LLVM context is in a `ThreadSafeContext`, but I have no need for that. My IR gen is single-threaded. IR gen (via `IRBuilder`) uses the context and module often, so I would need to do a whole lot of `withContextDo` with closures just to do things like create a new basic block.

2. The new cloning function requires a `ThreadSafeModule`, which can only be constructed by *moving* (i.e. consuming) a module. This prevents me from caching my C++ modules, which explodes the number I need to generate. Cloning should be read-only, so the module consuming here feels like it's just an implication of TSMs being used.

We can address both of these by:

1. Adding a `getContextUnlocked` to `ThreadSafeContext`, following what exists for `ThreadSafeModule`; this allows me to store a raw pointer for easy use with IR gen

2. Add an overload for `cloneToContext` which take a `const Module &` and use that as the implementation for the TSM version

@lhames for vis.

>From 8dde5ccf5561bdd81cc2f6a165cc2308857cbe78 Mon Sep 17 00:00:00 2001
From: jeaye <contact at jeaye.com>
Date: Mon, 14 Jul 2025 10:40:12 -0700
Subject: [PATCH] [ORC] Support non-moving cloneToContext

---
 .../ExecutionEngine/Orc/ThreadSafeModule.h    | 10 ++++
 .../ExecutionEngine/Orc/ThreadSafeModule.cpp  | 57 +++++++++++--------
 2 files changed, 43 insertions(+), 24 deletions(-)

diff --git a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
index 03165895b85d0..2e5076cc4d8d6 100644
--- a/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
+++ b/llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h
@@ -62,6 +62,12 @@ class ThreadSafeContext {
       return F((const LLVMContext *)nullptr);
   }
 
+  /// Get a raw pointer to the contained context without locking the context.
+  LLVMContext *getContextUnlocked() { return S->Ctx.get(); }
+
+  /// Get a raw pointer to the contained context without locking the context.
+  const LLVMContext *getContextUnlocked() const { return S->Ctx.get(); }
+
 private:
   std::shared_ptr<State> S;
 };
@@ -155,6 +161,10 @@ using GVModifier = std::function<void(GlobalValue &)>;
 
 /// Clones teh given module onto the given context.
 LLVM_ABI ThreadSafeModule
+cloneToContext(const Module &M, ThreadSafeContext TSCtx,
+               GVPredicate ShouldCloneDef = GVPredicate(),
+               GVModifier UpdateClonedDefSource = GVModifier());
+LLVM_ABI ThreadSafeModule
 cloneToContext(const ThreadSafeModule &TSMW, ThreadSafeContext TSCtx,
                GVPredicate ShouldCloneDef = GVPredicate(),
                GVModifier UpdateClonedDefSource = GVModifier());
diff --git a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
index 19c000e2472a8..39ad32d8e1ef7 100644
--- a/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
+++ b/llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp
@@ -14,53 +14,62 @@
 namespace llvm {
 namespace orc {
 
-ThreadSafeModule cloneToContext(const ThreadSafeModule &TSM,
+ThreadSafeModule cloneToContext(const Module &M,
                                 ThreadSafeContext TSCtx,
                                 GVPredicate ShouldCloneDef,
                                 GVModifier UpdateClonedDefSource) {
-  assert(TSM && "Can not clone null module");
-
   if (!ShouldCloneDef)
     ShouldCloneDef = [](const GlobalValue &) { return true; };
 
   // First copy the source module into a buffer.
   std::string ModuleName;
   SmallVector<char, 1> ClonedModuleBuffer;
-  TSM.withModuleDo([&](Module &M) {
-    ModuleName = M.getModuleIdentifier();
-    std::set<GlobalValue *> ClonedDefsInSrc;
-    ValueToValueMapTy VMap;
-    auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) {
-      if (ShouldCloneDef(*GV)) {
-        ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
-        return true;
-      }
-      return false;
-    });
+  ModuleName = M.getModuleIdentifier();
+  std::set<GlobalValue *> ClonedDefsInSrc;
+  ValueToValueMapTy VMap;
+  auto Tmp = CloneModule(M, VMap, [&](const GlobalValue *GV) {
+    if (ShouldCloneDef(*GV)) {
+      ClonedDefsInSrc.insert(const_cast<GlobalValue *>(GV));
+      return true;
+    }
+    return false;
+  });
 
-    if (UpdateClonedDefSource)
-      for (auto *GV : ClonedDefsInSrc)
-        UpdateClonedDefSource(*GV);
+  if (UpdateClonedDefSource)
+    for (auto *GV : ClonedDefsInSrc)
+      UpdateClonedDefSource(*GV);
 
-    BitcodeWriter BCWriter(ClonedModuleBuffer);
-    BCWriter.writeModule(*Tmp);
-    BCWriter.writeSymtab();
-    BCWriter.writeStrtab();
-  });
+  BitcodeWriter BCWriter(ClonedModuleBuffer);
+  BCWriter.writeModule(*Tmp);
+  BCWriter.writeSymtab();
+  BCWriter.writeStrtab();
 
   MemoryBufferRef ClonedModuleBufferRef(
       StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
       "cloned module buffer");
 
   // Then parse the buffer into the new Module.
-  auto M = TSCtx.withContextDo([&](LLVMContext *Ctx) {
+  auto R = TSCtx.withContextDo([&](LLVMContext *Ctx) {
     assert(Ctx && "No LLVMContext provided");
     auto TmpM = cantFail(parseBitcodeFile(ClonedModuleBufferRef, *Ctx));
     TmpM->setModuleIdentifier(ModuleName);
     return TmpM;
   });
 
-  return ThreadSafeModule(std::move(M), std::move(TSCtx));
+  return ThreadSafeModule(std::move(R), std::move(TSCtx));
+}
+
+ThreadSafeModule cloneToContext(const ThreadSafeModule &TSM,
+                                ThreadSafeContext TSCtx,
+                                GVPredicate ShouldCloneDef,
+                                GVModifier UpdateClonedDefSource) {
+  assert(TSM && "Can not clone null module");
+
+  ThreadSafeModule R;
+  TSM.withModuleDo([&](Module &M) {
+    R = cloneToContext(M, TSCtx, ShouldCloneDef, UpdateClonedDefSource);
+  });
+  return R;
 }
 
 ThreadSafeModule cloneToNewContext(const ThreadSafeModule &TSM,