[mlir:PDL] Fix bugs in PDLPatternModule merging

[rriddle]

Summary

• Constraints/Rewrites registered before a pattern was added were dropped
• Constraints/Rewrites may be registered multiple times (if different pattern sets depend on them)
• ModuleOp no longer has a terminator, so we shouldn't be removing the terminator from it

Depends On D114815

[jpienaar]

mlir/lib/IR/PatternMatch.cpp

185

And here you mean equivalent fn? E.g., not just exactly the same, but same functionality.

Could you add comment to document why we need to relax this from previous assert?

mlir/test/lib/Rewrite/TestPDLByteCode.cpp

104

But why is that a good thing? :) And duplicate == the duplicate will be ignored/will overwrite? (I don't see that documented on API)

[rriddle]

resolve comments

mlir/lib/IR/PatternMatch.cpp

 //===----------------------------------------------------------------------===//
 // PDLPatternModule
 //===----------------------------------------------------------------------===//
 
 void PDLPatternModule::mergeIn(PDLPatternModule &&other) {
   // Ignore the other module if it has no patterns.
   if (!other.pdlModule)
     return;
+
+  // Steal the functions of the other module.
+  for (auto &it : other.constraintFunctions)
+    registerConstraintFunction(it.first(), std::move(it.second));
+  for (auto &it : other.rewriteFunctions)
+    registerRewriteFunction(it.first(), std::move(it.second));
+
   // Steal the other state if we have no patterns.
   if (!pdlModule) {
-    constraintFunctions = std::move(other.constraintFunctions);
-    rewriteFunctions = std::move(other.rewriteFunctions);
     pdlModule = std::move(other.pdlModule);
     return;
   }
-  // Steal the functions of the other module.
-  for (auto &it : constraintFunctions)
-    registerConstraintFunction(it.first(), std::move(it.second));
-  for (auto &it : rewriteFunctions)
-    registerRewriteFunction(it.first(), std::move(it.second));
 
   // Merge the pattern operations from the other module into this one.
   Block *block = pdlModule->getBody();
-  block->getTerminator()->erase();
   block->getOperations().splice(block->end(),
                                 other.pdlModule->getBody()->getOperations());
 }
 
 //===----------------------------------------------------------------------===//
 // Function Registry
 
 void PDLPatternModule::registerConstraintFunction(
     StringRef name, PDLConstraintFunction constraintFn) {
-  auto it = constraintFunctions.try_emplace(name, std::move(constraintFn));
-  (void)it;
-  assert(it.second &&
-         "constraint with the given name has already been registered");
+  // TODO: Is it possible to diagnose when `name` is already registered to
+  // something other than `constraintFn`?
+  constraintFunctions.try_emplace(name, std::move(constraintFn));
 }
 
 void PDLPatternModule::registerRewriteFunction(StringRef name,
                                                PDLRewriteFunction rewriteFn) {
-  auto it = rewriteFunctions.try_emplace(name, std::move(rewriteFn));
-  (void)it;
-  assert(it.second && "native rewrite function with the given name has "
-                      "already been registered");
+  // TODO: Is it possible to diagnose when `name` is already registered to
+  // something other than `rewriteFn`?
+  rewriteFunctions.try_emplace(name, std::move(rewriteFn));
 }
 
 //===----------------------------------------------------------------------===//
 // RewriterBase
 //===----------------------------------------------------------------------===//
 
 RewriterBase::~RewriterBase() {
   // Out of line to provide a vtable anchor for the class.

mlir/test/lib/Rewrite/TestPDLByteCode.cpp

void runOnOperation() final {
     // patterns and one containing the operations to rewrite.
     ModuleOp patternModule = module.lookupSymbol<ModuleOp>(
         StringAttr::get(module->getContext(), "patterns"));
     ModuleOp irModule = module.lookupSymbol<ModuleOp>(
         StringAttr::get(module->getContext(), "ir"));
     if (!patternModule || !irModule)
       return;
 
+    RewritePatternSet patternList(module->getContext());
+
+    // Register ahead of time to test when functions are registered without a
+    // pattern.
+    patternList.getPDLPatterns().registerConstraintFunction(
+        "multi_entity_constraint", customMultiEntityConstraint);
+    patternList.getPDLPatterns().registerConstraintFunction(
+        "single_entity_constraint", customSingleEntityConstraint);
+
     // Process the pattern module.
     patternModule.getOperation()->remove();
     PDLPatternModule pdlPattern(patternModule);
+
+    // Note: This constraint was already registered, but we re-register here to
+    // ensure that duplication registration is allowed.
     pdlPattern.registerConstraintFunction("multi_entity_constraint",
                                           customMultiEntityConstraint);
-    pdlPattern.registerConstraintFunction("single_entity_constraint",
-                                          customSingleEntityConstraint);
     pdlPattern.registerConstraintFunction("multi_entity_var_constraint",
                                           customMultiEntityVariadicConstraint);
     pdlPattern.registerRewriteFunction("creator", customCreate);
     pdlPattern.registerRewriteFunction("var_creator",
                                        customVariadicResultCreate);
     pdlPattern.registerRewriteFunction("type_creator", customCreateType);
     pdlPattern.registerRewriteFunction("rewriter", customRewriter);
-
-    RewritePatternSet patternList(std::move(pdlPattern));
+    patternList.add(std::move(pdlPattern));
 
     // Invoke the pattern driver with the provided patterns.
     (void)applyPatternsAndFoldGreedily(irModule.getBodyRegion(),
                                        std::move(patternList));
   }
 };
 } // end anonymous namespace
 
 namespace mlir {
 namespace test {
 void registerTestPDLByteCodePass() { PassRegistration<TestPDLByteCodePass>(); }
 } // namespace test
 } // namespace mlir