[ORC] Add cloneExternalModuleToContext API.

cloneExternalModuleToContext can be used to clone an LLVM module onto a given
ThreadSafeContext. Callers of this function are responsible for ensuring
exclusive access to the source module and its LLVMContext.

Author: lhames

llvm/include/llvm/ExecutionEngine/Orc/ThreadSafeModule.h

@@ -159,6 +159,14 @@ cloneToContext(const ThreadSafeModule &TSMW, ThreadSafeContext TSCtx,
                GVPredicate ShouldCloneDef = GVPredicate(),
                GVModifier UpdateClonedDefSource = GVModifier());
 
+/// Clone the given module onto the given context.
+/// The caller is responsible for ensuring that the source module and its
+/// LLVMContext will not be concurrently accessed during the clone.
+LLVM_ABI ThreadSafeModule
+cloneExternalModuleToContext(const Module &M, ThreadSafeContext TSCtx,
+                             GVPredicate ShouldCloneDef = GVPredicate(),
+                             GVModifier UpdateClonedDefSource = GVModifier());
+
 /// Clones the given module on to a new context.
 LLVM_ABI ThreadSafeModule cloneToNewContext(
     const ThreadSafeModule &TSMW, GVPredicate ShouldCloneDef = GVPredicate(),

llvm/lib/ExecutionEngine/Orc/ThreadSafeModule.cpp

@@ -14,40 +14,39 @@
 namespace llvm {
 namespace orc {
 
-ThreadSafeModule cloneToContext(const ThreadSafeModule &TSM,
-                                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.
+static std::pair<std::string, SmallVector<char, 1>>
+serializeModule(const Module &M, GVPredicate ShouldCloneDef,
+                GVModifier UpdateClonedDefSource) {
   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;
-    });
-
-    if (UpdateClonedDefSource)
-      for (auto *GV : ClonedDefsInSrc)
-        UpdateClonedDefSource(*GV);
-
-    BitcodeWriter BCWriter(ClonedModuleBuffer);
-    BCWriter.writeModule(*Tmp);
-    BCWriter.writeSymtab();
-    BCWriter.writeStrtab();
+
+  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);
+
+  BitcodeWriter BCWriter(ClonedModuleBuffer);
+  BCWriter.writeModule(*Tmp);
+  BCWriter.writeSymtab();
+  BCWriter.writeStrtab();
+
+  return {std::move(ModuleName), std::move(ClonedModuleBuffer)};
+}
+
+ThreadSafeModule
+deserializeModule(std::string ModuleName,
+                  const SmallVector<char, 1> &ClonedModuleBuffer,
+                  ThreadSafeContext TSCtx) {
   MemoryBufferRef ClonedModuleBufferRef(
       StringRef(ClonedModuleBuffer.data(), ClonedModuleBuffer.size()),
       "cloned module buffer");
@@ -63,6 +62,40 @@ ThreadSafeModule cloneToContext(const ThreadSafeModule &TSM,
   return ThreadSafeModule(std::move(M), std::move(TSCtx));
 }
 
+ThreadSafeModule
+cloneExternalModuleToContext(const Module &M, ThreadSafeContext TSCtx,
+                             GVPredicate ShouldCloneDef,
+                             GVModifier UpdateClonedDefSource) {
+
+  if (!ShouldCloneDef)
+    ShouldCloneDef = [](const GlobalValue &) { return true; };
+
+  auto [ModuleName, ClonedModuleBuffer] = serializeModule(
+      M, std::move(ShouldCloneDef), std::move(UpdateClonedDefSource));
+
+  return deserializeModule(std::move(ModuleName), ClonedModuleBuffer,
+                           std::move(TSCtx));
+}
+
+ThreadSafeModule cloneToContext(const ThreadSafeModule &TSM,
+                                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.
+  auto [ModuleName, ClonedModuleBuffer] = TSM.withModuleDo([&](Module &M) {
+    return serializeModule(M, std::move(ShouldCloneDef),
+                           std::move(UpdateClonedDefSource));
+  });
+
+  return deserializeModule(std::move(ModuleName), ClonedModuleBuffer,
+                           std::move(TSCtx));
+}
+
 ThreadSafeModule cloneToNewContext(const ThreadSafeModule &TSM,
                                    GVPredicate ShouldCloneDef,
                                    GVModifier UpdateClonedDefSource) {

llvm/unittests/ExecutionEngine/Orc/ThreadSafeModuleTest.cpp

@@ -31,15 +31,21 @@ const llvm::StringRef FooSrc = R"(
   }
 )";
 
-static ThreadSafeModule parseModule(llvm::StringRef Source,
-                                    llvm::StringRef Name) {
-  auto Ctx = std::make_unique<LLVMContext>();
+static std::unique_ptr<Module>
+parseModuleRaw(llvm::StringRef Source, llvm::StringRef Name, LLVMContext &Ctx) {
   SMDiagnostic Err;
-  auto M = parseIR(MemoryBufferRef(Source, Name), Err, *Ctx);
+  auto M = parseIR(MemoryBufferRef(Source, Name), Err, Ctx);
   if (!M) {
     Err.print("Testcase source failed to parse: ", errs());
     exit(1);
   }
+  return M;
+}
+
+static ThreadSafeModule parseModule(llvm::StringRef Source,
+                                    llvm::StringRef Name) {
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = parseModuleRaw(Source, Name, *Ctx);
   return ThreadSafeModule(std::move(M), std::move(Ctx));
 }
 
@@ -128,6 +134,20 @@ TEST(ThreadSafeModuleTest, ConsumingModuleDo) {
   TSM.consumingModuleDo([](std::unique_ptr<Module> M) {});
 }
 
+TEST(ThreadSafeModuleTest, CloneExternalModuleToNewContext) {
+  auto Ctx = std::make_unique<LLVMContext>();
+  auto M = parseModuleRaw(FooSrc, "foo.ll", *Ctx);
+  auto TSCtx = ThreadSafeContext(std::make_unique<LLVMContext>());
+  auto TSM = cloneExternalModuleToContext(*M, TSCtx);
+  TSM.withModuleDo([&](Module &NewM) {
+    EXPECT_NE(&NewM.getContext(), Ctx.get());
+    TSCtx.withContextDo(
+        [&](LLVMContext *NewCtx) { EXPECT_EQ(&NewM.getContext(), NewCtx); });
+    EXPECT_FALSE(NewM.empty());
+    EXPECT_FALSE(verifyModule(NewM, &errs()));
+  });
+}
+
 TEST(ThreadSafeModuleTest, CloneToNewContext) {
   auto TSM1 = parseModule(FooSrc, "foo.ll");
   auto TSM2 = cloneToNewContext(TSM1);