Add static pass registration

Add static pass registration and change mlir-opt to use it. Future work is needed to refactor the registration for PassManager usage.

Change build targets to alwayslink to enforce registration.

PiperOrigin-RevId: 220390178

15 files changed, 146 insertions, 9 deletions

diff --git a/mlir/lib/Analysis/MemRefBoundCheck.cpp b/mlir/lib/Analysis/MemRefBoundCheck.cpp
index 0725cea7086..a7f0ebf4936 100644
--- a/mlir/lib/Analysis/MemRefBoundCheck.cpp
+++ b/mlir/lib/Analysis/MemRefBoundCheck.cpp
@@ -45,10 +45,14 @@ struct MemRefBoundCheck : public FunctionPass, StmtWalker<MemRefBoundCheck> {
   PassResult runOnCFGFunction(CFGFunction *f) override { return success(); }
 
   void visitOperationStmt(OperationStmt *opStmt);
+
+  static char passID;
 };
 
 } // end anonymous namespace
 
+char MemRefBoundCheck::passID = 0;
+
 FunctionPass *mlir::createMemRefBoundCheckPass() {
   return new MemRefBoundCheck();
 }
@@ -164,3 +168,7 @@ void MemRefBoundCheck::visitOperationStmt(OperationStmt *opStmt) {
 PassResult MemRefBoundCheck::runOnMLFunction(MLFunction *f) {
   return walk(f), success();
 }
+
+static PassRegistration<MemRefBoundCheck>
+    memRefBoundCheck("memref-bound-check",
+                     "Check memref accesses in an MLFunction");
diff --git a/mlir/lib/Analysis/MemRefDependenceCheck.cpp b/mlir/lib/Analysis/MemRefDependenceCheck.cpp
index 3ca669c5c85..7a620c1a3a8 100644
--- a/mlir/lib/Analysis/MemRefDependenceCheck.cpp
+++ b/mlir/lib/Analysis/MemRefDependenceCheck.cpp
@@ -51,10 +51,13 @@ struct MemRefDependenceCheck : public FunctionPass,
       loadsAndStores.push_back(opStmt);
     }
   }
+  static char passID;
 };
 
 } // end anonymous namespace
 
+char MemRefDependenceCheck::passID = 0;
+
 FunctionPass *mlir::createMemRefDependenceCheckPass() {
   return new MemRefDependenceCheck();
 }
@@ -132,3 +135,7 @@ PassResult MemRefDependenceCheck::runOnMLFunction(MLFunction *f) {
   checkDependences(loadsAndStores);
   return success();
 }
+
+static PassRegistration<MemRefDependenceCheck>
+    pass("memref-dependence-check",
+         "Checks dependences between all pairs of memref accesses.");
diff --git a/mlir/lib/Analysis/Pass.cpp b/mlir/lib/Analysis/Pass.cpp
index 1249c18c07e..ea9da5b0e80 100644
--- a/mlir/lib/Analysis/Pass.cpp
+++ b/mlir/lib/Analysis/Pass.cpp
@@ -23,6 +23,9 @@
 #include "mlir/IR/CFGFunction.h"
 #include "mlir/IR/MLFunction.h"
 #include "mlir/IR/Module.h"
+#include "mlir/Support/PassNameParser.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/ManagedStatic.h"
 using namespace mlir;
 
 /// Out of line virtual method to ensure vtables and metadata are emitted to a
@@ -51,3 +54,37 @@ PassResult FunctionPass::runOnFunction(Function *fn) {
 
   return success();
 }
+
+// TODO: The pass registry and pass name parsing should be moved out.
+static llvm::ManagedStatic<llvm::DenseMap<const void *, PassInfo>> passRegistry;
+
+void mlir::registerPass(StringRef arg, StringRef description,
+                        const void *passID,
+                        const PassAllocatorFunction &function) {
+  bool inserted = passRegistry
+                      ->insert(std::make_pair(
+                          passID, PassInfo(arg, description, passID, function)))
+                      .second;
+  assert(inserted && "Pass registered multiple times");
+  (void)inserted;
+}
+
+PassNameParser::PassNameParser(llvm::cl::Option &opt)
+    : llvm::cl::parser<const PassInfo *>(opt) {
+  for (const auto &kv : *passRegistry) {
+    addLiteralOption(kv.second.getPassArgument(), &kv.second,
+                     kv.second.getPassDescription());
+  }
+}
+
+void PassNameParser::printOptionInfo(const llvm::cl::Option &O,
+                                     size_t GlobalWidth) const {
+  PassNameParser *TP = const_cast<PassNameParser *>(this);
+  llvm::array_pod_sort(TP->Values.begin(), TP->Values.end(),
+                       [](const PassNameParser::OptionInfo *VT1,
+                          const PassNameParser::OptionInfo *VT2) {
+                         return VT1->Name.compare(VT2->Name);
+                       });
+  using llvm::cl::parser;
+  parser<const PassInfo *>::printOptionInfo(O, GlobalWidth);
+}
diff --git a/mlir/lib/Transforms/CFGFunctionViewGraph.cpp b/mlir/lib/Transforms/CFGFunctionViewGraph.cpp
index a75d26c1fbc..810264cb35e 100644
--- a/mlir/lib/Transforms/CFGFunctionViewGraph.cpp
+++ b/mlir/lib/Transforms/CFGFunctionViewGraph.cpp
@@ -74,13 +74,16 @@ void mlir::CFGFunction::viewGraph() const {
 
 namespace {
 struct PrintCFGPass : public FunctionPass {
-  PrintCFGPass(llvm::raw_ostream &os, bool shortNames, const llvm::Twine &title)
+  PrintCFGPass(llvm::raw_ostream &os = llvm::errs(), bool shortNames = false,
+               const llvm::Twine &title = "")
       : os(os), shortNames(shortNames), title(title) {}
   PassResult runOnCFGFunction(CFGFunction *function) override {
     mlir::writeGraph(os, function, shortNames, title);
     return success();
   }
 
+  static char passID;
+
 private:
   llvm::raw_ostream &os;
   bool shortNames;
@@ -88,8 +91,13 @@ private:
 };
 } // namespace
 
+char PrintCFGPass::passID = 0;
+
 FunctionPass *mlir::createPrintCFGGraphPass(llvm::raw_ostream &os,
                                             bool shortNames,
                                             const llvm::Twine &title) {
   return new PrintCFGPass(os, shortNames, title);
 }
+
+static PassRegistration<PrintCFGPass> pass("print-cfg-graph",
+                                           "Print CFG graph per function");
diff --git a/mlir/lib/Transforms/Canonicalizer.cpp b/mlir/lib/Transforms/Canonicalizer.cpp
index f34118ce21a..3a62f132459 100644
--- a/mlir/lib/Transforms/Canonicalizer.cpp
+++ b/mlir/lib/Transforms/Canonicalizer.cpp
@@ -35,9 +35,13 @@ namespace {
 /// Canonicalize operations in functions.
 struct Canonicalizer : public FunctionPass {
   PassResult runOnFunction(Function *fn) override;
+
+  static char passID;
 };
 } // end anonymous namespace
 
+char Canonicalizer::passID = 0;
+
 PassResult Canonicalizer::runOnFunction(Function *fn) {
   auto *context = fn->getContext();
   OwningPatternList patterns;
@@ -54,3 +58,6 @@ PassResult Canonicalizer::runOnFunction(Function *fn) {
 
 /// Create a Canonicalizer pass.
 FunctionPass *mlir::createCanonicalizerPass() { return new Canonicalizer(); }
+
+static PassRegistration<Canonicalizer> pass("canonicalize",
+                                            "Canonicalize operations");
diff --git a/mlir/lib/Transforms/ComposeAffineMaps.cpp b/mlir/lib/Transforms/ComposeAffineMaps.cpp
index af4a5d11521..61e2e8f2e83 100644
--- a/mlir/lib/Transforms/ComposeAffineMaps.cpp
+++ b/mlir/lib/Transforms/ComposeAffineMaps.cpp
@@ -50,10 +50,14 @@ struct ComposeAffineMaps : public FunctionPass, StmtWalker<ComposeAffineMaps> {
   void visitOperationStmt(OperationStmt *stmt);
   PassResult runOnMLFunction(MLFunction *f) override;
   using StmtWalker<ComposeAffineMaps>::walk;
+
+  static char passID;
 };
 
 } // end anonymous namespace
 
+char ComposeAffineMaps::passID = 0;
+
 FunctionPass *mlir::createComposeAffineMapsPass() {
   return new ComposeAffineMaps();
 }
@@ -92,3 +96,6 @@ PassResult ComposeAffineMaps::runOnMLFunction(MLFunction *f) {
   }
   return success();
 }
+
+static PassRegistration<ComposeAffineMaps> pass("compose-affine-maps",
+                                                "Compose affine maps");
diff --git a/mlir/lib/Transforms/ConstantFold.cpp b/mlir/lib/Transforms/ConstantFold.cpp
index 411d1caae29..9005c2bbf48 100644
--- a/mlir/lib/Transforms/ConstantFold.cpp
+++ b/mlir/lib/Transforms/ConstantFold.cpp
@@ -40,9 +40,13 @@ struct ConstantFold : public FunctionPass, StmtWalker<ConstantFold> {
   void visitForStmt(ForStmt *stmt);
   PassResult runOnCFGFunction(CFGFunction *f) override;
   PassResult runOnMLFunction(MLFunction *f) override;
+
+  static char passID;
 };
 } // end anonymous namespace
 
+char ConstantFold::passID = 0;
+
 /// Attempt to fold the specified operation, updating the IR to match.  If
 /// constants are found, we keep track of them in the existingConstants list.
 ///
@@ -174,3 +178,6 @@ PassResult ConstantFold::runOnMLFunction(MLFunction *f) {
 
 /// Creates a constant folding pass.
 FunctionPass *mlir::createConstantFoldPass() { return new ConstantFold(); }
+
+static PassRegistration<ConstantFold>
+    pass("constant-fold", "Constant fold operations in functions");
diff --git a/mlir/lib/Transforms/ConvertToCFG.cpp b/mlir/lib/Transforms/ConvertToCFG.cpp
index 52687da65ba..b36717d272f 100644
--- a/mlir/lib/Transforms/ConvertToCFG.cpp
+++ b/mlir/lib/Transforms/ConvertToCFG.cpp
@@ -70,6 +70,8 @@ public:
 
   PassResult runOnModule(Module *m) override;
 
+  static char passID;
+
 private:
   // Generates CFG functions for all ML functions in the module.
   void convertMLFunctions();
@@ -90,6 +92,8 @@ private:
 };
 } // end anonymous namespace
 
+char ModuleConverter::passID = 0;
+
 // 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.
@@ -163,3 +167,7 @@ void ModuleConverter::removeMLFunctions() {
 /// Function references are appropriately patched to refer to the newly
 /// generated CFG functions.
 ModulePass *mlir::createConvertToCFGPass() { return new ModuleConverter(); }
+
+static PassRegistration<ModuleConverter>
+    pass("convert-to-cfg",
+         "Convert all ML functions in the module to CFG ones");
diff --git a/mlir/lib/Transforms/LoopFusion.cpp b/mlir/lib/Transforms/LoopFusion.cpp
index d9cdf9d919b..ae4647e143d 100644
--- a/mlir/lib/Transforms/LoopFusion.cpp
+++ b/mlir/lib/Transforms/LoopFusion.cpp
@@ -45,6 +45,7 @@ struct LoopFusion : public FunctionPass {
   LoopFusion() {}
 
   PassResult runOnMLFunction(MLFunction *f) override;
+  static char passID;
 };
 
 // LoopCollector walks the statements in an MLFunction and builds a map from
@@ -75,6 +76,8 @@ public:
 
 } // end anonymous namespace
 
+char LoopFusion::passID = 0;
+
 FunctionPass *mlir::createLoopFusionPass() { return new LoopFusion; }
 
 // TODO(andydavis) Remove the following test code when more general loop
@@ -242,3 +245,5 @@ PassResult LoopFusion::runOnMLFunction(MLFunction *f) {
 
   return success();
 }
+
+static PassRegistration<LoopFusion> pass("loop-fusion", "Fuse loop nests");
diff --git a/mlir/lib/Transforms/LoopTiling.cpp b/mlir/lib/Transforms/LoopTiling.cpp
index bd66e337609..3bff008942c 100644
--- a/mlir/lib/Transforms/LoopTiling.cpp
+++ b/mlir/lib/Transforms/LoopTiling.cpp
@@ -42,10 +42,14 @@ namespace {
 struct LoopTiling : public FunctionPass {
   PassResult runOnMLFunction(MLFunction *f) override;
   constexpr static unsigned kDefaultTileSize = 32;
+
+  static char passID;
 };
 
 } // end anonymous namespace
 
+char LoopTiling::passID = 0;
+
 /// Creates a pass to perform loop tiling on all suitable loop nests of an
 /// MLFunction.
 FunctionPass *mlir::createLoopTilingPass() { return new LoopTiling(); }
@@ -238,3 +242,5 @@ PassResult LoopTiling::runOnMLFunction(MLFunction *f) {
   }
   return success();
 }
+
+static PassRegistration<LoopTiling> pass("loop-tile", "Tile loop nests");
diff --git a/mlir/lib/Transforms/LoopUnroll.cpp b/mlir/lib/Transforms/LoopUnroll.cpp
index 15c7014dc42..ae09098f9d5 100644
--- a/mlir/lib/Transforms/LoopUnroll.cpp
+++ b/mlir/lib/Transforms/LoopUnroll.cpp
@@ -56,22 +56,20 @@ struct LoopUnroll : public FunctionPass {
   Optional<unsigned> unrollFactor;
   Optional<bool> unrollFull;
 
-  explicit LoopUnroll(Optional<unsigned> unrollFactor,
-                      Optional<bool> unrollFull)
+  explicit LoopUnroll(Optional<unsigned> unrollFactor = None,
+                      Optional<bool> unrollFull = None)
       : unrollFactor(unrollFactor), unrollFull(unrollFull) {}
 
   PassResult runOnMLFunction(MLFunction *f) override;
 
   /// Unroll this for stmt. Returns false if nothing was done.
   bool runOnForStmt(ForStmt *forStmt);
+
+  static char passID;
 };
 } // end anonymous namespace
 
-FunctionPass *mlir::createLoopUnrollPass(int unrollFactor, int unrollFull) {
-  return new LoopUnroll(unrollFactor == -1 ? None
-                                           : Optional<unsigned>(unrollFactor),
-                        unrollFull == -1 ? None : Optional<bool>(unrollFull));
-}
+char LoopUnroll::passID = 0;
 
 PassResult LoopUnroll::runOnMLFunction(MLFunction *f) {
   // Gathers all innermost loops through a post order pruned walk.
@@ -286,3 +284,11 @@ bool mlir::loopUnrollByFactor(ForStmt *forStmt, uint64_t unrollFactor) {
 
   return true;
 }
+
+FunctionPass *mlir::createLoopUnrollPass(int unrollFactor, int unrollFull) {
+  return new LoopUnroll(unrollFactor == -1 ? None
+                                           : Optional<unsigned>(unrollFactor),
+                        unrollFull == -1 ? None : Optional<bool>(unrollFull));
+}
+
+static PassRegistration<LoopUnroll> pass("loop-unroll", "Unroll loops");
diff --git a/mlir/lib/Transforms/LoopUnrollAndJam.cpp b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
index f437b44ae26..ce6e939fae8 100644
--- a/mlir/lib/Transforms/LoopUnrollAndJam.cpp
+++ b/mlir/lib/Transforms/LoopUnrollAndJam.cpp
@@ -70,14 +70,18 @@ struct LoopUnrollAndJam : public FunctionPass {
   Optional<unsigned> unrollJamFactor;
   static const unsigned kDefaultUnrollJamFactor = 4;
 
-  explicit LoopUnrollAndJam(Optional<unsigned> unrollJamFactor)
+  explicit LoopUnrollAndJam(Optional<unsigned> unrollJamFactor = None)
       : unrollJamFactor(unrollJamFactor) {}
 
   PassResult runOnMLFunction(MLFunction *f) override;
   bool runOnForStmt(ForStmt *forStmt);
+
+  static char passID;
 };
 } // end anonymous namespace
 
+char LoopUnrollAndJam::passID = 0;
+
 FunctionPass *mlir::createLoopUnrollAndJamPass(int unrollJamFactor) {
   return new LoopUnrollAndJam(
       unrollJamFactor == -1 ? None : Optional<unsigned>(unrollJamFactor));
@@ -239,3 +243,6 @@ bool mlir::loopUnrollJamByFactor(ForStmt *forStmt, uint64_t unrollJamFactor) {
 
   return true;
 }
+
+static PassRegistration<LoopUnrollAndJam> pass("loop-unroll-jam",
+                                               "Unroll and jam loops");
diff --git a/mlir/lib/Transforms/PipelineDataTransfer.cpp b/mlir/lib/Transforms/PipelineDataTransfer.cpp
index c59e007e543..52052e09d7b 100644
--- a/mlir/lib/Transforms/PipelineDataTransfer.cpp
+++ b/mlir/lib/Transforms/PipelineDataTransfer.cpp
@@ -47,10 +47,14 @@ struct PipelineDataTransfer : public FunctionPass,
   // Collect all 'for' statements.
   void visitForStmt(ForStmt *forStmt) { forStmts.push_back(forStmt); }
   std::vector<ForStmt *> forStmts;
+
+  static char passID;
 };
 
 } // end anonymous namespace
 
+char PipelineDataTransfer::passID = 0;
+
 /// Creates a pass to pipeline explicit movement of data across levels of the
 /// memory hierarchy.
 FunctionPass *mlir::createPipelineDataTransferPass() {
@@ -306,3 +310,8 @@ PassResult PipelineDataTransfer::runOnForStmt(ForStmt *forStmt) {
 
   return success();
 }
+
+static PassRegistration<PipelineDataTransfer> pass(
+    "pipeline-data-transfer",
+    "Pipeline non-blocking data transfers between explicitly managed levels of "
+    "the memory hierarchy");
diff --git a/mlir/lib/Transforms/SimplifyAffineExpr.cpp b/mlir/lib/Transforms/SimplifyAffineExpr.cpp
index a412a83f66c..92d585f31bc 100644
--- a/mlir/lib/Transforms/SimplifyAffineExpr.cpp
+++ b/mlir/lib/Transforms/SimplifyAffineExpr.cpp
@@ -47,10 +47,14 @@ struct SimplifyAffineStructures : public FunctionPass,
 
   void visitIfStmt(IfStmt *ifStmt);
   void visitOperationStmt(OperationStmt *opStmt);
+
+  static char passID;
 };
 
 } // end anonymous namespace
 
+char SimplifyAffineStructures::passID = 0;
+
 FunctionPass *mlir::createSimplifyAffineStructuresPass() {
   return new SimplifyAffineStructures();
 }
@@ -83,3 +87,6 @@ PassResult SimplifyAffineStructures::runOnMLFunction(MLFunction *f) {
   walk(f);
   return success();
 }
+
+static PassRegistration<SimplifyAffineStructures>
+    pass("simplify-affine-structures", "Simplify affine expressions");
diff --git a/mlir/lib/Transforms/Vectorize.cpp b/mlir/lib/Transforms/Vectorize.cpp
index fa97b7025d4..63969af451f 100644
--- a/mlir/lib/Transforms/Vectorize.cpp
+++ b/mlir/lib/Transforms/Vectorize.cpp
@@ -199,10 +199,14 @@ struct Vectorize : public FunctionPass {
 
   // Thread-safe RAII contexts local to pass, BumpPtrAllocator freed on exit.
   MLFunctionMatcherContext MLContext;
+
+  static char passID;
 };
 
 } // end anonymous namespace
 
+char Vectorize::passID = 0;
+
 /////// TODO(ntv): Hoist to a VectorizationStrategy.cpp when appropriate. //////
 namespace {
 
@@ -669,3 +673,7 @@ PassResult Vectorize::runOnMLFunction(MLFunction *f) {
 }
 
 FunctionPass *mlir::createVectorizePass() { return new Vectorize(); }
+
+static PassRegistration<Vectorize>
+    pass("vectorize",
+         "Vectorize to a target independent n-D vector abstraction");