[mlir] Resolve TODO and use the pass argument instead of the TypeID for registration

[rriddle]

Summary

This removes the need to construct a pass during registration (to get the TypeID), and it also simplifies various pieces of code that interact with the pass registry.

[mehdi_amini]

mlir/lib/Pass/PassRegistry.cpp

113

If there is already an entry, can we check that the TypeID is matching?

[rriddle]

update

[mehdi_amini]

mlir/lib/Pass/PassRegistry.cpp

120

This isn't a sensitive code path, I'd make it a report_fatal_error and not guard it with ifndef NDEBUG

mlir/include/mlir/Pass/Pass.h

 class Pass {
 public:
   virtual ~Pass() = default;
 
   /// Returns the unique identifier that corresponds to this pass.
   TypeID getTypeID() const { return passID; }
 
   /// Returns the pass info for the specified pass class or null if unknown.
-  static const PassInfo *lookupPassInfo(TypeID passID);
-  template <typename PassT> static const PassInfo *lookupPassInfo() {
-    return lookupPassInfo(TypeID::get<PassT>());
-  }
+  static const PassInfo *lookupPassInfo(StringRef passArg);
 
-  /// Returns the pass info for this pass.
-  const PassInfo *lookupPassInfo() const { return lookupPassInfo(getTypeID()); }
+  /// Returns the pass info for this pass, or null if unknown.
+  const PassInfo *lookupPassInfo() const {
+    return lookupPassInfo(getArgument());
+  }
 
   /// Returns the derived pass name.
   virtual StringRef getName() const = 0;
 
   /// Register dependent dialects for the current pass.
   /// A pass is expected to register the dialects it will create entities for
   /// (Operations, Types, Attributes), other than dialect that exists in the
   /// input. For example, a pass that converts from Linalg to Affine would
   /// register the Affine dialect but does not need to register Linalg.
   virtual void getDependentDialects(DialectRegistry &registry) const {}
 
   /// Returns the command line argument used when registering this pass. Return
   /// an empty string if one does not exist.
-  virtual StringRef getArgument() const {
-    if (const PassInfo *passInfo = lookupPassInfo())
-      return passInfo->getPassArgument();
-    return "";
-  }
+  virtual StringRef getArgument() const { return ""; }
 
   /// Returns the name of the operation that this pass operates on, or None if
   /// this is a generic OperationPass.
   Optional<StringRef> getOpName() const { return opName; }
 
   //===--------------------------------------------------------------------===//
   // Options
   //===--------------------------------------------------------------------===//

mlir/include/mlir/Pass/PassRegistry.h

PassPipelineInfo(
       : PassRegistryEntry(arg, description, builder, optHandler) {}
 };
 
 /// A structure to represent the information for a derived pass class.
 class PassInfo : public PassRegistryEntry {
 public:
   /// PassInfo constructor should not be invoked directly, instead use
   /// PassRegistration or registerPass.
-  PassInfo(StringRef arg, StringRef description, TypeID passID,
+  PassInfo(StringRef arg, StringRef description,
            const PassAllocatorFunction &allocator);
 };
 
 //===----------------------------------------------------------------------===//
 // PassRegistration
 //===----------------------------------------------------------------------===//
 
 /// Register a specific dialect pipeline registry function with the system,

mlir/lib/Pass/PassRegistry.cpp

 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/MemoryBuffer.h"
 #include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 using namespace detail;
 
 /// Static mapping of all of the registered passes.
-static llvm::ManagedStatic<DenseMap<TypeID, PassInfo>> passRegistry;
+static llvm::ManagedStatic<llvm::StringMap<PassInfo>> passRegistry;
+
+/// A mapping of the above pass registry entries to the corresponding TypeID
+/// of the pass that they generate.
+static llvm::ManagedStatic<llvm::StringMap<TypeID>> passRegistryTypeIDs;
 
 /// Static mapping of all of the registered pass pipelines.
 static llvm::ManagedStatic<llvm::StringMap<PassPipelineInfo>>
     passPipelineRegistry;
 
 /// Utility to create a default registry function from a pass instance.
 static PassRegistryFunction
 buildDefaultRegistryFn(const PassAllocatorFunction &allocator) {
void mlir::registerPassPipeline(
   assert(inserted && "Pass pipeline registered multiple times");
   (void)inserted;
 }
 
 //===----------------------------------------------------------------------===//
 // PassInfo
 //===----------------------------------------------------------------------===//
 
-PassInfo::PassInfo(StringRef arg, StringRef description, TypeID passID,
+PassInfo::PassInfo(StringRef arg, StringRef description,
                    const PassAllocatorFunction &allocator)
     : PassRegistryEntry(
           arg, description, buildDefaultRegistryFn(allocator),
           // Use a temporary pass to provide an options instance.
           [=](function_ref<void(const PassOptions &)> optHandler) {
             optHandler(allocator()->passOptions);
           }) {}
 
 void mlir::registerPass(StringRef arg, StringRef description,
                         const PassAllocatorFunction &function) {
-  // TODO: We should use the 'arg' as the lookup key instead of the pass id.
-  TypeID passID = function()->getTypeID();
-  PassInfo passInfo(arg, description, passID, function);
-  passRegistry->try_emplace(passID, passInfo);
+  PassInfo passInfo(arg, description, function);
+  passRegistry->try_emplace(arg, passInfo);
+
+  // Verify that the registered pass has the same ID as any registered to this
+  // arg before it.
+  TypeID entryTypeID = function()->getTypeID();
+  auto it = passRegistryTypeIDs->try_emplace(arg, entryTypeID).first;
+  if (it->second != entryTypeID) {
+    llvm_unreachable("pass allocator creates a different pass than previously "
+                     "registered");
+  }
 }
 
-/// Returns the pass info for the specified pass class or null if unknown.
-const PassInfo *mlir::Pass::lookupPassInfo(TypeID passID) {
-  auto it = passRegistry->find(passID);
-  if (it == passRegistry->end())
-    return nullptr;
-  return &it->getSecond();
+/// Returns the pass info for the specified pass argument or null if unknown.
+const PassInfo *mlir::Pass::lookupPassInfo(StringRef passArg) {
+  auto it = passRegistry->find(passArg);
+  return it == passRegistry->end() ? nullptr : &it->second;
 }
 
 //===----------------------------------------------------------------------===//
 // PassOptions
 //===----------------------------------------------------------------------===//
 
 /// Out of line virtual function to provide home for the class.
 void detail::PassOptions::OptionBase::anchor() {}
TextualPipeline::resolvePipelineElement(PipelineElement &element,
   // Check to see if a pipeline was registered with this name.
   auto pipelineRegistryIt = passPipelineRegistry->find(element.name);
   if (pipelineRegistryIt != passPipelineRegistry->end()) {
     element.registryEntry = &pipelineRegistryIt->second;
     return success();
   }
 
   // If not, then this must be a specific pass name.
-  for (auto &passIt : *passRegistry) {
-    if (passIt.second.getPassArgument() == element.name) {
-      element.registryEntry = &passIt.second;
-      return success();
-    }
-  }
+  if ((element.registryEntry = Pass::lookupPassInfo(element.name)))
 
   // Emit an error for the unknown pass.
   auto *rawLoc = element.name.data();
   return errorHandler(rawLoc, "'" + element.name +
                                   "' does not refer to a "
                                   "registered pass or pass pipeline");
 }
 

mlir/test/lib/Pass/TestPassManager.cpp

 #include "mlir/Pass/PassManager.h"
 
 using namespace mlir;
 
 namespace {
 struct TestModulePass
     : public PassWrapper<TestModulePass, OperationPass<ModuleOp>> {
   void runOnOperation() final {}
+  StringRef getArgument() const final { return "test-module-pass"; }
 };
 struct TestFunctionPass : public PassWrapper<TestFunctionPass, FunctionPass> {
   void runOnFunction() final {}
+  StringRef getArgument() const final { return "test-function-pass"; }
 };
 class TestOptionsPass : public PassWrapper<TestOptionsPass, FunctionPass> {
 public:
   struct Options : public PassPipelineOptions<Options> {
     ListOption<int> listOption{*this, "list",
                                llvm::cl::MiscFlags::CommaSeparated,
                                llvm::cl::desc("Example list option")};
     ListOption<std::string> stringListOption{
         *this, "string-list", llvm::cl::MiscFlags::CommaSeparated,
         llvm::cl::desc("Example string list option")};
     Option<std::string> stringOption{*this, "string",
                                      llvm::cl::desc("Example string option")};
   };
   TestOptionsPass() = default;
   TestOptionsPass(const TestOptionsPass &) {}
   TestOptionsPass(const Options &options) {
     listOption = options.listOption;
     stringOption = options.stringOption;
     stringListOption = options.stringListOption;
   }
 
   void runOnFunction() final {}
+  StringRef getArgument() const final { return "test-options-pass"; }
 
   ListOption<int> listOption{*this, "list", llvm::cl::MiscFlags::CommaSeparated,
                              llvm::cl::desc("Example list option")};
   ListOption<std::string> stringListOption{
       *this, "string-list", llvm::cl::MiscFlags::CommaSeparated,
       llvm::cl::desc("Example string list option")};
   Option<std::string> stringOption{*this, "string",
                                    llvm::cl::desc("Example string option")};
 };
 
 /// A test pass that always aborts to enable testing the crash recovery
 /// mechanism of the pass manager.
 class TestCrashRecoveryPass
     : public PassWrapper<TestCrashRecoveryPass, OperationPass<>> {
   void runOnOperation() final { abort(); }
+  StringRef getArgument() const final { return "test-pass-crash"; }
 };
 
 /// A test pass that always fails to enable testing the failure recovery
 /// mechanisms of the pass manager.
 class TestFailurePass : public PassWrapper<TestFailurePass, OperationPass<>> {
   void runOnOperation() final { signalPassFailure(); }
 };