Make the implicit nesting behavior of the PassManager user-controllable and default to false

[mehdi_amini]

Summary

This is an error prone behavior, I frequently have ~20 min debugging sessions when I hit
an unexpected implicit nesting. This default makes the C++ API safer for users.

Depends On D90669

[rriddle]

Nice, this is really clean.

Can you add a unit test exercising the fatal error?

mlir/include/mlir/Pass/PassManager.h

51–53

This deserves some documentation.

[rriddle]

mlir/include/mlir/Pass/PassManager.h

152–153

Also, this comment needs to be updated in the previous revision.

[mehdi_amini]

Address comments

[silvas]

Are users expected to migrate to explicit mode?

[mehdi_amini]

They are encouraged to, but we don't consider the implicit behavior deprecated so they do how they see fit.

mlir/examples/toy/Ch5/toyc.cpp

if (enableOpt || isLoweringToAffine) {
     // the operations.
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
     optPM.addPass(mlir::toy::createShapeInferencePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
   }
 
   if (isLoweringToAffine) {
-    // Partially lower the toy dialect with a few cleanups afterwards.
-    pm.addPass(mlir::toy::createLowerToAffinePass());
-
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
+
+    // Partially lower the toy dialect with a few cleanups afterwards.
+    optPM.addPass(mlir::toy::createLowerToAffinePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
 
     // Add optimizations if enabled.
     if (enableOpt) {
       optPM.addPass(mlir::createLoopFusionPass());
       optPM.addPass(mlir::createMemRefDataFlowOptPass());
     }

mlir/examples/toy/Ch6/toyc.cpp

if (enableOpt || isLoweringToAffine) {
     // the operations.
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
     optPM.addPass(mlir::toy::createShapeInferencePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
   }
 
   if (isLoweringToAffine) {
-    // Partially lower the toy dialect with a few cleanups afterwards.
-    pm.addPass(mlir::toy::createLowerToAffinePass());
-
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
+
+    // Partially lower the toy dialect with a few cleanups afterwards.
+    optPM.addPass(mlir::toy::createLowerToAffinePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
 
     // Add optimizations if enabled.
     if (enableOpt) {
       optPM.addPass(mlir::createLoopFusionPass());
       optPM.addPass(mlir::createMemRefDataFlowOptPass());
     }

mlir/examples/toy/Ch7/toyc.cpp

if (enableOpt || isLoweringToAffine) {
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::toy::createShapeInferencePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
   }
 
   if (isLoweringToAffine) {
-    // Partially lower the toy dialect with a few cleanups afterwards.
-    pm.addPass(mlir::toy::createLowerToAffinePass());
-
     mlir::OpPassManager &optPM = pm.nest<mlir::FuncOp>();
+
+    // Partially lower the toy dialect with a few cleanups afterwards.
+    optPM.addPass(mlir::toy::createLowerToAffinePass());
     optPM.addPass(mlir::createCanonicalizerPass());
     optPM.addPass(mlir::createCSEPass());
 
     // Add optimizations if enabled.
     if (enableOpt) {
       optPM.addPass(mlir::createLoopFusionPass());
       optPM.addPass(mlir::createMemRefDataFlowOptPass());
     }

mlir/include/mlir/Pass/PassManager.h

 // OpPassManager
 //===----------------------------------------------------------------------===//
 
 /// This class represents a pass manager that runs passes on a specific
 /// operation type. This class is not constructed directly, but nested within
 /// other OpPassManagers or the top-level PassManager.
 class OpPassManager {
 public:
-  OpPassManager(Identifier name);
-  OpPassManager(StringRef name);
+  enum class Nesting { Implicit, Explicit };
+  OpPassManager(Identifier name, Nesting nesting);
+  OpPassManager(StringRef name, Nesting nesting);
   OpPassManager(OpPassManager &&rhs);
   OpPassManager(const OpPassManager &rhs);
   ~OpPassManager();
   OpPassManager &operator=(const OpPassManager &rhs);
 
   /// Iterator over the passes in this pass manager.
   using pass_iterator =
       llvm::pointee_iterator<MutableArrayRef<std::unique_ptr<Pass>>::iterator>;
enum class PassDisplayMode {
   // In this mode the results are displayed in a nested pipeline view that
   // mirrors the internal pass pipeline that is being executed in the pass
   // manager.
   Pipeline,
 };
 
 /// The main pass manager and pipeline builder.
 class PassManager : public OpPassManager {
 public:
-  PassManager(MLIRContext *ctx);
+  PassManager(MLIRContext *ctx, Nesting nesting = Nesting::Explicit);
   ~PassManager();
 
   /// Run the passes within this manager on the provided module.
   LLVM_NODISCARD
   LogicalResult run(ModuleOp module);
 
   /// Return an instance of the context.
   MLIRContext *getContext() const { return context; }

mlir/lib/Pass/Pass.cpp

 
 //===----------------------------------------------------------------------===//
 // OpPassManagerImpl
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace detail {
 struct OpPassManagerImpl {
-  OpPassManagerImpl(Identifier identifier)
-      : name(identifier), identifier(identifier) {}
-  OpPassManagerImpl(StringRef name) : name(name) {}
+  OpPassManagerImpl(Identifier identifier, OpPassManager::Nesting nesting)
+      : name(identifier), identifier(identifier), nesting(nesting) {}
+  OpPassManagerImpl(StringRef name, OpPassManager::Nesting nesting)
+      : name(name), nesting(nesting) {}
 
   /// Merge the passes of this pass manager into the one provided.
   void mergeInto(OpPassManagerImpl &rhs);
 
   /// Nest a new operation pass manager for the given operation kind under this
   /// pass manager.
   OpPassManager &nest(Identifier nestedName);
   OpPassManager &nest(StringRef nestedName);
struct OpPassManagerImpl {
   StringRef name;
 
   /// The cached identifier (internalized in the context) for the name of the
   /// operation that passes of this pass manager operate on.
   Optional<Identifier> identifier;
 
   /// The set of passes to run as part of this pass manager.
   std::vector<std::unique_ptr<Pass>> passes;
+
+  /// Control the implicit nesting of passes that mismatch the name set for this
+  /// OpPassManager.
+  OpPassManager::Nesting nesting;
 };
 } // end namespace detail
 } // end namespace mlir
 
 void OpPassManagerImpl::mergeInto(OpPassManagerImpl &rhs) {
   assert(name == rhs.name && "merging unrelated pass managers");
   for (auto &pass : passes)
     rhs.passes.push_back(std::move(pass));
   passes.clear();
 }
 
 OpPassManager &OpPassManagerImpl::nest(Identifier nestedName) {
-  OpPassManager nested(nestedName);
+  OpPassManager nested(nestedName, nesting);
   auto *adaptor = new OpToOpPassAdaptor(std::move(nested));
   addPass(std::unique_ptr<Pass>(adaptor));
   return adaptor->getPassManagers().front();
 }
 
 OpPassManager &OpPassManagerImpl::nest(StringRef nestedName) {
-  OpPassManager nested(nestedName);
+  OpPassManager nested(nestedName, nesting);
   auto *adaptor = new OpToOpPassAdaptor(std::move(nested));
   addPass(std::unique_ptr<Pass>(adaptor));
   return adaptor->getPassManagers().front();
 }
 
 void OpPassManagerImpl::addPass(std::unique_ptr<Pass> pass) {
   // If this pass runs on a different operation than this pass manager, then
-  // implicitly nest a pass manager for this operation.
+  // implicitly nest a pass manager for this operation if enabled.
   auto passOpName = pass->getOpName();
-  if (passOpName && passOpName != name)
+  if (passOpName && passOpName != name) {
+    if (nesting == OpPassManager::Nesting::Implicit)
       return nest(*passOpName).addPass(std::move(pass));
+    llvm::report_fatal_error(llvm::Twine("Can't add pass '") + pass->getName() +
+                             "' restricted to '" + *passOpName +
+                             "' on a PassManager intended to run on '" + name +
+                             "', did you intend to nest?");
+  }
 
   passes.emplace_back(std::move(pass));
 }
 
 void OpPassManagerImpl::coalesceAdjacentAdaptorPasses() {
   // Bail out early if there are no adaptor passes.
   if (llvm::none_of(passes, [](std::unique_ptr<Pass> &pass) {
         return isa<OpToOpPassAdaptor>(pass.get());
for (std::unique_ptr<Pass> &pass : oldPasses) {
     }
   }
 }
 
 //===----------------------------------------------------------------------===//
 // OpPassManager
 //===----------------------------------------------------------------------===//
 
-OpPassManager::OpPassManager(Identifier name)
-    : impl(new OpPassManagerImpl(name)) {}
-OpPassManager::OpPassManager(StringRef name)
-    : impl(new OpPassManagerImpl(name)) {}
+OpPassManager::OpPassManager(Identifier name, Nesting nesting)
+    : impl(new OpPassManagerImpl(name, nesting)) {}
+OpPassManager::OpPassManager(StringRef name, Nesting nesting)
+    : impl(new OpPassManagerImpl(name, nesting)) {}
 OpPassManager::OpPassManager(OpPassManager &&rhs) : impl(std::move(rhs.impl)) {}
 OpPassManager::OpPassManager(const OpPassManager &rhs) { *this = rhs; }
 OpPassManager &OpPassManager::operator=(const OpPassManager &rhs) {
-  impl.reset(new OpPassManagerImpl(rhs.impl->name));
+  impl.reset(new OpPassManagerImpl(rhs.impl->name, rhs.impl->nesting));
   for (auto &pass : rhs.impl->passes)
     impl->passes.emplace_back(pass->clone());
   return *this;
 }
 
 OpPassManager::~OpPassManager() {}
 
 OpPassManager::pass_iterator OpPassManager::begin() {
return module.emitError()
             "reproducer has been generated in '"
          << *crashReproducerFileName << "'";
 }
 
 //===----------------------------------------------------------------------===//
 // PassManager
 //===----------------------------------------------------------------------===//
 
-PassManager::PassManager(MLIRContext *ctx)
-    : OpPassManager(Identifier::get(ModuleOp::getOperationName(), ctx)),
+PassManager::PassManager(MLIRContext *ctx, Nesting nesting)
+    : OpPassManager(Identifier::get(ModuleOp::getOperationName(), ctx),
+                    nesting),
       context(ctx), passTiming(false), localReproducer(false),
       verifyPasses(true) {}
 
 PassManager::~PassManager() {}
 
 void PassManager::enableVerifier(bool enabled) { verifyPasses = enabled; }
 
 /// Run the passes within this manager on the provided module.

mlir/lib/Support/MlirOptMain.cpp

static LogicalResult performActions(raw_ostream &os, bool verifyDiagnostics,
 
   // Parse the input file and reset the context threading state.
   OwningModuleRef module(parseSourceFile(sourceMgr, context));
   context->enableMultithreading(wasThreadingEnabled);
   if (!module)
     return failure();
 
   // Apply any pass manager command line options.
-  PassManager pm(context);
+  PassManager pm(context, OpPassManager::Nesting::Implicit);
   pm.enableVerifier(verifyPasses);
   applyPassManagerCLOptions(pm);
 
   // Build the provided pipeline.
   if (failed(passPipeline.addToPipeline(pm)))
     return failure();
 
   // Run the pipeline.

mlir/test/lib/Transforms/TestConvVectorization.cpp

void TestConvVectorization::runOnOperation() {
   // supported as can be seen in splitFullAndPartialTransferPrecondition,
   // VectorTransforms.cpp
   vectorTransferPatterns.insert<VectorTransferFullPartialRewriter>(
       context, vectorTransformsOptions);
   applyPatternsAndFoldGreedily(module, std::move(vectorTransferPatterns));
 
   // Programmatic controlled lowering of linalg.copy and linalg.fill.
   PassManager pm(context);
-  pm.addPass(createConvertLinalgToLoopsPass());
+  pm.addNestedPass<FuncOp>(createConvertLinalgToLoopsPass());
   if (failed(pm.run(module)))
     llvm_unreachable("Unexpected failure in linalg to loops pass.");
 
   // Programmatic controlled lowering of vector.contract only.
   OwningRewritePatternList vectorContractLoweringPatterns;
   populateVectorContractLoweringPatterns(vectorContractLoweringPatterns,
                                          context, vectorTransformsOptions);
   applyPatternsAndFoldGreedily(module,

mlir/test/lib/Transforms/TestDynamicPipeline.cpp

 using namespace mlir;
 
 namespace {
 
 class TestDynamicPipelinePass
     : public PassWrapper<TestDynamicPipelinePass, OperationPass<>> {
 public:
   void getDependentDialects(DialectRegistry &registry) const override {
-    OpPassManager pm(ModuleOp::getOperationName());
+    OpPassManager pm(ModuleOp::getOperationName(),
+                     OpPassManager::Nesting::Implicit);
     parsePassPipeline(pipeline, pm, llvm::errs());
     pm.getDependentDialects(registry);
   }
 
   TestDynamicPipelinePass(){};
   TestDynamicPipelinePass(const TestDynamicPipelinePass &) {}
 
   void runOnOperation() override {
void runOnOperation() override {
 
     auto opName = symbolOp.getName();
     if (!opNames.empty() && !llvm::is_contained(opNames, opName)) {
       llvm::errs() << "dynamic-pipeline skip op name: " << opName << "\n";
       return;
     }
     if (!pm) {
       pm = std::make_unique<OpPassManager>(
-          getOperation()->getName().getIdentifier());
+          getOperation()->getName().getIdentifier(),
+          OpPassManager::Nesting::Implicit);
       parsePassPipeline(pipeline, *pm, llvm::errs());
     }
 
     // Check that running on the parent operation always immediately fails.
     if (runOnParent) {
       if (getOperation()->getParentOp())
         if (!failed(runPipeline(*pm, getOperation()->getParentOp())))
           signalPassFailure();

mlir/unittests/Pass/PassManagerTest.cpp

TEST(PassManagerTest, InvalidPass) {
   // check it later.
   std::unique_ptr<Diagnostic> diagnostic;
   context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
     diagnostic.reset(new Diagnostic(std::move(diag)));
   });
 
   // Instantiate and run our pass.
   PassManager pm(&context);
-  pm.addPass(std::make_unique<InvalidPass>());
+  pm.nest("invalid_op").addPass(std::make_unique<InvalidPass>());
   LogicalResult result = pm.run(module.get());
   EXPECT_TRUE(failed(result));
   ASSERT_TRUE(diagnostic.get() != nullptr);
   EXPECT_EQ(
       diagnostic->str(),
       "'invalid_op' op trying to schedule a pass on an unregistered operation");
+
+  // Check that adding the pass at the top-level triggers a fatal error.
+  ASSERT_DEATH(pm.addPass(std::make_unique<InvalidPass>()), "");
 }
 
 } // end namespace