[ORC] Fix the move-captured std::unique_ptr vs. std::function dilemma

[sgraenitz]

Summary

C++14 is there for a while now and the FIXME is getting confusing. There was an earlier attempt to fix this the way (I think) it was indended originally: https://github.com/llvm/llvm-project/commit/ce74c3b19f5b#diff-f27b9f96303108227f0119d85d6c3ddf (It was rolled back here: https://github.com/llvm/llvm-project/commit/6baaa4be7831#diff-f27b9f96303108227f0119d85d6c3ddf )

  ES->setDispatchMaterialization(
      [this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
+       auto Work = [MU = std::move(MU), &JD] { MU->doMaterialize(JD); };
-       // FIXME: Switch to move capture once we have c++14.
-       auto SharedMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
-       auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); };
        CompileThreads->async(std::move(Work));
      });

Unfortunately this doesn't work, because the lambda that captures the std::unique_ptr is being passed to ThreadPool::async() as a std::function: Holding the std::unique_ptr its copy constructor is deleted implicitly, but the standard requires the value of std::function to be copy-constructible.

What's the options? We certainly want to use ThreadPool, so we are bound to std::function. I think we don't want to go through a raw pointer capture by-value here, especially seeing that we pass a std::function next..

The current workaround creates and copies a std::shared_pointer. That's quite expensive. In a first step we can move-capture it in the Work lambda to avoid the copy. I think we don't get around the creation (assuming we keep a smart pointer), so why not receive it as a std::shared_ptr directly? The only disproportionate overhead I see is in the materializeOnCurrentThread case. I changed that to not calling the function at all.

What do you think?

[lhames]

The current workaround creates and copies a std::shared_pointer. That's quite expensive. In a first step we can move-capture it in the Work lambda to avoid the copy. I think we don't get around the creation (assuming we keep a smart pointer), so why not receive it as a std::shared_ptr directly? The only disproportionate overhead I see is in the materializeOnCurrentThread case. I changed that to not calling the function at all.

In the context of running a Materializer I don’t think we need to be worried about the shared_ptr construction overhead, but it would still be nice to tidy this up.

For our purposes it would be ideal if ThreadPool::async took an llvm::unique_function rather than a std::function. I wonder if that would work for other clients too? I think 99% of ThreadPool use cases would work equally well or better with a unique_function, and it looks like unique_function should be constructible from a std::function (with minor runtime overhead for execution) for the 1% of cases that are not supported.

[sgraenitz]

Indeed, that sounds like a better solution.

llvm/examples/SpeculativeJIT/SpeculativeJIT.cpp

SpeculativeJIT(
                      ConcurrentIRCompiler(std::move(JTMB))),
         S(Imps, *this->ES),
         SpeculateLayer(*this->ES, CompileLayer, S, Mangle, BlockFreqQuery()),
         CODLayer(*this->ES, SpeculateLayer, *this->LCTMgr,
                  std::move(ISMBuilder)) {
     MainJD.addGenerator(std::move(ProcessSymbolsGenerator));
     this->CODLayer.setImplMap(&Imps);
     this->ES->setDispatchMaterialization(
-
-        [this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
-          // FIXME: Switch to move capture once we have C++14.
-          auto SharedMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
-          auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); };
-          CompileThreads.async(std::move(Work));
+        [this](JITDylib &JD, std::shared_ptr<MaterializationUnit> MU) {
+          CompileThreads.async(
+              [MU = std::move(MU), &JD]() { MU->doMaterialize(JD); });
         });
     ExitOnErr(S.addSpeculationRuntime(MainJD, Mangle));
     LocalCXXRuntimeOverrides CXXRuntimeoverrides;
     ExitOnErr(CXXRuntimeoverrides.enable(MainJD, Mangle));
   }
 
   static std::unique_ptr<SectionMemoryManager> createMemMgr() {
     return std::make_unique<SectionMemoryManager>();

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

class ExecutionSession {
   friend class JITDylib;
 
 public:
   /// For reporting errors.
   using ErrorReporter = std::function<void(Error)>;
 
   /// For dispatching MaterializationUnit::materialize calls.
   using DispatchMaterializationFunction = std::function<void(
-      JITDylib &JD, std::unique_ptr<MaterializationUnit> MU)>;
+      JITDylib &JD, std::shared_ptr<MaterializationUnit> MU)>;
 
   /// Construct an ExecutionSession.
   ///
   /// SymbolStringPools may be shared between ExecutionSessions.
   ExecutionSession(std::shared_ptr<SymbolStringPool> SSP = nullptr);
 
   /// Add a symbol name to the SymbolStringPool and return a pointer to it.
   SymbolStringPtr intern(StringRef SymName) { return SSP->intern(SymName); }
ExecutionSession &setErrorReporter(ErrorReporter ReportError) {
     return *this;
   }
 
   /// Report a error for this execution session.
   ///
   /// Unhandled errors can be sent here to log them.
   void reportError(Error Err) { ReportError(std::move(Err)); }
 
-  /// Set the materialization dispatch function.
+  /// Set a materialization dispatch function. It can be used to delegate
+  /// compilation to different threads.
   ExecutionSession &setDispatchMaterialization(
       DispatchMaterializationFunction DispatchMaterialization) {
     this->DispatchMaterialization = std::move(DispatchMaterialization);
     return *this;
   }
 
   void legacyFailQuery(AsynchronousSymbolQuery &Q, Error Err);
 
public:
   /// Materialize the given unit.
   void dispatchMaterialization(JITDylib &JD,
                                std::unique_ptr<MaterializationUnit> MU) {
     LLVM_DEBUG({
       runSessionLocked([&]() {
         dbgs() << "Dispatching " << *MU << " for " << JD.getName() << "\n";
       });
     });
+    if (DispatchMaterialization) {
       DispatchMaterialization(JD, std::move(MU));
+    } else {
+      MU->doMaterialize(JD);
+    }
   }
 
   /// Dump the state of all the JITDylibs in this session.
   void dump(raw_ostream &OS);
 
 private:
   static void logErrorsToStdErr(Error Err) {
     logAllUnhandledErrors(std::move(Err), errs(), "JIT session error: ");
   }
 
-  static void
-  materializeOnCurrentThread(JITDylib &JD,
-                             std::unique_ptr<MaterializationUnit> MU) {
-    MU->doMaterialize(JD);
-  }
-
   void runOutstandingMUs();
 
   mutable std::recursive_mutex SessionMutex;
   std::shared_ptr<SymbolStringPool> SSP;
   VModuleKey LastKey = 0;
   ErrorReporter ReportError = logErrorsToStdErr;
-  DispatchMaterializationFunction DispatchMaterialization =
-      materializeOnCurrentThread;
+  DispatchMaterializationFunction DispatchMaterialization = nullptr;
 
   std::vector<std::unique_ptr<JITDylib>> JDs;
 
   // FIXME: Remove this (and runOutstandingMUs) once the linking layer works
   //        with callbacks from asynchronous queries.
   mutable std::recursive_mutex OutstandingMUsMutex;
   std::vector<std::pair<JITDylib *, std::unique_ptr<MaterializationUnit>>>
       OutstandingMUs;

llvm/lib/ExecutionEngine/Orc/LLJIT.cpp

else {
     CompileLayer = std::make_unique<IRCompileLayer>(
         *ES, ObjTransformLayer, std::move(*CompileFunction));
   }
 
   if (S.NumCompileThreads > 0) {
     CompileLayer->setCloneToNewContextOnEmit(true);
     CompileThreads = std::make_unique<ThreadPool>(S.NumCompileThreads);
     ES->setDispatchMaterialization(
-        [this](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
-          // FIXME: Switch to move capture once we have c++14.
-          auto SharedMU = std::shared_ptr<MaterializationUnit>(std::move(MU));
-          auto Work = [SharedMU, &JD]() { SharedMU->doMaterialize(JD); };
-          CompileThreads->async(std::move(Work));
+        [this](JITDylib &JD, std::shared_ptr<MaterializationUnit> MU) {
+          CompileThreads->async(
+              [MU = std::move(MU), &JD]() { MU->doMaterialize(JD); });
         });
   }
 }
 
 std::string LLJIT::mangle(StringRef UnmangledName) {
   std::string MangledName;
   {
     raw_string_ostream MangledNameStream(MangledName);

llvm/unittests/ExecutionEngine/Orc/CoreAPIsTest.cpp

TEST_F(CoreAPIsStandardTest, TestBasicWeakSymbolMaterialization) {
   EXPECT_TRUE(BarMaterialized) << "Bar was not materialized at all";
   EXPECT_TRUE(DuplicateBarDiscarded)
       << "Duplicate bar definition not discarded";
 }
 
 TEST_F(CoreAPIsStandardTest, DefineMaterializingSymbol) {
   bool ExpectNoMoreMaterialization = false;
   ES.setDispatchMaterialization(
-      [&](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
+      [&](JITDylib &JD, std::shared_ptr<MaterializationUnit> MU) {
         if (ExpectNoMoreMaterialization)
           ADD_FAILURE() << "Unexpected materialization";
         MU->doMaterialize(JD);
       });
 
   auto MU = std::make_unique<SimpleMaterializationUnit>(
       SymbolFlagsMap({{Foo, FooSym.getFlags()}}),
       [&](MaterializationResponsibility R) {
EXPECT_EQ(FooLookupResult.getFlags(), FooSym.getFlags())
       << "lookup returned incorrect flags";
 }
 
 TEST_F(CoreAPIsStandardTest, TestLookupWithThreadedMaterialization) {
 #if LLVM_ENABLE_THREADS
 
   std::thread MaterializationThread;
   ES.setDispatchMaterialization(
-      [&](JITDylib &JD, std::unique_ptr<MaterializationUnit> MU) {
+      [&](JITDylib &JD, std::shared_ptr<MaterializationUnit> MU) {
         MaterializationThread =
             std::thread([MU = std::move(MU), &JD] { MU->doMaterialize(JD); });
       });
 
   cantFail(JD.define(absoluteSymbols({{Foo, FooSym}})));
 
   auto FooLookupResult = cantFail(ES.lookup(makeJITDylibSearchOrder(&JD), Foo));