Multi-root PDL matching using upward traversals.

[sfuniak]

Summary

This is commit 4 of 4 for the multi-root matching in PDL, discussed in https://llvm.discourse.group/t/rfc-multi-root-pdl-patterns-for-kernel-matching/4148 (topic flagged for review).

This PR integrates the various components (root ordering algorithm, nondeterministic execution of PDL bytecode) to implement multi-root PDL matching. The main idea is for the pattern to specify mulitple candidate roots. The PDL-to-PDLInterp lowering selects one of these roots and "hangs" the pattern from this root, traversing the edges downwards (from operation to its operands) when possible and upwards (from values to its uses) when needed. The root is selected by invoking the optimal matching multiple times, once for each candidate root, and the connectors are determined form the optimal matching. The costs in the directed graph are equal to the number of upward edges that need to be traversed when connecting the given two candidate roots. It can be shown that, for this choice of the cost function, "hanging" the pattern an inner node is no better than from the optimal root.

The following three main additions were implemented as a part of this PR:

1. OperationPos predicate has been extended to allow tracing the operation accepting a value (the opposite of operation defining a value).
2. Predicate checking if two values are not equal - this is useful to ensure that we do not traverse the edge back downwards after we traversed it upwards.
3. Function for for building the cost graph among the candidate roots.
4. Updated buildPredicateList, building the predicates optimal branching has been determined.

Testing: unit tests (an integration test to follow once the stack of commits has landed)

[rriddle]

Made my way through most of it, will check the rest tomorrow.

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

454

Why remove the accessor but not the rootKind argument?

619–620

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

252

Does this revision not build directly? (i.e. are you intending to land all of the commits together?)

394

Can you use a different name than nominal? Confused me for a second when I looked at this (before I saw that EqualTo also corresponds to NE).

549–550

I'm a bit apprehensive about asserting this. I intend to support return values in constraints, and that would end up with an assert here. Can we have a method somewhere that provides this functionality in a fail-able way? (i.e. if the value isn't an OperationOp return None, otherwise return op.name()).

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

170

mlir/lib/Dialect/PDL/IR/PDL.cpp

43–44

This comment needs to be tweaked a bit.

[sfuniak]

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

454

Good point, will do.

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

252

It was supposed to build directly. I separated the functionality carefully. But I stacked the commits after-the-fact, which could be the reason why clang-tidy complains. Hopefully this error will go away when I update the diff.

394

Ok, I will just use is_true.

549–550

Sure, I can do that.

[sfuniak]

Migrated to pdl_interp.foreach and autodetecting roots

This update migrates the PDL-to-PDLInterp lowering to use the new ForEach construct. I have unified the generate methods and made the block passed to getValueAt mutable, so that it can insert the nested regions for each of the pdl_interp.foreach it encouters.

Finally, following the recommendations in the review, I change pdl.rewrite to accept a single optional root, which forces the root of the search tree. In the absence of this root, the lowering will choose the optimal root among the ones detected in the pattern.

[rriddle]

Nice!

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

587–590

Can you beef the documentation on what you mean by "search tree"? It isn't clear exactly what that is supposed to mean.

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

160

170

nit: Drop the const here

221

nit: Can you add a message here?

248

Can you switch this to early exit? i.e. the else branch is smaller, so you could do:

if (!operationPos->isUpward()) {

value = ...;
break;

}

257–258

Why is this calling the create directly? as opposed to using the builder.

262–263

builder.createBlock(...)?

338

Is the 2 a tuned value? or can we just use the default inline vector size?

339–343

nit: Use braces when the body is a nested control flow statement (the one exception to that though AFAIK is for (...) if (...))

394

nit: Use camelCase for variable names.

466

Is this still open?

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

170

unresolved?

335–340

383–390

442

530

580

nit; Assert messages should start with lower case

583

Is this guaranteed to be non-variadic?

778–791

[sfuniak]

Thank you for the review @rriddle. Will submit an update later today.

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

257–258

Just a MLIR rookie mistake, will fix.

338

I just copied whatever was there originally at line 348. Happy to use the default.

339–343

Thanks, I was not sure what the convention was.

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

170

Your suggested change did not compile: we have only comparison between two Optionals, not an Optional and a value. And because these are templates, implicit conversion from value to Optional is not invoked.

[sfuniak]

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

257–258

Actually, I remember now. ForEachOp::create is a custom function that creates the entry block with a loop variable (block argument). If I use a builder, then this will invoke OpTy::build, followed by Operation::create(state), which is not what I wanted.

Or is there another, more standard way to do this?

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

583

You are right, it can be variadic. I think the entire logic will still work, but we need to tweak the implementation of pdl_interp.get_users a bit. I will include that in an upcoming update.

[sfuniak]

Addressed review feedback.

[rriddle]

Did a quick scan, looks good to land and iterate!

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

587–590

unresolved?

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

257–258

We generally just use the build methods, create methods have generally been reserved for top-level things and are generally left-over from before FuncOp/ModuleOp were operations. Can we remove ForEachOp::create and use a build method instead?

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

545–546

Can you move these inner methods to separate functions? buildPredicateList has become quite large.

mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir

494

Can you add tests for when this is variadic?

[sfuniak]

Final review feedback; also fixed the connected component checking code.

[sfuniak]

Rebased

[Mogball]

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

337–339

Are there going to be more cases down the line?

383

Seeing a struct def inside a loop is... weird to me. I'm not sure if LLVM's coding standards say much about where structs defined inside functions should be, but it would make more sense to me at the top of the function body.

[sfuniak]

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

337–339

No, but you need a cast to pdl::ResultOp or pdl::ResultsOp to extract the parent. The TypeSwitch approach seems cleaner.

383

No problem, I will move it.

[sfuniak]

Updated code to simplified pdl_interp.get_users.

[sfuniak]

Use pdl_interp.get_value for upward traversal of value ranges.

[sfuniak]

Switch from pdl_interp.get_value to pdl_interp.extract.

[sfuniak]

Rebased.

[mehdi_amini]

This broke the gcc 5.4 build:

/var/lib/buildkite-agent/builds/buildkite-588bd64db9-bjm5q-1/mlir/mlir-core/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp:212:36: error: cannot call member function 'void {anonymous}::PatternLowering::generate(mlir::pdl_to_pdl_interp::BoolNode*, mlir::Block*&, mlir::Value)' without object
           [&](auto *derivedNode) { generate(derivedNode, currentBlock, val); })
                                    ^
/var/lib/buildkite-agent/builds/buildkite-588bd64db9-bjm5q-1/mlir/mlir-core/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp: In instantiation of '{anonymous}::PatternLowering::generateMatcher(mlir::pdl_to_pdl_interp::MatcherNode&, mlir::Region&)::<lambda(auto:16*)> [with auto:16 = mlir::pdl_to_pdl_interp::SwitchNode]':
/var/lib/buildkite-agent/builds/buildkite-588bd64db9-bjm5q-1/mlir/mlir-core/llvm/include/llvm/ADT/TypeSwitch.h:169:13:   required from 'llvm::TypeSwitch<T, void>& llvm::TypeSwitch<T, void>::Case(CallableT&&) [with CaseT = mlir::pdl_to_pdl_interp::SwitchNode; CallableT = {anonymous}::PatternLowering::generateMatcher(mlir::pdl_to_pdl_interp::MatcherNode&, mlir::Region&)::<lambda(auto:16*)>&; T = mlir::pdl_to_pdl_interp::MatcherNode*]'
/var/lib/buildkite-agent/builds/buildkite-588bd64db9-bjm5q-1/mlir/mlir-core/llvm/include/llvm/ADT/TypeSwitch.h:46:49:   required from 'DerivedT& llvm::detail::TypeSwitchBase<DerivedT, T>::Case(CallableT&&) [with CaseT = mlir::pdl_to_pdl_interp::BoolNode; CaseT2 = mlir::pdl_to_pdl_interp::SwitchNode; CaseTs = {}; CallableT = {anonymous}::PatternLowering::generateMatcher(mlir::pdl_to_pdl_interp::MatcherNode&, mlir::Region&)::<lambda(auto:16*)>; DerivedT = llvm::TypeSwitch<mlir::pdl_to_pdl_interp::MatcherNode*, void>; T = mlir::pdl_to_pdl_interp::MatcherNode*]'
/var/lib/buildkite-agent/builds/buildkite-588bd64db9-bjm5q-1/mlir/mlir-core/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp:212:79:   required from here
/var/lib/buildkite-agent/builds/buildkite-588bd64db9-bjm5q-1/mlir/mlir-core/mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp:212:36: error: cannot call member function 'void {anonymous}::PatternLowering::generate(mlir::pdl_to_pdl_interp::SwitchNode*, mlir::Block*, mlir::Value)' without object

[sfuniak]

@mehdi_amini thanks for bringing this up; I will take a look.

[rriddle]

In D108550#3156533, @sfuniak wrote:

@mehdi_amini thanks for bringing this up; I will take a look.

This looks like the gcc5 bug for auto lambdas that capture this, you likely just need to use this-> for member calls.

mlir/include/mlir/Dialect/PDL/IR/PDLOps.td

pdl.pattern : benefit(1) {
       %root = pdl.operation "foo.op"(%inputOperand) -> (%resultType)
       pdl.rewrite %root {
         pdl.replace %root with (%inputOperand)
       }
     }
     ```
   }];
 
-  let arguments = (ins OptionalAttr<StrAttr>:$rootKind,
-                       Confined<I16Attr, [IntNonNegative]>:$benefit,
+  let arguments = (ins Confined<I16Attr, [IntNonNegative]>:$benefit,
                        OptionalAttr<SymbolNameAttr>:$sym_name);
   let regions = (region SizedRegion<1>:$body);
   let assemblyFormat = [{
-    ($sym_name^)? `:` `benefit` `(` $benefit `)`
-    (`,` `root` `(` $rootKind^ `)`)? attr-dict-with-keyword $body
+    ($sym_name^)? `:` `benefit` `(` $benefit `)` attr-dict-with-keyword $body
   }];
 
   let builders = [
-    OpBuilder<(ins CArg<"Optional<StringRef>", "llvm::None">:$rootKind,
-      CArg<"Optional<uint16_t>", "1">:$benefit,
+    OpBuilder<(ins CArg<"Optional<uint16_t>", "1">:$benefit,
                    CArg<"Optional<StringRef>", "llvm::None">:$name)>,
   ];
   let extraClassDeclaration = [{
     //===------------------------------------------------------------------===//
     // SymbolOpInterface Methods
     //===------------------------------------------------------------------===//
 
     /// A PatternOp may optionally define a symbol.
     bool isOptionalSymbol() { return true; }
 
     /// Returns the rewrite operation of this pattern.
     RewriteOp getRewriter();
-
-    /// Return the root operation kind that this pattern matches, or None if
-    /// there isn't a specific root.
-    Optional<StringRef> getRootKind();
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::ReplaceOp
 //===----------------------------------------------------------------------===//
 
 def PDL_ReplaceOp : PDL_Op<"replace", [

 }
 
 //===----------------------------------------------------------------------===//
 // pdl::RewriteOp
 //===----------------------------------------------------------------------===//
 
 def PDL_RewriteOp : PDL_Op<"rewrite", [
      Terminator, HasParent<"pdl::PatternOp">, NoTerminator, NoRegionArguments,
-     SingleBlock
+     SingleBlock, AttrSizedOperandSegments
   ]> {
   let summary = "Specify the rewrite of a matched pattern";
   let description = [{
     `pdl.rewrite` operations terminate the region of a `pdl.pattern` and specify
-    the main rewrite of a `pdl.pattern`, on the specified root operation. The
+    the main rewrite of a `pdl.pattern`, on the optional root operation. The
     rewrite is specified either via a string name (`name`) to a native
     rewrite function, or via the region body. The rewrite region, if specified,
     must contain a single block. If the rewrite is external it functions
     similarly to `pdl.apply_native_rewrite`, and takes a set of constant
     parameters and a set of additional positional values defined within the
     matcher as arguments. If the rewrite is external, the root operation is
-    passed to the native function as the first argument.
+    passed to the native function as the leading arguments. The root operation,
+    if provided, specifies which operation is used as the root of the search
+    tree. If it is omitted, the pdl_interp lowering will automatically select
+    the best root of the search tree among all operations in the pattern.
 
     Example:
 
     ```mlir
     // Specify an external rewrite function:
     pdl.rewrite %root with "myExternalRewriter"(%value : !pdl.value)
 
-    // Specify the rewrite inline using PDL:
+    // Specify a rewrite inline using PDL with the given root:
     pdl.rewrite %root {
       %op = pdl.operation "foo.op"(%arg0, %arg1)
       pdl.replace %root with %op
     }
+
+    // Specify a rewrite inline using PDL, automatically selecting root:
+    pdl.rewrite {
+      %op1 = pdl.operation "foo.op"(%arg0, %arg1)
+      %op2 = pdl.operation "bar.op"(%arg0, %arg1)
+      pdl.replace %root1 with %op1
+      pdl.replace %root2 with %op2
+    }
     ```
   }];
 
-  let arguments = (ins PDL_Operation:$root,
+  let arguments = (ins Optional<PDL_Operation>:$root,
                        OptionalAttr<StrAttr>:$name,
                        Variadic<PDL_AnyType>:$externalArgs,
                        OptionalAttr<ArrayAttr>:$externalConstParams);
   let regions = (region AnyRegion:$body);
   let assemblyFormat = [{
-    $root (`with` $name^ ($externalConstParams^)?
+    ($root^)? (`with` $name^ ($externalConstParams^)?
                              (`(` $externalArgs^ `:` type($externalArgs) `)`)?)?
               ($body^)?
     attr-dict-with-keyword
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::TypeOp
 //===----------------------------------------------------------------------===//
 

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

public:
   /// within the module.
   void lower(ModuleOp module);
 
 private:
   using ValueMap = llvm::ScopedHashTable<Position *, Value>;
   using ValueMapScope = llvm::ScopedHashTableScope<Position *, Value>;
 
   /// Generate interpreter operations for the tree rooted at the given matcher
-  /// node.
-  Block *generateMatcher(MatcherNode &node);
+  /// node, in the specified region.
+  Block *generateMatcher(MatcherNode &node, Region &region);
 
-  /// Get or create an access to the provided positional value within the
-  /// current block.
-  Value getValueAt(Block *cur, Position *pos);
-
-  /// Create an interpreter predicate operation, branching to the provided true
-  /// and false destinations.
-  void generatePredicate(Block *currentBlock, Qualifier *question,
-                         Qualifier *answer, Value val, Block *trueDest,
-                         Block *falseDest);
-
-  /// Create an interpreter switch predicate operation, with a provided default
-  /// and several case destinations.
-  void generateSwitch(SwitchNode *switchNode, Block *currentBlock,
-                      Qualifier *question, Value val, Block *defaultDest);
-
-  /// Create the interpreter operations to record a successful pattern match.
-  void generateRecordMatch(Block *currentBlock, Block *nextBlock,
-                           pdl::PatternOp pattern);
+  /// Get or create an access to the provided positional value in the current
+  /// block. This operation may mutate the provided block pointer if nested
+  /// regions (i.e., pdl_interp.iterate) are required.
+  Value getValueAt(Block *&currentBlock, Position *pos);
+
+  /// Create the interpreter predicate operations. This operation may mutate the
+  /// provided current block pointer if nested regions (iterates) are required.
+  void generate(BoolNode *boolNode, Block *&currentBlock, Value val);
+
+  /// Create the interpreter switch / predicate operations, with several case
+  /// destinations. This operation never mutates the provided current block
+  /// pointer, because the switch operation does not need Values beyond `val`.
+  void generate(SwitchNode *switchNode, Block *currentBlock, Value val);
+
+  /// Create the interpreter operations to record a successful pattern match
+  /// using the contained root operation. This operation may mutate the current
+  /// block pointer if nested regions (i.e., pdl_interp.iterate) are required.
+  void generate(SuccessNode *successNode, Block *&currentBlock);
 
   /// Generate a rewriter function for the given pattern operation, and returns
   /// a reference to that function.
   SymbolRefAttr generateRewriter(pdl::PatternOp pattern,
                                  SmallVectorImpl<Position *> &usedMatchValues);
 
   /// Generate the rewriter code for the given operation.
   void generateRewriter(pdl::ApplyNativeRewriteOp rewriteOp,
void PatternLowering::lower(ModuleOp module) {
   // Insert the root operation, i.e. argument to the matcher, at the root
   // position.
   Block *matcherEntryBlock = matcherFunc.addEntryBlock();
   values.insert(predicateBuilder.getRoot(), matcherEntryBlock->getArgument(0));
 
   // Generate a root matcher node from the provided PDL module.
   std::unique_ptr<MatcherNode> root = MatcherNode::generateMatcherTree(
       module, predicateBuilder, valueToPosition);
-  Block *firstMatcherBlock = generateMatcher(*root);
+  Block *firstMatcherBlock = generateMatcher(*root, matcherFunc.getBody());
+  assert(failureBlockStack.empty() && "Failed to empty the stack");
 
   // After generation, merged the first matched block into the entry.
   matcherEntryBlock->getOperations().splice(matcherEntryBlock->end(),
                                             firstMatcherBlock->getOperations());
   firstMatcherBlock->erase();
 }
 
-Block *PatternLowering::generateMatcher(MatcherNode &node) {
+Block *PatternLowering::generateMatcher(MatcherNode &node, Region &region) {
   // Push a new scope for the values used by this matcher.
-  Block *block = matcherFunc.addBlock();
+  Block *const block = &region.emplaceBlock();
   ValueMapScope scope(values);
 
   // If this is the return node, simply insert the corresponding interpreter
   // finalize.
   if (isa<ExitNode>(node)) {
     builder.setInsertionPointToEnd(block);
     builder.create<pdl_interp::FinalizeOp>(matcherFunc.getLoc());
     return block;
   }
 
-  // If this node contains a position, get the corresponding value for this
-  // block.
-  Position *position = node.getPosition();
-  Value val = position ? getValueAt(block, position) : Value();
-
   // Get the next block in the match sequence.
+  // This is intentionally executed first, before we get the value for the
+  // position associated with the node, so that we preserve an "there exist"
+  // semantics: if getting a value requires an upward traversal (going from a
+  // value to its consumers), we want to perform the check on all the consumers
+  // before we pass control to the failure node.
   std::unique_ptr<MatcherNode> &failureNode = node.getFailureNode();
-  Block *nextBlock;
+  Block *failureBlock;
   if (failureNode) {
-    nextBlock = generateMatcher(*failureNode);
-    failureBlockStack.push_back(nextBlock);
+    failureBlock = generateMatcher(*failureNode, region);
+    failureBlockStack.push_back(failureBlock);
   } else {
     assert(!failureBlockStack.empty() && "expected valid failure block");
-    nextBlock = failureBlockStack.back();
+    failureBlock = failureBlockStack.back();
   }
 
+  // If this node contains a position, get the corresponding value for this
+  // block.
+  Block *currentBlock = block;
+  Position *position = node.getPosition();
+  Value val = position ? getValueAt(currentBlock, position) : Value();
+
   // If this value corresponds to an operation, record that we are going to use
   // its location as part of a fused location.
   bool isOperationValue = val && val.getType().isa<pdl::OperationType>();
   if (isOperationValue)
     locOps.insert(val);
 
-  // Generate code for a boolean predicate node.
-  if (auto *boolNode = dyn_cast<BoolNode>(&node)) {
-    auto *child = generateMatcher(*boolNode->getSuccessNode());
-    generatePredicate(block, node.getQuestion(), boolNode->getAnswer(), val,
-                      child, nextBlock);
-
-    // Generate code for a switch node.
-  } else if (auto *switchNode = dyn_cast<SwitchNode>(&node)) {
-    generateSwitch(switchNode, block, node.getQuestion(), val, nextBlock);
-
-    // Generate code for a success node.
-  } else if (auto *successNode = dyn_cast<SuccessNode>(&node)) {
-    generateRecordMatch(block, nextBlock, successNode->getPattern());
+  // Dispatch to the correct method based on derived node type.
+  TypeSwitch<MatcherNode *>(&node)
+      .Case<BoolNode, SwitchNode>(
+          [&](auto *derivedNode) { generate(derivedNode, currentBlock, val); })
+      .Case([&](SuccessNode *successNode) {
+        generate(successNode, currentBlock);
+      });
+
+  // Pop all the failure blocks that were inserted due to nesting of
+  // pdl_interp.iterate.
+  while (failureBlockStack.back() != failureBlock) {
+    failureBlockStack.pop_back();
+    assert(!failureBlockStack.empty());
   }
 
+  // Pop the new failure block.
   if (failureNode)
     failureBlockStack.pop_back();
+
   if (isOperationValue)
     locOps.remove(val);
+
   return block;
 }
 
-Value PatternLowering::getValueAt(Block *cur, Position *pos) {
+Value PatternLowering::getValueAt(Block *&currentBlock, Position *pos) {
   if (Value val = values.lookup(pos))
     return val;
 
   // Get the value for the parent position.
-  Value parentVal = getValueAt(cur, pos->getParent());
+  Value parentVal = getValueAt(currentBlock, pos->getParent());
 
   // TODO: Use a location from the position.
   Location loc = parentVal.getLoc();
-  builder.setInsertionPointToEnd(cur);
+  builder.setInsertionPointToEnd(currentBlock);
   Value value;
   switch (pos->getKind()) {
-  case Predicates::OperationPos:
+  case Predicates::OperationPos: {
+    auto *operationPos = cast<OperationPosition>(pos);
+    if (operationPos->isUpward()) {
+      // The first operation retrieves the users.
+      Type operationTy = builder.getType<pdl::OperationType>();
+      value = builder.create<pdl_interp::GetUsersOp>(
+          loc, pdl::RangeType::get(operationTy), parentVal,
+          operationPos->getIndex());
+
+      // The second operation iterates over them.
+      assert(!failureBlockStack.empty() && "expected valid failure block");
+      auto foreach = pdl_interp::ForEachOp::create(builder, loc, value,
+                                                   failureBlockStack.back());
+      value = foreach.getLoopVariable();
+
+      // The third operation is the continuation to the next iterate.
+      Block *continueBlock = &foreach.region().emplaceBlock();
+      builder.setInsertionPointToEnd(continueBlock);
+      builder.create<pdl_interp::ContinueOp>(loc);
+      failureBlockStack.push_back(continueBlock);
+
+      // Update the current block pointer.
+      currentBlock = &foreach.region().front();
+      builder.setInsertionPointToEnd(currentBlock);
+    } else {
       value = builder.create<pdl_interp::GetDefiningOpOp>(
           loc, builder.getType<pdl::OperationType>(), parentVal);
+    }
     break;
+  }
   case Predicates::OperandPos: {
     auto *operandPos = cast<OperandPosition>(pos);
     value = builder.create<pdl_interp::GetOperandOp>(
         loc, builder.getType<pdl::ValueType>(), parentVal,
         operandPos->getOperandNumber());
     break;
   }
   case Predicates::OperandGroupPos: {
value = builder.create<pdl_interp::GetResultsOp>(
         loc, resPos->isVariadic() ? pdl::RangeType::get(valueTy) : valueTy,
         parentVal, resPos->getResultGroupNumber());
     break;
   }
   default:
     llvm_unreachable("Generating unknown Position getter");
     break;
   }
+
   values.insert(pos, value);
   return value;
 }
 
-void PatternLowering::generatePredicate(Block *currentBlock,
-                                        Qualifier *question, Qualifier *answer,
-                                        Value val, Block *trueDest,
-                                        Block *falseDest) {
-  builder.setInsertionPointToEnd(currentBlock);
+void PatternLowering::generate(BoolNode *boolNode, Block *&currentBlock,
+                               Value val) {
   Location loc = val.getLoc();
+  Qualifier *question = boolNode->getQuestion();
+  Qualifier *answer = boolNode->getAnswer();
+  Region *region = currentBlock->getParent();
+
+  // Execute the getValue queries first, so that we create success
+  // matcher in the correct (possibly nested) region.
+  SmallVector<Value, 2> args;
+  if (auto *equalToQuestion = dyn_cast<EqualToQuestion>(question))
+    args = {getValueAt(currentBlock, equalToQuestion->getValue())};
+  else if (auto *cstQuestion = dyn_cast<ConstraintQuestion>(question))
+    for (Position *position : std::get<1>(cstQuestion->getValue()))
+      args.push_back(getValueAt(currentBlock, position));
+
+  // Generate the matcher in the current (potentially nested) region
+  // and get the failure successor.
+  Block *success = generateMatcher(*boolNode->getSuccessNode(), *region);
+  Block *failure = failureBlockStack.back();
+
+  // Finally, create the predicate.
+  builder.setInsertionPointToEnd(currentBlock);
   Predicates::Kind kind = question->getKind();
   switch (kind) {
   case Predicates::IsNotNullQuestion:
-    builder.create<pdl_interp::IsNotNullOp>(loc, val, trueDest, falseDest);
+    builder.create<pdl_interp::IsNotNullOp>(loc, val, success, failure);
     break;
   case Predicates::OperationNameQuestion: {
     auto *opNameAnswer = cast<OperationNameAnswer>(answer);
     builder.create<pdl_interp::CheckOperationNameOp>(
-        loc, val, opNameAnswer->getValue().getStringRef(), trueDest, falseDest);
+        loc, val, opNameAnswer->getValue().getStringRef(), success, failure);
     break;
   }
   case Predicates::TypeQuestion: {
     auto *ans = cast<TypeAnswer>(answer);
     if (val.getType().isa<pdl::RangeType>())
       builder.create<pdl_interp::CheckTypesOp>(
-          loc, val, ans->getValue().cast<ArrayAttr>(), trueDest, falseDest);
+          loc, val, ans->getValue().cast<ArrayAttr>(), success, failure);
     else
       builder.create<pdl_interp::CheckTypeOp>(
-          loc, val, ans->getValue().cast<TypeAttr>(), trueDest, falseDest);
+          loc, val, ans->getValue().cast<TypeAttr>(), success, failure);
     break;
   }
   case Predicates::AttributeQuestion: {
     auto *ans = cast<AttributeAnswer>(answer);
     builder.create<pdl_interp::CheckAttributeOp>(loc, val, ans->getValue(),
-                                                 trueDest, falseDest);
+                                                 success, failure);
     break;
   }
   case Predicates::OperandCountAtLeastQuestion:
   case Predicates::OperandCountQuestion:
     builder.create<pdl_interp::CheckOperandCountOp>(
         loc, val, cast<UnsignedAnswer>(answer)->getValue(),
         /*compareAtLeast=*/kind == Predicates::OperandCountAtLeastQuestion,
-        trueDest, falseDest);
+        success, failure);
     break;
   case Predicates::ResultCountAtLeastQuestion:
   case Predicates::ResultCountQuestion:
     builder.create<pdl_interp::CheckResultCountOp>(
         loc, val, cast<UnsignedAnswer>(answer)->getValue(),
         /*compareAtLeast=*/kind == Predicates::ResultCountAtLeastQuestion,
-        trueDest, falseDest);
+        success, failure);
     break;
   case Predicates::EqualToQuestion: {
-    auto *equalToQuestion = cast<EqualToQuestion>(question);
-    builder.create<pdl_interp::AreEqualOp>(
-        loc, val, getValueAt(currentBlock, equalToQuestion->getValue()),
-        trueDest, falseDest);
+    bool true_answer = isa<TrueAnswer>(answer);
+    builder.create<pdl_interp::AreEqualOp>(loc, val, args.front(),
+                                           true_answer ? success : failure,
+                                           true_answer ? failure : success);
     break;
   }
   case Predicates::ConstraintQuestion: {
-    auto *cstQuestion = cast<ConstraintQuestion>(question);
-    SmallVector<Value, 2> args;
-    for (Position *position : std::get<1>(cstQuestion->getValue()))
-      args.push_back(getValueAt(currentBlock, position));
+    auto value = cast<ConstraintQuestion>(question)->getValue();
     builder.create<pdl_interp::ApplyConstraintOp>(
-        loc, std::get<0>(cstQuestion->getValue()), args,
-        std::get<2>(cstQuestion->getValue()).cast<ArrayAttr>(), trueDest,
-        falseDest);
+        loc, std::get<0>(value), args, std::get<2>(value).cast<ArrayAttr>(),
+        success, failure);
     break;
   }
   default:
     llvm_unreachable("Generating unknown Predicate operation");
   }
 }
 
 template <typename OpT, typename PredT, typename ValT = typename PredT::KeyTy>
 static void createSwitchOp(Value val, Block *defaultDest, OpBuilder &builder,
                            llvm::MapVector<Qualifier *, Block *> &dests) {
   std::vector<ValT> values;
   std::vector<Block *> blocks;
   values.reserve(dests.size());
   blocks.reserve(dests.size());
   for (const auto &it : dests) {
     blocks.push_back(it.second);
     values.push_back(cast<PredT>(it.first)->getValue());
   }
   builder.create<OpT>(val.getLoc(), val, values, defaultDest, blocks);
 }
 
-void PatternLowering::generateSwitch(SwitchNode *switchNode,
-                                     Block *currentBlock, Qualifier *question,
-                                     Value val, Block *defaultDest) {
+void PatternLowering::generate(SwitchNode *switchNode, Block *currentBlock,
+                               Value val) {
+  Qualifier *question = switchNode->getQuestion();
+  Region *region = currentBlock->getParent();
+  Block *defaultDest = failureBlockStack.back();
+
   // If the switch question is not an exact answer, i.e. for the `at_least`
   // cases, we generate a special block sequence.
   Predicates::Kind kind = question->getKind();
   if (kind == Predicates::OperandCountAtLeastQuestion ||
       kind == Predicates::ResultCountAtLeastQuestion) {
     // Order the children such that the cases are in reverse numerical order.
     SmallVector<unsigned> sortedChildren = llvm::to_vector<16>(
         llvm::seq<unsigned>(0, switchNode->getChildren().size()));
if (kind == Predicates::OperandCountAtLeastQuestion ||
     //   goto failure
     // if (child2.match())
     //   ...
     //
     // failure:
     //   ...
     //
     failureBlockStack.push_back(defaultDest);
+    Location loc = val.getLoc();
     for (unsigned idx : sortedChildren) {
       auto &child = switchNode->getChild(idx);
-      Block *childBlock = generateMatcher(*child.second);
+      // TODO - can the matcher nest? probably yes, make sure this still works
+      Block *childBlock = generateMatcher(*child.second, *region);
       Block *predicateBlock = builder.createBlock(childBlock);
-      generatePredicate(predicateBlock, question, child.first, val, childBlock,
-                        defaultDest);
+      builder.setInsertionPointToEnd(predicateBlock);
+      unsigned ans = cast<UnsignedAnswer>(child.first)->getValue();
+      switch (kind) {
+      case Predicates::OperandCountAtLeastQuestion:
+        builder.create<pdl_interp::CheckOperandCountOp>(
+            loc, val, ans, /*compareAtLeast=*/true, childBlock, defaultDest);
+        break;
+      case Predicates::ResultCountAtLeastQuestion:
+        builder.create<pdl_interp::CheckResultCountOp>(
+            loc, val, ans, /*compareAtLeast=*/true, childBlock, defaultDest);
+        break;
+      default:
+        llvm_unreachable("Generating invalid AtLeast operation");
+      }
       failureBlockStack.back() = predicateBlock;
     }
     Block *firstPredicateBlock = failureBlockStack.pop_back_val();
     currentBlock->getOperations().splice(currentBlock->end(),
                                          firstPredicateBlock->getOperations());
     firstPredicateBlock->erase();
     return;
   }
 
   // Otherwise, generate each of the children and generate an interpreter
   // switch.
   llvm::MapVector<Qualifier *, Block *> children;
   for (auto &it : switchNode->getChildren())
-    children.insert({it.first, generateMatcher(*it.second)});
+    children.insert({it.first, generateMatcher(*it.second, *region)});
   builder.setInsertionPointToEnd(currentBlock);
 
   switch (question->getKind()) {
   case Predicates::OperandCountQuestion:
     return createSwitchOp<pdl_interp::SwitchOperandCountOp, UnsignedAnswer,
                           int32_t>(val, defaultDest, builder, children);
   case Predicates::ResultCountQuestion:
     return createSwitchOp<pdl_interp::SwitchResultCountOp, UnsignedAnswer,
void PatternLowering::generate(SwitchNode *switchNode, Block *currentBlock,
   case Predicates::AttributeQuestion:
     return createSwitchOp<pdl_interp::SwitchAttributeOp, AttributeAnswer>(
         val, defaultDest, builder, children);
   default:
     llvm_unreachable("Generating unknown switch predicate.");
   }
 }
 
-void PatternLowering::generateRecordMatch(Block *currentBlock, Block *nextBlock,
-                                          pdl::PatternOp pattern) {
+void PatternLowering::generate(SuccessNode *successNode, Block *&currentBlock) {
+  pdl::PatternOp pattern = successNode->getPattern();
+  Value root = successNode->getRoot();
+
   // Generate a rewriter for the pattern this success node represents, and track
   // any values used from the match region.
   SmallVector<Position *, 8> usedMatchValues;
   SymbolRefAttr rewriterFuncRef = generateRewriter(pattern, usedMatchValues);
 
   // Process any values used in the rewrite that are defined in the match.
   std::vector<Value> mappedMatchValues;
   mappedMatchValues.reserve(usedMatchValues.size());
   for (Position *position : usedMatchValues)
     mappedMatchValues.push_back(getValueAt(currentBlock, position));
 
   // Collect the set of operations generated by the rewriter.
   SmallVector<StringRef, 4> generatedOps;
   for (auto op : pattern.getRewriter().body().getOps<pdl::OperationOp>())
     generatedOps.push_back(*op.name());
   ArrayAttr generatedOpsAttr;
   if (!generatedOps.empty())
     generatedOpsAttr = builder.getStrArrayAttr(generatedOps);
 
   // Grab the root kind if present.
   StringAttr rootKindAttr;
-  if (Optional<StringRef> rootKind = pattern.getRootKind())
+  if (pdl::OperationOp rootOp = root.getDefiningOp<pdl::OperationOp>())
+    if (Optional<StringRef> rootKind = rootOp.name())
       rootKindAttr = builder.getStringAttr(*rootKind);
 
   builder.setInsertionPointToEnd(currentBlock);
   builder.create<pdl_interp::RecordMatchOp>(
       pattern.getLoc(), mappedMatchValues, locOps.getArrayRef(),
       rewriterFuncRef, rootKindAttr, generatedOpsAttr, pattern.benefitAttr(),
-      nextBlock);
+      failureBlockStack.back());
 }
 
 SymbolRefAttr PatternLowering::generateRewriter(
     pdl::PatternOp pattern, SmallVectorImpl<Position *> &usedMatchValues) {
   FuncOp rewriterFunc =
       FuncOp::create(pattern.getLoc(), "pdl_generated_rewriter",
                      builder.getFunctionType(llvm::None, llvm::None));
   rewriterSymbolTable.insert(rewriterFunc);
auto mapRewriteValue = [&](Value oldValue) {
     usedMatchValues.push_back(inputPos);
     return newValue = rewriterFunc.front().addArgument(oldValue.getType());
   };
 
   // If this is a custom rewriter, simply dispatch to the registered rewrite
   // method.
   pdl::RewriteOp rewriter = pattern.getRewriter();
   if (StringAttr rewriteName = rewriter.nameAttr()) {
+    SmallVector<Value> args;
+    if (rewriter.root())
+      args.push_back(mapRewriteValue(rewriter.root()));
     auto mappedArgs = llvm::map_range(rewriter.externalArgs(), mapRewriteValue);
-    SmallVector<Value, 4> args(1, mapRewriteValue(rewriter.root()));
     args.append(mappedArgs.begin(), mappedArgs.end());
     builder.create<pdl_interp::ApplyRewriteOp>(
         rewriter.getLoc(), /*resultTypes=*/TypeRange(), rewriteName, args,
         rewriter.externalConstParamsAttr());
   } else {
     // Otherwise this is a dag rewriter defined using PDL operations.
     for (Operation &rewriteOp : *rewriter.getBody()) {
       llvm::TypeSwitch<Operation *>(&rewriteOp)

mlir/lib/Conversion/PDLToPDLInterp/Predicate.h

enum Kind : unsigned {
   OperandCountQuestion,
   ResultCountAtLeastQuestion,
   ResultCountQuestion,
   EqualToQuestion,
   ConstraintQuestion,
 
   // Answers.
   AttributeAnswer,
-  TrueAnswer,
+  FalseAnswer,
   OperationNameAnswer,
+  TrueAnswer,
   TypeAnswer,
   UnsignedAnswer,
 };
 } // end namespace Predicates
 
 /// Base class for all predicates, used to allow efficient pointer comparison.
 template <typename ConcreteT, typename BaseT, typename Key,
           Predicates::Kind Kind>
struct OperandGroupPosition
   bool isVariadic() const { return std::get<2>(key); }
 };
 
 //===----------------------------------------------------------------------===//
 // OperationPosition
 
 /// An operation position describes an operation node in the IR. Other position
 /// kinds are formed with respect to an operation position.
-struct OperationPosition : public PredicateBase<OperationPosition, Position,
-                                                std::pair<Position *, unsigned>,
-                                                Predicates::OperationPos> {
+struct OperationPosition
+    : public PredicateBase<OperationPosition, Position,
+                           std::tuple<Position *, Optional<unsigned>, unsigned>,
   explicit OperationPosition(const KeyTy &key) : Base(key) {
-    parent = key.first;
+    parent = std::get<0>(key);
+  }
+
+  /// Returns a hash suitable for the given keytype.
+  static llvm::hash_code hashKey(const KeyTy &key) {
+    return llvm::hash_value(key);
   }
 
   /// Gets the root position.
   static OperationPosition *getRoot(StorageUniquer &uniquer) {
-    return Base::get(uniquer, nullptr, 0);
+    return Base::get(uniquer, nullptr, llvm::None, 0);
   }
+
   /// Gets an operation position with the given parent.
-  static OperationPosition *get(StorageUniquer &uniquer, Position *parent) {
-    return Base::get(uniquer, parent, parent->getOperationDepth() + 1);
+  /// The optional operand index indicates a downward traversal from operand to
+  /// the operation.
+  static OperationPosition *get(StorageUniquer &uniquer, Position *parent,
+                                Optional<unsigned> operand = llvm::None) {
+    return Base::get(uniquer, parent, operand, parent->getOperationDepth() + 1);
   }
 
+  /// Returns if this operation position is upward, accepting an input.
+  bool isUpward() const { return std::get<1>(key).hasValue(); }
+
+  /// Returns the operand index for an upward operation position.
+  unsigned getIndex() const { return *std::get<1>(key); }
+
   /// Returns the depth of this position.
-  unsigned getDepth() const { return key.second; }
+  unsigned getDepth() const { return std::get<2>(key); }
 
   /// Returns if this operation position corresponds to the root.
   bool isRoot() const { return getDepth() == 0; }
 };
 
 //===----------------------------------------------------------------------===//
 // ResultPosition
 

 };
 
 /// An Answer representing a boolean `true` value.
 struct TrueAnswer
     : PredicateBase<TrueAnswer, Qualifier, void, Predicates::TrueAnswer> {
   using Base::Base;
 };
 
+/// An Answer representing a boolean 'false' value.
+struct FalseAnswer
+    : PredicateBase<FalseAnswer, Qualifier, void, Predicates::FalseAnswer> {
+  using Base::Base;
+};
+
 /// An Answer representing a `Type` value. The value is stored as either a
 /// TypeAttr, or an ArrayAttr of TypeAttr.
 struct TypeAnswer : public PredicateBase<TypeAnswer, Qualifier, Attribute,
                                          Predicates::TypeAnswer> {
   using Base::Base;
 };
 
 /// An Answer representing an unsigned value.
PredicateUniquer() {
     registerParametricStorageType<ResultGroupPosition>();
     registerParametricStorageType<TypePosition>();
 
     // Register the types of Questions with the uniquer.
     registerParametricStorageType<AttributeAnswer>();
     registerParametricStorageType<OperationNameAnswer>();
     registerParametricStorageType<TypeAnswer>();
     registerParametricStorageType<UnsignedAnswer>();
+    registerSingletonStorageType<FalseAnswer>();
     registerSingletonStorageType<TrueAnswer>();
 
     // Register the types of Answers with the uniquer.
     registerParametricStorageType<ConstraintQuestion>();
     registerParametricStorageType<EqualToQuestion>();
     registerSingletonStorageType<AttributeQuestion>();
     registerSingletonStorageType<IsNotNullQuestion>();
     registerSingletonStorageType<OperandCountQuestion>();
public:
 
   /// Returns the parent position defining the value held by the given operand.
   OperationPosition *getOperandDefiningOp(Position *p) {
     assert((isa<OperandPosition, OperandGroupPosition>(p)) &&
            "expected operand position");
     return OperationPosition::get(uniquer, p);
   }
 
+  /// Returns the operation position accepting the value at the given position.
+  OperationPosition *getOperandAcceptingOp(Position *p, unsigned operand) {
+    assert((isa<OperandPosition, OperandGroupPosition, ResultPosition,
+                ResultGroupPosition>(p)) &&
+           "expected result position");
+    return OperationPosition::get(uniquer, p, operand);
+  }
+
   /// Returns an attribute position for an attribute of the given operation.
   Position *getAttribute(OperationPosition *p, StringRef name) {
     return AttributePosition::get(uniquer, p, Identifier::get(name, ctx));
   }
 
   /// Returns an operand position for an operand of the given operation.
   Position *getOperand(OperationPosition *p, unsigned operand) {
     return OperandPosition::get(uniquer, p, operand);
public:
   using Predicate = std::pair<Qualifier *, Qualifier *>;
 
   /// Create a predicate comparing an attribute to a known value.
   Predicate getAttributeConstraint(Attribute attr) {
     return {AttributeQuestion::get(uniquer),
             AttributeAnswer::get(uniquer, attr)};
   }
 
-  /// Create a predicate comparing two values.
+  /// Create a predicate checking if two values are equal.
   Predicate getEqualTo(Position *pos) {
     return {EqualToQuestion::get(uniquer, pos), TrueAnswer::get(uniquer)};
   }
 
+  /// Create a predicate checking if two values are not equal.
+  Predicate getNotEqualTo(Position *pos) {
+    return {EqualToQuestion::get(uniquer, pos), FalseAnswer::get(uniquer)};
+  }
+
   /// Create a predicate that applies a generic constraint.
   Predicate getConstraint(StringRef name, ArrayRef<Position *> pos,
                           Attribute params) {
     return {
         ConstraintQuestion::get(uniquer, std::make_tuple(name, pos, params)),
         TrueAnswer::get(uniquer)};
   }
 

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.h

 
 //===----------------------------------------------------------------------===//
 // SuccessNode
 
 /// A SuccessNode denotes that a given high level pattern has successfully been
 /// matched. This does not terminate the matcher, as there may be multiple
 /// successful matches.
 struct SuccessNode : public MatcherNode {
-  explicit SuccessNode(pdl::PatternOp pattern,
+  explicit SuccessNode(pdl::PatternOp pattern, Value root,
                        std::unique_ptr<MatcherNode> failureNode);
 
   /// Returns if the given matcher node is an instance of this class, used to
   /// support type casting.
   static bool classof(const MatcherNode *node) {
     return node->getMatcherTypeID() == TypeID::get<SuccessNode>();
   }
 
   /// Return the high level pattern operation that is matched with this node.
   pdl::PatternOp getPattern() const { return pattern; }
 
+  /// Return the chosen root of the pattern.
+  Value getRoot() const { return root; }
+
 private:
   /// The high level pattern operation that was successfully matched with this
   /// node.
   pdl::PatternOp pattern;
+
+  /// The chosen root of the pattern.
+  Value root;
 };
 
 //===----------------------------------------------------------------------===//
 // SwitchNode
 
 /// A SwitchNode denotes a question with multiple potential results. These nodes
 /// branch to a specific node based on the result of the question.
 struct SwitchNode : public MatcherNode {

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

 //===- PredicateTree.cpp - Predicate tree merging -------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "PredicateTree.h"
+#include "mlir/Conversion/PDLToPDLInterp/RootOrdering.h"
 #include "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/Dialect/PDLInterp/IR/PDLInterp.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
+#include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/TypeSwitch.h"
+#include "llvm/Support/Debug.h"
+#include <queue>
+
+#define DEBUG_TYPE "pdl-predicate-tree"
 
 using namespace mlir;
 using namespace mlir::pdl_to_pdl_interp;
 
 //===----------------------------------------------------------------------===//
 // Predicate List Building
 //===----------------------------------------------------------------------===//
 
TypeSwitch<Operation *>(val.getDefiningOp())
         getTreePredicates(predList, op.parent(), builder, inputs,
                           (Position *)parentPos);
       });
 }
 
 static void getTreePredicates(std::vector<PositionalPredicate> &predList,
                               Value val, PredicateBuilder &builder,
                               DenseMap<Value, Position *> &inputs,
-                              OperationPosition *pos) {
+                              OperationPosition *pos,
+                              Optional<unsigned> ignoreOperand = llvm::None) {
   assert(val.getType().isa<pdl::OperationType>() && "expected operation");
   pdl::OperationOp op = cast<pdl::OperationOp>(val.getDefiningOp());
   OperationPosition *opPos = cast<OperationPosition>(pos);
 
   // Ensure getDefiningOp returns a non-null operation.
   if (!opPos->isRoot())
     predList.emplace_back(pos, builder.getIsNotNull());
 
if (operands.size() == 1 && operands[0].getType().isa<pdl::RangeType>()) {
     getTreePredicates(predList, operands.front(), builder, inputs,
                       builder.getAllOperands(opPos));
   } else {
     bool foundVariableLength = false;
     for (auto operandIt : llvm::enumerate(operands)) {
       bool isVariadic = operandIt.value().getType().isa<pdl::RangeType>();
       foundVariableLength |= isVariadic;
 
+      // Ignore the specified operand, usually because this position was
+      // visited in an upward traversal via a nondeterministic choice.
+      if (ignoreOperand && *ignoreOperand == operandIt.index())
+        continue;
+
       Position *pos =
           foundVariableLength
               ? builder.getOperandGroup(opPos, operandIt.index(), isVariadic)
               : builder.getOperand(opPos, operandIt.index());
       getTreePredicates(predList, operandIt.value(), builder, inputs, pos);
     }
   }
   /// Results.
TypeSwitch<Operation *>(&op)
           getConstraintPredicates(constraintOp, predList, builder, inputs);
         })
         .Case<pdl::ResultOp, pdl::ResultsOp>([&](auto resultOp) {
           getResultPredicates(resultOp, predList, builder, inputs);
         });
   }
 }
 
+namespace {
+
+/// An op accepting a value at an optional index.
+struct OpIndex {
+  Value parent;
+  Optional<unsigned> index;
+};
+
+/// The parent and operand index of each operation for each root, stored
+/// as a nested map [root][operation].
+using ParentMaps = DenseMap<Value, DenseMap<Value, OpIndex>>;
+
+} // namespace
+
+/// Given a pattern, determines the set of roots present in this pattern.
+/// These are the operations whose results are not consumed by other operations.
+static SmallVector<Value> detectRoots(pdl::PatternOp pattern) {
+  // First, collect all the operations that are used as operands
+  // to other operations. These are not roots by default.
+  DenseSet<Value> used;
+  for (auto operationOp : pattern.body().getOps<pdl::OperationOp>())
+    for (Value operand : operationOp.operands())
+      TypeSwitch<Operation *>(operand.getDefiningOp())
+          .Case<pdl::ResultOp, pdl::ResultsOp>([&used](auto resultOp) {
+            used.insert(resultOp.parent());
+          });
+
+  // Remove the specified root from the use set, so that we can
+  // always select it as a root, even if it is used by other operations.
+  if (Value root = pattern.getRewriter().root())
+    used.erase(root);
+
+  // Finally, collect all the unused operations.
+  SmallVector<Value> roots;
+  for (Value operationOp : pattern.body().getOps<pdl::OperationOp>())
+    if (!used.contains(operationOp))
+      roots.push_back(operationOp);
+
+  return roots;
+}
+
+/// Given a list of candidate roots, builds the cost graph for connecting them.
+/// The graph is formed by traversing the DAG of operations starting from each
+/// root and marking the depth of each connector value (operand). Then we join
+/// the candidate roots based on the common connector values, taking the one
+/// with the minimum depth. Along the way, we compute, for each candidate root,
+/// a mapping from each operation (in the DAG underneath this root) to its
+/// parent operation and the corresponding operand index.
+static void buildCostGraph(ArrayRef<Value> roots, RootOrderingGraph &graph,
+                           ParentMaps &parentMaps) {
+
+  // A root of a value and its depth (distance from root to the value).
+  struct RootDepth {
+    Value root;
+    unsigned depth = 0;
+  };
+
+  // Map from candidate connector values to their roots and depths. Using a
+  // small vector with 1 entry because most values belong to a single root.
+  llvm::MapVector<Value, SmallVector<RootDepth, 1>> connectorsRootsDepths;
+
+  // Perform a breadth-first traversal of the op DAG rooted at each root.
+  for (Value root : roots) {
+    // The entry of a queue. The entry consists of the following items:
+    // * the value in the DAG underneath the root;
+    // * the parent of the value;
+    // * the operand index of the value in its parent;
+    // * the depth of the visited value.
+    struct Entry {
+      Value value;
+      Value parent;
+      Optional<unsigned> index;
+      unsigned depth;
+      Entry(Value value, Value parent, Optional<unsigned> index, unsigned depth)
+          : value(value), parent(parent), index(index), depth(depth) {}
+    };
+
+    // The queue of visited values. A value may be present multiple times in
+    // the queue, for multiple parents. We only accept the first occurrence,
+    // which is guaranteed to have the lowest depth.
+    std::queue<Entry> toVisit;
+    toVisit.emplace(root, Value(), 0, 0);
+
+    // The map from value to its parent for the current root.
+    DenseMap<Value, OpIndex> &parentMap = parentMaps[root];
+
+    while (!toVisit.empty()) {
+      Entry entry = toVisit.front();
+      toVisit.pop();
+      // Skip if already visited.
+      if (!parentMap.insert({entry.value, {entry.parent, entry.index}}).second)
+        continue;
+
+      // Mark the root and depth of the value.
+      connectorsRootsDepths[entry.value].push_back({root, entry.depth});
+
+      // Traverse the operands of an operation and result ops.
+      // We intentionally do not traverse attributes and types, because those
+      // are expensive to join on.
+      TypeSwitch<Operation *>(entry.value.getDefiningOp())
+          .Case<pdl::OperationOp>([&](auto operationOp) {
+            for (auto p : llvm::enumerate(operationOp.operands()))
+              toVisit.emplace(p.value(), entry.value, p.index(),
+                              entry.depth + 1);
+          })
+          .Case<pdl::ResultOp, pdl::ResultsOp>([&](auto resultOp) {
+            toVisit.emplace(resultOp.parent(), entry.value, resultOp.index(),
+                            entry.depth);
+          });
+    }
+  }
+
+  // Now build the cost graph.
+  // This is simply a minimum over all depths for the target root.
+  unsigned nextID = 0;
+  for (const auto &connectorRootsDepths : connectorsRootsDepths) {
+    Value value = connectorRootsDepths.first;
+    ArrayRef<RootDepth> rootsDepths = connectorRootsDepths.second;
+    // If there is only one root for this value, this will not trigger
+    // any edges in the cost graph (a perf optimization).
+    if (rootsDepths.size() == 1)
+      continue;
+
+    for (const RootDepth &p : rootsDepths) {
+      for (const RootDepth &q : rootsDepths) {
+        if (&p == &q)
+          continue;
+        // insert or retrieve the property of edge from p to q
+        RootOrderingCost &cost = graph[q.root][p.root];
+        if (!cost.connector /* new edge */ || cost.cost.first > q.depth) {
+          if (!cost.connector)
+            cost.cost.second = nextID++;
+          cost.cost.first = q.depth;
+          cost.connector = value;
+        }
+      }
+    }
+  }
+
+  assert((llvm::hasSingleElement(roots) || graph.size() == roots.size()) &&
+         "the pattern contains a candidate root disconnected from the others");
+}
+
 /// Given a pattern operation, build the set of matcher predicates necessary to
 /// match this pattern.
-static void buildPredicateList(pdl::PatternOp pattern,
+static Value buildPredicateList(pdl::PatternOp pattern,
                                 PredicateBuilder &builder,
                                 std::vector<PositionalPredicate> &predList,
                                 DenseMap<Value, Position *> &valueToPosition) {
-  getTreePredicates(predList, pattern.getRewriter().root(), builder,
-                    valueToPosition, builder.getRoot());
+  SmallVector<Value> roots = detectRoots(pattern);
+
+  // Build the root ordering graph and compute the parent maps.
+  RootOrderingGraph graph;
+  ParentMaps parentMaps;
+  buildCostGraph(roots, graph, parentMaps);
+  LLVM_DEBUG({
+    llvm::dbgs() << "Graph:\n";
+    for (auto &target : graph) {
+      llvm::dbgs() << "  * " << target.first << "\n";
+      for (auto &source : target.second) {
+        RootOrderingCost c = source.second;
+        llvm::dbgs() << "      <- " << source.first << ": " << c.cost.first
+                     << ":" << c.cost.second << " via " << c.connector.getLoc()
+                     << "\n";
+      }
+    }
+  });
+
+  // Solve the optimal branching problem for each candidate root, or use the
+  // provided one.
+  Value bestRoot = pattern.getRewriter().root();
+  OptimalBranching::EdgeList bestEdges;
+  if (!bestRoot) {
+    unsigned bestCost = 0;
+    LLVM_DEBUG(llvm::dbgs() << "Candidate roots:\n");
+    for (Value root : roots) {
+      OptimalBranching solver(graph, root);
+      unsigned cost = solver.solve();
+      LLVM_DEBUG(llvm::dbgs() << "  * " << root << ": " << cost << "\n");
+      if (!bestRoot || bestCost > cost) {
+        bestCost = cost;
+        bestRoot = root;
+        bestEdges = solver.preOrderTraversal(roots);
+      }
+    }
+  } else {
+    OptimalBranching solver(graph, bestRoot);
+    solver.solve();
+    bestEdges = solver.preOrderTraversal(roots);
+  }
+
+  LLVM_DEBUG(llvm::dbgs() << "Calling key getTreePredicates:\n");
+  LLVM_DEBUG(llvm::dbgs() << "  * Value: " << bestRoot << "\n");
+
+  // The best root is the starting point for the traversal. Get the tree
+  // predicates for the DAG rooted at bestRoot.
+  getTreePredicates(predList, bestRoot, builder, valueToPosition,
+                    builder.getRoot());
+
+  // Traverse the selected optimal branching. For all edges in order, traverse
+  // up starting from the connector, until the candidate root is reached, and
+  // call getTreePredicates at every node along the way.
+  for (const std::pair<Value, Value> &edge : bestEdges) {
+    Value target = edge.first;
+    Value source = edge.second;
+
+    // Check if we already visited the target root. This happens in two cases:
+    // 1) the initial root (bestRoot);
+    // 2) a root that is dominated by (contained in the subtree rooted at) an
+    //    already visited root.
+    if (valueToPosition.count(target))
+      continue;
+
+    // Determine the connector.
+    Value connector = graph[target][source].connector;
+    assert(connector && "Invalid edge");
+    LLVM_DEBUG(llvm::dbgs() << "  * Connector: " << connector.getLoc() << "\n");
+    DenseMap<Value, OpIndex> parentMap = parentMaps.lookup(target);
+    Position *pos = valueToPosition.lookup(connector);
+    assert(pos && "The value has not been traversed yet");
+    bool first = true;
+
+    // Traverse from the connector upwards towards the target root.
+    for (Value value = connector; value != target;) {
+      OpIndex opIndex = parentMap.lookup(value);
+      assert(opIndex.parent && "Missing parent");
+      value = opIndex.parent;
+
+      // Visit the node.
+      TypeSwitch<Operation *>(value.getDefiningOp())
+          .Case<pdl::OperationOp>([&](auto operationOp) {
+            LLVM_DEBUG(llvm::dbgs() << "  * Value: " << value << "\n");
+            OperationPosition *opPos =
+                builder.getOperandAcceptingOp(pos, *opIndex.index);
+
+            // Guard against traversing back to where we came from.
+            if (first) {
+              Position *parent = pos->getParent();
+              predList.emplace_back(opPos, builder.getNotEqualTo(parent));
+              first = false;
+            }
+
+            // Guard against duplicate upward visits. These are not possible,
+            // because if this value was already visited, it would have been
+            // cheaper to start the traversal at this value rather than at the
+            // `connector`, violating the optimality of our spanning tree.
+            bool inserted = valueToPosition.try_emplace(value, opPos).second;
+            assert(inserted && "Duplicate upward visit");
+
+            // Obtain the tree predicates at the current value
+            getTreePredicates(predList, value, builder, valueToPosition, opPos,
+                              opIndex.index);
+
+            // Update the position
+            pos = opPos;
+          })
+          .Case<pdl::ResultOp>([&](auto resultOp) {
+            // Traverse up an individual result.
+            auto *opPos = dyn_cast<OperationPosition>(pos);
+            assert(opPos && "Operations and results must be interleaved");
+            pos = builder.getResult(opPos, *opIndex.index);
+          })
+          .Case<pdl::ResultsOp>([&](auto resultOp) {
+            // Traverse up a group of results.
+            auto *opPos = dyn_cast<OperationPosition>(pos);
+            assert(opPos && "Operations and results must be interleaved");
+            if (opIndex.index)
+              pos = builder.getResultGroup(opPos, opIndex.index,
+                                           /*isVariadic=*/false);
+            else
+              pos = builder.getAllResults(opPos);
+          });
+    }
+  }
+
   getNonTreePredicates(pattern, predList, builder, valueToPosition);
+
+  return bestRoot;
 }
 
 //===----------------------------------------------------------------------===//
 // Pattern Predicate Tree Merging
 //===----------------------------------------------------------------------===//
 
 namespace {
 
struct OrderedPredicateDenseInfo {
   static unsigned getHashValue(const OrderedPredicate &p) {
     return llvm::hash_combine(p.position, p.question);
   }
 };
 
 /// This class wraps a set of ordered predicates that are used within a specific
 /// pattern operation.
 struct OrderedPredicateList {
-  OrderedPredicateList(pdl::PatternOp pattern) : pattern(pattern) {}
+  OrderedPredicateList(pdl::PatternOp pattern, Value root)
+      : pattern(pattern), root(root) {}
 
   pdl::PatternOp pattern;
+  Value root;
   DenseSet<OrderedPredicate *> predicates;
 };
 } // end anonymous namespace
 
 /// Returns true if the given matcher refers to the same predicate as the given
 /// ordered predicate. This means that the position and questions of the two
 /// match.
 static bool isSamePredicate(MatcherNode *node, OrderedPredicate *predicate) {

 /// and then either diverge from the CFG or reach the end of a branch and start
 /// creating new nodes.
 static void propagatePattern(std::unique_ptr<MatcherNode> &node,
                              OrderedPredicateList &list,
                              std::vector<OrderedPredicate *>::iterator current,
                              std::vector<OrderedPredicate *>::iterator end) {
   if (current == end) {
     // We've hit the end of a pattern, so create a successful result node.
-    node = std::make_unique<SuccessNode>(list.pattern, std::move(node));
+    node =
+        std::make_unique<SuccessNode>(list.pattern, list.root, std::move(node));
 
     // If the pattern doesn't contain this predicate, ignore it.
   } else if (list.predicates.find(*current) == list.predicates.end()) {
     propagatePattern(node, list, std::next(current), end);
 
     // If the current matcher node is invalid, create a new one for this
     // position and continue propagation.
   } else if (!node) {
static void insertExitNode(std::unique_ptr<MatcherNode> *root) {
   *root = std::make_unique<ExitNode>();
 }
 
 /// Given a module containing PDL pattern operations, generate a matcher tree
 /// using the patterns within the given module and return the root matcher node.
 std::unique_ptr<MatcherNode>
 MatcherNode::generateMatcherTree(ModuleOp module, PredicateBuilder &builder,
                                  DenseMap<Value, Position *> &valueToPosition) {
-  // Collect the set of predicates contained within the pattern operations of
-  // the module.
-  SmallVector<std::pair<pdl::PatternOp, std::vector<PositionalPredicate>>, 16>
-      patternsAndPredicates;
+  // The set of predicates contained within the pattern operations of the
+  // module.
+  struct PatternPredicates {
+    /// A pattern.
+    pdl::PatternOp pattern;
+
+    /// A root of the pattern chosen among the candidate roots in pdl.rewrite.
+    Value root;
+
+    /// The extracted predicates for this pattern and root.
+    std::vector<PositionalPredicate> predicates;
+
+    PatternPredicates(pdl::PatternOp pattern, Value root,
+                      std::vector<PositionalPredicate> predicates)
+        : pattern(pattern), root(root), predicates(std::move(predicates)) {}
+  };
+
+  SmallVector<PatternPredicates, 16> patternsAndPredicates;
   for (pdl::PatternOp pattern : module.getOps<pdl::PatternOp>()) {
     std::vector<PositionalPredicate> predicateList;
+    Value root =
         buildPredicateList(pattern, builder, predicateList, valueToPosition);
-    patternsAndPredicates.emplace_back(pattern, std::move(predicateList));
+    patternsAndPredicates.emplace_back(pattern, root, std::move(predicateList));
   }
 
   // Associate a pattern result with each unique predicate.
   DenseSet<OrderedPredicate, OrderedPredicateDenseInfo> uniqued;
   for (auto &patternAndPredList : patternsAndPredicates) {
-    for (auto &predicate : patternAndPredList.second) {
+    for (auto &predicate : patternAndPredList.predicates) {
       auto it = uniqued.insert(predicate);
-      it.first->patternToAnswer.try_emplace(patternAndPredList.first,
+      it.first->patternToAnswer.try_emplace(patternAndPredList.pattern,
                                             predicate.answer);
     }
   }
 
   // Associate each pattern to a set of its ordered predicates for later lookup.
   std::vector<OrderedPredicateList> lists;
   lists.reserve(patternsAndPredicates.size());
   for (auto &patternAndPredList : patternsAndPredicates) {
-    OrderedPredicateList list(patternAndPredList.first);
-    for (auto &predicate : patternAndPredList.second) {
+    OrderedPredicateList list(patternAndPredList.pattern,
+                              patternAndPredList.root);
+    for (auto &predicate : patternAndPredList.predicates) {
       OrderedPredicate *orderedPredicate = &*uniqued.find(predicate);
       list.predicates.insert(orderedPredicate);
 
       // Increment the primary sum for each reference to a particular predicate.
       ++orderedPredicate->primary;
     }
     lists.push_back(std::move(list));
   }
BoolNode::BoolNode(Position *position, Qualifier *question, Qualifier *answer,
     : MatcherNode(TypeID::get<BoolNode>(), position, question,
                   std::move(failureNode)),
       answer(answer), successNode(std::move(successNode)) {}
 
 //===----------------------------------------------------------------------===//
 // SuccessNode
 //===----------------------------------------------------------------------===//
 
-SuccessNode::SuccessNode(pdl::PatternOp pattern,
+SuccessNode::SuccessNode(pdl::PatternOp pattern, Value root,
                          std::unique_ptr<MatcherNode> failureNode)
     : MatcherNode(TypeID::get<SuccessNode>(), /*position=*/nullptr,
                   /*question=*/nullptr, std::move(failureNode)),
-      pattern(pattern) {}
+      pattern(pattern), root(root) {}
 
 //===----------------------------------------------------------------------===//
 // SwitchNode
 //===----------------------------------------------------------------------===//
 
 SwitchNode::SwitchNode(Position *position, Qualifier *question)
     : MatcherNode(TypeID::get<SwitchNode>(), position, question) {}

mlir/lib/Dialect/PDL/IR/PDL.cpp

 //===----------------------------------------------------------------------===//
 
 /// Returns true if the given operation is used by a "binding" pdl operation
 /// within the main matcher body of a `pdl.pattern`.
 static bool hasBindingUseInMatcher(Operation *op, Block *matcherBlock) {
   for (OpOperand &use : op->getUses()) {
     Operation *user = use.getOwner();
     if (user->getBlock() != matcherBlock)
       continue;
-    if (isa<AttributeOp, OperandOp, OperandsOp, OperationOp>(user))
-      return true;
-    // Only the first operand of RewriteOp may be bound to, i.e. the root
-    // operation of the pattern.
-    if (isa<RewriteOp>(user) && use.getOperandNumber() == 0)
+    if (isa<AttributeOp, OperandOp, OperandsOp, OperationOp, RewriteOp>(user))
       return true;
     // A result by itself is not binding, it must also be bound.
     if (isa<ResultOp, ResultsOp>(user) &&
         hasBindingUseInMatcher(user, matcherBlock))
       return true;
   }
   return false;
 }
if (!isa_and_nonnull<PDLDialect>(op->getDialect())) {
       return WalkResult::interrupt();
     }
     return WalkResult::advance();
   });
   return failure(result.wasInterrupted());
 }
 
 void PatternOp::build(OpBuilder &builder, OperationState &state,
-                      Optional<StringRef> rootKind, Optional<uint16_t> benefit,
-                      Optional<StringRef> name) {
-  build(builder, state,
-        rootKind ? builder.getStringAttr(*rootKind) : StringAttr(),
-        builder.getI16IntegerAttr(benefit ? *benefit : 0),
+                      Optional<uint16_t> benefit, Optional<StringRef> name) {
+  build(builder, state, builder.getI16IntegerAttr(benefit ? *benefit : 0),
         name ? builder.getStringAttr(*name) : StringAttr());
   state.regions[0]->emplaceBlock();
 }
 
 /// Returns the rewrite operation of this pattern.
 RewriteOp PatternOp::getRewriter() {
   return cast<RewriteOp>(body().front().getTerminator());
 }
 
-/// Return the root operation kind that this pattern matches, or None if
-/// there isn't a specific root.
-Optional<StringRef> PatternOp::getRootKind() {
-  OperationOp rootOp = cast<OperationOp>(getRewriter().root().getDefiningOp());
-  return rootOp.name();
-}
-
 //===----------------------------------------------------------------------===//
 // pdl::ReplaceOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ReplaceOp op) {
   if (op.replOperation() && !op.replValues().empty())
     return op.emitOpError() << "expected no replacement values to be provided"
                                " when the replacement operation is present";

mlir/test/Conversion/PDLToPDLInterp/pdl-to-pdl-interp-matcher.mlir

module @predicate_ordering  {
   }
 
   pdl.pattern : benefit(1) {
     %resultType = pdl.type
     %apply = pdl.operation -> (%resultType : !pdl.type)
     pdl.rewrite %apply with "rewriter"
   }
 }
+
+
+// -----
+
+// CHECK-LABEL: module @multi_root
+module @multi_root  {
+  // Check the lowering of a simple two-root pattern.
+  // This checks that we correctly generate the pdl_interp.choose_op operation
+  // and tie the break between %root1 and %root2 in favor of %root1.
+
+  // CHECK: func @matcher(%[[ROOT1:.*]]: !pdl.operation)
+  // CHECK-DAG: %[[VAL1:.*]] = pdl_interp.get_operand 0 of %[[ROOT1]]
+  // CHECK-DAG: %[[OP1:.*]] = pdl_interp.get_defining_op of %[[VAL1]]
+  // CHECK-DAG: %[[OPS:.*]] = pdl_interp.get_users of %[[VAL1]] : !pdl.value at 0
+  // CHECK-DAG: pdl_interp.foreach %[[ROOT2:.*]] : !pdl.operation in %[[OPS]]
+  // CHECK-DAG: %[[VAL2:.*]] = pdl_interp.get_operand 1 of %[[ROOT2]]
+  // CHECK-DAG: %[[OP2:.*]] = pdl_interp.get_defining_op of %[[VAL2]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[OP1]] : !pdl.operation -> ^{{.*}}, ^[[FAIL:.*]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[OP2]] : !pdl.operation
+  // CHECK-DAG: pdl_interp.is_not_null %[[VAL1]] : !pdl.value
+  // CHECK-DAG: pdl_interp.is_not_null %[[VAL2]] : !pdl.value
+  // CHECK-DAG: pdl_interp.is_not_null %[[ROOT2]] : !pdl.operation
+  // CHECK-DAG: pdl_interp.are_equal %[[ROOT2]], %[[ROOT1]] : !pdl.operation -> ^[[FAIL]]
+
+  pdl.pattern @rewrite_multi_root : benefit(1) {
+    %input1 = pdl.operand
+    %input2 = pdl.operand
+    %type = pdl.type
+    %op1 = pdl.operation(%input1 : !pdl.value) -> (%type : !pdl.type)
+    %val1 = pdl.result 0 of %op1
+    %root1 = pdl.operation(%val1 : !pdl.value)
+    %op2 = pdl.operation(%input2 : !pdl.value) -> (%type : !pdl.type)
+    %val2 = pdl.result 0 of %op2
+    %root2 = pdl.operation(%val1, %val2 : !pdl.value, !pdl.value)
+    pdl.rewrite %root1 with "rewriter"(%root2 : !pdl.operation)
+  }
+}
+
+
+// -----
+
+// CHECK-LABEL: module @overlapping_roots
+module @overlapping_roots  {
+  // Check the lowering of a degenerate two-root pattern, where one root
+  // is in the subtree rooted at another.
+
+  // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
+  // CHECK-DAG: %[[VAL:.*]] = pdl_interp.get_operand 0 of %[[ROOT]]
+  // CHECK-DAG: %[[OP:.*]] = pdl_interp.get_defining_op of %[[VAL]]
+  // CHECK-DAG: %[[INPUT1:.*]] = pdl_interp.get_operand 0 of %[[OP]]
+  // CHECK-DAG: %[[INPUT2:.*]] = pdl_interp.get_operand 1 of %[[OP]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[VAL]] : !pdl.value
+  // CHECK-DAG: pdl_interp.is_not_null %[[OP]] : !pdl.operation
+  // CHECK-DAG: pdl_interp.is_not_null %[[INPUT1]] : !pdl.value
+  // CHECK-DAG: pdl_interp.is_not_null %[[INPUT2]] : !pdl.value
+
+  pdl.pattern @rewrite_overlapping_roots : benefit(1) {
+    %input1 = pdl.operand
+    %input2 = pdl.operand
+    %type = pdl.type
+    %op = pdl.operation(%input1, %input2 : !pdl.value, !pdl.value) -> (%type : !pdl.type)
+    %val = pdl.result 0 of %op
+    %root = pdl.operation(%val : !pdl.value)
+    pdl.rewrite with "rewriter"(%root : !pdl.operation)
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module @force_overlapped_root
+module @force_overlapped_root  {
+  // Check the lowering of a degenerate two-root pattern, where one root
+  // is in the subtree rooted at another, and we are forced to use this
+  // root as the root of the search tree.
+
+  // CHECK: func @matcher(%[[OP:.*]]: !pdl.operation)
+  // CHECK-DAG: %[[VAL:.*]] = pdl_interp.get_result 0 of %[[OP]]
+  // CHECK-DAG: pdl_interp.check_operand_count of %[[OP]] is 2
+  // CHECK-DAG: pdl_interp.check_result_count of %[[OP]] is 1
+  // CHECK-DAG: %[[INPUT2:.*]] = pdl_interp.get_operand 1 of %[[OP]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[INPUT2]] : !pdl.value
+  // CHECK-DAG: %[[INPUT1:.*]] = pdl_interp.get_operand 0 of %[[OP]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[INPUT1]] : !pdl.value
+  // CHECK-DAG: %[[ROOTS:.*]] = pdl_interp.get_users of %[[VAL]] : !pdl.value at 0
+  // CHECK-DAG: pdl_interp.foreach %[[ROOT:.*]] : !pdl.operation in %[[ROOTS]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[ROOT]] : !pdl.operation
+  // CHECK-DAG: pdl_interp.check_operand_count of %[[ROOT]] is 1
+
+  pdl.pattern @rewrite_forced_overlapped_root : benefit(1) {
+    %input1 = pdl.operand
+    %input2 = pdl.operand
+    %type = pdl.type
+    %op = pdl.operation(%input1, %input2 : !pdl.value, !pdl.value) -> (%type : !pdl.type)
+    %val = pdl.result 0 of %op
+    %root = pdl.operation(%val : !pdl.value)
+    pdl.rewrite %op with "rewriter"(%root : !pdl.operation)
+  }
+}

mlir/test/Dialect/PDL/invalid.mlir

 //===----------------------------------------------------------------------===//
 // pdl::RewriteOp
 //===----------------------------------------------------------------------===//
 
 pdl.pattern : benefit(1) {
   %op = pdl.operation "foo.op"
 
   // expected-error@below {{expected rewrite region to be non-empty if external name is not specified}}
-  "pdl.rewrite"(%op) ({}) : (!pdl.operation) -> ()
+  "pdl.rewrite"(%op) ({}) {
+    operand_segment_sizes = dense<[1,0]> : vector<2xi32>
+  } : (!pdl.operation) -> ()
 }
 
 // -----
 
 pdl.pattern : benefit(1) {
   %op = pdl.operation "foo.op"
 
   // expected-error@below {{expected no external arguments when the rewrite is specified inline}}
   "pdl.rewrite"(%op, %op) ({
     ^bb1:
-  }) : (!pdl.operation, !pdl.operation) -> ()
+  }) {
+    operand_segment_sizes = dense<1> : vector<2xi32>
+  }: (!pdl.operation, !pdl.operation) -> ()
 }
 
 // -----
 
 pdl.pattern : benefit(1) {
   %op = pdl.operation "foo.op"
 
   // expected-error@below {{expected no external constant parameters when the rewrite is specified inline}}
   "pdl.rewrite"(%op) ({
     ^bb1:
-  }) {externalConstParams = []} : (!pdl.operation) -> ()
+  }) {
+    operand_segment_sizes = dense<[1,0]> : vector<2xi32>,
+    externalConstParams = []} : (!pdl.operation) -> ()
 }
 
 // -----
 
 pdl.pattern : benefit(1) {
   %op = pdl.operation "foo.op"
 
   // expected-error@below {{expected rewrite region to be empty when rewrite is external}}
   "pdl.rewrite"(%op) ({
     ^bb1:
-  }) {name = "foo"} : (!pdl.operation) -> ()
+  }) {
+    name = "foo",
+    operand_segment_sizes = dense<[1,0]> : vector<2xi32>
+  } : (!pdl.operation) -> ()
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // pdl::TypeOp
 //===----------------------------------------------------------------------===//
 

mlir/test/Dialect/PDL/ops.mlir

 pdl.pattern @rewrite_with_args_and_params : benefit(1) {
   %input = pdl.operand
   %root = pdl.operation(%input : !pdl.value)
   pdl.rewrite %root with "rewriter"["I am param"](%input : !pdl.value)
 }
 
 // -----
 
+pdl.pattern @rewrite_multi_root_optimal : benefit(2) {
+  %input1 = pdl.operand
+  %input2 = pdl.operand
+  %type = pdl.type
+  %op1 = pdl.operation(%input1 : !pdl.value) -> (%type : !pdl.type)
+  %val1 = pdl.result 0 of %op1
+  %root1 = pdl.operation(%val1 : !pdl.value)
+  %op2 = pdl.operation(%input2 : !pdl.value) -> (%type : !pdl.type)
+  %val2 = pdl.result 0 of %op2
+  %root2 = pdl.operation(%val1, %val2 : !pdl.value, !pdl.value)
+  pdl.rewrite with "rewriter"["I am param"](%root1, %root2 : !pdl.operation, !pdl.operation)
+}
+
+// -----
+
+pdl.pattern @rewrite_multi_root_forced : benefit(2) {
+  %input1 = pdl.operand
+  %input2 = pdl.operand
+  %type = pdl.type
+  %op1 = pdl.operation(%input1 : !pdl.value) -> (%type : !pdl.type)
+  %val1 = pdl.result 0 of %op1
+  %root1 = pdl.operation(%val1 : !pdl.value)
+  %op2 = pdl.operation(%input2 : !pdl.value) -> (%type : !pdl.type)
+  %val2 = pdl.result 0 of %op2
+  %root2 = pdl.operation(%val1, %val2 : !pdl.value, !pdl.value)
+  pdl.rewrite %root1 with "rewriter"["I am param"](%root2 : !pdl.operation)
+}
+
+// -----
+
 // Check that the result type of an operation within a rewrite can be inferred
 // from a pdl.replace.
 pdl.pattern @infer_type_from_operation_replace : benefit(1) {
   %type1 = pdl.type : i32
   %type2 = pdl.type
   %root = pdl.operation -> (%type1, %type2 : !pdl.type, !pdl.type)
   pdl.rewrite %root {
     %type3 = pdl.type