Update MLIRContext to allow injecting an external ThreadPool (NFC)

[mehdi_amini]

Summary

The context can be created with threading disabled, to avoid creating a thread pool
that may be destroyed when injecting another one later.

[jpienaar]

Seems good

mlir/lib/IR/MLIRContext.cpp

595

This interaction was confusing to me initially, could you add a comment here? May also be good to document that threadPool != nullptr doesn't mean threading is enabled.

[mehdi_amini]

Improve doc

[stellaraccident]

mlir/lib/IR/MLIRContext.cpp

606

So, if I want to bring my own thread pool and avoid thrashing things by having the context spin up a bunch of threads just to reset them, is the sequence:

llvm::ThreadPool myAwesomeThreadPool;
MLIRContext ctx(registry, Threading::disabled);
ctx.setThreadPool(myAwesomeThreadPool);
ctx.enableMultithreading();

I feel that a slight improvement might be to call ctx.enableMultithreading() from setThreadPool.

Also, since the comment says This method requires that multithreading is disabled for this context prior to the call., maybe add an assert() on that here?

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

57

We generally capitalize the values of enums .

mlir/lib/IR/MLIRContext.cpp

264–272

This comment looks off, can you reword the first part?

600

Should you wrap this in the previous if?

[mehdi_amini]

Address comments:

• More docs, in particular include an extended version of Stella's example in the class doc.
• Enable multi-threading when setting an external pool
• Add some asserts around the invariants for threadPool/ownedThreadPool.

mlir/include/mlir/IR/MLIRContext.h

 /// MLIRContext is the top-level object for a collection of MLIR operations. It
 /// holds immortal uniqued objects like types, and the tables used to unique
 /// them.
 ///
 /// MLIRContext gets a redundant "MLIR" prefix because otherwise it ends up with
 /// a very generic name ("Context") and because it is uncommon for clients to
 /// interact with it.
 ///
+/// The context wrap some multi-threading facilities, and in particular by
+/// default it will implicitly create a thread pool.
+/// This can be undesirable if multiple context exists at the same time or if a
+/// process will be long-lived and create and destroy contexts.
+/// To control better thread spawning, an externally owned ThreadPool can be
+/// injected in the context. For example:
+///
+///  llvm::ThreadPool myThreadPool;
+///  while (auto *request = nextCompilationRequests()) {
+///    MLIRContext ctx(registry, MLIRContext::Threading::DISABLED);
+///    ctx.setThreadPool(myThreadPool);
+///    processRequest(request, cxt);
+///  }
+///
 class MLIRContext {
 public:
+  enum class Threading { DISABLED, ENABLED };
   /// Create a new Context.
-  explicit MLIRContext();
-  explicit MLIRContext(const DialectRegistry &registry);
+  explicit MLIRContext(Threading multithreading = Threading::ENABLED);
+  explicit MLIRContext(const DialectRegistry &registry,
+                       Threading multithreading = Threading::ENABLED);
   ~MLIRContext();
 
   /// Return information about all IR dialects loaded in the context.
   std::vector<Dialect *> getLoadedDialects();
 
   /// Return the dialect registry associated with this context.
   const DialectRegistry &getDialectRegistry();
 
public:
   bool isMultithreadingEnabled();
 
   /// Set the flag specifying if multi-threading is disabled by the context.
   void disableMultithreading(bool disable = true);
   void enableMultithreading(bool enable = true) {
     disableMultithreading(!enable);
   }
 
-  /// Return the thread pool owned by this context. This method requires that
+  /// Set a new thread pool to be used in this context. This method requires
+  /// that multithreading is disabled for this context prior to the call. This
+  /// allows to share a thread pool across multiple contexts, as well as
+  /// decoupling the lifetime of the threads from the contexts. The thread pool
+  /// must outlive the context. Multi-threading will be enabled as part of this
+  /// method.
+  void setThreadPool(llvm::ThreadPool &pool);
+
+  /// Return the thread pool used by this context. This method requires that
   /// multithreading be enabled within the context, and should generally not be
   /// used directly. Users should instead prefer the threading utilities within
   /// Threading.h.
   llvm::ThreadPool &getThreadPool();
 
   /// Return true if we should attach the operation to diagnostics emitted via
   /// Operation::emit.
   bool shouldPrintOpOnDiagnostic();

mlir/lib/IR/MLIRContext.cpp

#endif
 
   /// If the current stack trace should be attached when emitting diagnostics.
   bool printStackTraceOnDiagnostic = false;
 
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 
-  /// The thread pool to use when processing MLIR tasks in parallel.
-  llvm::Optional<llvm::ThreadPool> threadPool;
+  /// This points to the ThreadPool used when processing MLIR tasks in parallel.
+  /// It can't be nullptr when multi-threading is enabled. Otherwise if
+  /// multi-threading is disabled, and the threadpool wasn't externally provided
+  /// using `setThreadPool`, this will be nullptr.
+  llvm::ThreadPool *threadPool = nullptr;
+
+  /// In case where the thread pool is owned by the context, this ensures
+  /// destruction with the context.
+  std::unique_ptr<llvm::ThreadPool> ownedThreadPool;
 
   /// This is a list of dialects that are created referring to this context.
   /// The MLIRContext owns the objects.
   DenseMap<StringRef, std::unique_ptr<Dialect>> loadedDialects;
   DialectRegistry dialectsRegistry;
 
   /// This is a mapping from operation name to AbstractOperation for registered
   /// operations.
#endif
   /// Cached Attribute Instances.
   BoolAttr falseAttr, trueAttr;
   UnitAttr unitAttr;
   UnknownLoc unknownLocAttr;
   DictionaryAttr emptyDictionaryAttr;
   StringAttr emptyStringAttr;
 
 public:
-  MLIRContextImpl() : identifiers(identifierAllocator) {
-    if (threadingIsEnabled)
-      threadPool.emplace();
+  MLIRContextImpl(bool threadingIsEnabled)
+      : threadingIsEnabled(threadingIsEnabled),
+        identifiers(identifierAllocator) {
+    if (threadingIsEnabled) {
+      ownedThreadPool = std::make_unique<llvm::ThreadPool>();
+      threadPool = ownedThreadPool.get();
+    }
   }
   ~MLIRContextImpl() {
     for (auto typeMapping : registeredTypes)
       typeMapping.second->~AbstractType();
     for (auto attrMapping : registeredAttributes)
       attrMapping.second->~AbstractAttribute();
   }
 };
 } // end namespace mlir
 
-MLIRContext::MLIRContext() : MLIRContext(DialectRegistry()) {}
+MLIRContext::MLIRContext(Threading setting)
+    : MLIRContext(DialectRegistry(), setting) {}
 
-MLIRContext::MLIRContext(const DialectRegistry &registry)
-    : impl(new MLIRContextImpl) {
+MLIRContext::MLIRContext(const DialectRegistry &registry, Threading setting)
+    : impl(new MLIRContextImpl(setting == Threading::ENABLED)) {
   // Initialize values based on the command line flags if they were provided.
   if (clOptions.isConstructed()) {
     disableMultithreading(clOptions->disableThreading);
     printOpOnDiagnostic(clOptions->printOpOnDiagnostic);
     printStackTraceOnDiagnostic(clOptions->printStackTraceOnDiagnostic);
   }
 
   // Pre-populate the registry.
void MLIRContext::disableMultithreading(bool disable) {
 
   // Update the threading mode for each of the uniquers.
   impl->affineUniquer.disableMultithreading(disable);
   impl->attributeUniquer.disableMultithreading(disable);
   impl->typeUniquer.disableMultithreading(disable);
 
   // Destroy thread pool (stop all threads) if it is no longer needed, or create
   // a new one if multithreading was re-enabled.
-  if (!impl->threadingIsEnabled)
-    impl->threadPool.reset();
-  else if (!impl->threadPool.hasValue())
-    impl->threadPool.emplace();
+  if (disable) {
+    // If the thread pool is owned, explicitly set it to nullptr to avoid
+    // keeping a dangling pointer around. If the thread pool is externally
+    // owned, we don't do anything.
+    if (impl->ownedThreadPool) {
+      assert(impl->threadPool);
+      impl->threadPool = nullptr;
+      impl->ownedThreadPool.reset();
+    }
+  } else if (!impl->threadPool) {
+    // The thread pool isn't externally provided.
+    assert(!impl->ownedThreadPool);
+    impl->ownedThreadPool = std::make_unique<llvm::ThreadPool>();
+    impl->threadPool = impl->ownedThreadPool.get();
+  }
+}
+
+void MLIRContext::setThreadPool(llvm::ThreadPool &pool) {
+  assert(!isMultithreadingEnabled() &&
+         "expected multi-threading to be disabled when setting a ThreadPool");
+  impl->threadPool = &pool;
+  impl->ownedThreadPool.reset();
+  enableMultithreading();
 }
 
 llvm::ThreadPool &MLIRContext::getThreadPool() {
   assert(isMultithreadingEnabled() &&
          "expected multi-threading to be enabled within the context");
+  assert(impl->threadPool &&
+         "multi-threading is enabled but threadpool not set");
   return *impl->threadPool;
 }
 
 void MLIRContext::enterMultiThreadedExecution() {
 #ifndef NDEBUG
   ++impl->multiThreadedExecutionContext;
 #endif
 }