[mlir] Destroy MLIRContext thread pool when disable multi threading

[ezhulenev]

Fix typo

[rriddle]

mlir/lib/IR/MLIRContext.cpp

265

Can you use Optional here instead?

579–585

[ezhulenev]

Use Optional to store a thread pool

[ezhulenev]

Small fix

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::ThreadPool threadPool;
+  std::unique_ptr<llvm::ThreadPool> threadPool;
 
   /// 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) {}
+  MLIRContextImpl()
+      : threadPool(std::make_unique<llvm::ThreadPool>()),
+        identifiers(identifierAllocator) {}
   ~MLIRContextImpl() {
     for (auto typeMapping : registeredTypes)
       typeMapping.second->~AbstractType();
     for (auto attrMapping : registeredAttributes)
       attrMapping.second->~AbstractAttribute();
   }
 };
 } // end namespace mlir

 /// Set the flag specifying if multi-threading is disabled by the context.
 void MLIRContext::disableMultithreading(bool disable) {
   impl->threadingIsEnabled = !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.
+  if (!impl->threadingIsEnabled)
+    impl->threadPool.reset();
+
+  // Create new thread pool if multi threading was re-enabled.
+  if (impl->threadingIsEnabled && impl->threadPool == nullptr)
+    impl->threadPool = std::make_unique<llvm::ThreadPool>();
 }
 
 llvm::ThreadPool &MLIRContext::getThreadPool() {
   assert(isMultithreadingEnabled() &&
          "expected multi-threading to be enabled within the context");
-  return impl->threadPool;
+  return *impl->threadPool;
 }
 
 void MLIRContext::enterMultiThreadedExecution() {
 #ifndef NDEBUG
   ++impl->multiThreadedExecutionContext;
 #endif
 }
 void MLIRContext::exitMultiThreadedExecution() {