[Mlir-commits] [mlir] a32300a - Make the --mlir-disable-threading command line option overrides the C++ API usage

Mehdi Amini llvmlistbot at llvm.org
Tue Sep 14 20:21:06 PDT 2021


Author: Mehdi Amini
Date: 2021-09-15T03:20:48Z
New Revision: a32300a68f6c94b7b275e3560ed31e9174cec5ad

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

LOG: Make the --mlir-disable-threading command line option overrides the C++ API usage

This seems in-line with the intent and how we build tools around it.
Update the description for the flag accordingly.
Also use an injected thread pool in MLIROptMain, now we will create
threads up-front and reuse them across split buffers.

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

Added: 
    

Modified: 
    mlir/include/mlir/IR/MLIRContext.h
    mlir/lib/IR/MLIRContext.cpp
    mlir/lib/Support/MlirOptMain.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h
index 196c6ef441a6a..6d45fbce7b48f 100644
--- a/mlir/include/mlir/IR/MLIRContext.h
+++ b/mlir/include/mlir/IR/MLIRContext.h
@@ -129,6 +129,8 @@ class MLIRContext {
   bool isMultithreadingEnabled();
 
   /// Set the flag specifying if multi-threading is disabled by the context.
+  /// The command line debugging flag `--mlir-disable-threading` is overriding
+  /// this call and making it a no-op!
   void disableMultithreading(bool disable = true);
   void enableMultithreading(bool enable = true) {
     disableMultithreading(!enable);
@@ -140,6 +142,9 @@ class MLIRContext {
   /// decoupling the lifetime of the threads from the contexts. The thread pool
   /// must outlive the context. Multi-threading will be enabled as part of this
   /// method.
+  /// The command line debugging flag `--mlir-disable-threading` will still
+  /// prevent threading from being enabled and threading won't be enabled after
+  /// this call in this case.
   void setThreadPool(llvm::ThreadPool &pool);
 
   /// Return the thread pool used by this context. This method requires that

diff  --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp
index 7e4ec8261a2fc..dd1fa945fb029 100644
--- a/mlir/lib/IR/MLIRContext.cpp
+++ b/mlir/lib/IR/MLIRContext.cpp
@@ -57,7 +57,8 @@ namespace {
 struct MLIRContextOptions {
   llvm::cl::opt<bool> disableThreading{
       "mlir-disable-threading",
-      llvm::cl::desc("Disabling multi-threading within MLIR")};
+      llvm::cl::desc("Disable multi-threading within MLIR, overrides any "
+                     "further call to MLIRContext::enableMultiThreading()")};
 
   llvm::cl::opt<bool> printOpOnDiagnostic{
       "mlir-print-op-on-diagnostic",
@@ -74,6 +75,14 @@ struct MLIRContextOptions {
 
 static llvm::ManagedStatic<MLIRContextOptions> clOptions;
 
+static bool isThreadingGloballyDisabled() {
+#if LLVM_ENABLE_THREADS != 0
+  return clOptions.isConstructed() && clOptions->disableThreading;
+#else
+  return true;
+#endif
+}
+
 /// Register a set of useful command-line options that can be used to configure
 /// various flags within the MLIRContext. These flags are used when constructing
 /// an MLIR context for initialization.
@@ -362,10 +371,10 @@ MLIRContext::MLIRContext(Threading setting)
     : MLIRContext(DialectRegistry(), setting) {}
 
 MLIRContext::MLIRContext(const DialectRegistry &registry, Threading setting)
-    : impl(new MLIRContextImpl(setting == Threading::ENABLED)) {
+    : impl(new MLIRContextImpl(setting == Threading::ENABLED &&
+                               !isThreadingGloballyDisabled())) {
   // Initialize values based on the command line flags if they were provided.
   if (clOptions.isConstructed()) {
-    disableMultithreading(clOptions->disableThreading);
     printOpOnDiagnostic(clOptions->printOpOnDiagnostic);
     printStackTraceOnDiagnostic(clOptions->printStackTraceOnDiagnostic);
   }
@@ -582,6 +591,11 @@ bool MLIRContext::isMultithreadingEnabled() {
 
 /// Set the flag specifying if multi-threading is disabled by the context.
 void MLIRContext::disableMultithreading(bool disable) {
+  // This API can be overridden by the global debugging flag
+  // --mlir-disable-threading
+  if (isThreadingGloballyDisabled())
+    return;
+
   impl->threadingIsEnabled = !disable;
 
   // Update the threading mode for each of the uniquers.

diff  --git a/mlir/lib/Support/MlirOptMain.cpp b/mlir/lib/Support/MlirOptMain.cpp
index 7f4b6e1e99cb9..6da9f993eeafb 100644
--- a/mlir/lib/Support/MlirOptMain.cpp
+++ b/mlir/lib/Support/MlirOptMain.cpp
@@ -32,6 +32,7 @@
 #include "llvm/Support/Regex.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/StringSaver.h"
+#include "llvm/Support/ThreadPool.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 using namespace mlir;
@@ -93,19 +94,22 @@ static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
 
 /// 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,
-                                   const PassPipelineCLParser &passPipeline,
-                                   DialectRegistry &registry) {
+static LogicalResult
+processBuffer(raw_ostream &os, std::unique_ptr<MemoryBuffer> ownedBuffer,
+              bool verifyDiagnostics, bool verifyPasses,
+              bool allowUnregisteredDialects, bool preloadDialectsInContext,
+              const PassPipelineCLParser &passPipeline,
+              DialectRegistry &registry, 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);
+
   // Parse the input file.
-  MLIRContext context(registry);
   if (preloadDialectsInContext)
     context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(allowUnregisteredDialects);
@@ -143,20 +147,24 @@ LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 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;
   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, passPipeline,
-                               registry);
+                               preloadDialectsInContext, passPipeline, registry,
+                               threadPool);
         },
         outputStream);
 
   return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
                        verifyPasses, allowUnregisteredDialects,
-                       preloadDialectsInContext, passPipeline, registry);
+                       preloadDialectsInContext, passPipeline, registry,
+                       threadPool);
 }
 
 LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,


        


More information about the Mlir-commits mailing list