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

[Mehdi Amini]

Author: Mehdi Amini
Date: 2022-01-08T02:26:41Z
New Revision: 3e13c4c37c8c0379b15c509ca483503c649164ed

URL: https://github.com/llvm/llvm-project/commit/3e13c4c37c8c0379b15c509ca483503c649164ed
DIFF: https://github.com/llvm/llvm-project/commit/3e13c4c37c8c0379b15c509ca483503c649164ed.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

Reland 45adf608024 after build fixes

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..5c3f7a7732403 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),