[mlir][PDL] Add support for native constraints with results

[martin-luecke]

Summary

This adds support for native PDL (and PDLL) C++ constraints to return results.

This is useful for situations where a pattern checks for certain constraints of multiple interdependent attributes and computes a new attribute value based on them. Currently, for such an example it is required to escape to C++ during matching to perform the check and after a successful match again escape to native C++ to perform the computation during the rewriting part of the pattern.
With this work we can do the computation in C++ during matching and use the result in the rewriting part of the pattern. Effectively this enables a choice in the trade-off of memory consumption during matching vs recomputation of values.

This is an example of a situation where this is useful: We have two operations with certain attributes that have interdependent constraints. For instance attr_foo: one_of [0, 2, 4, 8], attr_bar: one_of [0, 2, 4, 8] and attr_foo == attr_bar. The pattern should only match if all conditions are true. The new operation should be created with a new attribute which is computed from the two matched attributes e.g. attr_baz = attr_foo * attr_bar. For the check we already escape to native C++ and have all values at hand so it makes sense to directly compute the new attribute value as well:

PDLL
Constraint checkAndCompute(attr0: Attr, attr1: Attr) -> Attr;

Pattern example with benefit(1) { 
    let foo = op<test.foo>() {attr = attr_foo : Attr}; 
    let bar = op<test.bar>(foo) {attr = attr_bar : Attr}; 
    let attr_baz = checkAndCompute(attr_foo, attr_bar); 
    rewrite bar with { 
        let baz = op<test.baz> {attr=attr_baz}; 
        replace bar with baz; 
    }; 
}

To achieve this the following notable changes were necessary:
PDLL

• Remove check in PDLL parser that prevented native constraints from returning results

PDL

• Change PDL definition of pdl.apply_native_constraint to allow variadic results

PDL_interp

• Change PDL_interp definition of pdl_interp.apply_constraint to allow variadic results

PDLToPDLInterp Pass:
The input to the pass is an arbitrary number of PDL patterns. The pass collects the predicates
that are required to match all of the pdl patterns and establishes an ordering that allows
creation of a single efficient matcher function to match all of them. Values that are matched and
possibly used in the rewriting part of a pattern are represented as positions. This allows fusion
and thus reusing a single position for multiple matching patterns.
Accordingly, we introduce ConstraintPosition, which records the type and index of the result of the constraint.
The problem is for the corresponding value to be used in the rewriting part of a pattern it has
to be an input to the pdl_interp.record_match operation, which is generated early during the pass
such that its surrounding block can be referred to by branching operations. In consequence the value has
to be materialized after the original pdl.apply_native_constraint has been deleted but before
we get the chance to generate the corresponding pdl_interp.apply_constraint operation.
We solve this by emitting a placeholder value when a ConstraintPosition is evaluated.
These placeholder values (due to fusion there may be multiple for one constraint result) are
replaced later when the actual pdl_interp.apply_constraint operation is created.

Bytecode generator and interpreter:
Constraint functions which return results have a different type compared to existing constraint functions.
They have the same type as native rewrite functions and hence are registered as rewrite functions.
Other options:

• Have a specific registration bucket for constraints with results
• Change the type for all native constraints to include (possibly empty) results and store all natvie constraints in the same registration bucket. (It is open how this interacts with the convenience PDL Constraint Builders, possibly something for a follow-up PR)

[Mogball]

Awesome work. Thanks for the patch! The high-level concerns are:

• whether a use-before-def is possible due to the PredicateTree attempting to access a ConstraintPosition before the relevant ConstraintQuestion is emitted
• changing native constraint functions to return results instead of clobbering rewrite functions to achieve the same effect

In the future, it would be great if the patch could be split up into smaller pieces to make it easier to review. For example, this patch could have been

1. Add results to pdl.apply_native_constraint and pdl_interp.apply_constraint and implement the lowering in PDLToPDLInterp
2. Add registration for native constraints with results and support in bytecode lowering and the interpreter.
3. Add PDLL support.

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

44

Please add an example of apply_native_constraint that returns results

50

mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td

108

mlir/include/mlir/IR/PatternMatch.h

1585

It does make more sense to me to give constraint functions results by default, and current constraints will have "empty results". This means there will be a single consistent representation of constraint functions.

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

381

There is an awkward problem here that the results of a native constraint could be accessed before the constraint is invoked, resulting in an SSA value use before definition error in this pass. Essentially, the assumption the pattern fuser makes is that all Positions can be accessed in any order, and PDLInterp and the interpreter rely on null propagation to make this work.

This may be working now because the results of the constraints are only accessed in the rewriter. On the other hand, it's possible the predicate tree will never access a ConstraintPosition value before the ConstraintQuestion itself is encoded, because structurally it is impossible for a pattern to use a native constraint result without posing the ConstraintQuestion first, meaning across patterns it's impossible for the number of constraints on the ConstraintPosition to be greater than the number of ConstraintQuestion of the same kind. If the number is equal, the PredicateTree *may* be stable, so they won't get swapped. I'm not 100% sure, however.

In any case, I think it makes more sense to map the ConstraintQuestion to the pdl_interp::ApplyConstraint operation and then retrieve the relevant results here, instead of relying on creating "placeholder" values.

mlir/lib/Conversion/PDLToPDLInterp/Predicate.h

478

Do the types need to be encoded?

I guess there is nothing stopping a user from writing a native constraint that dynamically returns different results, but this would be impossible to construct from PDLL.

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

279

280

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

98

Please also document this restriction in the op definition

mlir/lib/Rewrite/ByteCode.cpp

784

please use llvm::report_fatal_error

818

isa<T>(...) is the preferred form

1458

This clobbering of using rewrite functions to contain constraints with results seems problematic. At this point I don't remember why rewrite functions are allowed to fail in the first place. In any case, please add results to regular constraint functions instead.

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

125

please add a test where the result of a pdl.apply_native_constraint is used inside the matcher

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

 //===----------------------------------------------------------------------===//
 
 def PDL_ApplyNativeConstraintOp
     : PDL_Op<"apply_native_constraint", [HasParent<"pdl::PatternOp">]> {
   let summary = "Apply a native constraint to a set of provided entities";
   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.
+    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)
     ```
   }];
 
   let arguments = (ins StrAttr:$name, Variadic<PDL_AnyType>:$args);
-  let assemblyFormat = "$name `(` $args `:` type($args) `)` attr-dict";
+  let results = (outs Variadic<PDL_AnyType>:$results);
+  let assemblyFormat = "$name `(` $args `:` type($args) `)` (`:`  type($results)^ )? attr-dict";
   let hasVerifier = 1;
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::ApplyNativeRewriteOp
 //===----------------------------------------------------------------------===//
 
 def PDL_ApplyNativeRewriteOp

mlir/include/mlir/Dialect/PDLInterp/IR/PDLInterpOps.td

 // pdl_interp::ApplyConstraintOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_ApplyConstraintOp : PDLInterp_PredicateOp<"apply_constraint"> {
   let summary = "Apply a constraint to a set of positional values";
   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,
-    otherwise the false destination is taken.
+    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.
 
     Example:
 
     ```mlir
     // Apply `myConstraint` to the entities defined by `input`, `attr`, and
     // `op`.
     pdl_interp.apply_constraint "myConstraint"(%input, %attr, %op : !pdl.value, !pdl.attribute, !pdl.operation) -> ^matchDest, ^failureDest
     ```
   }];
 
   let arguments = (ins StrAttr:$name, Variadic<PDL_AnyType>:$args);
+  let results = (outs Variadic<PDL_AnyType>:$results);
   let assemblyFormat = [{
-    $name `(` $args `:` type($args) `)` attr-dict `->` successors
+    $name `(` $args `:` type($args) `)` (`:` type($results)^)? attr-dict `->` successors
   }];
 }
 
 //===----------------------------------------------------------------------===//
 // pdl_interp::ApplyRewriteOp
 //===----------------------------------------------------------------------===//
 
 def PDLInterp_ApplyRewriteOp : PDLInterp_Op<"apply_rewrite"> {

mlir/include/mlir/IR/PatternMatch.h

public:
   template <typename ConstraintFnT>
   void registerConstraintFunction(StringRef name,
                                   ConstraintFnT &&constraintFn) {
     registerConstraintFunction(name,
                                detail::pdl_function_builder::buildConstraintFn(
                                    std::forward<ConstraintFnT>(constraintFn)));
   }
 
+  /// Register a constraint function that produces results with PDL. A
+  /// constraint function with results uses the same registry as rewrite
+  /// functions. It may be specified as follows:
+  ///
+  ///   * `LogicalResult (PatternRewriter &, PDLResultList &,
+  ///   ArrayRef<PDLValue>)`
+  ///
+  ///   The arguments of the constraint function are passed via the low-level
+  ///   PDLValue form, and the results are manually appended to the given result
+  ///   list.
+  void registerConstraintFunctionWithResults(StringRef name,
+                                             PDLRewriteFunction constraintFn);
+
   /// Register a rewrite function with PDL. A rewrite function may be specified
   /// in one of two ways:
   ///
   ///   * `void (PatternRewriter &, PDLResultList &, ArrayRef<PDLValue>)`
   ///
   ///   In this overload the arguments of the constraint function are passed via
   ///   the low-level PDLValue form, and the results are manually appended to
   ///   the given result list.

mlir/lib/Conversion/PDLToPDLInterp/PDLToPDLInterp.cpp

private:
 
   /// The set of operation values whose whose location will be used for newly
   /// generated operations.
   SetVector<Value> locOps;
 
   /// A mapping between pattern operations and the corresponding configuration
   /// set.
   DenseMap<Operation *, PDLPatternConfigSet *> *configMap;
+
+  /// A mapping from a constraint question and result index that together
+  /// refer to a value created by a constraint to the temporary placeholder
+  /// values created for them.
+  std::multimap<std::pair<ConstraintQuestion *, unsigned>, Value>
+      constraintResultMap;
 };
 } // namespace
 
 PatternLowering::PatternLowering(
     pdl_interp::FuncOp matcherFunc, ModuleOp rewriterModule,
     DenseMap<Operation *, PDLPatternConfigSet *> *configMap)
     : builder(matcherFunc.getContext()), matcherFunc(matcherFunc),
       rewriterModule(rewriterModule), rewriterSymbolTable(rewriterModule),
case Predicates::TypeLiteralPos: {
     Attribute rawTypeAttr = typePos->getValue();
     if (TypeAttr typeAttr = dyn_cast<TypeAttr>(rawTypeAttr))
       value = builder.create<pdl_interp::CreateTypeOp>(loc, typeAttr);
     else
       value = builder.create<pdl_interp::CreateTypesOp>(
           loc, cast<ArrayAttr>(rawTypeAttr));
     break;
   }
+  case Predicates::ConstraintResultPos: {
+    // The corresponding pdl.ApplyNativeConstraint op has already been deleted
+    // and the new pdl_interp.ApplyConstraint has not been created yet. To
+    // enable referring to results created by this operation we build a
+    // placeholder value that will be replaced when the actual
+    // pdl_interp.ApplyConstraint operation is created.
+    auto *constrResPos = cast<ConstraintPosition>(pos);
+    Value placeholderValue = builder.create<pdl_interp::CreateAttributeOp>(
+        loc, StringAttr::get(builder.getContext(), "placeholder"));
+    constraintResultMap.insert(
+        {{constrResPos->getQuestion(), constrResPos->getIndex()},
+         placeholderValue});
+    value = placeholderValue;
+    break;
+  }
   default:
     llvm_unreachable("Generating unknown Position getter");
     break;
   }
 
   values.insert(pos, value);
   return value;
 }
case Predicates::EqualToQuestion: {
     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);
-    builder.create<pdl_interp::ApplyConstraintOp>(loc, cstQuestion->getName(),
-                                                  args, success, failure);
+    auto applyConstraintOp = builder.create<pdl_interp::ApplyConstraintOp>(
+        loc, cstQuestion->getResultTypes(), cstQuestion->getName(), args,
+        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 = constraintResultMap.equal_range(substitutionKey);
+      std::for_each(range.first, range.second, [&](const auto &elem) {
+        Value placeholder = elem.second;
+        placeholder.replaceAllUsesWith(result.value());
+        placeholder.getDefiningOp()->erase();
+      });
+      constraintResultMap.erase(substitutionKey);
+    }
     break;
   }
   default:
     llvm_unreachable("Generating unknown Predicate operation");
   }
 }
 
 template <typename OpT, typename PredT, typename ValT = typename PredT::KeyTy>

mlir/lib/Conversion/PDLToPDLInterp/Predicate.h

 namespace Predicates {
 /// An enumeration of the kinds of predicates.
 enum Kind : unsigned {
   /// Positions, ordered by decreasing priority.
   OperationPos,
   OperandPos,
   OperandGroupPos,
   AttributePos,
+  ConstraintResultPos,
   ResultPos,
   ResultGroupPos,
   TypePos,
   AttributeLiteralPos,
   TypeLiteralPos,
   UsersPos,
   ForEachPos,
 

 /// A position describing a literal attribute.
 struct AttributeLiteralPosition
     : public PredicateBase<AttributeLiteralPosition, Position, Attribute,
                            Predicates::AttributeLiteralPos> {
   using PredicateBase::PredicateBase;
 };
 
 //===----------------------------------------------------------------------===//
+// 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; }
+};
+
+//===----------------------------------------------------------------------===//
 // ForEachPosition
 
 /// A position describing an iterative choice of an operation.
 struct ForEachPosition : public PredicateBase<ForEachPosition, Position,
                                               std::pair<Position *, unsigned>,
                                               Predicates::ForEachPos> {
   explicit ForEachPosition(const KeyTy &key) : Base(key) { parent = key.first; }
 

 //===----------------------------------------------------------------------===//
 // Questions
 
 /// Compare an `Attribute` to a constant value.
 struct AttributeQuestion
     : public PredicateBase<AttributeQuestion, Qualifier, void,
                            Predicates::AttributeQuestion> {};
 
-/// Apply a parameterized constraint to multiple position values.
+/// Apply a parameterized constraint to multiple position values and possibly
+/// produce results.
 struct ConstraintQuestion
-    : public PredicateBase<ConstraintQuestion, Qualifier,
-                           std::tuple<StringRef, ArrayRef<Position *>>,
-                           Predicates::ConstraintQuestion> {
+    : public PredicateBase<
+          ConstraintQuestion, Qualifier,
+          std::tuple<StringRef, ArrayRef<Position *>, ArrayRef<Type>>,
   using Base::Base;
 
   /// Return the name of the constraint.
   StringRef getName() const { return std::get<0>(key); }
 
   /// Return the arguments of the constraint.
   ArrayRef<Position *> getArgs() const { return std::get<1>(key); }
 
+  /// Return the result types of the constraint.
+  ArrayRef<Type> getResultTypes() const { return std::get<2>(key); }
+
   /// Construct an instance with the given storage allocator.
   static ConstraintQuestion *construct(StorageUniquer::StorageAllocator &alloc,
                                        KeyTy key) {
     return Base::construct(alloc, KeyTy{alloc.copyInto(std::get<0>(key)),
-                                        alloc.copyInto(std::get<1>(key))});
+                                        alloc.copyInto(std::get<1>(key)),
+                                        alloc.copyInto(std::get<2>(key))});
   }
 };
 
 /// Compare the equality of two values.
 struct EqualToQuestion
     : public PredicateBase<EqualToQuestion, Qualifier, Position *,
                            Predicates::EqualToQuestion> {
   using Base::Base;

 /// This class provides a storage uniquer that is used to allocate predicate
 /// instances.
 class PredicateUniquer : public StorageUniquer {
 public:
   PredicateUniquer() {
     // Register the types of Positions with the uniquer.
     registerParametricStorageType<AttributePosition>();
     registerParametricStorageType<AttributeLiteralPosition>();
+    registerParametricStorageType<ConstraintPosition>();
     registerParametricStorageType<ForEachPosition>();
     registerParametricStorageType<OperandPosition>();
     registerParametricStorageType<OperandGroupPosition>();
     registerParametricStorageType<OperationPosition>();
     registerParametricStorageType<ResultPosition>();
     registerParametricStorageType<ResultGroupPosition>();
     registerParametricStorageType<TypePosition>();
     registerParametricStorageType<TypeLiteralPosition>();
public:
   }
 
   /// Returns the operation position equivalent to the given position.
   OperationPosition *getPassthroughOp(Position *p) {
     assert((isa<ForEachPosition>(p)) && "expected users position");
     return OperationPosition::get(uniquer, p);
   }
 
+  // Returns a position for a new value created by a constraint.
+  ConstraintPosition *getConstraintPosition(ConstraintQuestion *q,
+                                            unsigned index) {
+    return ConstraintPosition::get(uniquer, std::make_pair(q, index));
+  }
+
   /// Returns an attribute position for an attribute of the given operation.
   Position *getAttribute(OperationPosition *p, StringRef name) {
     return AttributePosition::get(uniquer, p, StringAttr::get(ctx, name));
   }
 
   /// Returns an attribute position for the given attribute.
   Position *getAttributeLiteral(Attribute attr) {
     return AttributeLiteralPosition::get(uniquer, attr);
public:
   }
 
   /// 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) {
-    return {ConstraintQuestion::get(uniquer, std::make_tuple(name, pos)),
+  Predicate getConstraint(StringRef name, ArrayRef<Position *> args,
+                          ArrayRef<Type> resultTypes) {
+    return {ConstraintQuestion::get(uniquer,
+                                    std::make_tuple(name, args, resultTypes)),
             TrueAnswer::get(uniquer)};
   }
 
   /// Create a predicate comparing a value with null.
   Predicate getIsNotNull() {
     return {IsNotNullQuestion::get(uniquer), TrueAnswer::get(uniquer)};
   }
 

mlir/lib/Conversion/PDLToPDLInterp/PredicateTree.cpp

static void getConstraintPredicates(pdl::ApplyNativeConstraintOp op,
   std::vector<Position *> allPositions;
   allPositions.reserve(arguments.size());
   for (Value arg : arguments)
     allPositions.push_back(inputs.lookup(arg));
 
   // Push the constraint to the furthest position.
   Position *pos = *std::max_element(allPositions.begin(), allPositions.end(),
                                     comparePosDepth);
-  PredicateBuilder::Predicate pred =
-      builder.getConstraint(op.getName(), allPositions);
+  ResultRange results = op.getResults();
+  PredicateBuilder::Predicate pred = builder.getConstraint(
+      op.getName(), allPositions, SmallVector<Type>(results.getTypes()));
+
+  // for each result register a position so it can be used later
+  for (auto result : llvm::enumerate(results)) {
+    ConstraintQuestion *q = cast<ConstraintQuestion>(pred.first);
+    ConstraintPosition *pos = builder.getConstraintPosition(q, result.index());
+    inputs[result.value()] = pos;
+  }
   predList.emplace_back(pos, pred);
 }
 
 static void getResultPredicates(pdl::ResultOp op,
                                 std::vector<PositionalPredicate> &predList,
                                 PredicateBuilder &builder,
                                 DenseMap<Value, Position *> &inputs) {
   Position *&resultPos = inputs[op];

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

 
 //===----------------------------------------------------------------------===//
 // pdl::ApplyNativeConstraintOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ApplyNativeConstraintOp::verify() {
   if (getNumOperands() == 0)
     return emitOpError("expected at least one argument");
+  if (llvm::any_of(getResults(), [](OpResult result) {
+        return result.getType().isa<OperationType>();
+      })) {
+    return emitOpError(
+        "returning an operation from a constraint is not supported");
+  }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // pdl::ApplyNativeRewriteOp
 //===----------------------------------------------------------------------===//
 
 LogicalResult ApplyNativeRewriteOp::verify() {

mlir/lib/IR/PatternMatch.cpp

void PDLPatternModule::registerConstraintFunction(
     StringRef name, PDLConstraintFunction constraintFn) {
   // TODO: Is it possible to diagnose when `name` is already registered to
   // a function that is not equivalent to `constraintFn`?
   // Allow existing mappings in the case multiple patterns depend on the same
   // constraint.
   constraintFunctions.try_emplace(name, std::move(constraintFn));
 }
 
+void PDLPatternModule::registerConstraintFunctionWithResults(
+    StringRef name, PDLRewriteFunction constraintFn) {
+  // TODO: Is it possible to diagnose when `name` is already registered to
+  // a function that is not equivalent to `rewriteFn`?
+  // Allow existing mappings in the case multiple patterns depend on the same
+  // rewrite.
+  registerRewriteFunction(name, std::move(constraintFn));
+}
+
 void PDLPatternModule::registerRewriteFunction(StringRef name,
                                                PDLRewriteFunction rewriteFn) {
   // TODO: Is it possible to diagnose when `name` is already registered to
   // a function that is not equivalent to `rewriteFn`?
   // Allow existing mappings in the case multiple patterns depend on the same
   // rewrite.
   rewriteFunctions.try_emplace(name, std::move(rewriteFn));
 }

mlir/lib/Rewrite/ByteCode.cpp

TypeSwitch<Operation *>(op)
           [&](auto interpOp) { this->generate(interpOp, writer); })
       .Default([](Operation *) {
         llvm_unreachable("unknown `pdl_interp` operation");
       });
 }
 
 void Generator::generate(pdl_interp::ApplyConstraintOp op,
                          ByteCodeWriter &writer) {
-  assert(constraintToMemIndex.count(op.getName()) &&
-         "expected index for constraint function");
+  // Constraints that should return a value have to be registered as rewrites.
+  // If a constraint and a rewrite of similar name are registered the
+  // constraint takes precedence
+  ResultRange results = op.getResults();
+  if (results.size() == 0 && constraintToMemIndex.count(op.getName()) != 0) {
     writer.append(OpCode::ApplyConstraint, constraintToMemIndex[op.getName()]);
+  } else if (results.size() > 0 &&
+             externalRewriterToMemIndex.count(op.getName()) != 0) {
+    writer.append(OpCode::ApplyConstraint,
+                  externalRewriterToMemIndex[op.getName()]);
+  } else {
+    assert(true && "expected index for constraint function, make sure it is "
+                   "registered properly. Note that native constraints with "
+                   "results have to be registered using "
+                   "PDLPatternModule::registerConstraintFunctionWithResults.");
+  }
   writer.appendPDLValueList(op.getArgs());
+  writer.append(ByteCodeField(results.size()));
+  for (Value result : results) {
+    // We record the expected kind of the result, so that we can provide extra
+    // verification of the native rewrite function and handle the failure case
+    // of constraints accordingly.
+    writer.appendPDLValueKind(result);
+
+    // Range results also need to append the range storage index.
+    if (result.getType().isa<pdl::RangeType>())
+      writer.append(getRangeStorageIndex(result));
+    writer.append(result);
+  }
   writer.append(op.getSuccessors());
 }
 void Generator::generate(pdl_interp::ApplyRewriteOp op,
                          ByteCodeWriter &writer) {
   assert(externalRewriterToMemIndex.count(op.getName()) &&
          "expected index for rewrite function");
   writer.append(OpCode::ApplyRewrite, externalRewriterToMemIndex[op.getName()]);
   writer.appendPDLValueList(op.getArgs());
 
   ResultRange results = op.getResults();
   writer.append(ByteCodeField(results.size()));
   for (Value result : results) {
-    // In debug mode we also record the expected kind of the result, so that we
+    // We record the expected kind of the result, so that we
     // can provide extra verification of the native rewrite function.
-#ifndef NDEBUG
     writer.appendPDLValueKind(result);
-#endif
 
     // Range results also need to append the range storage index.
-    if (isa<pdl::RangeType>(result.getType()))
+    if (result.getType().isa<pdl::RangeType>())
       writer.append(getRangeStorageIndex(result));
     writer.append(result);
   }
 }
 void Generator::generate(pdl_interp::AreEqualOp op, ByteCodeWriter &writer) {
   Value lhs = op.getLhs();
   if (isa<pdl::RangeType>(lhs.getType())) {
     writer.append(OpCode::AreRangesEqual);
void PDLByteCode::initializeMutableState(PDLByteCodeMutableState &state) const {
   for (const PDLByteCodePattern &pattern : patterns)
     state.currentPatternBenefits.push_back(pattern.getBenefit());
 }
 
 //===----------------------------------------------------------------------===//
 // ByteCode Execution
 
 namespace {
+/// This class is an instantiation of the PDLResultList that provides access to
+/// the returned results. This API is not on `PDLResultList` to avoid
+/// overexposing access to information specific solely to the ByteCode.
+class ByteCodeRewriteResultList : public PDLResultList {
+public:
+  ByteCodeRewriteResultList(unsigned maxNumResults)
+      : PDLResultList(maxNumResults) {}
+
+  /// Return the list of PDL results.
+  MutableArrayRef<PDLValue> getResults() { return results; }
+
+  /// Return the type ranges allocated by this list.
+  MutableArrayRef<llvm::OwningArrayRef<Type>> getAllocatedTypeRanges() {
+    return allocatedTypeRanges;
+  }
+
+  /// Return the value ranges allocated by this list.
+  MutableArrayRef<llvm::OwningArrayRef<Value>> getAllocatedValueRanges() {
+    return allocatedValueRanges;
+  }
+};
+
 /// This class provides support for executing a bytecode stream.
 class ByteCodeExecutor {
 public:
   ByteCodeExecutor(
       const ByteCodeField *curCodeIt, MutableArrayRef<const void *> memory,
       MutableArrayRef<llvm::OwningArrayRef<Operation *>> opRangeMemory,
       MutableArrayRef<TypeRange> typeRangeMemory,
       std::vector<llvm::OwningArrayRef<Type>> &allocatedTypeRangeMemory,
private:
   void executeReplaceOp(PatternRewriter &rewriter);
   void executeSwitchAttribute();
   void executeSwitchOperandCount();
   void executeSwitchOperationName();
   void executeSwitchResultCount();
   void executeSwitchType();
   void executeSwitchTypes();
 
+  // Helper for native constraints and rewrites
+  void processNativeFunResults(ByteCodeRewriteResultList &results,
+                               unsigned numResults,
+                               LogicalResult &rewriteResult);
+
   /// Pushes a code iterator to the stack.
   void pushCodeIt(const ByteCodeField *it) { resumeCodeIt.push_back(it); }
 
   /// Pops a code iterator from the stack, returning true on success.
   void popCodeIt() {
     assert(!resumeCodeIt.empty() && "attempt to pop code off empty stack");
     curCodeIt = resumeCodeIt.back();
     resumeCodeIt.pop_back();
private:
   std::enable_if_t<llvm::is_detected<has_pointer_traits, T>::value, T>
   readInline() {
     const void *pointer;
     std::memcpy(&pointer, curCodeIt, sizeof(const void *));
     curCodeIt += sizeof(const void *) / sizeof(ByteCodeField);
     return T::getFromOpaquePointer(pointer);
   }
 
+  /// Skip a list of values on the bytecode buffer.
+  void skip(size_t skipN) { curCodeIt += skipN; }
+
   /// Jump to a specific successor based on a predicate value.
   void selectJump(bool isTrue) { selectJump(size_t(isTrue ? 0 : 1)); }
   /// Jump to a specific successor based on a destination index.
   void selectJump(size_t destIndex) {
     curCodeIt = &code[read<ByteCodeAddr>(destIndex * 2)];
   }
 
   /// Handle a switch operation with the provided value and cases.
private:
   /// References to ByteCode data necessary for execution.
   ArrayRef<const void *> uniquedMemory;
   ArrayRef<ByteCodeField> code;
   ArrayRef<PatternBenefit> currentPatternBenefits;
   ArrayRef<PDLByteCodePattern> patterns;
   ArrayRef<PDLConstraintFunction> constraintFunctions;
   ArrayRef<PDLRewriteFunction> rewriteFunctions;
 };
-
-/// This class is an instantiation of the PDLResultList that provides access to
-/// the returned results. This API is not on `PDLResultList` to avoid
-/// overexposing access to information specific solely to the ByteCode.
-class ByteCodeRewriteResultList : public PDLResultList {
-public:
-  ByteCodeRewriteResultList(unsigned maxNumResults)
-      : PDLResultList(maxNumResults) {}
-
-  /// Return the list of PDL results.
-  MutableArrayRef<PDLValue> getResults() { return results; }
-
-  /// Return the type ranges allocated by this list.
-  MutableArrayRef<llvm::OwningArrayRef<Type>> getAllocatedTypeRanges() {
-    return allocatedTypeRanges;
-  }
-
-  /// Return the value ranges allocated by this list.
-  MutableArrayRef<llvm::OwningArrayRef<Value>> getAllocatedValueRanges() {
-    return allocatedValueRanges;
-  }
-};
 } // namespace
 
 void ByteCodeExecutor::executeApplyConstraint(PatternRewriter &rewriter) {
   LLVM_DEBUG(llvm::dbgs() << "Executing ApplyConstraint:\n");
-  const PDLConstraintFunction &constraintFn = constraintFunctions[read()];
+  ByteCodeField fun_idx = read();
   SmallVector<PDLValue, 16> args;
   readList<PDLValue>(args);
 
   LLVM_DEBUG({
     llvm::dbgs() << "  * Arguments: ";
     llvm::interleaveComma(args, llvm::dbgs());
+    llvm::dbgs() << "\n";
   });
 
-  // Invoke the constraint and jump to the proper destination.
-  selectJump(succeeded(constraintFn(rewriter, args)));
+  ByteCodeField numResults = read();
+  if (numResults == 0) {
+    const PDLConstraintFunction &constraintFn = constraintFunctions[fun_idx];
+    LogicalResult rewriteResult = constraintFn(rewriter, args);
+    // Depending on the constraint jump to the proper destination.
+    selectJump(succeeded(rewriteResult));
+    return;
+  }
+  const PDLRewriteFunction &constraintFn = rewriteFunctions[fun_idx];
+  ByteCodeRewriteResultList results(numResults);
+  LogicalResult rewriteResult = constraintFn(rewriter, results, args);
+  ArrayRef<PDLValue> constraintResults = results.getResults();
+  LLVM_DEBUG({
+    if (succeeded(rewriteResult)) {
+      llvm::dbgs() << "  * Constraint succeeded\n";
+      llvm::dbgs() << "  * Results: ";
+      llvm::interleaveComma(constraintResults, llvm::dbgs());
+      llvm::dbgs() << "\n";
+    } else {
+      llvm::dbgs() << "  * Constraint failed\n";
+    }
+  });
+  assert((failed(rewriteResult) || constraintResults.size() == numResults) &&
+         "native PDL rewrite function succeeded but returned "
+         "unexpected number of results");
+  processNativeFunResults(results, numResults, rewriteResult);
+
+  // Depending on the constraint jump to the proper destination.
+  selectJump(succeeded(rewriteResult));
 }
 
 LogicalResult ByteCodeExecutor::executeApplyRewrite(PatternRewriter &rewriter) {
   LLVM_DEBUG(llvm::dbgs() << "Executing ApplyRewrite:\n");
   const PDLRewriteFunction &rewriteFn = rewriteFunctions[read()];
   SmallVector<PDLValue, 16> args;
   readList<PDLValue>(args);
 
   LLVM_DEBUG({
     llvm::dbgs() << "  * Arguments: ";
     llvm::interleaveComma(args, llvm::dbgs());
   });
 
   // Execute the rewrite function.
   ByteCodeField numResults = read();
   ByteCodeRewriteResultList results(numResults);
   LogicalResult rewriteResult = rewriteFn(rewriter, results, args);
 
   assert(results.getResults().size() == numResults &&
          "native PDL rewrite function returned unexpected number of results");
 
-  // Store the results in the bytecode memory.
-  for (PDLValue &result : results.getResults()) {
-    LLVM_DEBUG(llvm::dbgs() << "  * Result: " << result << "\n");
+  processNativeFunResults(results, numResults, rewriteResult);
 
-// In debug mode we also verify the expected kind of the result.
-#ifndef NDEBUG
-    assert(result.getKind() == read<PDLValue::Kind>() &&
-           "native PDL rewrite function returned an unexpected type of result");
-#endif
+  if (failed(rewriteResult)) {
+    LLVM_DEBUG(llvm::dbgs() << "  - Failed");
+    return failure();
+  }
+  return success();
+}
+
+void ByteCodeExecutor::processNativeFunResults(
+    ByteCodeRewriteResultList &results, unsigned numResults,
+    LogicalResult &rewriteResult) {
+  // Store the results in the bytecode memory or handle missing results on
+  // failure.
+  for (unsigned resultIdx = 0; resultIdx < numResults; resultIdx++) {
+    PDLValue::Kind resultKind = read<PDLValue::Kind>();
+
+    // Skip the according number of values on the buffer on failure and exit
+    // early as there are no results to process.
+    if (failed(rewriteResult)) {
+      if (resultKind == PDLValue::Kind::TypeRange ||
+          resultKind == PDLValue::Kind::ValueRange) {
+        skip(2);
+      } else {
+        skip(1);
+      }
+      return;
+    }
+    PDLValue result = results.getResults()[resultIdx];
+    LLVM_DEBUG(llvm::dbgs() << "  * Result: " << result << "\n");
+    assert(result.getKind() == resultKind &&
+           "native PDL rewrite function returned an unexpected type of "
+           "result");
 
     // If the result is a range, we need to copy it over to the bytecodes
     // range memory.
     if (std::optional<TypeRange> typeRange = result.dyn_cast<TypeRange>()) {
       unsigned rangeIndex = read();
       typeRangeMemory[rangeIndex] = *typeRange;
       memory[read()] = &typeRangeMemory[rangeIndex];
     } else if (std::optional<ValueRange> valueRange =
                    result.dyn_cast<ValueRange>()) {
       unsigned rangeIndex = read();
       valueRangeMemory[rangeIndex] = *valueRange;
       memory[read()] = &valueRangeMemory[rangeIndex];
     } else {
       memory[read()] = result.getAsOpaquePointer();
     }
   }
 
   // Copy over any underlying storage allocated for result ranges.
   for (auto &it : results.getAllocatedTypeRanges())
     allocatedTypeRangeMemory.push_back(std::move(it));
   for (auto &it : results.getAllocatedValueRanges())
     allocatedValueRangeMemory.push_back(std::move(it));
-
-  // Process the result of the rewrite.
-  if (failed(rewriteResult)) {
-    LLVM_DEBUG(llvm::dbgs() << "  - Failed");
-    return failure();
-  }
-  return success();
 }
 
 void ByteCodeExecutor::executeAreEqual() {
   LLVM_DEBUG(llvm::dbgs() << "Executing AreEqual:\n");
   const void *lhs = read<const void *>();
   const void *rhs = read<const void *>();
 
   LLVM_DEBUG(llvm::dbgs() << "  * " << lhs << " == " << rhs << "\n");

mlir/lib/Tools/PDLL/Parser/Parser.cpp

if (curToken.isString()) {
     return emitError(name.getLoc(),
                      "external declarations must be declared in global scope");
   } else if (curToken.is(Token::error)) {
     return failure();
   }
   if (failed(parseToken(Token::semicolon,
                         "expected `;` after native declaration")))
     return failure();
-  // TODO: PDL should be able to support constraint results in certain
-  // situations, we should revise this.
-  if (std::is_same<ast::UserConstraintDecl, T>::value && !results.empty()) {
-    return emitError(
-        "native Constraints currently do not support returning results");
-  }
   return T::createNative(ctx, name, arguments, results, optCodeStr, resultType);
 }
 
 LogicalResult Parser::parseUserConstraintOrRewriteSignature(
     SmallVectorImpl<ast::VariableDecl *> &arguments,
     SmallVectorImpl<ast::VariableDecl *> &results,
     ast::DeclScope *&argumentScope, ast::Type &resultType) {
   // Parse the argument list of the decl.

mlir/python/mlir/dialects/_pdl_ops_ext.py

 
 
 class ApplyNativeConstraintOp:
     """Specialization for PDL apply native constraint op class."""
 
     def __init__(
         self,
         name: Union[str, StringAttr],
+        results: Sequence[Type] = None,
         args: Optional[Sequence[Union[OpView, Operation, Value]]] = None,
         *,
         loc=None,
         ip=None,
     ):
         if args is None:
             args = []
+        if results is None:
+            results = []
         args = _get_values(args)
-        super().__init__(name, args, loc=loc, ip=ip)
+        results = _get_values(results)
+        super().__init__(results, name, args, loc=loc, ip=ip)
 
 
 class ApplyNativeRewriteOp:
     """Specialization for PDL apply native rewrite op class."""
 
     def __init__(
         self,
         results: Sequence[Type],

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

pdl.pattern : benefit(1) {
 
     pdl.apply_native_constraint "multi_constraint"(%input0, %input1, %result0 : !pdl.value, !pdl.value, !pdl.value)
     rewrite %root with "rewriter"
   }
 }
 
 // -----
 
+// CHECK-LABEL: module @constraint_with_result
+module @constraint_with_result {
+  // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
+  // CHECK: %[[ATTR:.*]] = pdl_interp.apply_constraint "check_op_and_get_attr_constr"(%[[ROOT]]
+  // CHECK: pdl_interp.record_match @rewriters::@pdl_generated_rewriter(%[[ROOT]], %[[ATTR]] : !pdl.operation, !pdl.attribute)
+  pdl.pattern : benefit(1) {
+    %root = operation
+    %attr = pdl.apply_native_constraint "check_op_and_get_attr_constr"(%root : !pdl.operation) : !pdl.attribute
+    rewrite %root with "rewriter"(%attr : !pdl.attribute)
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module @constraint_with_unused_result
+module @constraint_with_unused_result {
+  // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
+  // CHECK: %[[ATTR:.*]] = pdl_interp.apply_constraint "check_op_and_get_attr_constr"(%[[ROOT]]
+  // CHECK: pdl_interp.record_match @rewriters::@pdl_generated_rewriter(%[[ROOT]] : !pdl.operation)
+  pdl.pattern : benefit(1) {
+    %root = operation
+    %attr = pdl.apply_native_constraint "check_op_and_get_attr_constr"(%root : !pdl.operation) : !pdl.attribute
+    rewrite %root with "rewriter"
+  }
+}
+
+// -----
+
+// CHECK-LABEL: module @constraint_with_result_multiple
+module @constraint_with_result_multiple {
+  // check that native constraints work as expected even when multiple identical constraints are fused  
+
+  // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
+  // CHECK: %[[ATTR:.*]] = pdl_interp.apply_constraint "check_op_and_get_attr_constr"(%[[ROOT]]
+  // CHECK-NOT: pdl_interp.apply_constraint "check_op_and_get_attr_constr"
+  // CHECK: pdl_interp.record_match @rewriters::@pdl_generated_rewriter_0(%[[ROOT]], %[[ATTR]]  : !pdl.operation, !pdl.attribute)
+  // CHECK: pdl_interp.record_match @rewriters::@pdl_generated_rewriter(%[[ROOT]], %[[ATTR]] : !pdl.operation, !pdl.attribute)
+  pdl.pattern : benefit(1) {
+    %root = operation
+    %attr = pdl.apply_native_constraint "check_op_and_get_attr_constr"(%root : !pdl.operation) : !pdl.attribute
+    rewrite %root with "rewriter"(%attr : !pdl.attribute)
+  }
+  pdl.pattern : benefit(1) {
+    %root = operation
+    %attr = pdl.apply_native_constraint "check_op_and_get_attr_constr"(%root : !pdl.operation) : !pdl.attribute
+    rewrite %root with "rewriter"(%attr : !pdl.attribute)
+  }
+}
+
+// -----
+
 // CHECK-LABEL: module @inputs
 module @inputs {
   // CHECK: func @matcher(%[[ROOT:.*]]: !pdl.operation)
   // CHECK-DAG: pdl_interp.check_operand_count of %[[ROOT]] is 2
 
   // Get the input and check the type.
   // CHECK-DAG:   %[[INPUT:.*]] = pdl_interp.get_operand 0 of %[[ROOT]]
   // CHECK-DAG:   pdl_interp.is_not_null %[[INPUT]] : !pdl.value

mlir/test/Dialect/PDL/ops.mlir

pdl.pattern @apply_rewrite_with_no_results : benefit(1) {
   %root = operation
   rewrite %root {
     apply_native_rewrite "NativeRewrite"(%root : !pdl.operation)
   }
 }
 
 // -----
 
+pdl.pattern @apply_constraint_with_no_results : benefit(1) {
+  %root = operation
+  apply_native_constraint "NativeConstraint"(%root : !pdl.operation)
+  rewrite %root with "rewriter"
+}
+
+// -----
+
+pdl.pattern @apply_constraint_with_results : benefit(1) {
+  %root = operation
+  %attr = apply_native_constraint "NativeConstraint"(%root : !pdl.operation) : !pdl.attribute
+  rewrite %root {
+    apply_native_rewrite "NativeRewrite"(%attr : !pdl.attribute)
+  }
+}
+
+// -----
+
 pdl.pattern @attribute_with_dict : benefit(1) {
   %root = operation
   rewrite %root {
     %attr = attribute = {some_unit_attr} attributes {pdl.special_attribute}
     apply_native_rewrite "NativeRewrite"(%attr : !pdl.attribute)
   }
 }
 

mlir/test/Rewrite/pdl-bytecode.mlir

 // CHECK: "test.replaced_by_pattern"
 module @ir attributes { test.apply_constraint_2 } {
   "test.failure_op"() { test_attr } : () -> ()
   "test.success_op"() : () -> (i32, i64)
 }
 
 // -----
 
+// Test returning an attribute from a native constraint
+module @patterns {
+  pdl_interp.func @matcher(%root : !pdl.operation) {
+    pdl_interp.check_operation_name of %root is "test.success_op" -> ^pat, ^end
+
+  ^pat:
+    %attr = pdl_interp.apply_constraint "op_constr_return_attr"(%root : !pdl.operation) : !pdl.attribute -> ^pat2, ^end
+
+  ^pat2:
+    pdl_interp.record_match @rewriters::@success(%root, %attr : !pdl.operation, !pdl.attribute) : benefit(1), loc([%root]) -> ^end
+
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    pdl_interp.func @success(%root : !pdl.operation, %attr : !pdl.attribute) {
+      %op = pdl_interp.create_operation "test.replaced_by_pattern" {"attr" = %attr}
+      pdl_interp.erase %root
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.apply_constraint_3
+// CHECK-NOT: "test.replaced_by_pattern"
+// CHECK: "test.replaced_by_pattern"() {attr = "test.success"}
+module @ir attributes { test.apply_constraint_3 } {
+  "test.failure_op"() : () -> ()
+  "test.success_op"() : () -> ()
+}
+
+// -----
+
+// Test returning a type from a native constraint.
+module @patterns {
+  pdl_interp.func @matcher(%root : !pdl.operation) {
+    pdl_interp.check_operation_name of %root is "test.success_op" -> ^pat, ^end
+
+  ^pat:
+    %new_type = pdl_interp.apply_constraint "op_constr_return_type"(%root : !pdl.operation) : !pdl.type -> ^pat2, ^end
+
+  ^pat2:
+    pdl_interp.record_match @rewriters::@success(%root, %new_type : !pdl.operation, !pdl.type) : benefit(1), loc([%root]) -> ^end
+
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    pdl_interp.func @success(%root : !pdl.operation, %new_type : !pdl.type) {
+      %op = pdl_interp.create_operation "test.replaced_by_pattern" -> (%new_type : !pdl.type)
+      pdl_interp.erase %root
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.apply_constraint_4
+// CHECK-NOT: "test.replaced_by_pattern"
+// CHECK: "test.replaced_by_pattern"() : () -> f32
+module @ir attributes { test.apply_constraint_4 } {
+  "test.failure_op"() : () -> ()
+  "test.success_op"() : () -> ()
+}
+
+// -----
+
+// Test success and failure cases of native constraints with pdl.range results.
+module @patterns {
+  pdl_interp.func @matcher(%root : !pdl.operation) {
+    pdl_interp.check_operation_name of %root is "test.success_op" -> ^pat, ^end
+  
+  ^pat:
+    %num_results = pdl_interp.create_attribute 2 : i32
+    %types = pdl_interp.apply_constraint "op_constr_return_type_range"(%root, %num_results : !pdl.operation, !pdl.attribute) : !pdl.range<type> -> ^pat1, ^end
+
+  ^pat1:
+    pdl_interp.record_match @rewriters::@success(%root, %types : !pdl.operation, !pdl.range<type>) : benefit(1), loc([%root]) -> ^end
+
+  ^end:
+    pdl_interp.finalize
+  }
+
+  module @rewriters {
+    pdl_interp.func @success(%root : !pdl.operation, %types : !pdl.range<type>) {
+      %op = pdl_interp.create_operation "test.replaced_by_pattern" -> (%types : !pdl.range<type>)
+      pdl_interp.erase %root
+      pdl_interp.finalize
+    }
+  }
+}
+
+// CHECK-LABEL: test.apply_constraint_5
+// CHECK-NOT: "test.replaced_by_pattern"
+// CHECK: "test.replaced_by_pattern"() : () -> (f32, f32)
+module @ir attributes { test.apply_constraint_5 } {
+  "test.failure_op"() : () -> ()
+  "test.success_op"() : () -> ()
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // pdl_interp::ApplyRewriteOp
 //===----------------------------------------------------------------------===//
 
 module @patterns {
   pdl_interp.func @matcher(%root : !pdl.operation) {
     pdl_interp.check_operation_name of %root is "test.op" -> ^pat, ^end
 

mlir/test/lib/Rewrite/TestPDLByteCode.cpp

 }
 static LogicalResult customMultiEntityVariadicConstraint(
     PatternRewriter &rewriter, ValueRange operandValues, TypeRange typeValues) {
   if (operandValues.size() != 2 || typeValues.size() != 2)
     return failure();
   return success();
 }
 
+// Custom constraint that returns a value if the op is named test.success_op
+static LogicalResult customValueResultConstraint(PatternRewriter &rewriter,
+                                                 PDLResultList &results,
+                                                 ArrayRef<PDLValue> args) {
+  auto *op = args[0].cast<Operation *>();
+  if (op->getName().getStringRef() == "test.success_op") {
+    StringAttr customAttr = rewriter.getStringAttr("test.success");
+    results.push_back(customAttr);
+    return success();
+  }
+  return failure();
+}
+
+// Custom constraint that returns a type if the op is named test.success_op
+static LogicalResult customTypeResultConstraint(PatternRewriter &rewriter,
+                                                PDLResultList &results,
+                                                ArrayRef<PDLValue> args) {
+  auto *op = args[0].cast<Operation *>();
+  if (op->getName().getStringRef() == "test.success_op") {
+    results.push_back(rewriter.getF32Type());
+    return success();
+  }
+  return failure();
+}
+
+// Custom constraint that returns a type range of variable length if the op is
+// named test.success_op
+static LogicalResult customTypeRangeResultConstraint(PatternRewriter &rewriter,
+                                                     PDLResultList &results,
+                                                     ArrayRef<PDLValue> args) {
+  auto *op = args[0].cast<Operation *>();
+  int numTypes = args[1].cast<Attribute>().cast<IntegerAttr>().getInt();
+
+  if (op->getName().getStringRef() == "test.success_op") {
+    SmallVector<Type> types;
+    for (int i = 0; i < numTypes; i++) {
+      types.push_back(rewriter.getF32Type());
+    }
+    results.push_back(TypeRange(types));
+    return success();
+  }
+  return failure();
+}
+
 // Custom creator invoked from PDL.
 static Operation *customCreate(PatternRewriter &rewriter, Operation *op) {
   return rewriter.create(OperationState(op->getLoc(), "test.success"));
 }
 static auto customVariadicResultCreate(PatternRewriter &rewriter,
                                        Operation *root) {
   return std::make_pair(root->getOperands(), root->getOperands().getTypes());
 }
void runOnOperation() final {
     // will be ignored). This tests that we support separating the registration
     // of library functions from the construction of patterns, and also that we
     // allow multiple patterns to depend on the same library functions (without
     // asserting/crashing).
     pdlPattern.registerConstraintFunction("multi_entity_constraint",
                                           customMultiEntityConstraint);
     pdlPattern.registerConstraintFunction("multi_entity_var_constraint",
                                           customMultiEntityVariadicConstraint);
+    pdlPattern.registerConstraintFunctionWithResults(
+        "op_constr_return_attr", customValueResultConstraint);
+    pdlPattern.registerConstraintFunctionWithResults(
+        "op_constr_return_type", customTypeResultConstraint);
+    pdlPattern.registerConstraintFunctionWithResults(
+        "op_constr_return_type_range", customTypeRangeResultConstraint);
     pdlPattern.registerRewriteFunction("creator", customCreate);
     pdlPattern.registerRewriteFunction("var_creator",
                                        customVariadicResultCreate);
     pdlPattern.registerRewriteFunction("type_creator", customCreateType);
     pdlPattern.registerRewriteFunction("str_creator", customCreateStrAttr);
     pdlPattern.registerRewriteFunction("rewriter", customRewriter);
     patternList.add(std::move(pdlPattern));
 

mlir/test/mlir-pdll/Parser/constraint-failure.pdll

Pattern {
   // CHECK: external declarations must be declared in global scope
   Constraint ExternalConstraint();
 }
 
 // -----
 
 // CHECK: expected `;` after native declaration
 Constraint Foo() [{}]
-
-// -----
-
-// CHECK: native Constraints currently do not support returning results
-Constraint Foo() -> Op;

mlir/test/mlir-pdll/Parser/constraint.pdll

 // RUN: mlir-pdll %s -I %S -split-input-file | FileCheck %s
 
 // CHECK:  Module
 // CHECK:  `-UserConstraintDecl {{.*}} Name<Foo> ResultType<Tuple<>>
 Constraint Foo();
 
 // -----
 
 // CHECK:  Module
 // CHECK:  `-UserConstraintDecl {{.*}} Name<Foo> ResultType<Tuple<>> Code< /* Native Code */ >
 Constraint Foo() [{ /* Native Code */ }];
 
 // -----
 
+// Test that native constraints support returning results.
+
+// CHECK:  Module
+// CHECK:  `-UserConstraintDecl {{.*}} Name<Foo> ResultType<Attr>
+Constraint Foo() -> Attr;
+
+// -----
+
 // CHECK: Module
 // CHECK: `-UserConstraintDecl {{.*}} Name<Foo> ResultType<Value>
 // CHECK:   `Inputs`
 // CHECK:     `-VariableDecl {{.*}} Name<arg> Type<Value>
 // CHECK:   `Results`
 // CHECK:     `-VariableDecl {{.*}} Name<> Type<Value>
 // CHECK:   `-CompoundStmt {{.*}}
 // CHECK:     `-ReturnStmt {{.*}}