Enable `Pass::initialize()` to fail by returning a LogicalResult

[rriddle]

mlir/include/mlir/Pass/Pass.h

173–175

Can you document the meaning of the result, both here and in website doc?

[mehdi_amini]

Add doc

[rriddle]

LG. Thanks!

mlir/docs/PassManagement.md

 
 In certain situations, a Pass may contain state that is constructed dynamically,
 but is potentially expensive to recompute in successive runs of the Pass. One
 such example is when using [`PDL`-based](Dialects/PDLOps.md)
 [patterns](PatternRewriter.md), which are compiled into a bytecode during
 runtime. In these situations, a pass may override the following hook to
 initialize this heavy state:
 
-*   `void initialize(MLIRContext *context)`
+*   `LogicalResult initialize(MLIRContext *context)`
 
 This hook is executed once per run of a full pass pipeline, meaning that it does
 not have access to the state available during a `runOnOperation` call. More
 concretely, all necessary accesses to an `MLIRContext` should be driven via the
 provided `context` parameter, and methods that utilize "per-run" state such as
 `getContext`/`getOperation`/`getAnalysis`/etc. must not be used.
+In case of an error during initialization, the pass is expected to emit an error
+diagnostic and return a `failure()` which will abort the pass pipeline execution.
 
 ## Analysis Management
 
 An important concept, along with transformation passes, are analyses. These are
 conceptually similar to transformation passes, except that they compute
 information on a specific operation without modifying it. In MLIR, analyses are
 not passes but free-standing classes that are computed lazily on-demand and
 cached to avoid unnecessary recomputation. An analysis in MLIR must adhere to

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.
-  virtual void initialize(MLIRContext *context) {}
+  /// Returns a LogicalResult to indicate failure, in which case the pass
+  /// pipeline won't execute.
+  virtual LogicalResult initialize(MLIRContext *context) { return success(); }
 
   /// Schedule an arbitrary pass pipeline on the provided operation.
   /// This can be invoke any time in a pass to dynamic schedule more passes.
   /// The provided operation must be the current one or one nested below.
   LogicalResult runPipeline(OpPassManager &pipeline, Operation *op) {
     return passState->pipelineExecutor(pipeline, op);
   }
 

mlir/include/mlir/Pass/PassManager.h

public:
 
   /// Return the current nesting mode.
   Nesting getNesting();
 
 private:
   /// Initialize all of the passes within this pass manager with the given
   /// initialization generation. The initialization generation is used to detect
   /// if a pass manager has already been initialized.
-  void initialize(MLIRContext *context, unsigned newInitGeneration);
+  LogicalResult initialize(MLIRContext *context, unsigned newInitGeneration);
 
   /// A pointer to an internal implementation instance.
   std::unique_ptr<detail::OpPassManagerImpl> impl;
 
   /// Allow access to initialize.
   friend detail::OpToOpPassAdaptor;
 
   /// Allow access to the constructor.

mlir/lib/Pass/Pass.cpp

 void OpPassManager::getDependentDialects(DialectRegistry &dialects) const {
   registerDialectsForPipeline(*this, dialects);
 }
 
 void OpPassManager::setNesting(Nesting nesting) { impl->nesting = nesting; }
 
 OpPassManager::Nesting OpPassManager::getNesting() { return impl->nesting; }
 
-void OpPassManager::initialize(MLIRContext *context,
+LogicalResult OpPassManager::initialize(MLIRContext *context,
                                         unsigned newInitGeneration) {
   if (impl->initializationGeneration == newInitGeneration)
-    return;
+    return success();
   impl->initializationGeneration = newInitGeneration;
   for (Pass &pass : getPasses()) {
     // If this pass isn't an adaptor, directly initialize it.
     auto *adaptor = dyn_cast<OpToOpPassAdaptor>(&pass);
     if (!adaptor) {
-      pass.initialize(context);
+      if (failed(pass.initialize(context)))
+        return failure();
       continue;
     }
 
     // Otherwise, initialize each of the adaptors pass managers.
     for (OpPassManager &adaptorPM : adaptor->getPassManagers())
-      adaptorPM.initialize(context, newInitGeneration);
+      if (failed(adaptorPM.initialize(context, newInitGeneration)))
+        return failure();
   }
+  return success();
 }
 
 //===----------------------------------------------------------------------===//
 // OpToOpPassAdaptor
 //===----------------------------------------------------------------------===//
 
 LogicalResult OpToOpPassAdaptor::run(Pass *pass, Operation *op,
                                      AnalysisManager am, bool verifyPasses,
auto dynamic_pipeline_callback = [&](OpPassManager &pipeline,
     if (!op->isAncestor(root))
       return root->emitOpError()
              << "Trying to schedule a dynamic pipeline on an "
                 "operation that isn't "
                 "nested under the current operation the pass is processing";
     assert(pipeline.getOpName() == root->getName().getStringRef());
 
     // Initialize the user provided pipeline and execute the pipeline.
-    pipeline.initialize(root->getContext(), parentInitGeneration);
+    if (failed(pipeline.initialize(root->getContext(), parentInitGeneration)))
+      return failure();
     AnalysisManager nestedAm = root == op ? am : am.nest(root);
     return OpToOpPassAdaptor::runPipeline(pipeline.getPasses(), root, nestedAm,
                                           verifyPasses, parentInitGeneration,
                                           pi, &parentInfo);
   };
   pass->passState.emplace(op, am, dynamic_pipeline_callback);
 
   // Instrument before the pass has run.
LogicalResult PassManager::run(Operation *op) {
   getDependentDialects(dependentDialects);
   context->appendDialectRegistry(dependentDialects);
   for (StringRef name : dependentDialects.getDialectNames())
     context->getOrLoadDialect(name);
 
   // 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);
+    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();

mlir/lib/Transforms/Canonicalizer.cpp

 
 using namespace mlir;
 
 namespace {
 /// Canonicalize operations in nested regions.
 struct Canonicalizer : public CanonicalizerBase<Canonicalizer> {
   /// Initialize the canonicalizer by building the set of patterns used during
   /// execution.
-  void initialize(MLIRContext *context) override {
+  LogicalResult initialize(MLIRContext *context) override {
     OwningRewritePatternList owningPatterns;
     for (auto *op : context->getRegisteredOperations())
       op->getCanonicalizationPatterns(owningPatterns, context);
     patterns = std::move(owningPatterns);
+    return success();
   }
   void runOnOperation() override {
     (void)applyPatternsAndFoldGreedily(getOperation()->getRegions(), patterns);
   }
 
   FrozenRewritePatternList patterns;
 };
 } // end anonymous namespace
 
 /// Create a Canonicalizer pass.
 std::unique_ptr<Pass> mlir::createCanonicalizerPass() {
   return std::make_unique<Canonicalizer>();
 }