Add PassResult and have passes return PassResult to indicate failure/success.

For FunctionPass's for passes that want to stop upon error encountered.

PiperOrigin-RevId: 213058651

mlir/include/mlir/Transforms/Pass.h

@@ -18,39 +18,61 @@
 #ifndef MLIR_PASS_H
 #define MLIR_PASS_H
 
+#include "llvm/Support/Compiler.h"
+
 namespace mlir {
 class CFGFunction;
 class MLFunction;
 class Module;
 
+// Values that can be used by to signal success/failure. This can be implicitly
+// converted to/from boolean values, with false representing success and true
+// failure.
+struct LLVM_NODISCARD PassResult {
+  enum ResultEnum { Success, Failure } value;
+  PassResult(ResultEnum v) : value(v) {}
+  operator bool() const { return value == Failure; }
+};
+
+static PassResult success() { return PassResult::Success; }
+static PassResult failure() { return PassResult::Failure; }
+
 class Pass {
 public:
   virtual ~Pass() = default;
-  virtual void runOnModule(Module *m) = 0;
+  virtual PassResult runOnModule(Module *m) = 0;
 };
 
 class ModulePass : public Pass {
 public:
-  virtual void runOnModule(Module *m) override = 0;
+  virtual PassResult runOnModule(Module *m) override = 0;
 };
 
 class FunctionPass : public Pass {
 public:
-  virtual void runOnCFGFunction(CFGFunction *f) = 0;
-  virtual void runOnMLFunction(MLFunction *f) = 0;
-  virtual void runOnModule(Module *m) override;
+  virtual PassResult runOnCFGFunction(CFGFunction *f) = 0;
+  virtual PassResult runOnMLFunction(MLFunction *f) = 0;
+
+  // Iterates over all functions in a module, halting upon failure.
+  virtual PassResult runOnModule(Module *m) override;
 };
 
 class CFGFunctionPass : public FunctionPass {
 public:
-  virtual void runOnMLFunction(MLFunction *f) override {}
-  virtual void runOnCFGFunction(CFGFunction *f) override = 0;
+  virtual PassResult runOnMLFunction(MLFunction *f) override {
+    // Skip over MLFunction.
+    return success();
+  }
+  virtual PassResult runOnCFGFunction(CFGFunction *f) override = 0;
 };
 
 class MLFunctionPass : public FunctionPass {
 public:
-  virtual void runOnCFGFunction(CFGFunction *f) override {}
-  virtual void runOnMLFunction(MLFunction *f) override = 0;
+  virtual PassResult runOnCFGFunction(CFGFunction *f) override {
+    // Skip over CFGFunction.
+    return success();
+  }
+  virtual PassResult runOnMLFunction(MLFunction *f) override = 0;
 };
 
 } // end namespace mlir

mlir/lib/Transforms/CFGFunctionViewGraph.cpp

@@ -101,8 +101,9 @@ namespace {
 struct PrintCFGPass : public CFGFunctionPass {
   PrintCFGPass(llvm::raw_ostream &os, bool shortNames, const llvm::Twine &title)
       : os(os), shortNames(shortNames), title(title) {}
-  void runOnCFGFunction(CFGFunction *function) override {
+  PassResult runOnCFGFunction(CFGFunction *function) override {
     mlir::writeGraph(os, function, shortNames, title);
+    return success();
   }
 
 private:

mlir/lib/Transforms/ConvertToCFG.cpp

@@ -68,7 +68,7 @@ class ModuleConverter : public ModulePass {
 public:
   explicit ModuleConverter() {}
 
-  void runOnModule(Module *m) override;
+  PassResult runOnModule(Module *m) override;
 
 private:
   // Generates CFG functions for all ML functions in the module.
@@ -93,10 +93,11 @@ private:
 // Iterates over all functions in the module generating CFG functions
 // equivalent to ML functions and replacing references to ML functions
 // with references to the generated ML functions.
-void ModuleConverter::runOnModule(Module *m) {
+PassResult ModuleConverter::runOnModule(Module *m) {
   module = m;
   convertMLFunctions();
   replaceReferences();
+  return success();
 }
 
 void ModuleConverter::convertMLFunctions() {

mlir/lib/Transforms/LoopUnroll.cpp

@@ -59,7 +59,7 @@ struct LoopUnroll : public MLFunctionPass {
                       Optional<bool> unrollFull)
       : unrollFactor(unrollFactor), unrollFull(unrollFull) {}
 
-  void runOnMLFunction(MLFunction *f) override;
+  PassResult runOnMLFunction(MLFunction *f) override;
   /// Unroll this for stmt. Returns false if nothing was done.
   bool runOnForStmt(ForStmt *forStmt);
   bool loopUnrollFull(ForStmt *forStmt);
@@ -73,7 +73,7 @@ MLFunctionPass *mlir::createLoopUnrollPass(int unrollFactor, int unrollFull) {
                         unrollFull == -1 ? None : Optional<bool>(unrollFull));
 }
 
-void LoopUnroll::runOnMLFunction(MLFunction *f) {
+PassResult LoopUnroll::runOnMLFunction(MLFunction *f) {
   // Gathers all innermost loops through a post order pruned walk.
   class InnermostLoopGatherer : public StmtWalker<InnermostLoopGatherer, bool> {
   public:
@@ -141,7 +141,7 @@ void LoopUnroll::runOnMLFunction(MLFunction *f) {
     auto &loops = slg.loops;
     for (auto *forStmt : loops)
       loopUnrollFull(forStmt);
-    return;
+    return success();
   }
 
   InnermostLoopGatherer ilg;
@@ -149,6 +149,7 @@ void LoopUnroll::runOnMLFunction(MLFunction *f) {
   auto &loops = ilg.loops;
   for (auto *forStmt : loops)
     runOnForStmt(forStmt);
+  return success();
 }
 
 /// Unroll a for stmt. Default unroll factor is 4.

mlir/lib/Transforms/LoopUnrollAndJam.cpp

@@ -72,7 +72,7 @@ struct LoopUnrollAndJam : public MLFunctionPass {
   explicit LoopUnrollAndJam(Optional<unsigned> unrollJamFactor)
       : unrollJamFactor(unrollJamFactor) {}
 
-  void runOnMLFunction(MLFunction *f) override;
+  PassResult runOnMLFunction(MLFunction *f) override;
   bool runOnForStmt(ForStmt *forStmt);
   bool loopUnrollJamByFactor(ForStmt *forStmt, uint64_t unrollJamFactor);
 };
@@ -83,15 +83,16 @@ MLFunctionPass *mlir::createLoopUnrollAndJamPass(int unrollJamFactor) {
       unrollJamFactor == -1 ? None : Optional<unsigned>(unrollJamFactor));
 }
 
-void LoopUnrollAndJam::runOnMLFunction(MLFunction *f) {
+PassResult LoopUnrollAndJam::runOnMLFunction(MLFunction *f) {
   // Currently, just the outermost loop from the first loop nest is
   // unroll-and-jammed by this pass. However, runOnForStmt can be called on any
   // for Stmt.
   if (!isa<ForStmt>(f->begin()))
-    return;
+    return success();
 
   auto *forStmt = cast<ForStmt>(f->begin());
   runOnForStmt(forStmt);
+  return success();
 }
 
 /// Unroll and jam a 'for' stmt. Default unroll jam factor is

mlir/lib/Transforms/Pass.cpp

@@ -27,12 +27,15 @@
 using namespace mlir;
 
 /// Function passes walk a module and look at each function with their
-/// corresponding hooks.
-void FunctionPass::runOnModule(Module *m) {
+/// corresponding hooks and terminates upon error encountered.
+PassResult FunctionPass::runOnModule(Module *m) {
   for (auto &fn : *m) {
     if (auto *mlFunc = dyn_cast<MLFunction>(&fn))
-      runOnMLFunction(mlFunc);
+      if (runOnMLFunction(mlFunc))
+        return failure();
     if (auto *cfgFunc = dyn_cast<CFGFunction>(&fn))
-      runOnCFGFunction(cfgFunc);
+      if (runOnCFGFunction(cfgFunc))
+        return failure();
   }
+  return success();
 }

mlir/lib/Transforms/SimplifyAffineExpr.cpp

@@ -39,10 +39,10 @@ namespace {
 struct SimplifyAffineExpr : public FunctionPass {
   explicit SimplifyAffineExpr() {}
 
-  void runOnMLFunction(MLFunction *f);
+  PassResult runOnMLFunction(MLFunction *f);
   // Does nothing on CFG functions for now. No reusable walkers/visitors exist
   // for this yet? TODO(someone).
-  void runOnCFGFunction(CFGFunction *f) {}
+  PassResult runOnCFGFunction(CFGFunction *f) { return success(); }
 };
 
 } // end anonymous namespace
@@ -55,7 +55,7 @@ AffineMap *MutableAffineMap::getAffineMap() {
   return AffineMap::get(numDims, numSymbols, results, rangeSizes, context);
 }
 
-void SimplifyAffineExpr::runOnMLFunction(MLFunction *f) {
+PassResult SimplifyAffineExpr::runOnMLFunction(MLFunction *f) {
   struct MapSimplifier : public StmtWalker<MapSimplifier> {
     MLIRContext *context;
     MapSimplifier(MLIRContext *context) : context(context) {}
@@ -74,4 +74,5 @@ void SimplifyAffineExpr::runOnMLFunction(MLFunction *f) {
 
   MapSimplifier v(f->getContext());
   v.walkPostOrder(f);
+  return success();
 }

mlir/tools/mlir-opt/mlir-opt.cpp

@@ -186,8 +186,10 @@ static OptResult performActions(SourceMgr &sourceMgr, MLIRContext *context) {
       break;
     }
 
-    pass->runOnModule(module.get());
+    PassResult result = pass->runOnModule(module.get());
     delete pass;
+    if (result)
+      return OptFailure;
 
     // Verify that the result of the pass is still valid.
     module->verify();