[mlir][PassManager] Only reinitialize the pass manager if the context registry changes

[rriddle]

Summary

This prevents needless reinitialization for clients that want to reuse a pass manager multiple times. A new getRegisryHash function is exposed by the context to give a rough indicator of when the context registry has changed.

mlir/include/mlir/IR/MLIRContext.h

public:
   /// not been loaded for this namespace/TypeID yet, use the provided ctor to
   /// create one on the fly and load it. Returns a pointer to the dialect owned
   /// by the context.
   /// The use of this method is in general discouraged in favor of
   /// 'getOrLoadDialect<DialectClass>()'.
   Dialect *getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
                             function_ref<std::unique_ptr<Dialect>()> ctor);
 
+  /// Returns a hash of the registry of the context that may be used to give
+  /// a rough indicator of if the state of the context registry has changed. The
+  /// context registry correlates to loaded dialects and their entities
+  /// (attributes, operations, types, etc.).
+  llvm::hash_code getRegistryHash();
+
 private:
   const std::unique_ptr<MLIRContextImpl> impl;
 
   MLIRContext(const MLIRContext &) = delete;
   void operator=(const MLIRContext &) = delete;
 };
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/Pass/PassManager.h

private:
   Optional<PassDisplayMode> passStatisticsMode;
 
   /// A manager for pass instrumentations.
   std::unique_ptr<PassInstrumentor> instrumentor;
 
   /// An optional factory to use when generating a crash reproducer if valid.
   ReproducerStreamFactory crashReproducerStreamFactory;
 
+  /// A hash key used to detect when reinitialization is necessary.
+  llvm::hash_code initializationKey;
+
   /// Flag that specifies if pass timing is enabled.
   bool passTiming : 1;
 
   /// Flag that specifies if the generated crash reproducer should be local.
   bool localReproducer : 1;
 
   /// A flag that indicates if the IR should be verified in between passes.
   bool verifyPasses : 1;

mlir/lib/IR/MLIRContext.cpp

#endif
   // Abort if dialect with namespace has already been registered.
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
 
   return dialect.get();
 }
 
+llvm::hash_code MLIRContext::getRegistryHash() {
+  llvm::hash_code hash(0);
+  // Factor in number of loaded dialects, attributes, operations, types.
+  hash = llvm::hash_combine(hash, impl->loadedDialects.size());
+  hash = llvm::hash_combine(hash, impl->registeredAttributes.size());
+  hash = llvm::hash_combine(hash, impl->registeredOperations.size());
+  hash = llvm::hash_combine(hash, impl->registeredTypes.size());
+  return hash;
+}
+
 bool MLIRContext::allowsUnregisteredDialects() {
   return impl->allowUnregisteredDialects;
 }
 
 void MLIRContext::allowUnregisteredDialects(bool allowing) {
   impl->allowUnregisteredDialects = allowing;
 }
 

mlir/lib/Pass/Pass.cpp

 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
 PassManager::PassManager(MLIRContext *ctx, Nesting nesting,
                          StringRef operationName)
     : OpPassManager(Identifier::get(operationName, ctx), nesting), context(ctx),
+      initializationKey(DenseMapInfo<llvm::hash_code>::getTombstoneKey()),
       passTiming(false), localReproducer(false), verifyPasses(true) {}
 
 PassManager::~PassManager() {}
 
 void PassManager::enableVerifier(bool enabled) { verifyPasses = enabled; }
 
 /// Run the passes within this manager on the provided operation.
 LogicalResult PassManager::run(Operation *op) {
   MLIRContext *context = getContext();
   assert(op->getName().getIdentifier() == getOpName(*context) &&
          "operation has a different name than the PassManager");
 
   // Before running, make sure to coalesce any adjacent pass adaptors in the
   // pipeline.
   getImpl().coalesceAdjacentAdaptorPasses();
 
   // Register all dialects for the current pipeline.
   DialectRegistry dependentDialects;
   getDependentDialects(dependentDialects);
   dependentDialects.loadAll(context);
 
   // Initialize all of the passes within the pass manager with a new generation.
+  llvm::hash_code newInitKey = context->getRegistryHash();
+  if (newInitKey != initializationKey) {
     initialize(context, impl->initializationGeneration + 1);
+    initializationKey = newInitKey;
+  }
 
   // Construct a top level analysis manager for the pipeline.
   ModuleAnalysisManager am(op, instrumentor.get());
 
   // Notify the context that we start running a pipeline for book keeping.
   context->enterMultiThreadedExecution();
 
   // If reproducer generation is enabled, run the pass manager with crash