Xilinx
diff --git a/‎mlir/include/mlir/Dialect/PDL/IR/PDLOps.td
Lines changed: 4 additions & 2 deletions b/‎mlir/include/mlir/Dialect/PDL/IR/PDLOps.td
Lines changed: 4 additions & 2 deletions
diff --git a/‎mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td
Lines changed: 5 additions & 2 deletions b/‎mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td
Lines changed: 5 additions & 2 deletions
diff --git a/‎mlir/include/mlir/IR/PDLPatternMatch.h.inc
Lines changed: 11 additions & 7 deletions b/‎mlir/include/mlir/IR/PDLPatternMatch.h.inc
Lines changed: 11 additions & 7 deletions
diff --git a/‎mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
Lines changed: 23 additions & 37 deletions b/‎mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp
Lines changed: 23 additions & 37 deletions
diff --git a/‎mlir/lib/Conversion/PDLToPDLInterp/Predicate.h
Lines changed: 22 additions & 0 deletions b/‎mlir/lib/Conversion/PDLToPDLInterp/Predicate.h
Lines changed: 22 additions & 0 deletions
diff --git a/‎mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp
Lines changed: 51 additions & 45 deletions b/‎mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp
Lines changed: 51 additions & 45 deletions
@@ -35,13 +35,15 @@ def PDL_ApplyNativeConstraintOp
   let description = [{
     `pdl.apply_native_constraint` operations apply a native C++ constraint, that
     has been registered externally with the consumer of PDL, to a given set of
-    entities and optionally return a number of values..
+    entities and optionally return a number of values.
 
     Example:
 
     ```mlir
     // Apply `myConstraint` to the entities defined by `input`, `attr`, and `op`.
     pdl.apply_native_constraint "myConstraint"(%input, %attr, %op : !pdl.value, !pdl.attribute, !pdl.operation)
+    // Apply constraint `with_result` to `root`. This constraint returns an attribute.
+    %attr = pdl.apply_native_constraint "with_result"(%root : !pdl.operation) : !pdl.attribute
     ```
   }];
 
@@ -50,7 +52,7 @@ def PDL_ApplyNativeConstraintOp
                        DefaultValuedAttr<BoolAttr, "false">:$isNegated);
   let results = (outs Variadic<PDL_AnyType>:$results);
   let assemblyFormat = [{
-    $name (`(` $args^ `:` type($args) `)`)? (`:` type($results)^)? attr-dict
+    $name `(` $args `:` type($args) `)` (`:`  type($results)^ )? attr-dict
   }];
   let hasVerifier = 1;
 }
 
@@ -88,7 +88,9 @@ def PDLInterp_ApplyConstraintOp : PDLInterp_PredicateOp<"apply_constraint"> {
   let description = [{
     `pdl_interp.apply_constraint` operations apply a generic constraint, that
     has been registered with the interpreter, with a given set of positional
-    values. On success, this operation branches to the true destination,
+    values.
+    The constraint function may return any number of results.
+    On success, this operation branches to the true destination,
     otherwise the false destination is taken. This behavior can be reversed
     by setting the attribute `isNegated` to true.
 
@@ -106,7 +108,8 @@ def PDLInterp_ApplyConstraintOp : PDLInterp_PredicateOp<"apply_constraint"> {
                        DefaultValuedAttr<BoolAttr, "false">:$isNegated);
   let results = (outs Variadic<PDL_AnyType>:$results);
   let assemblyFormat = [{
-    $name (`(` $args^ `:` type($args) `)`)? (`:` type($results)^)? attr-dict `->` successors
+    $name `(` $args `:` type($args) `)` (`:` type($results)^)? attr-dict 
+    `->` successors
   }];
 }
 
 
@@ -318,8 +318,9 @@ protected:
 /// A generic PDL pattern constraint function. This function applies a
 /// constraint to a given set of opaque PDLValue entities. Returns success if
 /// the constraint successfully held, failure otherwise.
-using PDLConstraintFunction =
-    std::function<LogicalResult(PatternRewriter &, ArrayRef<PDLValue>)>;
+using PDLConstraintFunction = std::function<LogicalResult(
+    PatternRewriter &, PDLResultList &, ArrayRef<PDLValue>)>;
+
 /// A native PDL rewrite function. This function performs a rewrite on the
 /// given set of values. Any results from this rewrite that should be passed
 /// back to PDL should be added to the provided result list. This method is only
@@ -726,7 +727,7 @@ std::enable_if_t<
     PDLConstraintFunction>
 buildConstraintFn(ConstraintFnT &&constraintFn) {
   return [constraintFn = std::forward<ConstraintFnT>(constraintFn)](
-             PatternRewriter &rewriter,
+             PatternRewriter &rewriter, PDLResultList &,
              ArrayRef<PDLValue> values) -> LogicalResult {
     auto argIndices = std::make_index_sequence<
         llvm::function_traits<ConstraintFnT>::num_args - 1>();
@@ -842,10 +843,13 @@ public:
   /// Register a constraint function with PDL. A constraint function may be
   /// specified in one of two ways:
   ///
-  ///   * `LogicalResult (PatternRewriter &, ArrayRef<PDLValue>)`
+  ///   * `LogicalResult (PatternRewriter &,
+  ///                     PDLResultList &,
+  ///                     ArrayRef<PDLValue>)`
   ///
   ///   In this overload the arguments of the constraint function are passed via
-  ///   the low-level PDLValue form.
+  ///   the low-level PDLValue form, and the results are manually appended to
+  ///   the given result list.
   ///
   ///   * `LogicalResult (PatternRewriter &, ValueTs... values)`
   ///
@@ -963,8 +967,8 @@ public:
   }
 };
 class PDLResultList {};
-using PDLConstraintFunction =
-    std::function<LogicalResult(PatternRewriter &, ArrayRef<PDLValue>)>;
+using PDLConstraintFunction = std::function<LogicalResult(
+    PatternRewriter &, PDLResultList &, ArrayRef<PDLValue>)>;
 using PDLRewriteFunction = std::function<LogicalResult(
     PatternRewriter &, PDLResultList &, ArrayRef<PDLValue>)>;
 
 
@@ -50,7 +50,8 @@ struct PatternLowering {
 
   /// Generate interpreter operations for the tree rooted at the given matcher
   /// node, in the specified region.
-  Block *generateMatcher(MatcherNode &node, Region &region);
+  Block *generateMatcher(MatcherNode &node, Region &region,
+                         Block *block = nullptr);
 
   /// Get or create an access to the provided positional value in the current
   /// block. This operation may mutate the provided block pointer if nested
@@ -149,9 +150,9 @@ struct PatternLowering {
   /// set.
   DenseMap<Operation *, PDLPatternConfigSet *> *configMap;
 
-  /// A mapping between constraint questions that refer to values created by
-  /// constraints and the temporary placeholder values created for them.
-  std::multimap<std::pair<ConstraintQuestion *, unsigned>, Value> substitutions;
+  /// A mapping from a constraint question to the ApplyConstraintOp
+  /// that implements it.
+  DenseMap<ConstraintQuestion *, pdl_interp::ApplyConstraintOp> constraintOpMap;
 };
 } // namespace
 
@@ -186,9 +187,11 @@ void PatternLowering::lower(ModuleOp module) {
   firstMatcherBlock->erase();
 }
 
-Block *PatternLowering::generateMatcher(MatcherNode &node, Region &region) {
+Block *PatternLowering::generateMatcher(MatcherNode &node, Region &region,
+                                        Block *block) {
   // Push a new scope for the values used by this matcher.
-  Block *block = &region.emplaceBlock();
+  if (!block)
+    block = &region.emplaceBlock();
   ValueMapScope scope(values);
 
   // If this is the return node, simply insert the corresponding interpreter
@@ -369,18 +372,12 @@ Value PatternLowering::getValueAt(Block *&currentBlock, Position *pos) {
     break;
   }
   case Predicates::ConstraintResultPos: {
-    // At this point in time the corresponding pdl.ApplyNativeConstraint op has
-    // been deleted and the new pdl_interp.ApplyConstraint has not been created
-    // yet. To enable use of results created by these operations we build a
-    // placeholder value that will be replaced when the actual
-    // pdl_interp.ApplyConstraint operation is created.
+    // Due to the order of traversal, the ApplyConstraintOp has already been
+    // created and we can find it in constraintOpMap.
     auto *constrResPos = cast<ConstraintPosition>(pos);
-    Value placeholderValue = builder.create<pdl_interp::CreateAttributeOp>(
-        loc, StringAttr::get(builder.getContext(), "placeholder"));
-    substitutions.insert(
-        {{constrResPos->getQuestion(), constrResPos->getIndex()},
-         placeholderValue});
-    value = placeholderValue;
+    auto i = constraintOpMap.find(constrResPos->getQuestion());
+    assert(i != constraintOpMap.end());
+    value = i->second->getResult(constrResPos->getIndex());
     break;
   }
   default:
@@ -409,12 +406,11 @@ void PatternLowering::generate(BoolNode *boolNode, Block *&currentBlock,
       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);
+  // Generate a new block as success successor and get the failure successor.
+  Block *success = &region->emplaceBlock();
   Block *failure = failureBlockStack.back();
 
-  // Finally, create the predicate.
+  // Create the predicate.
   builder.setInsertionPointToEnd(currentBlock);
   Predicates::Kind kind = question->getKind();
   switch (kind) {
@@ -469,27 +465,17 @@ void PatternLowering::generate(BoolNode *boolNode, Block *&currentBlock,
     auto applyConstraintOp = builder.create<pdl_interp::ApplyConstraintOp>(
         loc, cstQuestion->getResultTypes(), cstQuestion->getName(), args,
         cstQuestion->getIsNegated(), success, failure);
-    // Replace the generated placeholders with the results of the constraint and
-    // erase them
-    for (auto result : llvm::enumerate(applyConstraintOp.getResults())) {
-      std::pair<ConstraintQuestion *, unsigned> substitutionKey = {
-          cstQuestion, result.index()};
-      // Check if there are substitutions to perform. If the result is never
-      // used or multiple calls to the same constraint have been merged,
-      // no substitutions will have been generated for this specific op.
-      auto range = substitutions.equal_range(substitutionKey);
-      std::for_each(range.first, range.second, [&](const auto &elem) {
-        Value placeholder = elem.second;
-        placeholder.replaceAllUsesWith(result.value());
-        placeholder.getDefiningOp()->erase();
-      });
-      substitutions.erase(substitutionKey);
-    }
+
+    constraintOpMap.insert({cstQuestion, applyConstraintOp});
     break;
   }
   default:
     llvm_unreachable("Generating unknown Predicate operation");
   }
+
+  // Generate the matcher in the current (potentially nested) region.
+  // This might use the results of the current predicate.
+  generateMatcher(*boolNode->getSuccessNode(), *region, success);
 }
 
 template <typename OpT, typename PredT, typename ValT = typename PredT::KeyTy>
 
@@ -299,6 +299,28 @@ struct OperationPosition : public PredicateBase<OperationPosition, Position,
   bool isOperandDefiningOp() const;
 };
 
+//===----------------------------------------------------------------------===//
+// ConstraintPosition
+
+struct ConstraintQuestion;
+
+/// A position describing the result of a native constraint. It saves the
+/// corresponding ConstraintQuestion and result index to enable referring
+/// back to them
+struct ConstraintPosition
+    : public PredicateBase<ConstraintPosition, Position,
+                           std::pair<ConstraintQuestion *, unsigned>,
+                           Predicates::ConstraintResultPos> {
+  using PredicateBase::PredicateBase;
+
+  /// Returns the ConstraintQuestion to enable keeping track of the native
+  /// constraint this position stems from.
+  ConstraintQuestion *getQuestion() const { return key.first; }
+
+  // Returns the result index of this position
+  unsigned getIndex() const { return key.second; }
+};
+
 //===----------------------------------------------------------------------===//
 // ResultPosition
 
 
@@ -15,10 +15,10 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "llvm/ADT/MapVector.h"
+#include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include <queue>
-#include "llvm/ADT/SmallPtrSet.h"
 
 #define DEBUG_TYPE "pdl-predicate-tree"
 
@@ -50,14 +50,15 @@ static void getTreePredicates(std::vector<PositionalPredicate> &predList,
                               DenseMap<Value, Position *> &inputs,
                               AttributePosition *pos) {
   assert(isa<pdl::AttributeType>(val.getType()) && "expected attribute type");
-  pdl::AttributeOp attr = cast<pdl::AttributeOp>(val.getDefiningOp());
   predList.emplace_back(pos, builder.getIsNotNull());
 
-  // If the attribute has a type or value, add a constraint.
-  if (Value type = attr.getValueType())
-    getTreePredicates(predList, type, builder, inputs, builder.getType(pos));
-  else if (Attribute value = attr.getValueAttr())
-    predList.emplace_back(pos, builder.getAttributeConstraint(value));
+  if (auto attr = dyn_cast<pdl::AttributeOp>(val.getDefiningOp())) {
+    // If the attribute has a type or value, add a constraint.
+    if (Value type = attr.getValueType())
+      getTreePredicates(predList, type, builder, inputs, builder.getType(pos));
+    else if (Attribute value = attr.getValueAttr())
+      predList.emplace_back(pos, builder.getAttributeConstraint(value));
+  }
 }
 
 /// Collect all of the predicates for the given operand position.
@@ -265,34 +266,32 @@ static void getConstraintPredicates(pdl::ApplyNativeConstraintOp op,
                                     DenseMap<Value, Position *> &inputs) {
   OperandRange arguments = op.getArgs();
 
-  Position *pos = nullptr;
   std::vector<Position *> allPositions;
+  allPositions.reserve(arguments.size());
+  for (Value arg : arguments)
+    allPositions.push_back(inputs.lookup(arg));
 
-  // If this constraint has no arguments, this means it has no dependencies, and
-  // the same applies to all results
-  if (arguments.empty()) {
-    pos = builder.getRoot();
-  } else {
-    allPositions.reserve(arguments.size());
-    for (Value arg : arguments)
-      allPositions.push_back(inputs.lookup(arg));
-
-    // Push the constraint to the furthest position.
-    pos = *std::max_element(allPositions.begin(), allPositions.end(),
-                            comparePosDepth);
-  }
-  assert(pos && "Must have a non-null value");
-
+  // Push the constraint to the furthest position.
+  Position *pos = *std::max_element(allPositions.begin(), allPositions.end(),
+                                    comparePosDepth);
   ResultRange results = op.getResults();
   PredicateBuilder::Predicate pred = builder.getConstraint(
       op.getName(), allPositions, SmallVector<Type>(results.getTypes()),
       op.getIsNegated());
 
-  // for each result register a position so it can be used later
-  for (auto result : llvm::enumerate(results)) {
+  // For each result register a position so it can be used later
+  for (auto [i, result] : llvm::enumerate(results)) {
     ConstraintQuestion *q = cast<ConstraintQuestion>(pred.first);
-    ConstraintPosition *pos = builder.getConstraintPosition(q, result.index());
-    inputs[result.value()] = pos;
+    ConstraintPosition *pos = builder.getConstraintPosition(q, i);
+    auto [it, inserted] = inputs.insert({result, pos});
+    // If this is an input value that has been visited in the tree, add a
+    // constraint to ensure that both instances refer to the same value.
+    if (!inserted) {
+      auto minMaxPositions =
+          std::minmax<Position *>(pos, it->second, comparePosDepth);
+      predList.emplace_back(minMaxPositions.second,
+                            builder.getEqualTo(minMaxPositions.first));
+    }
   }
   predList.emplace_back(pos, pred);
 }
@@ -896,15 +895,14 @@ static void insertExitNode(std::unique_ptr<MatcherNode> *root) {
 }
 
 /// Sorts the range begin/end with the partial order given by cmp.
-/// cmp must be a partial ordering.
 template <typename Iterator, typename Compare>
-void stableTopologicalSort(Iterator begin, Iterator end, Compare cmp) {
+static void stableTopologicalSort(Iterator begin, Iterator end, Compare cmp) {
   while (begin != end) {
     // Cannot compute sortBeforeOthers in the predicate of stable_partition
     // because stable_partition will not keep the [begin, end) range intact
     // while it runs.
     llvm::SmallPtrSet<typename Iterator::value_type, 16> sortBeforeOthers;
-    for(auto i = begin; i != end; ++i) {
+    for (auto i = begin; i != end; ++i) {
       if (std::none_of(begin, end, [&](auto const &b) { return cmp(b, *i); }))
         sortBeforeOthers.insert(*i);
     }
@@ -917,6 +915,28 @@ void stableTopologicalSort(Iterator begin, Iterator end, Compare cmp) {
   }
 }
 
+/// Returns true if 'b' depends on a result of 'a'.
+static bool dependsOn(OrderedPredicate *a, OrderedPredicate *b) {
+  auto *cqa = dyn_cast<ConstraintQuestion>(a->question);
+  if (!cqa)
+    return false;
+
+  auto positionDependsOnA = [&](Position *p) {
+    auto *cp = dyn_cast<ConstraintPosition>(p);
+    return cp && cp->getQuestion() == cqa;
+  };
+
+  if (auto *cqb = dyn_cast<ConstraintQuestion>(b->question)) {
+    // Does any argument of b use a?
+    return llvm::any_of(cqb->getArgs(), positionDependsOnA);
+  }
+  if (auto *equalTo = dyn_cast<EqualToQuestion>(b->question)) {
+    return positionDependsOnA(b->position) ||
+           positionDependsOnA(equalTo->getValue());
+  }
+  return positionDependsOnA(b->position);
+}
+
 /// 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>
@@ -999,21 +1019,7 @@ MatcherNode::generateMatcherTree(ModuleOp module, PredicateBuilder &builder,
 
   // Mostly keep the now established order, but also ensure that
   // ConstraintQuestions come after the results they use.
-  stableTopologicalSort(ordered.begin(), ordered.end(),
-                        [](OrderedPredicate *a, OrderedPredicate *b) {
-                          auto *cqa = dyn_cast<ConstraintQuestion>(a->question);
-                          auto *cqb = dyn_cast<ConstraintQuestion>(b->question);
-                          if (cqa && cqb) {
-                            // Does any argument of b use a? Then b must be
-                            // sorted after a.
-                            return llvm::any_of(
-                                cqb->getArgs(), [&](Position *p) {
-                                  auto *cp = dyn_cast<ConstraintPosition>(p);
-                                  return cp && cp->getQuestion() == cqa;
-                                });
-                          }
-                          return false;
-                        });
+  stableTopologicalSort(ordered.begin(), ordered.end(), dependsOn);
 
   // Build the matchers for each of the pattern predicate lists.
   std::unique_ptr<MatcherNode> root;