Give PassInstrumentor a SmartMutex to lock access to the held instrumentations.

PiperOrigin-RevId: 239031524

mlir/g3doc/WritingAPass.md

@@ -367,9 +367,9 @@ that counts the number of times DominanceInfo is computed:
 
 ```c++
 struct DominanceCounterInstrumentation : public PassInstrumentation {
-  std::atomic<int> &count;
+  unsigned &count;
 
-  DominanceCounterInstrumentation(std::atomic<int> &count) : count(count) {}
+  DominanceCounterInstrumentation(unsigned &count) : count(count) {}
   void runAfterAnalysis(llvm::StringRef, AnalysisID *id,
                         const llvm::Any &) override {
     if (id == AnalysisID::getID<DominanceInfo>())
@@ -380,7 +380,7 @@ struct DominanceCounterInstrumentation : public PassInstrumentation {
 PassManager pm;
 
 // Add the instrumentation to the pass manager.
-std::atomic<int> domInfoCount;
+unsigned domInfoCount;
 pm.addInstrumentation(new DominanceCounterInstrumentation(domInfoCount));
 
 // Run the pass manager on a module.

mlir/include/mlir/Pass/PassInstrumentation.h

@@ -26,6 +26,10 @@ namespace mlir {
 struct AnalysisID;
 class Pass;
 
+namespace detail {
+struct PassInstrumentorImpl;
+} // end namespace detail
+
 /// PassInstrumentation provdes several entry points into the pass manager
 /// infrastructure. Instrumentations should be added directly to a PassManager
 /// before running a pipeline.
@@ -66,51 +70,61 @@ public:
 /// their respective call backs.
 class PassInstrumentor {
 public:
+  PassInstrumentor();
+  PassInstrumentor(PassInstrumentor &&) = delete;
+  PassInstrumentor(const PassInstrumentor &) = delete;
+  ~PassInstrumentor();
+
   /// See PassInstrumentation::runBeforePass for details.
   template <typename IRUnitT> void runBeforePass(Pass *pass, IRUnitT *ir) {
-    llvm::Any irAny(ir);
-    for (auto &instr : instrumentations)
-      instr->runBeforePass(pass, irAny);
+    runBeforePass(pass, llvm::Any(ir));
   }
 
   /// See PassInstrumentation::runAfterPass for details.
   template <typename IRUnitT> void runAfterPass(Pass *pass, IRUnitT *ir) {
-    llvm::Any irAny(ir);
-    for (auto &instr : llvm::reverse(instrumentations))
-      instr->runAfterPass(pass, irAny);
+    runAfterPass(pass, llvm::Any(ir));
   }
 
   /// See PassInstrumentation::runAfterPassFailed for details.
   template <typename IRUnitT> void runAfterPassFailed(Pass *pass, IRUnitT *ir) {
-    llvm::Any irAny(ir);
-    for (auto &instr : llvm::reverse(instrumentations))
-      instr->runAfterPassFailed(pass, irAny);
+    runAfterPassFailed(pass, llvm::Any(ir));
   }
 
   /// See PassInstrumentation::runBeforeAnalysis for details.
   template <typename IRUnitT>
   void runBeforeAnalysis(llvm::StringRef name, AnalysisID *id, IRUnitT *ir) {
-    llvm::Any irAny(ir);
-    for (auto &instr : instrumentations)
-      instr->runBeforeAnalysis(name, id, irAny);
+    runBeforeAnalysis(name, id, llvm::Any(ir));
   }
 
   /// See PassInstrumentation::runAfterAnalysis for details.
   template <typename IRUnitT>
   void runAfterAnalysis(llvm::StringRef name, AnalysisID *id, IRUnitT *ir) {
-    llvm::Any irAny(ir);
-    for (auto &instr : llvm::reverse(instrumentations))
-      instr->runAfterAnalysis(name, id, irAny);
+    runAfterAnalysis(name, id, llvm::Any(ir));
   }
 
   /// Add the given instrumentation to the collection. This takes ownership over
   /// the given pointer.
-  void addInstrumentation(PassInstrumentation *pi) {
-    instrumentations.emplace_back(pi);
-  }
+  void addInstrumentation(PassInstrumentation *pi);
 
 private:
-  std::vector<std::unique_ptr<PassInstrumentation>> instrumentations;
+  /// See PassInstrumentation::runBeforePass for details.
+  void runBeforePass(Pass *pass, const llvm::Any &ir);
+
+  /// See PassInstrumentation::runAfterPass for details.
+  void runAfterPass(Pass *pass, const llvm::Any &ir);
+
+  /// See PassInstrumentation::runAfterPassFailed for details.
+  void runAfterPassFailed(Pass *pass, const llvm::Any &ir);
+
+  /// See PassInstrumentation::runBeforeAnalysis for details.
+  void runBeforeAnalysis(llvm::StringRef name, AnalysisID *id,
+                         const llvm::Any &ir);
+
+  /// See PassInstrumentation::runAfterAnalysis for details.
+  void runAfterAnalysis(llvm::StringRef name, AnalysisID *id,
+                        const llvm::Any &ir);
+
+  std::unique_ptr<detail::PassInstrumentorImpl> impl;
 };
 
 } // end namespace mlir

mlir/lib/Pass/Pass.cpp

@@ -23,6 +23,7 @@
 #include "PassDetail.h"
 #include "mlir/IR/Module.h"
 #include "mlir/Pass/PassManager.h"
+#include "llvm/Support/Mutex.h"
 
 using namespace mlir;
 using namespace mlir::detail;
@@ -279,3 +280,66 @@ void ModuleAnalysisManager::invalidate(const detail::PreservedAnalyses &pa) {
 //===----------------------------------------------------------------------===//
 
 PassInstrumentation::~PassInstrumentation() {}
+
+//===----------------------------------------------------------------------===//
+// PassInstrumentor
+//===----------------------------------------------------------------------===//
+
+namespace mlir {
+namespace detail {
+struct PassInstrumentorImpl {
+  /// Mutex to keep instrumentation access thread-safe.
+  llvm::sys::SmartMutex<true> mutex;
+
+  /// Set of registered instrumentations.
+  std::vector<std::unique_ptr<PassInstrumentation>> instrumentations;
+};
+} // end namespace detail
+} // end namespace mlir
+
+PassInstrumentor::PassInstrumentor() : impl(new PassInstrumentorImpl()) {}
+PassInstrumentor::~PassInstrumentor() {}
+
+/// See PassInstrumentation::runBeforePass for details.
+void PassInstrumentor::runBeforePass(Pass *pass, const llvm::Any &ir) {
+  llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+  for (auto &instr : impl->instrumentations)
+    instr->runBeforePass(pass, ir);
+}
+
+/// See PassInstrumentation::runAfterPass for details.
+void PassInstrumentor::runAfterPass(Pass *pass, const llvm::Any &ir) {
+  llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+  for (auto &instr : llvm::reverse(impl->instrumentations))
+    instr->runAfterPass(pass, ir);
+}
+
+/// See PassInstrumentation::runAfterPassFailed for details.
+void PassInstrumentor::runAfterPassFailed(Pass *pass, const llvm::Any &ir) {
+  llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+  for (auto &instr : llvm::reverse(impl->instrumentations))
+    instr->runAfterPassFailed(pass, ir);
+}
+
+/// See PassInstrumentation::runBeforeAnalysis for details.
+void PassInstrumentor::runBeforeAnalysis(llvm::StringRef name, AnalysisID *id,
+                                         const llvm::Any &ir) {
+  llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+  for (auto &instr : impl->instrumentations)
+    instr->runBeforeAnalysis(name, id, ir);
+}
+
+/// See PassInstrumentation::runAfterAnalysis for details.
+void PassInstrumentor::runAfterAnalysis(llvm::StringRef name, AnalysisID *id,
+                                        const llvm::Any &ir) {
+  llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+  for (auto &instr : llvm::reverse(impl->instrumentations))
+    instr->runAfterAnalysis(name, id, ir);
+}
+
+/// Add the given instrumentation to the collection. This takes ownership over
+/// the given pointer.
+void PassInstrumentor::addInstrumentation(PassInstrumentation *pi) {
+  llvm::sys::SmartScopedLock<true> instrumentationLock(impl->mutex);
+  impl->instrumentations.emplace_back(pi);
+}