Avoid creating a ThreadPool in MlirOptMain when `--mlir-disable-threading` option...

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

Reland 45adf608024 after build fixes

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

diff --git a/mlir/lib/Support/MlirOptMain.cpp b/mlir/lib/Support/MlirOptMain.cpp
index 9a8b21d..5c3f7a7 100644 (file)
--- 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),