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?

622–623

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).

548–549

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

171

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

37–38

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.

548–549

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–594

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(...)?

337

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

338–342

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

394–397

nit: Use camelCase for variable names.

466

Is this still open?

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

171

unresolved?

336–341

384–391

443

593

643

nit; Assert messages should start with lower case

646

Is this guaranteed to be non-variadic?

797–810

[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.

337

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

338–342

Thanks, I was not sure what the convention was.

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

171

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

646

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–594

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

608–609

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

338–340

Are there going to be more cases down the line?

384

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

338–340

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

384

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 the starting point in the pattern for the subgraph
+    isomorphism search. Pattern matching will proceed from this node downward
+    (towards the defining operation) or upward (towards the users) until all
+    the operations in the pattern have been matched. If the root is omitted,
+    the pdl_interp lowering will automatically select the best root of the
+    pdl.rewrite among all the 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 *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() && "unable to locate failure block");
   }
 
+  // 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()) {
+      // Standard (downward) traversal which directly follows the defining op.
       value = builder.create<pdl_interp::GetDefiningOpOp>(
           loc, builder.getType<pdl::OperationType>(), parentVal);
       break;
+    }
+
+    // The first operation retrieves the users.
+    value = builder.create<pdl_interp::GetUsersOp>(loc, parentVal,
+                                                   operationPos->getIndex());
+
+    // The second operation iterates over them.
+    assert(!failureBlockStack.empty() && "expected valid failure block");
+    auto foreach = builder.create<pdl_interp::ForEachOp>(
+        loc, value, failureBlockStack.back(), /*initLoop=*/true);
+    value = foreach.getLoopVariable();
+
+    // The third operation is the continuation to the next iterate.
+    Block *continueBlock = builder.createBlock(&foreach.region());
+    builder.create<pdl_interp::ContinueOp>(loc);
+    failureBlockStack.push_back(continueBlock);
+
+    // Update the current block pointer.
+    currentBlock = &foreach.region().front();
+    builder.setInsertionPointToEnd(currentBlock);
+    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> 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 trueAnswer = isa<TrueAnswer>(answer);
+    builder.create<pdl_interp::AreEqualOp>(loc, val, args.front(),
+                                           trueAnswer ? success : failure,
+                                           trueAnswer ? 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);
+      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>,
+  static constexpr unsigned kDown = std::numeric_limits<unsigned>::max();
+
   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, kDown, 0);
   }
-  /// Gets an operation position with the given parent.
+
+  /// Gets an downward operation position with the given parent.
   static OperationPosition *get(StorageUniquer &uniquer, Position *parent) {
-    return Base::get(uniquer, parent, parent->getOperationDepth() + 1);
+    return Base::get(uniquer, parent, kDown, parent->getOperationDepth() + 1);
   }
 
+  /// Gets an upward operation position with the given parent and operand.
+  static OperationPosition *get(StorageUniquer &uniquer, Position *parent,
+                                Optional<unsigned> operand) {
+    return Base::get(uniquer, parent, operand, parent->getOperationDepth() + 1);
+  }
+
+  /// Returns the operand index for an upward operation position.
+  Optional<unsigned> getIndex() const { return std::get<1>(key); }
+
+  /// Returns if this operation position is upward, accepting an input.
+  bool isUpward() const { return getIndex().getValueOr(0) != kDown; }
+
   /// 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 position of operation using the value at the given index.
+  OperationPosition *getUsersOp(Position *p, Optional<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 "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 an iterative 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 {
+      Entry(Value value, Value parent, Optional<unsigned> index, unsigned depth)
+          : value(value), parent(parent), index(index), depth(depth) {}
+
+      Value value;
+      Value parent;
+      Optional<unsigned> index;
+      unsigned 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) {
+            OperandRange operands = operationOp.operands();
+            // Special case when we pass all the operands in one range.
+            // For those, the index is empty.
+            if (operands.size() == 1 &&
+                operands[0].getType().isa<pdl::RangeType>()) {
+              toVisit.emplace(operands[0], entry.value, llvm::None,
+                              entry.depth + 1);
+              return;
+            }
+
+            // Default case: visit all the operands.
+            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");
+}
+
+/// Visit a node during upward traversal.
+void visitUpward(std::vector<PositionalPredicate> &predList, OpIndex opIndex,
+                 PredicateBuilder &builder,
+                 DenseMap<Value, Position *> &valueToPosition, Position *&pos,
+                 bool &first) {
+  Value value = opIndex.parent;
+  TypeSwitch<Operation *>(value.getDefiningOp())
+      .Case<pdl::OperationOp>([&](auto operationOp) {
+        LLVM_DEBUG(llvm::dbgs() << "  * Value: " << value << "\n");
+        OperationPosition *opPos = builder.getUsersOp(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");
+        bool isVariadic = value.getType().isa<pdl::RangeType>();
+        if (opIndex.index)
+          pos = builder.getResultGroup(opPos, opIndex.index, isVariadic);
+        else
+          pos = builder.getAllResults(opPos);
+      });
+}
+
 /// 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");
+      visitUpward(predList, opIndex, builder, valueToPosition, pos, first);
+      value = opIndex.parent;
+    }
+  }
+
   getNonTreePredicates(pattern, predList, builder, valueToPosition);
+
+  return bestRoot;
 }
 
 //===----------------------------------------------------------------------===//
 // Pattern Predicate Tree Merging
 //===----------------------------------------------------------------------===//
 
 namespace {
 
struct OrderedPredicate {
   /// The first and second order benefit sums.
   /// The primary sum is the number of occurrences of this predicate among all
   /// of the patterns.
   unsigned primary = 0;
   /// The secondary sum is a squared summation of the primary sum of all of the
   /// predicates within each pattern that contains this predicate. This allows
   /// for favoring predicates that are more commonly shared within a pattern, as
   /// opposed to those shared across patterns.
   unsigned secondary = 0;
 
   /// A map between a pattern operation and the answer to the predicate question
   /// within that pattern.
   DenseMap<Operation *, Qualifier *> patternToAnswer;
 
   /// Returns true if this predicate is ordered before `rhs`, based on the cost
   /// model.
   bool operator<(const OrderedPredicate &rhs) const {
     // Sort by:
     // * higher first and secondary order sums
     // * lower depth
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 {
+    PatternPredicates(pdl::PatternOp pattern, Value root,
+                      std::vector<PositionalPredicate> predicates)
+        : pattern(pattern), root(root), predicates(std::move(predicates)) {}
+
+    /// 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;
+  };
+
+  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

 //===- PDL.cpp - Pattern Descriptor Language Dialect ----------------------===//
 //
 // 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 "mlir/Dialect/PDL/IR/PDL.h"
 #include "mlir/Dialect/PDL/IR/PDLOps.h"
 #include "mlir/Dialect/PDL/IR/PDLTypes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
-#include "llvm/ADT/StringSwitch.h"
+#include "llvm/ADT/DenseSet.h"
+#include "llvm/ADT/TypeSwitch.h"
 
 using namespace mlir;
 using namespace mlir::pdl;
 
 #include "mlir/Dialect/PDL/IR/PDLOpsDialect.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // PDLDialect
 //===----------------------------------------------------------------------===//
 
 void PDLDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/PDL/IR/PDLOps.cpp.inc"
       >();
   registerTypes();
 }
 
 //===----------------------------------------------------------------------===//
 // PDL Operations
 //===----------------------------------------------------------------------===//
 
-/// 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)
-      return true;
+/// Returns true if the given operation is used by a "binding" pdl operation.
+static bool hasBindingUse(Operation *op) {
+  for (Operation *user : op->getUsers())
     // A result by itself is not binding, it must also be bound.
-    if (isa<ResultOp, ResultsOp>(user) &&
-        hasBindingUseInMatcher(user, matcherBlock))
+    if (!isa<ResultOp, ResultsOp>(user) || hasBindingUse(user))
       return true;
-  }
   return false;
 }
 
-/// Returns success if the given operation is used by a "binding" pdl operation
-/// within the main matcher body of a `pdl.pattern`. On failure, emits an error
-/// with the given context message.
-static LogicalResult
-verifyHasBindingUseInMatcher(Operation *op,
-                             StringRef bindableContextStr = "`pdl.operation`") {
-  // If the pattern is not a pattern, there is nothing to do.
+/// Returns success if the given operation is not in the main matcher body or
+/// is used by a "binding" operation. On failure, emits an error.
+static LogicalResult verifyHasBindingUse(Operation *op) {
+  // If the parent is not a pattern, there is nothing to do.
   if (!isa<PatternOp>(op->getParentOp()))
     return success();
-  if (hasBindingUseInMatcher(op, op->getBlock()))
+  if (hasBindingUse(op))
     return success();
-  return op->emitOpError()
-         << "expected a bindable (i.e. " << bindableContextStr
-         << ") user when defined in the matcher body of a `pdl.pattern`";
+  return op->emitOpError(
+      "expected a bindable user when defined in the matcher body of a "
+      "`pdl.pattern`");
+}
+
+/// Visits all the pdl.operand(s), pdl.result(s), and pdl.operation(s)
+/// connected to the given operation.
+static void visit(Operation *op, DenseSet<Operation *> &visited) {
+  // If the parent is not a pattern, there is nothing to do.
+  if (!isa<PatternOp>(op->getParentOp()) || isa<RewriteOp>(op))
+    return;
+
+  // Ignore if already visited.
+  if (visited.contains(op))
+    return;
+
+  // Mark as visited.
+  visited.insert(op);
+
+  // Traverse the operands / parent.
+  TypeSwitch<Operation *>(op)
+      .Case<OperationOp>([&visited](auto operation) {
+        for (Value operand : operation.operands())
+          visit(operand.getDefiningOp(), visited);
+      })
+      .Case<ResultOp, ResultsOp>([&visited](auto result) {
+        visit(result.parent().getDefiningOp(), visited);
+      });
+
+  // Traverse the users.
+  for (Operation *user : op->getUsers())
+    visit(user, visited);
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::ApplyNativeConstraintOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(ApplyNativeConstraintOp op) {
   if (op.getNumOperands() == 0)
static LogicalResult verify(AttributeOp op) {
   Value attrType = op.type();
   Optional<Attribute> attrValue = op.value();
 
   if (!attrValue && isa<RewriteOp>(op->getParentOp()))
     return op.emitOpError("expected constant value when specified within a "
                           "`pdl.rewrite`");
   if (attrValue && attrType)
     return op.emitOpError("expected only one of [`type`, `value`] to be set");
-  return verifyHasBindingUseInMatcher(op);
+  return verifyHasBindingUse(op);
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::OperandOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(OperandOp op) {
-  return verifyHasBindingUseInMatcher(op);
-}
+static LogicalResult verify(OperandOp op) { return verifyHasBindingUse(op); }
 
 //===----------------------------------------------------------------------===//
 // pdl::OperandsOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(OperandsOp op) {
-  return verifyHasBindingUseInMatcher(op);
-}
+static LogicalResult verify(OperandsOp op) { return verifyHasBindingUse(op); }
 
 //===----------------------------------------------------------------------===//
 // pdl::OperationOp
 //===----------------------------------------------------------------------===//
 
 static ParseResult parseOperationOpAttributes(
     OpAsmParser &p, SmallVectorImpl<OpAsmParser::OperandType> &attrOperands,
     ArrayAttr &attrNamesAttr) {
static LogicalResult verify(OperationOp op) {
 
   // If the operation is within a rewrite body and doesn't have type inference,
   // ensure that the result types can be resolved.
   if (isWithinRewrite && !op.hasTypeInference()) {
     if (failed(verifyResultTypesAreInferrable(op, op.types())))
       return failure();
   }
 
-  return verifyHasBindingUseInMatcher(op, "`pdl.operation` or `pdl.rewrite`");
+  return verifyHasBindingUse(op);
 }
 
 bool OperationOp::hasTypeInference() {
   Optional<StringRef> opName = name();
   if (!opName)
     return false;
 
   OperationName name(*opName, getContext());
   if (const AbstractOperation *op = name.getAbstractOperation())
     return op->getInterface<InferTypeOpInterface>();
   return false;
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::PatternOp
 //===----------------------------------------------------------------------===//
 
 static LogicalResult verify(PatternOp pattern) {
   Region &body = pattern.body();
-  auto *term = body.front().getTerminator();
-  if (!isa<RewriteOp>(term)) {
+  Operation *term = body.front().getTerminator();
+  auto rewrite_op = dyn_cast<RewriteOp>(term);
+  if (!rewrite_op) {
     return pattern.emitOpError("expected body to terminate with `pdl.rewrite`")
         .attachNote(term->getLoc())
         .append("see terminator defined here");
   }
 
-  // Check that all values defined in the top-level pattern are referenced at
-  // least once in the source tree.
+  // Check that all values defined in the top-level pattern belong to the PDL
+  // dialect.
   WalkResult result = body.walk([&](Operation *op) -> WalkResult {
     if (!isa_and_nonnull<PDLDialect>(op->getDialect())) {
       pattern
           .emitOpError("expected only `pdl` operations within the pattern body")
           .attachNote(op->getLoc())
           .append("see non-`pdl` operation defined here");
       return WalkResult::interrupt();
     }
     return WalkResult::advance();
   });
-  return failure(result.wasInterrupted());
+  if (result.wasInterrupted())
+    return failure();
+
+  // Check that there is at least one operation.
+  if (body.front().getOps<OperationOp>().empty())
+    return pattern.emitOpError(
+        "the pattern must contain at least one `pdl.operation`");
+
+  // Determine if the operations within the pdl.pattern form a connected
+  // component. This is determined by starting the search from the first
+  // operand/result/operation and visiting their users / parents / operands.
+  // We limit our attention to operations that have a user in pdl.rewrite,
+  // those that do not will be detected via other means (expected bindable
+  // user).
+  bool first = true;
+  DenseSet<Operation *> visited;
+  for (Operation &op : body.front()) {
+    // The following are the operations forming the connected component.
+    if (!isa<OperandOp, OperandsOp, ResultOp, ResultsOp, OperationOp>(op))
+      continue;
+
+    // Determine if the operation has a user in `pdl.rewrite`.
+    bool hasUserInRewrite = false;
+    for (Operation *user : op.getUsers()) {
+      Region *region = user->getParentRegion();
+      if (isa<RewriteOp>(user) ||
+          (region && isa<RewriteOp>(region->getParentOp()))) {
+        hasUserInRewrite = true;
+        break;
+      }
+    }
+
+    // If the operation does not have a user in `pdl.rewrite`, ignore it.
+    if (!hasUserInRewrite)
+      continue;
+
+    if (first) {
+      // For the first operation, invoke visit.
+      visit(&op, visited);
+      first = false;
+    } else if (!visited.count(&op)) {
+      // For the subsequent operations, check if already visited.
+      return pattern
+          .emitOpError("the operations must form a connected component")
+          .attachNote(op.getLoc())
+          .append("see a disconnected value / operation here");
+    }
+  }
+
+  return success();
 }
 
 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";
static LogicalResult verify(RewriteOp op) {
 
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::TypeOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(TypeOp op) {
-  return verifyHasBindingUseInMatcher(
-      op, "`pdl.attribute`, `pdl.operand`, or `pdl.operation`");
-}
+static LogicalResult verify(TypeOp op) { return verifyHasBindingUse(op); }
 
 //===----------------------------------------------------------------------===//
 // pdl::TypesOp
 //===----------------------------------------------------------------------===//
 
-static LogicalResult verify(TypesOp op) {
-  return verifyHasBindingUseInMatcher(op, "`pdl.operands`, or `pdl.operation`");
-}
+static LogicalResult verify(TypesOp op) { return verifyHasBindingUse(op); }
 
 //===----------------------------------------------------------------------===//
 // TableGen'd op method definitions
 //===----------------------------------------------------------------------===//
 
 #define GET_OP_CLASSES
 #include "mlir/Dialect/PDL/IR/PDLOps.cpp.inc"

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

pdl.pattern : benefit(1) {
     pdl.rewrite %root with "rewriter"
   }
 }
 
 
 // -----
 
 // CHECK-LABEL: module @predicate_ordering
-module @predicate_ordering  {
+module @predicate_ordering {
   // Check that the result is checked for null first, before applying the
   // constraint. The null check is prevalent in both patterns, so should be
   // prioritized first.
 
   // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
   // CHECK:   %[[RESULT:.*]] = pdl_interp.get_result 0 of %[[ROOT]]
   // CHECK-NEXT: pdl_interp.is_not_null %[[RESULT]]
   // CHECK:   %[[RESULT_TYPE:.*]] = pdl_interp.get_value_type of %[[RESULT]]
   // CHECK: pdl_interp.apply_constraint "typeConstraint" [](%[[RESULT_TYPE]]
 
   pdl.pattern : benefit(1) {
     %resultType = pdl.type
     pdl.apply_native_constraint "typeConstraint"[](%resultType : !pdl.type)
     %root = pdl.operation -> (%resultType : !pdl.type)
     pdl.rewrite %root with "rewriter"
   }
 
   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 0 of %[[VAL1]] : !pdl.value
+  // 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(%[[ROOT:.*]]: !pdl.operation)
+  // CHECK-DAG: %[[VAL:.*]] = pdl_interp.get_result 0 of %[[ROOT]]
+  // CHECK-DAG: pdl_interp.check_operand_count of %[[ROOT]] is 2
+  // CHECK-DAG: pdl_interp.check_result_count of %[[ROOT]] is 1
+  // CHECK-DAG: %[[INPUT2:.*]] = pdl_interp.get_operand 1 of %[[ROOT]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[INPUT2]] : !pdl.value
+  // CHECK-DAG: %[[INPUT1:.*]] = pdl_interp.get_operand 0 of %[[ROOT]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[INPUT1]] : !pdl.value
+  // CHECK-DAG: %[[OPS:.*]] = pdl_interp.get_users 0 of %[[VAL]] : !pdl.value
+  // CHECK-DAG: pdl_interp.foreach %[[OP:.*]] : !pdl.operation in %[[OPS]]
+  // CHECK-DAG:   pdl_interp.is_not_null %[[OP]] : !pdl.operation
+  // CHECK-DAG:   pdl_interp.check_operand_count of %[[OP]] is 1
+
+  pdl.pattern @rewrite_forced_overlapped_root : benefit(1) {
+    %input1 = pdl.operand
+    %input2 = pdl.operand
+    %type = pdl.type
+    %root = pdl.operation(%input1, %input2 : !pdl.value, !pdl.value) -> (%type : !pdl.type)
+    %val = pdl.result 0 of %root
+    %op = pdl.operation(%val : !pdl.value)
+    pdl.rewrite %root with "rewriter"(%op : !pdl.operation)
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module @variadic_results_all
+module @variadic_results_all {
+  // Check the correct lowering when using all results of an operation
+  // and passing it them as operands to another operation.
+
+  // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
+  // CHECK-DAG: pdl_interp.check_operand_count of %[[ROOT]] is 0
+  // CHECK-DAG: %[[VALS:.*]] = pdl_interp.get_results of %[[ROOT]] : !pdl.range<value>
+  // CHECK-DAG: %[[OPS:.*]] = pdl_interp.get_users of %[[VALS]] : !pdl.range<value>
+  // CHECK-DAG: pdl_interp.foreach %[[OP:.*]] : !pdl.operation in %[[OPS]]
+  // CHECK-DAG:   pdl_interp.is_not_null %[[OP]]
+  // CHECK-DAG:   pdl_interp.check_result_count of %[[OP]] is 0
+  pdl.pattern @variadic_results_all : benefit(1) {
+    %types = pdl.types
+    %root = pdl.operation -> (%types : !pdl.range<type>)
+    %vals = pdl.results of %root
+    %op = pdl.operation(%vals : !pdl.range<value>)
+    pdl.rewrite %root with "rewriter"(%op : !pdl.operation)
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module @variadic_results_at
+module @variadic_results_at {
+  // Check the correct lowering when using selected results of an operation
+  // and passing it them as an operand to another operation.
+
+  // CHECK: func @matcher(%[[ROOT1:.*]]: !pdl.operation)
+  // CHECK-DAG: %[[VALS:.*]] = pdl_interp.get_operands 0 of %[[ROOT1]] : !pdl.range<value>
+  // CHECK-DAG: %[[OP:.*]] = pdl_interp.get_defining_op of %[[VALS]] : !pdl.range<value>
+  // CHECK-DAG: pdl_interp.is_not_null %[[OP]] : !pdl.operation
+  // CHECK-DAG: pdl_interp.check_operand_count of %[[ROOT1]] is at_least 1
+  // CHECK-DAG: pdl_interp.check_result_count of %[[ROOT1]] is 0
+  // CHECK-DAG: %[[VAL:.*]] = pdl_interp.get_operands 1 of %[[ROOT1]] : !pdl.value
+  // CHECK-DAG: pdl_interp.is_not_null %[[VAL]]
+  // CHECK-DAG: pdl_interp.is_not_null %[[VALS]]
+  // CHECK-DAG: %[[ROOTS2:.*]] = pdl_interp.get_users 1 of %[[VALS]] : !pdl.range<value>
+  // CHECK-DAG: pdl_interp.foreach %[[ROOT2:.*]] : !pdl.operation in %[[ROOTS2]] {
+  // CHECK-DAG:   pdl_interp.is_not_null %[[ROOT2]]
+  // CHECK-DAG:   pdl_interp.check_operand_count of %[[ROOT2]] is at_least 1
+  // CHECK-DAG:   pdl_interp.check_result_count of %[[ROOT2]] is 0
+  // CHECK-DAG:   pdl_interp.check_operand_count of %[[OP]] is 0
+  // CHECK-DAG:   pdl_interp.check_result_count of %[[OP]] is at_least 1
+  pdl.pattern @variadic_results_at : benefit(1) {
+    %type = pdl.type
+    %types = pdl.types
+    %val = pdl.operand
+    %op = pdl.operation -> (%types, %type : !pdl.range<type>, !pdl.type)
+    %vals = pdl.results 0 of %op -> !pdl.range<value>
+    %root1 = pdl.operation(%vals, %val : !pdl.range<value>, !pdl.value)
+    %root2 = pdl.operation(%val, %vals : !pdl.value, !pdl.range<value>)
+    pdl.rewrite with "rewriter"(%root1, %root2 : !pdl.operation, !pdl.operation)
+  }
+}

mlir/test/Dialect/PDL/invalid.mlir

pdl.rewrite %op {
     // expected-error@below {{expected constant value when specified within a `pdl.rewrite`}}
     %attr = pdl.attribute
   }
 }
 
 // -----
 
 pdl.pattern : benefit(1) {
-  // expected-error@below {{expected a bindable (i.e. `pdl.operation`) user when defined in the matcher body of a `pdl.pattern`}}
+  // expected-error@below {{expected a bindable user when defined in the matcher body of a `pdl.pattern`}}
   %unused = pdl.attribute
 
   %op = pdl.operation "foo.op"
   pdl.rewrite %op with "rewriter"
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // pdl::OperandOp
 //===----------------------------------------------------------------------===//
 
 pdl.pattern : benefit(1) {
-  // expected-error@below {{expected a bindable (i.e. `pdl.operation`) user when defined in the matcher body of a `pdl.pattern`}}
+  // expected-error@below {{expected a bindable user when defined in the matcher body of a `pdl.pattern`}}
   %unused = pdl.operand
 
   %op = pdl.operation "foo.op"
   pdl.rewrite %op with "rewriter"
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // pdl::OperandsOp
 //===----------------------------------------------------------------------===//
 
 pdl.pattern : benefit(1) {
-  // expected-error@below {{expected a bindable (i.e. `pdl.operation`) user when defined in the matcher body of a `pdl.pattern`}}
+  // expected-error@below {{expected a bindable user when defined in the matcher body of a `pdl.pattern`}}
   %unused = pdl.operands
 
   %op = pdl.operation "foo.op"
   pdl.rewrite %op with "rewriter"
 }
 
 // -----
 
pdl.rewrite %op {
     // expected-note@below {{result type #0 was not constrained}}
     %newOp = pdl.operation "foo.op" -> (%type : !pdl.type)
   }
 }
 
 // -----
 
 pdl.pattern : benefit(1) {
-  // expected-error@below {{expected a bindable (i.e. `pdl.operation` or `pdl.rewrite`) user when defined in the matcher body of a `pdl.pattern`}}
+  // expected-error@below {{expected a bindable user when defined in the matcher body of a `pdl.pattern`}}
   %unused = pdl.operation "foo.op"
 
   %op = pdl.operation "foo.op"
   pdl.rewrite %op with "rewriter"
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // pdl::PatternOp
 //===----------------------------------------------------------------------===//
 
 // expected-error@below {{expected body to terminate with `pdl.rewrite`}}
 pdl.pattern : benefit(1) {
   // expected-note@below {{see terminator defined here}}
   return
 }
 
 // -----
 
+// expected-error@below {{the pattern must contain at least one `pdl.operation`}}
+pdl.pattern : benefit(1) {
+  pdl.rewrite with "foo"
+}
+
+// -----
 // expected-error@below {{expected only `pdl` operations within the pattern body}}
 pdl.pattern : benefit(1) {
   // expected-note@below {{see non-`pdl` operation defined here}}
   "test.foo.other_op"() : () -> ()
 
   %root = pdl.operation "foo.op"
   pdl.rewrite %root with "foo"
 }
 
 // -----
+// expected-error@below {{the operations must form a connected component}}
+pdl.pattern : benefit(1) {
+  %op1 = pdl.operation "foo.op"
+  %op2 = pdl.operation "bar.op"
+  // expected-note@below {{see a disconnected value / operation here}}
+  %val = pdl.result 0 of %op2
+  pdl.rewrite %op1 with "foo"(%val : !pdl.value)
+}
+
+// -----
+// expected-error@below {{the operations must form a connected component}}
+pdl.pattern : benefit(1) {
+  %type = pdl.type
+  %op1 = pdl.operation "foo.op" -> (%type : !pdl.type)
+  %val = pdl.result 0 of %op1
+  %op2 = pdl.operation "bar.op"(%val : !pdl.value)
+  // expected-note@below {{see a disconnected value / operation here}}
+  %op3 = pdl.operation "baz.op"
+  pdl.rewrite {
+    pdl.erase %op1
+    pdl.erase %op2
+    pdl.erase %op3
+  }
+}
+
+// -----
 
 pdl.pattern : benefit(1) {
   %type = pdl.type : i32
   %root = pdl.operation "foo.op" -> (%type : !pdl.type)
   pdl.rewrite %root {
     %newOp = pdl.operation "foo.op" -> (%type : !pdl.type)
     %newResult = pdl.result 0 of %newOp
 

 //===----------------------------------------------------------------------===//
 // 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
 //===----------------------------------------------------------------------===//
 
 pdl.pattern : benefit(1) {
-  // expected-error@below {{expected a bindable (i.e. `pdl.attribute`, `pdl.operand`, or `pdl.operation`) user when defined in the matcher body of a `pdl.pattern`}}
+  // expected-error@below {{expected a bindable user when defined in the matcher body of a `pdl.pattern`}}
   %unused = pdl.type
 
   %op = pdl.operation "foo.op"
   pdl.rewrite %op with "rewriter"
 }
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // pdl::TypesOp
 //===----------------------------------------------------------------------===//
 
 pdl.pattern : benefit(1) {
-  // expected-error@below {{expected a bindable (i.e. `pdl.operands`, or `pdl.operation`) user when defined in the matcher body of a `pdl.pattern`}}
+  // expected-error@below {{expected a bindable user when defined in the matcher body of a `pdl.pattern`}}
   %unused = pdl.types
 
   %op = pdl.operation "foo.op"
   pdl.rewrite %op with "rewriter"
 }

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