[Mlir-commits] [mlir] 45adf60 - Avoid creating a ThreadPool in MlirOptMain when `--mlir-disable-threading` option is set

[Mehdi Amini]

Author: Mehdi Amini
Date: 2022-01-08T02:18:19Z
New Revision: 45adf608024655d3594993d041af9602da41780f

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

LOG: Avoid creating a ThreadPool in MlirOptMain when `--mlir-disable-threading` option is set

a32300a changed it to create a ThreadPool eagerly so that it gets reused
across buffers, however it also made it so that we create a ThreadPool
early even if we're not gonna use it later because of the command line
option `--mlir-disable-threading` is provided.

Fix #53056

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

Added: 
    

Modified: 
    mlir/lib/Support/MlirOptMain.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Support/MlirOptMain.cpp b/mlir/lib/Support/MlirOptMain.cpp
index 9a8b21d37f254..94c7563b1b149 100644
--- a/mlir/lib/Support/MlirOptMain.cpp
+++ b/mlir/lib/Support/MlirOptMain.cpp
@@ -94,7 +94,7 @@ processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
               bool verifyDiagnostics, bool verifyPasses,
               bool allowUnregisteredDialects, bool preloadDialectsInContext,
               PassPipelineFn passManagerSetupFn, DialectRegistry &registry,
-              llvm::ThreadPool &threadPool) {
+              llvm::ThreadPool *threadPool) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
@@ -102,7 +102,8 @@ processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
   // Create a context just for the current buffer. Disable threading on creation
   // since we'll inject the thread-pool separately.
   MLIRContext context(registry, MLIRContext::Threading::DISABLED);
-  context.setThreadPool(threadPool);
+  if (threadPool)
+    context.setThreadPool(*threadPool);
 
   // Parse the input file.
   if (preloadDialectsInContext)
@@ -144,7 +145,15 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
   // up into small pieces and checks each independently.
   // We use an explicit threadpool to avoid creating and joining/destroying
   // threads for each of the split.
-  llvm::ThreadPool threadPool;
+  ThreadPool *threadPool = nullptr;
+  // Create a temporary context for the sake of checking if
+  // --mlir-disable-threading was passed on the command line.
+  // We use the thread-pool this context is creating, and avoid
+  // creating any thread when disabled.
+  MLIRContext threadPoolCtx;
+  if (threadPoolCtx.isMultithreadingEnabled())
+    threadPool = &threadPoolCtx.getThreadPool();
+
   if (splitInputFile)
     return splitAndProcessBuffer(
         std::move(buffer),
@@ -152,14 +161,14 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
           return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
                                verifyPasses, allowUnregisteredDialects,
                                preloadDialectsInContext, passManagerSetupFn,
-                               registry, threadPool);
+                               registry, threadPool.get());
         },
         outputStream);
 
   return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
                        verifyPasses, allowUnregisteredDialects,
                        preloadDialectsInContext, passManagerSetupFn, registry,
-                       threadPool);
+                       threadPool.get());
 }
 
 LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,