[Mlir-commits] [mlir] fe3c425 - [mlir] Destroy MLIRContext thread pool when disable multi threading

[Eugene Zhulenev]

Author: Eugene Zhulenev
Date: 2021-06-28T13:26:56-07:00
New Revision: fe3c425ae01389eb865d7d979e1ae0c53cc92740

URL: https://github.com/llvm/llvm-project/commit/fe3c425ae01389eb865d7d979e1ae0c53cc92740
DIFF: https://github.com/llvm/llvm-project/commit/fe3c425ae01389eb865d7d979e1ae0c53cc92740.diff

LOG: [mlir] Destroy MLIRContext thread pool when disable multi threading

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D105057

Added: 
    

Modified: 
    mlir/lib/IR/MLIRContext.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 1ae3e6c21cc51..ddb909949cdfc 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -262,7 +262,7 @@ class MLIRContextImpl {
   //===--------------------------------------------------------------------===//
 
   /// The thread pool to use when processing MLIR tasks in parallel.
-  llvm::ThreadPool threadPool;
+  llvm::Optional<llvm::ThreadPool> threadPool;
 
   /// This is a list of dialects that are created referring to this context.
   /// The MLIRContext owns the objects.
@@ -334,7 +334,10 @@ class MLIRContextImpl {
   StringAttr emptyStringAttr;
 
 public:
-  MLIRContextImpl() : identifiers(identifierAllocator) {}
+  MLIRContextImpl() : identifiers(identifierAllocator) {
+    if (threadingIsEnabled)
+      threadPool.emplace();
+  }
   ~MLIRContextImpl() {
     for (auto typeMapping : registeredTypes)
       typeMapping.second->~AbstractType();
@@ -573,12 +576,19 @@ void MLIRContext::disableMultithreading(bool disable) {
   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();
 }
 
 llvm::ThreadPool &MLIRContext::getThreadPool() {
   assert(isMultithreadingEnabled() &&
          "expected multi-threading to be enabled within the context");
-  return impl->threadPool;
+  return *impl->threadPool;
 }
 
 void MLIRContext::enterMultiThreadedExecution() {