Simplify FunctionPass to eliminate the CFGFunctionPass/MLFunctionPass

distinction.  FunctionPasses can now choose to get called on all functions, or
have the driver split CFG/ML Functions up for them.  NFC.

PiperOrigin-RevId: 218775885

12 files changed, 72 insertions, 67 deletions

diff --git a/mlir/include/mlir/Transforms/CFGFunctionViewGraph.h b/mlir/include/mlir/Transforms/CFGFunctionViewGraph.h
index 728ef0e7e8f..1a5895e423f 100644
--- a/mlir/include/mlir/Transforms/CFGFunctionViewGraph.h
+++ b/mlir/include/mlir/Transforms/CFGFunctionViewGraph.h
@@ -41,9 +41,9 @@ llvm::raw_ostream &writeGraph(llvm::raw_ostream &os,
                               bool shortNames = false, const Twine &title = "");
 
 /// Creates a pass to print CFG graphs.
-CFGFunctionPass *createPrintCFGGraphPass(llvm::raw_ostream &os = llvm::errs(),
-                                         bool shortNames = false,
-                                         const llvm::Twine &title = "");
+FunctionPass *createPrintCFGGraphPass(llvm::raw_ostream &os = llvm::errs(),
+                                      bool shortNames = false,
+                                      const llvm::Twine &title = "");
 
 } // end namespace mlir
 
diff --git a/mlir/include/mlir/Transforms/Pass.h b/mlir/include/mlir/Transforms/Pass.h
index f18bc65be69..cd7b7027907 100644
--- a/mlir/include/mlir/Transforms/Pass.h
+++ b/mlir/include/mlir/Transforms/Pass.h
@@ -21,6 +21,7 @@
 #include "llvm/Support/Compiler.h"
 
 namespace mlir {
+class Function;
 class CFGFunction;
 class MLFunction;
 class Module;
@@ -41,38 +42,43 @@ public:
 
   static PassResult success() { return PassResult::Success; }
   static PassResult failure() { return PassResult::Failure; }
+
+private:
+  /// Out of line virtual method to ensure vtables and metadata are emitted to a
+  /// single .o file.
+  virtual void anchor();
 };
 
 class ModulePass : public Pass {
 public:
   virtual PassResult runOnModule(Module *m) override = 0;
+
+private:
+  /// Out of line virtual method to ensure vtables and metadata are emitted to a
+  /// single .o file.
+  virtual void anchor();
 };
 
+/// FunctionPass's are run on every function in a module, and multiple functions
+/// may be optimized concurrently by different instances of the function pass.
+/// By subclassing this, your pass promises only to look at the function psased
+/// in to it, it isn't allowed to inspect or modify other functions in the
+/// module.
 class FunctionPass : public Pass {
 public:
-  virtual PassResult runOnCFGFunction(CFGFunction *f) = 0;
-  virtual PassResult runOnMLFunction(MLFunction *f) = 0;
+  /// Implement this function to be run on every function in the module.  If you
+  /// do not implement this, the default implementation will dispatch to
+  /// runOnCFGFunction or runOnMLFunction.
+  virtual PassResult runOnFunction(Function *fn);
 
-  // Iterates over all functions in a module, halting upon failure.
-  virtual PassResult runOnModule(Module *m) override;
-};
+  /// Implement this function if you want to see CFGFunction's specifically.
+  virtual PassResult runOnCFGFunction(CFGFunction *fn) { return success(); }
 
-class CFGFunctionPass : public FunctionPass {
-public:
-  virtual PassResult runOnMLFunction(MLFunction *f) override {
-    // Skip over MLFunction.
-    return success();
-  }
-  virtual PassResult runOnCFGFunction(CFGFunction *f) override = 0;
-};
+  /// Implement this function if you want to see MLFunction's specifically.
+  virtual PassResult runOnMLFunction(MLFunction *fn) { return success(); }
 
-class MLFunctionPass : public FunctionPass {
-public:
-  virtual PassResult runOnCFGFunction(CFGFunction *f) override {
-    // Skip over CFGFunction.
-    return success();
-  }
-  virtual PassResult runOnMLFunction(MLFunction *f) override = 0;
+  // Iterates over all functions in a module, halting upon failure.
+  virtual PassResult runOnModule(Module *m) override;
 };
 
 } // end namespace mlir
diff --git a/mlir/include/mlir/Transforms/Passes.h b/mlir/include/mlir/Transforms/Passes.h
index 1c34e36a57e..31c03c3cccb 100644
--- a/mlir/include/mlir/Transforms/Passes.h
+++ b/mlir/include/mlir/Transforms/Passes.h
@@ -28,7 +28,6 @@
 namespace mlir {
 
 class FunctionPass;
-class MLFunctionPass;
 class ModulePass;
 
 /// Creates a constant folding pass.
@@ -39,27 +38,26 @@ FunctionPass *createCanonicalizerPass();
 
 /// Creates a pass to vectorize loops, operations and data types using a
 /// target-independent, n-D virtual vector abstraction.
-MLFunctionPass *createVectorizePass();
+FunctionPass *createVectorizePass();
 
 /// Creates a loop unrolling pass. Default option or command-line options take
 /// effect if -1 is passed as parameter.
-MLFunctionPass *createLoopUnrollPass(int unrollFactor = -1,
-                                     int unrollFull = -1);
+FunctionPass *createLoopUnrollPass(int unrollFactor = -1, int unrollFull = -1);
 
 /// Creates a loop unroll jam pass to unroll jam by the specified factor. A
 /// factor of -1 lets the pass use the default factor or the one on the command
 /// line if provided.
-MLFunctionPass *createLoopUnrollAndJamPass(int unrollJamFactor = -1);
+FunctionPass *createLoopUnrollAndJamPass(int unrollJamFactor = -1);
 
 /// Creates an simplification pass for affine structures.
 FunctionPass *createSimplifyAffineStructuresPass();
 
 /// Creates a pass to pipeline explicit movement of data across levels of the
 /// memory hierarchy.
-MLFunctionPass *createPipelineDataTransferPass();
+FunctionPass *createPipelineDataTransferPass();
 
 /// Creates a pass which composes all affine maps applied to loads and stores.
-MLFunctionPass *createComposeAffineMapsPass();
+FunctionPass *createComposeAffineMapsPass();
 
 /// Replaces all ML functions in the module with equivalent CFG functions.
 /// Function references are appropriately patched to refer to the newly
diff --git a/mlir/lib/Transforms/CFGFunctionViewGraph.cpp b/mlir/lib/Transforms/CFGFunctionViewGraph.cpp
index dea1fc90af2..a75d26c1fbc 100644
--- a/mlir/lib/Transforms/CFGFunctionViewGraph.cpp
+++ b/mlir/lib/Transforms/CFGFunctionViewGraph.cpp
@@ -73,7 +73,7 @@ void mlir::CFGFunction::viewGraph() const {
 }
 
 namespace {
-struct PrintCFGPass : public CFGFunctionPass {
+struct PrintCFGPass : public FunctionPass {
   PrintCFGPass(llvm::raw_ostream &os, bool shortNames, const llvm::Twine &title)
       : os(os), shortNames(shortNames), title(title) {}
   PassResult runOnCFGFunction(CFGFunction *function) override {
@@ -88,8 +88,8 @@ private:
 };
 } // namespace
 
-CFGFunctionPass *mlir::createPrintCFGGraphPass(llvm::raw_ostream &os,
-                                               bool shortNames,
-                                               const llvm::Twine &title) {
+FunctionPass *mlir::createPrintCFGGraphPass(llvm::raw_ostream &os,
+                                            bool shortNames,
+                                            const llvm::Twine &title) {
   return new PrintCFGPass(os, shortNames, title);
 }
diff --git a/mlir/lib/Transforms/Canonicalizer.cpp b/mlir/lib/Transforms/Canonicalizer.cpp
index b7bfd785101..4f308171467 100644
--- a/mlir/lib/Transforms/Canonicalizer.cpp
+++ b/mlir/lib/Transforms/Canonicalizer.cpp
@@ -34,21 +34,10 @@ namespace {
 
 /// Canonicalize operations in functions.
 struct Canonicalizer : public FunctionPass {
-  PassResult runOnCFGFunction(CFGFunction *f) override;
-  PassResult runOnMLFunction(MLFunction *f) override;
-  PassResult runOnFunction(Function *fn);
+  PassResult runOnFunction(Function *fn) override;
 };
 } // end anonymous namespace
 
-
-PassResult Canonicalizer::runOnCFGFunction(CFGFunction *fn) {
-  return runOnFunction(fn);
-}
-
-PassResult Canonicalizer::runOnMLFunction(MLFunction *fn) {
-  return runOnFunction(fn);
-}
-
 PassResult Canonicalizer::runOnFunction(Function *fn) {
   auto *context = fn->getContext();
   OwningPatternList patterns;
diff --git a/mlir/lib/Transforms/ComposeAffineMaps.cpp b/mlir/lib/Transforms/ComposeAffineMaps.cpp
index 8f7919d0760..ced9cbcb86a 100644
--- a/mlir/lib/Transforms/ComposeAffineMaps.cpp
+++ b/mlir/lib/Transforms/ComposeAffineMaps.cpp
@@ -41,21 +41,20 @@ namespace {
 // also AffineApplyOps. After forward subtituting its results, AffineApplyOps
 // with no remaining uses are collected and erased after the walk.
 // TODO(andydavis) Remove this when Chris adds instruction combiner pass.
-struct ComposeAffineMaps : public MLFunctionPass,
-                           StmtWalker<ComposeAffineMaps> {
+struct ComposeAffineMaps : public FunctionPass, StmtWalker<ComposeAffineMaps> {
   std::vector<OperationStmt *> affineApplyOpsToErase;
 
   explicit ComposeAffineMaps() {}
   using StmtListType = llvm::iplist<Statement>;
   void walk(StmtListType::iterator Start, StmtListType::iterator End);
   void visitOperationStmt(OperationStmt *stmt);
-  PassResult runOnMLFunction(MLFunction *f);
+  PassResult runOnMLFunction(MLFunction *f) override;
   using StmtWalker<ComposeAffineMaps>::walk;
 };
 
 } // end anonymous namespace
 
-MLFunctionPass *mlir::createComposeAffineMapsPass() {
+FunctionPass *mlir::createComposeAffineMapsPass() {
   return new ComposeAffineMaps();
 }
 
diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp
index 3f4b673f273..b8a5cf9d0a1 100644
--- a/mlir/lib/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Transforms/LoopUnroll.cpp
@@ -52,7 +52,7 @@ namespace {
 /// full unroll threshold was specified, in which case, fully unrolls all loops
 /// with trip count less than the specified threshold. The latter is for testing
 /// purposes, especially for testing outer loop unrolling.
-struct LoopUnroll : public MLFunctionPass {
+struct LoopUnroll : public FunctionPass {
   Optional<unsigned> unrollFactor;
   Optional<bool> unrollFull;
 
@@ -61,12 +61,13 @@ struct LoopUnroll : public MLFunctionPass {
       : unrollFactor(unrollFactor), unrollFull(unrollFull) {}
 
   PassResult runOnMLFunction(MLFunction *f) override;
+
   /// Unroll this for stmt. Returns false if nothing was done.
   bool runOnForStmt(ForStmt *forStmt);
 };
 } // end anonymous namespace
 
-MLFunctionPass *mlir::createLoopUnrollPass(int unrollFactor, int unrollFull) {
+FunctionPass *mlir::createLoopUnrollPass(int unrollFactor, int unrollFull) {
   return new LoopUnroll(unrollFactor == -1 ? None
                                            : Optional<unsigned>(unrollFactor),
                         unrollFull == -1 ? None : Optional<bool>(unrollFull));
diff --git a/mlir/lib/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
index a4a69ecdc5c..b05063232ed 100644
--- a/mlir/lib/Transforms/LoopUnrollAndJam.cpp
+++ b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
@@ -66,7 +66,7 @@ static llvm::cl::opt<unsigned>
 namespace {
 /// Loop unroll jam pass. Currently, this just unroll jams the first
 /// outer loop in an MLFunction.
-struct LoopUnrollAndJam : public MLFunctionPass {
+struct LoopUnrollAndJam : public FunctionPass {
   Optional<unsigned> unrollJamFactor;
   static const unsigned kDefaultUnrollJamFactor = 4;
 
@@ -78,7 +78,7 @@ struct LoopUnrollAndJam : public MLFunctionPass {
 };
 } // end anonymous namespace
 
-MLFunctionPass *mlir::createLoopUnrollAndJamPass(int unrollJamFactor) {
+FunctionPass *mlir::createLoopUnrollAndJamPass(int unrollJamFactor) {
   return new LoopUnrollAndJam(
       unrollJamFactor == -1 ? None : Optional<unsigned>(unrollJamFactor));
 }
diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp
index 91f3b845a98..d96d65b5fb7 100644
--- a/mlir/lib/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Transforms/PipelineDataTransfer.cpp
@@ -39,7 +39,7 @@ using namespace mlir;
 
 namespace {
 
-struct PipelineDataTransfer : public MLFunctionPass,
+struct PipelineDataTransfer : public FunctionPass,
                               StmtWalker<PipelineDataTransfer> {
   PassResult runOnMLFunction(MLFunction *f) override;
   PassResult runOnForStmt(ForStmt *forStmt);
@@ -53,7 +53,7 @@ struct PipelineDataTransfer : public MLFunctionPass,
 
 /// Creates a pass to pipeline explicit movement of data across levels of the
 /// memory hierarchy.
-MLFunctionPass *mlir::createPipelineDataTransferPass() {
+FunctionPass *mlir::createPipelineDataTransferPass() {
   return new PipelineDataTransfer();
 }
 
diff --git a/mlir/lib/Transforms/SimplifyAffineExpr.cpp b/mlir/lib/Transforms/SimplifyAffineExpr.cpp
index 0a36387caaa..f4a23e6b153 100644
--- a/mlir/lib/Transforms/SimplifyAffineExpr.cpp
+++ b/mlir/lib/Transforms/SimplifyAffineExpr.cpp
@@ -41,10 +41,10 @@ struct SimplifyAffineStructures : public FunctionPass,
                                   StmtWalker<SimplifyAffineStructures> {
   explicit SimplifyAffineStructures() {}
 
-  PassResult runOnMLFunction(MLFunction *f);
+  PassResult runOnMLFunction(MLFunction *f) override;
   // Does nothing on CFG functions for now. No reusable walkers/visitors exist
   // for this yet? TODO(someone).
-  PassResult runOnCFGFunction(CFGFunction *f) { return success(); }
+  PassResult runOnCFGFunction(CFGFunction *f) override { return success(); }
 
   void visitIfStmt(IfStmt *ifStmt);
   void visitOperationStmt(OperationStmt *opStmt);
diff --git a/mlir/lib/Transforms/Utils/Pass.cpp b/mlir/lib/Transforms/Utils/Pass.cpp
index 8b1110798bd..e4edee31900 100644
--- a/mlir/lib/Transforms/Utils/Pass.cpp
+++ b/mlir/lib/Transforms/Utils/Pass.cpp
@@ -23,19 +23,31 @@
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/MLFunction.h"
 #include "mlir/IR/Module.h"
-
 using namespace mlir;
 
+/// Out of line virtual method to ensure vtables and metadata are emitted to a
+/// single .o file.
+void Pass::anchor() {}
+
+/// Out of line virtual method to ensure vtables and metadata are emitted to a
+/// single .o file.
+void ModulePass::anchor() {}
+
 /// Function passes walk a module and look at each function with their
 /// corresponding hooks and terminates upon error encountered.
 PassResult FunctionPass::runOnModule(Module *m) {
   for (auto &fn : *m) {
-    if (auto *mlFunc = dyn_cast<MLFunction>(&fn))
-      if (runOnMLFunction(mlFunc))
-        return failure();
-    if (auto *cfgFunc = dyn_cast<CFGFunction>(&fn))
-      if (runOnCFGFunction(cfgFunc))
-        return failure();
+    if (runOnFunction(&fn))
+      return failure();
   }
   return success();
 }
+
+PassResult FunctionPass::runOnFunction(Function *fn) {
+  if (auto *mlFunc = dyn_cast<MLFunction>(fn))
+    return runOnMLFunction(mlFunc);
+  if (auto *cfgFunc = dyn_cast<CFGFunction>(fn))
+    return runOnCFGFunction(cfgFunc);
+
+  return success();
+}
diff --git a/mlir/lib/Transforms/Vectorize.cpp b/mlir/lib/Transforms/Vectorize.cpp
index 37c4f3abb59..f38fb8a117b 100644
--- a/mlir/lib/Transforms/Vectorize.cpp
+++ b/mlir/lib/Transforms/Vectorize.cpp
@@ -39,7 +39,7 @@ static cl::list<unsigned> clVirtualVectorSize(
 
 namespace {
 
-struct Vectorize : public MLFunctionPass {
+struct Vectorize : public FunctionPass {
   PassResult runOnMLFunction(MLFunction *f) override;
 
   // Thread-safe RAII contexts local to pass, BumpPtrAllocator freed on exit.
@@ -73,4 +73,4 @@ PassResult Vectorize::runOnMLFunction(MLFunction *f) {
   return PassResult::Success;
 }
 
-MLFunctionPass *mlir::createVectorizePass() { return new Vectorize(); }
+FunctionPass *mlir::createVectorizePass() { return new Vectorize(); }