[mlir:PDLL] Allow complex constraints on Rewrite arguments/results

[rriddle]

Summary

The documentation already has examples of this, and it allows for
using nicer C++ types when defining native Rewrites.

mlir/include/mlir/Tools/PDLL/AST/Nodes.h

createImpl(Context &ctx, const Name &name, ArrayRef<VariableDecl *> inputs,
              const CompoundStmt *body, Type resultType);
 
   UserConstraintDecl(const Name &name, unsigned numInputs,
                      bool hasNativeInputTypes, unsigned numResults,
                      Optional<StringRef> codeBlock, const CompoundStmt *body,
                      Type resultType)
       : Base(name.getLoc(), &name), numInputs(numInputs),
         numResults(numResults), codeBlock(codeBlock), constraintBody(body),
-        resultType(resultType) {}
+        resultType(resultType), hasNativeInputTypes(hasNativeInputTypes) {}
 
   /// The number of inputs to this constraint.
   unsigned numInputs;
 
   /// The number of explicit results to this constraint.
   unsigned numResults;
 
   /// The optional code block of this constraint.

mlir/include/mlir/Tools/PDLL/Parser/CodeComplete.h

public:
   /// Signal code completion for a member access into the given operation type.
   virtual void codeCompleteOperationMemberAccess(ast::OperationType opType);
 
   /// Signal code completion for a member access into the given operation type.
   virtual void codeCompleteOperationAttributeName(StringRef opName) {}
 
   /// Signal code completion for a constraint name with an optional decl scope.
   /// `currentType` is the current type of the variable that will use the
-  /// constraint, or nullptr if a type is unknown. `allowNonCoreConstraints`
-  /// indicates if user defined constraints are allowed in the completion
-  /// results. `allowInlineTypeConstraints` enables inline type constraints for
-  /// Attr/Value/ValueRange.
+  /// constraint, or nullptr if a type is unknown. `allowInlineTypeConstraints`
+  /// enables inline type constraints for Attr/Value/ValueRange.
   virtual void codeCompleteConstraintName(ast::Type currentType,
-                                          bool allowNonCoreConstraints,
                                           bool allowInlineTypeConstraints,
                                           const ast::DeclScope *scope);
 
   /// Signal code completion for a dialect name.
   virtual void codeCompleteDialectName() {}
 
   /// Signal code completion for an operation name in the given dialect.
   virtual void codeCompleteOperationName(StringRef dialectName) {}

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

 CodeCompleteContext::~CodeCompleteContext() = default;
 
 void CodeCompleteContext::codeCompleteTupleMemberAccess(
     ast::TupleType tupleType) {}
 void CodeCompleteContext::codeCompleteOperationMemberAccess(
     ast::OperationType opType) {}
 
 void CodeCompleteContext::codeCompleteConstraintName(
-    ast::Type currentType, bool allowNonCoreConstraints,
-    bool allowInlineTypeConstraints, const ast::DeclScope *scope) {}
+    ast::Type currentType, bool allowInlineTypeConstraints,
+    const ast::DeclScope *scope) {}

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

private:
   FailureOr<ast::Expr *> parseTypeConstraintExpr();
 
   /// Try to parse a single reference to a constraint. `typeConstraint` is the
   /// location of a previously parsed type constraint for the entity that will
   /// be constrained by the parsed constraint. `existingConstraints` are any
   /// existing constraints that have already been parsed for the same entity
   /// that will be constrained by this constraint. `allowInlineTypeConstraints`
   /// allows the use of inline Type constraints, e.g. `Value<valueType: Type>`.
-  /// If `allowNonCoreConstraints` is true, then complex (e.g. user defined
-  /// constraints) may be used with the variable.
   FailureOr<ast::ConstraintRef>
   parseConstraint(Optional<SMRange> &typeConstraint,
                   ArrayRef<ast::ConstraintRef> existingConstraints,
-                  bool allowInlineTypeConstraints,
-                  bool allowNonCoreConstraints);
+                  bool allowInlineTypeConstraints);
 
   /// Try to parse the constraint for a UserConstraintDecl/UserRewriteDecl
   /// argument or result variable. The constraints for these variables do not
   /// allow inline type constraints, and only permit a single constraint.
   FailureOr<ast::ConstraintRef> parseArgOrResultConstraint();
 
   //===--------------------------------------------------------------------===//
   // Exprs
private:
   FailureOr<ast::VariableDecl *>
   createArgOrResultVariableDecl(StringRef name, SMRange loc,
                                 const ast::ConstraintRef &constraint);
 
   /// Validate the constraints used to constraint a variable decl.
   /// `inferredType` is the type of the variable inferred by the constraints
   /// within the list, and is updated to the most refined type as determined by
   /// the constraints. Returns success if the constraint list is valid, failure
-  /// otherwise. If `allowNonCoreConstraints` is true, then complex (e.g. user
-  /// defined constraints) may be used with the variable.
+  /// otherwise.
   LogicalResult
   validateVariableConstraints(ArrayRef<ast::ConstraintRef> constraints,
-                              ast::Type &inferredType,
-                              bool allowNonCoreConstraints = true);
+                              ast::Type &inferredType);
   /// Validate a single reference to a constraint. `inferredType` contains the
   /// currently inferred variabled type and is refined within the type defined
   /// by the constraint. Returns success if the constraint is valid, failure
-  /// otherwise. If `allowNonCoreConstraints` is true, then complex (e.g. user
-  /// defined constraints) may be used with the variable.
+  /// otherwise.
   LogicalResult validateVariableConstraint(const ast::ConstraintRef &ref,
-                                           ast::Type &inferredType,
-                                           bool allowNonCoreConstraints = true);
+                                           ast::Type &inferredType);
   LogicalResult validateTypeConstraintExpr(const ast::Expr *typeExpr);
   LogicalResult validateTypeRangeConstraintExpr(const ast::Expr *typeExpr);
 
   //===--------------------------------------------------------------------===//
   // Exprs
 
   FailureOr<ast::CallExpr *>
   createCallExpr(SMRange loc, ast::Expr *parentExpr,
private:
 
   /// The set of various code completion methods. Every completion method
   /// returns `failure` to stop the parsing process after providing completion
   /// results.
 
   LogicalResult codeCompleteMemberAccess(ast::Expr *parentExpr);
   LogicalResult codeCompleteAttributeName(Optional<StringRef> opName);
   LogicalResult codeCompleteConstraintName(ast::Type inferredType,
-                                           bool allowNonCoreConstraints,
                                            bool allowInlineTypeConstraints);
   LogicalResult codeCompleteDialectName();
   LogicalResult codeCompleteOperationName(StringRef dialectName);
   LogicalResult codeCompletePatternMetadata();
   LogicalResult codeCompleteIncludeFilename(StringRef curPath);
 
   void codeCompleteCallSignature(ast::Node *parent, unsigned currentNumArgs);
   void codeCompleteOperationOperandsSignature(Optional<StringRef> opName,
FailureOr<ast::VariableDecl *> Parser::parseArgumentDecl() {
 
   return createArgOrResultVariableDecl(name, nameLoc, *cst);
 }
 
 FailureOr<ast::VariableDecl *> Parser::parseResultDecl(unsigned resultNum) {
   // Check to see if this result is named.
   if (curToken.is(Token::identifier) || curToken.isDependentKeyword()) {
     // Check to see if this name actually refers to a Constraint.
-    ast::Decl *existingDecl = curDeclScope->lookup(curToken.getSpelling());
-    if (isa_and_nonnull<ast::ConstraintDecl>(existingDecl)) {
-      // If yes, and this is a Rewrite, give a nice error message as non-Core
-      // constraints are not supported on Rewrite results.
-      if (parserContext == ParserContext::Rewrite) {
-        return emitError(
-            "`Rewrite` results are only permitted to use core constraints, "
-            "such as `Attr`, `Op`, `Type`, `TypeRange`, `Value`, `ValueRange`");
-      }
-
-      // Otherwise, parse this as an unnamed result variable.
-    } else {
+    if (!curDeclScope->lookup<ast::ConstraintDecl>(curToken.getSpelling())) {
       // If it wasn't a constraint, parse the result similarly to a variable. If
       // there is already an existing decl, we will emit an error when defining
       // this variable later.
       StringRef name = curToken.getSpelling();
       SMRange nameLoc = curToken.getLoc();
       consumeToken();
 
       if (failed(parseToken(Token::colon,
return defineVariableDecl(name, nameLoc, type, /*initExpr=*/nullptr,
                             constraints);
 }
 
 LogicalResult Parser::parseVariableDeclConstraintList(
     SmallVectorImpl<ast::ConstraintRef> &constraints) {
   Optional<SMRange> typeConstraint;
   auto parseSingleConstraint = [&] {
     FailureOr<ast::ConstraintRef> constraint = parseConstraint(
-        typeConstraint, constraints, /*allowInlineTypeConstraints=*/true,
-        /*allowNonCoreConstraints=*/true);
+        typeConstraint, constraints, /*allowInlineTypeConstraints=*/true);
     if (failed(constraint))
       return failure();
     constraints.push_back(*constraint);
     return success();
   };
 
   // Check to see if this is a single constraint, or a list.
   if (!consumeIf(Token::l_square))
     return parseSingleConstraint();
 
   do {
     if (failed(parseSingleConstraint()))
       return failure();
   } while (consumeIf(Token::comma));
   return parseToken(Token::r_square, "expected `]` after constraint list");
 }
 
 FailureOr<ast::ConstraintRef>
 Parser::parseConstraint(Optional<SMRange> &typeConstraint,
                         ArrayRef<ast::ConstraintRef> existingConstraints,
-                        bool allowInlineTypeConstraints,
-                        bool allowNonCoreConstraints) {
+                        bool allowInlineTypeConstraints) {
   auto parseTypeConstraint = [&](ast::Expr *&typeExpr) -> LogicalResult {
     if (!allowInlineTypeConstraints) {
       return emitError(
           curToken.getLoc(),
           "inline `Attr`, `Value`, and `ValueRange` type constraints are not "
           "permitted on arguments or results");
     }
     if (typeConstraint)
case Token::identifier: {
     return emitErrorAndNote(
         loc, "invalid reference to non-constraint", cstDecl->getLoc(),
         "see the definition of `" + constraintName + "` here");
   }
     // Handle single entity constraint code completion.
   case Token::code_complete: {
     // Try to infer the current type for use by code completion.
     ast::Type inferredType;
-    if (failed(validateVariableConstraints(existingConstraints, inferredType,
-                                           allowNonCoreConstraints)))
+    if (failed(validateVariableConstraints(existingConstraints, inferredType)))
       return failure();
 
-    return codeCompleteConstraintName(inferredType, allowNonCoreConstraints,
-                                      allowInlineTypeConstraints);
+    return codeCompleteConstraintName(inferredType, allowInlineTypeConstraints);
   }
   default:
     break;
   }
   return emitError(loc, "expected identifier constraint");
 }
 
 FailureOr<ast::ConstraintRef> Parser::parseArgOrResultConstraint() {
-  // Constraint arguments may apply more complex constraints via the arguments.
-  bool allowNonCoreConstraints = parserContext == ParserContext::Constraint;
-
   Optional<SMRange> typeConstraint;
   return parseConstraint(typeConstraint, /*existingConstraints=*/llvm::None,
-                         /*allowInlineTypeConstraints=*/false,
-                         allowNonCoreConstraints);
+                         /*allowInlineTypeConstraints=*/false);
 }
 
 //===----------------------------------------------------------------------===//
 // Exprs
 
 FailureOr<ast::Expr *> Parser::parseExpr() {
   if (curToken.is(Token::underscore))
     return parseUnderscoreExpr();
if (failed(varDecl))
     return failure();
 
   return *varDecl;
 }
 
 FailureOr<ast::VariableDecl *>
 Parser::createArgOrResultVariableDecl(StringRef name, SMRange loc,
                                       const ast::ConstraintRef &constraint) {
-  // Constraint arguments may apply more complex constraints via the arguments.
-  bool allowNonCoreConstraints = parserContext == ParserContext::Constraint;
   ast::Type argType;
-  if (failed(validateVariableConstraint(constraint, argType,
-                                        allowNonCoreConstraints)))
+  if (failed(validateVariableConstraint(constraint, argType)))
     return failure();
   return defineVariableDecl(name, loc, argType, constraint);
 }
 
 LogicalResult
 Parser::validateVariableConstraints(ArrayRef<ast::ConstraintRef> constraints,
-                                    ast::Type &inferredType,
-                                    bool allowNonCoreConstraints) {
+                                    ast::Type &inferredType) {
   for (const ast::ConstraintRef &ref : constraints)
-    if (failed(validateVariableConstraint(ref, inferredType,
-                                          allowNonCoreConstraints)))
+    if (failed(validateVariableConstraint(ref, inferredType)))
       return failure();
   return success();
 }
 
 LogicalResult Parser::validateVariableConstraint(const ast::ConstraintRef &ref,
-                                                 ast::Type &inferredType,
-                                                 bool allowNonCoreConstraints) {
+                                                 ast::Type &inferredType) {
   ast::Type constraintType;
   if (const auto *cst = dyn_cast<ast::AttrConstraintDecl>(ref.constraint)) {
     if (const ast::Expr *typeExpr = cst->getTypeExpr()) {
       if (failed(validateTypeConstraintExpr(typeExpr)))
         return failure();
     }
     constraintType = ast::AttributeType::get(ctx);
   } else if (const auto *cst =
constraintType = valueTy;
                  dyn_cast<ast::ValueRangeConstraintDecl>(ref.constraint)) {
     if (const ast::Expr *typeExpr = cst->getTypeExpr()) {
       if (failed(validateTypeRangeConstraintExpr(typeExpr)))
         return failure();
     }
     constraintType = valueRangeTy;
   } else if (const auto *cst =
                  dyn_cast<ast::UserConstraintDecl>(ref.constraint)) {
-    if (!allowNonCoreConstraints) {
-      return emitError(ref.referenceLoc,
-                       "`Rewrite` arguments and results are only permitted to "
-                       "use core constraints, such as `Attr`, `Op`, `Type`, "
-                       "`TypeRange`, `Value`, `ValueRange`");
-    }
-
     ArrayRef<ast::VariableDecl *> inputs = cst->getInputs();
     if (inputs.size() != 1) {
       return emitErrorAndNote(ref.referenceLoc,
                               "`Constraint`s applied via a variable constraint "
                               "list must take a single input, but got " +
                                   Twine(inputs.size()),
                               cst->getLoc(),
                               "see definition of constraint here");

 LogicalResult Parser::codeCompleteAttributeName(Optional<StringRef> opName) {
   if (opName)
     codeCompleteContext->codeCompleteOperationAttributeName(*opName);
   return failure();
 }
 
 LogicalResult
 Parser::codeCompleteConstraintName(ast::Type inferredType,
-                                   bool allowNonCoreConstraints,
                                    bool allowInlineTypeConstraints) {
   codeCompleteContext->codeCompleteConstraintName(
-      inferredType, allowNonCoreConstraints, allowInlineTypeConstraints,
-      curDeclScope);
+      inferredType, allowInlineTypeConstraints, curDeclScope);
   return failure();
 }
 
 LogicalResult Parser::codeCompleteDialectName() {
   codeCompleteContext->codeCompleteDialectName();
   return failure();
 }
 

mlir/lib/Tools/mlir-pdll-lsp-server/PDLLServer.cpp

for (const ods::Attribute &attr : odsOp->getAttributes()) {
                         constraint.getCppClass())
               .str()};
       item.insertTextFormat = lsp::InsertTextFormat::PlainText;
       completionList.items.emplace_back(item);
     }
   }
 
   void codeCompleteConstraintName(ast::Type currentType,
-                                  bool allowNonCoreConstraints,
                                   bool allowInlineTypeConstraints,
                                   const ast::DeclScope *scope) final {
     auto addCoreConstraint = [&](StringRef constraint, StringRef mlirType,
                                  StringRef snippetText = "") {
       lsp::CompletionItem item;
       item.label = constraint.str();
       item.kind = lsp::CompletionItemKind::Class;
       item.detail = (constraint + " constraint").str();
if (allowInlineTypeConstraints) {
         addCoreConstraint("ValueRange<type>", "mlir::ValueRange",
                           "ValueRange<$1>");
     }
 
     // If a scope was provided, check it for potential constraints.
     while (scope) {
       for (const ast::Decl *decl : scope->getDecls()) {
         if (const auto *cst = dyn_cast<ast::UserConstraintDecl>(decl)) {
-          if (!allowNonCoreConstraints)
-            continue;
-
           lsp::CompletionItem item;
           item.label = cst->getName().getName().str();
           item.kind = lsp::CompletionItemKind::Interface;
           item.sortText = "2_" + item.label;
 
           // Skip constraints that are not single-arg. We currently only
           // complete variable constraints.
           if (cst->getInputs().size() != 1)

mlir/test/mlir-pdll/CodeGen/CPP/general.pdll

-// RUN: mlir-pdll %s -I %S -split-input-file -x cpp | FileCheck %s
+// RUN: mlir-pdll %s -I %S -I %S/../../../../include -split-input-file -x cpp | FileCheck %s
 
 // Check that we generate a wrapper pattern for each PDL pattern. Also
 // add in a pattern awkwardly named the same as our generated patterns to
 // check that we handle overlap.
 
 // CHECK: struct GeneratedPDLLPattern0 : ::mlir::PDLPatternModule {
 // CHECK:  template <typename... ConfigsT>
 // CHECK:  : ::mlir::PDLPatternModule(::mlir::parseSourceString<::mlir::ModuleOp>(

 Pattern NamedPattern => erase op<test.op2>;
 Pattern GeneratedPDLLPattern1 => erase op<>;
 Pattern => erase op<test.op3>;
 
 // -----
 
 // Check the generation of native constraints and rewrites.
 
+#include "include/ods.td"
+
 // CHECK:      static ::mlir::LogicalResult TestCstPDLFn(::mlir::PatternRewriter &rewriter,
 // CHECK-SAME:     ::mlir::Attribute attr, ::mlir::Operation * op, ::mlir::Type type,
 // CHECK-SAME:     ::mlir::Value value, ::mlir::TypeRange typeRange, ::mlir::ValueRange valueRange) {
 // CHECK-NEXT:   return success();
 // CHECK: }
 
 // CHECK-NOT: TestUnusedCst
 
 // CHECK: static void TestRewritePDLFn(::mlir::PatternRewriter &rewriter,
 // CHECK-SAME:     ::mlir::Attribute attr, ::mlir::Operation * op, ::mlir::Type type,
 // CHECK-SAME:     ::mlir::Value value, ::mlir::TypeRange typeRange, ::mlir::ValueRange valueRange) {
 // CHECK: foo;
 // CHECK: }
 
+// CHECK: TestAttrInterface TestRewriteODSPDLFn(::mlir::PatternRewriter &rewriter, TestAttrInterface attr) {
 // CHECK: static ::mlir::Attribute TestRewriteSinglePDLFn(::mlir::PatternRewriter &rewriter) {
 // CHECK: std::tuple<::mlir::Attribute, ::mlir::Type> TestRewriteTuplePDLFn(::mlir::PatternRewriter &rewriter) {
 
 // CHECK-NOT: TestUnusedRewrite
 
 // CHECK: struct TestCstAndRewrite : ::mlir::PDLPatternModule {
 // CHECK:   registerConstraintFunction("TestCst", TestCstPDLFn);
 // CHECK:   registerRewriteFunction("TestRewrite", TestRewritePDLFn);
 
 Constraint TestCst(attr: Attr, op: Op, type: Type, value: Value, typeRange: TypeRange, valueRange: ValueRange) [{
   return success();
 }];
 Constraint TestUnusedCst() [{ return success(); }];
 
 Rewrite TestRewrite(attr: Attr, op: Op, type: Type, value: Value, typeRange: TypeRange, valueRange: ValueRange) [{ foo; }];
+Rewrite TestRewriteODS(attr: TestAttrInterface) -> TestAttrInterface [{}];
 Rewrite TestRewriteSingle() -> Attr [{}];
 Rewrite TestRewriteTuple() -> (Attr, Type) [{}];
 Rewrite TestUnusedRewrite(op: Op) [{}];
 
 Pattern TestCstAndRewrite {
   let root = op<>(operand: Value, operands: ValueRange) -> (type: Type, types: TypeRange);
   TestCst(attr<"true">, root, type, operand, types, operands);
   rewrite root with {
     TestRewrite(attr<"true">, root, type, operand, types, operands);
+    TestRewriteODS(attr<"true">);
     TestRewriteSingle();
     TestRewriteTuple();
     erase root;
   };
 }

mlir/test/mlir-pdll/CodeGen/CPP/include/ods.td

+include "mlir/IR/OpBase.td"
+
+def TestAttrInterface : AttrInterface<"TestAttrInterface">;

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

 
 // -----
 
 // CHECK: inline `Attr`, `Value`, and `ValueRange` type constraints are not permitted on arguments or results
 Rewrite Foo(arg: Value<type>){}
 
 // -----
 
-Constraint ValueConstraint(value: Value);
-
-// CHECK: arguments and results are only permitted to use core constraints, such as `Attr`, `Op`, `Type`, `TypeRange`, `Value`, `ValueRange`
-Rewrite Foo(arg: ValueConstraint);
-
-// -----
-
 // CHECK: expected `)` to end argument list
 Rewrite Foo(arg: Value{}
 
 // -----
 
 //===----------------------------------------------------------------------===//
 // Results
 //===----------------------------------------------------------------------===//

 
 // -----
 
 // CHECK: inline `Attr`, `Value`, and `ValueRange` type constraints are not permitted on arguments or results
 Rewrite Foo() -> Value<type>){}
 
 // -----
 
-Constraint ValueConstraint(value: Value);
-
-// CHECK: results are only permitted to use core constraints, such as `Attr`, `Op`, `Type`, `TypeRange`, `Value`, `ValueRange`
-Rewrite Foo() -> ValueConstraint;
-
-// -----
-
 //===----------------------------------------------------------------------===//
 // Native Rewrites
 //===----------------------------------------------------------------------===//
 
 Pattern {
   // CHECK: external declarations must be declared in global scope
   Rewrite ExternalConstraint();
 }
 
 // -----
 
 // CHECK: expected `;` after native declaration
 Rewrite Foo() [{}]