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

[mehdi_amini]

Summary

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

[rriddle]

mlir/lib/Support/MlirOptMain.cpp

149–155

This feels not ideal, is there a way to check statically that threading is enabled? I thought we had something on MLIRContext for that. If not, could we add it?

[mehdi_amini]

mlir/lib/Support/MlirOptMain.cpp

149–155

By "statically" you mean "without an MLIRContext instance"?

We don't have this I believe. The use of --mlir-disable-threading is private to the context implementation and not exposed publicly right now.

[jdd]

mlir/lib/Support/MlirOptMain.cpp

149–155

There is isThreadingGloballyDisabled() in MLIRContext.cpp. You could make that public and use it.

[rriddle]

mlir/lib/Support/MlirOptMain.cpp

149–155

Ah, I was thinking of isThreadingGloballyDisabled in MLIRContext.cpp. Maybe we could expose that in some way?

[mehdi_amini]

mlir/lib/Support/MlirOptMain.cpp

149–155

Well, of course, but I was a bit reluctant to make public the access to cl::opt things that I was seeing as conceptually private :)

[rriddle]

LGTM in general, if we intend to do this more than once (either in another tool, or in user code) would be nice to expose something general for it.

mlir/lib/Support/MlirOptMain.cpp

149–155

True true, this just seemed like something we might need to do elsewhere. Either way, what you have is fine for local needs.

mlir/lib/Support/MlirOptMain.cpp

 
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
 static LogicalResult
 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());
 
   // 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)
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(allowUnregisteredDialects);
   if (verifyDiagnostics)
     context.printOpOnDiagnostic(false);
   context.getDebugActionManager().registerActionHandler<DebugCounter>();
LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,
                                 bool preloadDialectsInContext) {
   // The split-input-file mode is a very specific mode that slices the file
   // 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),
         [&](std::unique_ptr<MemoryBuffer> chunkBuffer, raw_ostream &os) {
           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,
                                 std::unique_ptr<MemoryBuffer> buffer,
                                 const PassPipelineCLParser &passPipeline,
                                 DialectRegistry &registry, bool splitInputFile,
                                 bool verifyDiagnostics, bool verifyPasses,
                                 bool allowUnregisteredDialects,