Clarify the invariant of the MLIR pass pipeline around `Pass::initialize()`

[mehdi_amini]

Summary

This method should not load new dialect or affect the context itself.

[rriddle]

mlir/include/mlir/Pass/Pass.h

185

Typo here geD

[bondhugula]

mlir/lib/Pass/Pass.cpp

823

Nit: bookkeeping

mlir/include/mlir/Pass/Pass.h

protected:
 
   /// The polymorphic API that runs the pass over the currently held operation.
   virtual void runOnOperation() = 0;
 
   /// Initialize any complex state necessary for running this pass. This hook
   /// should not rely on any state accessible during the execution of a pass.
   /// For example, `getContext`/`getOperation`/`getAnalysis`/etc. should not be
   /// invoked within this hook.
-  /// Returns a LogicalResult to indicate failure, in which case the pass
-  /// pipeline won't execute.
+  /// This method is invoked after all dependent dialects for the pipeline are
+  /// loaded, and is not allowed to load any further dialects (override the
+  /// `geDependentDialects()` for this purpose instead). Returns a LogicalResult
+  /// to indicate failure, in which case the pass pipeline won't execute.
   virtual LogicalResult initialize(MLIRContext *context) { return success(); }
 
   /// Indicate if the current pass can be scheduled on the given operation type.
   /// This is useful for generic operation passes to add restrictions on the
   /// operations they operate on.
   virtual bool canScheduleOn(RegisteredOperationName opName) const = 0;
 
   /// Schedule an arbitrary pass pipeline on the provided operation.

mlir/lib/Pass/Pass.cpp

LogicalResult PassManager::run(Operation *op) {
   context->appendDialectRegistry(dependentDialects);
   for (StringRef name : dependentDialects.getDialectNames())
     context->getOrLoadDialect(name);
 
   // Before running, make sure to finalize the pipeline pass list.
   if (failed(getImpl().finalizePassList(context)))
     return failure();
 
+  // Notify the context that we start running a pipeline for book keeping.
+  context->enterMultiThreadedExecution();
+
   // Initialize all of the passes within the pass manager with a new generation.
   llvm::hash_code newInitKey = context->getRegistryHash();
   if (newInitKey != initializationKey) {
     if (failed(initialize(context, impl->initializationGeneration + 1)))
       return failure();
     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
   // handling enabled.
   LogicalResult result =
       crashReproGenerator ? runWithCrashRecovery(op, am) : runPasses(op, am);
 
   // Notify the context that the run is done.
   context->exitMultiThreadedExecution();