[MLIR][PDL] Add PDLL support for negated native constraints

[martin-luecke]

Summary

This commit enables the expression of negated native constraints in PDLL:

If a constraint is prefixed with "not" it is parsed as a negated constraint and hence the attribute isNegated of the emitted pdl.apply_native_constraint operation is set to true.
In first instance this is only supported for the calling of external native C++ constraints and generation of PDL patterns.

Previously, negating a native constraint would have been handled by creating an additional native call, e.g.

PDLL
Constraint checkA(input: Attr);
Constarint checkNotA(input: Attr);

or by including an explicit additional operand for negation, e.g.
Constraint checkA(input: Attr, negated: Attr);

With this a constraint can simply be negated by prefixing it with not. e.g.

PDLL
Constraint simpleConstraint(op: Op);

Pattern example {
    let inputOp = op<test.bar>() ->(type: Type);
    let root = op<test.foo>(inputOp.0) -> ();
    not simpleConstraint(inputOp);
    simpleConstraint(root);
    erase root;
}

Depends on D153871

[Mogball]

LGTM but please wait for @rriddle to review

mlir/lib/Tools/PDLL/CodeGen/MLIRGen.cpp

592

can you guard this on std::is_same_v<PDLOpT, ApplyConstantOp> and then use setIsNegated function directly?

mlir/test/mlir-pdll/Parser/expr.pdll

63

I feel like not TestConstraint is a more intuitive syntax here

[martin-luecke]

Address comments:

• setting attribute directly rather than through the general interface
• syntax for negation changed from ! to not

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

 //===----------------------------------------------------------------------===//
 
 /// This expression represents a call to a decl, such as a
 /// UserConstraintDecl/UserRewriteDecl.
 class CallExpr final : public Node::NodeBase<CallExpr, Expr>,
                        private llvm::TrailingObjects<CallExpr, Expr *> {
 public:
   static CallExpr *create(Context &ctx, SMRange loc, Expr *callable,
-                          ArrayRef<Expr *> arguments, Type resultType);
+                          ArrayRef<Expr *> arguments, Type resultType,
+                          bool isNegated = false);
 
   /// Return the callable of this call.
   Expr *getCallableExpr() const { return callable; }
 
   /// Return the arguments of this call.
   MutableArrayRef<Expr *> getArguments() {
     return {getTrailingObjects<Expr *>(), numArgs};
   }
   ArrayRef<Expr *> getArguments() const {
     return const_cast<CallExpr *>(this)->getArguments();
   }
 
+  /// Returns whether the result of this call is to be negated.
+  bool getIsNegated() const { return isNegated; }
+
 private:
-  CallExpr(SMRange loc, Type type, Expr *callable, unsigned numArgs)
-      : Base(loc, type), callable(callable), numArgs(numArgs) {}
+  CallExpr(SMRange loc, Type type, Expr *callable, unsigned numArgs,
+           bool isNegated)
+      : Base(loc, type), callable(callable), numArgs(numArgs),
+        isNegated(isNegated) {}
 
   /// The callable of this call.
   Expr *callable;
 
   /// The number of arguments of the call.
   unsigned numArgs;
 
   /// TrailingObject utilities.
   friend llvm::TrailingObjects<CallExpr, Expr *>;
+
+  // Is the result of this call to be negated.
+  bool isNegated;
 };
 
 //===----------------------------------------------------------------------===//
 // DeclRefExpr
 //===----------------------------------------------------------------------===//
 
 /// This expression represents a reference to a Decl node.
 class DeclRefExpr : public Node::NodeBase<DeclRefExpr, Expr> {

mlir/lib/Tools/PDLL/AST/NodePrinter.cpp

 
 void NodePrinter::printImpl(const AttributeExpr *expr) {
   os << "AttributeExpr " << expr << " Value<\"" << expr->getValue() << "\">\n";
 }
 
 void NodePrinter::printImpl(const CallExpr *expr) {
   os << "CallExpr " << expr << " Type<";
   print(expr->getType());
-  os << ">\n";
+  os << ">";
+  if (expr->getIsNegated())
+    os << " Negated";
+  os << "\n";
   printChildren(expr->getCallableExpr());
   printChildren("Arguments", expr->getArguments());
 }
 
 void NodePrinter::printImpl(const DeclRefExpr *expr) {
   os << "DeclRefExpr " << expr << " Type<";
   print(expr->getType());
   os << ">\n";

mlir/lib/Tools/PDLL/AST/Nodes.cpp

return new (ctx.getAllocator().Allocate<AttributeExpr>())
       AttributeExpr(ctx, loc, copyStringWithNull(ctx, value));
 }
 
 //===----------------------------------------------------------------------===//
 // CallExpr
 //===----------------------------------------------------------------------===//
 
 CallExpr *CallExpr::create(Context &ctx, SMRange loc, Expr *callable,
-                           ArrayRef<Expr *> arguments, Type resultType) {
+                           ArrayRef<Expr *> arguments, Type resultType,
+                           bool isNegated) {
   unsigned allocSize = CallExpr::totalSizeToAlloc<Expr *>(arguments.size());
   void *rawData = ctx.getAllocator().Allocate(allocSize, alignof(CallExpr));
 
-  CallExpr *expr =
-      new (rawData) CallExpr(loc, resultType, callable, arguments.size());
+  CallExpr *expr = new (rawData)
+      CallExpr(loc, resultType, callable, arguments.size(), isNegated);
   std::uninitialized_copy(arguments.begin(), arguments.end(),
                           expr->getArguments().begin());
   return expr;
 }
 
 //===----------------------------------------------------------------------===//
 // DeclRefExpr
 //===----------------------------------------------------------------------===//

mlir/lib/Tools/PDLL/CodeGen/MLIRGen.cpp

private:
   SmallVector<Value> genExprImpl(const ast::DeclRefExpr *expr);
   Value genExprImpl(const ast::MemberAccessExpr *expr);
   Value genExprImpl(const ast::OperationExpr *expr);
   Value genExprImpl(const ast::RangeExpr *expr);
   SmallVector<Value> genExprImpl(const ast::TupleExpr *expr);
   Value genExprImpl(const ast::TypeExpr *expr);
 
   SmallVector<Value> genConstraintCall(const ast::UserConstraintDecl *decl,
-                                       Location loc, ValueRange inputs);
+                                       Location loc, ValueRange inputs,
+                                       bool isNegated = false);
   SmallVector<Value> genRewriteCall(const ast::UserRewriteDecl *decl,
                                     Location loc, ValueRange inputs);
   template <typename PDLOpT, typename T>
   SmallVector<Value> genConstraintOrRewriteCall(const T *decl, Location loc,
-                                                ValueRange inputs);
+                                                ValueRange inputs,
+                                                bool isNegated = false);
 
   //===--------------------------------------------------------------------===//
   // Fields
   //===--------------------------------------------------------------------===//
 
   /// The MLIR builder used for building the resultant IR.
   OpBuilder builder;
 
SmallVector<Value> CodeGen::genExprImpl(const ast::CallExpr *expr) {
 
   // Resolve the callable expression of this call.
   auto *callableExpr = dyn_cast<ast::DeclRefExpr>(expr->getCallableExpr());
   assert(callableExpr && "unhandled CallExpr callable");
 
   // Generate the PDL based on the type of callable.
   const ast::Decl *callable = callableExpr->getDecl();
   if (const auto *decl = dyn_cast<ast::UserConstraintDecl>(callable))
-    return genConstraintCall(decl, loc, arguments);
+    return genConstraintCall(decl, loc, arguments, expr->getIsNegated());
   if (const auto *decl = dyn_cast<ast::UserRewriteDecl>(callable))
     return genRewriteCall(decl, loc, arguments);
   llvm_unreachable("unhandled CallExpr callable");
 }
 
 SmallVector<Value> CodeGen::genExprImpl(const ast::DeclRefExpr *expr) {
   if (const auto *varDecl = dyn_cast<ast::VariableDecl>(expr->getDecl()))
     return genVar(varDecl);
Value CodeGen::genExprImpl(const ast::TypeExpr *expr) {
   assert(type && "invalid MLIR type data");
   return builder.create<pdl::TypeOp>(genLoc(expr->getLoc()),
                                      builder.getType<pdl::TypeType>(),
                                      TypeAttr::get(type));
 }
 
 SmallVector<Value>
 CodeGen::genConstraintCall(const ast::UserConstraintDecl *decl, Location loc,
-                           ValueRange inputs) {
+                           ValueRange inputs, bool isNegated) {
   // Apply any constraints defined on the arguments to the input values.
   for (auto it : llvm::zip(decl->getInputs(), inputs))
     applyVarConstraints(std::get<0>(it), std::get<1>(it));
 
   // Generate the constraint call.
   SmallVector<Value> results =
-      genConstraintOrRewriteCall<pdl::ApplyNativeConstraintOp>(decl, loc,
-                                                               inputs);
+      genConstraintOrRewriteCall<pdl::ApplyNativeConstraintOp>(
+          decl, loc, inputs, isNegated);
 
   // Apply any constraints defined on the results of the constraint.
   for (auto it : llvm::zip(decl->getResults(), results))
     applyVarConstraints(std::get<0>(it), std::get<1>(it));
   return results;
 }
 
 SmallVector<Value> CodeGen::genRewriteCall(const ast::UserRewriteDecl *decl,
                                            Location loc, ValueRange inputs) {
   return genConstraintOrRewriteCall<pdl::ApplyNativeRewriteOp>(decl, loc,
                                                                inputs);
 }
 
 template <typename PDLOpT, typename T>
-SmallVector<Value> CodeGen::genConstraintOrRewriteCall(const T *decl,
-                                                       Location loc,
-                                                       ValueRange inputs) {
+SmallVector<Value>
+CodeGen::genConstraintOrRewriteCall(const T *decl, Location loc,
+                                    ValueRange inputs, bool isNegated) {
   const ast::CompoundStmt *cstBody = decl->getBody();
 
   // If the decl doesn't have a statement body, it is a native decl.
   if (!cstBody) {
     ast::Type declResultType = decl->getResultType();
     SmallVector<Type> resultTypes;
     if (ast::TupleType tupleType = declResultType.dyn_cast<ast::TupleType>()) {
       for (ast::Type type : tupleType.getElementTypes())
         resultTypes.push_back(genType(type));
     } else {
       resultTypes.push_back(genType(declResultType));
     }
-    Operation *pdlOp = builder.create<PDLOpT>(
+    PDLOpT pdlOp = builder.create<PDLOpT>(
         loc, resultTypes, decl->getName().getName(), inputs);
+    if (isNegated && std::is_same_v<PDLOpT, pdl::ApplyNativeConstraintOp>)
+      cast<pdl::ApplyNativeConstraintOp>(pdlOp).setIsNegated(true);
     return pdlOp->getResults();
   }
 
   // Otherwise, this is a PDLL decl.
   VariableMapTy::ScopeTy varScope(variables);
 
   // Map the inputs of the call to the decl arguments.
   // Note: This is only valid because we do not support recursion, meaning

mlir/lib/Tools/PDLL/Parser/Lexer.h

enum Kind {
     kw_type,
     KW_DEPENDENT_END,
 
     /// General keywords.
     kw_Attr,
     kw_erase,
     kw_let,
     kw_Constraint,
+    kw_not,
     kw_Op,
     kw_OpName,
     kw_Pattern,
     kw_replace,
     kw_return,
     kw_rewrite,
     kw_Rewrite,
     kw_Type,

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

Token Lexer::lexIdentifier(const char *tokStart) {
   // Check to see if this identifier is a keyword.
   StringRef str(tokStart, curPtr - tokStart);
   Token::Kind kind = StringSwitch<Token::Kind>(str)
                          .Case("attr", Token::kw_attr)
                          .Case("Attr", Token::kw_Attr)
                          .Case("erase", Token::kw_erase)
                          .Case("let", Token::kw_let)
                          .Case("Constraint", Token::kw_Constraint)
+                         .Case("not", Token::kw_not)
                          .Case("op", Token::kw_op)
                          .Case("Op", Token::kw_Op)
                          .Case("OpName", Token::kw_OpName)
                          .Case("Pattern", Token::kw_Pattern)
                          .Case("replace", Token::kw_replace)
                          .Case("return", Token::kw_return)
                          .Case("rewrite", Token::kw_rewrite)
                          .Case("Rewrite", Token::kw_Rewrite)

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

private:
 
   //===--------------------------------------------------------------------===//
   // Exprs
 
   FailureOr<ast::Expr *> parseExpr();
 
   /// Identifier expressions.
   FailureOr<ast::Expr *> parseAttributeExpr();
-  FailureOr<ast::Expr *> parseCallExpr(ast::Expr *parentExpr);
+  FailureOr<ast::Expr *> parseCallExpr(ast::Expr *parentExpr,
+                                       bool isNegated = false);
   FailureOr<ast::Expr *> parseDeclRefExpr(StringRef name, SMRange loc);
   FailureOr<ast::Expr *> parseIdentifierExpr();
   FailureOr<ast::Expr *> parseInlineConstraintLambdaExpr();
   FailureOr<ast::Expr *> parseInlineRewriteLambdaExpr();
   FailureOr<ast::Expr *> parseMemberAccessExpr(ast::Expr *parentExpr);
+  FailureOr<ast::Expr *> parseNegatedExpr();
   FailureOr<ast::OpNameDecl *> parseOperationName(bool allowEmptyName = false);
   FailureOr<ast::OpNameDecl *> parseWrappedOperationName(bool allowEmptyName);
   FailureOr<ast::Expr *>
   parseOperationExpr(OpResultTypeContext inputResultTypeContext =
                          OpResultTypeContext::Explicit);
   FailureOr<ast::Expr *> parseTupleExpr();
   FailureOr<ast::Expr *> parseTypeExpr();
   FailureOr<ast::Expr *> parseUnderscoreExpr();
private:
   LogicalResult validateTypeConstraintExpr(const ast::Expr *typeExpr);
   LogicalResult validateTypeRangeConstraintExpr(const ast::Expr *typeExpr);
 
   //===--------------------------------------------------------------------===//
   // Exprs
 
   FailureOr<ast::CallExpr *>
   createCallExpr(SMRange loc, ast::Expr *parentExpr,
-                 MutableArrayRef<ast::Expr *> arguments);
+                 MutableArrayRef<ast::Expr *> arguments,
+                 bool isNegated = false);
   FailureOr<ast::DeclRefExpr *> createDeclRefExpr(SMRange loc, ast::Decl *decl);
   FailureOr<ast::DeclRefExpr *>
   createInlineVariableExpr(ast::Type type, StringRef name, SMRange loc,
                            ArrayRef<ast::ConstraintRef> constraints);
   FailureOr<ast::MemberAccessExpr *>
   createMemberAccessExpr(ast::Expr *parentExpr, StringRef name, SMRange loc);
 
   /// Validate the member access `name` into the given parent expression. On
FailureOr<ast::Expr *> Parser::parseExpr() {
   FailureOr<ast::Expr *> lhsExpr;
   switch (curToken.getKind()) {
   case Token::kw_attr:
     lhsExpr = parseAttributeExpr();
     break;
   case Token::kw_Constraint:
     lhsExpr = parseInlineConstraintLambdaExpr();
     break;
+  case Token::kw_not:
+    lhsExpr = parseNegatedExpr();
+    break;
   case Token::identifier:
     lhsExpr = parseIdentifierExpr();
     break;
   case Token::kw_op:
     lhsExpr = parseOperationExpr();
     break;
   case Token::kw_Rewrite:
     lhsExpr = parseInlineRewriteLambdaExpr();
FailureOr<ast::Expr *> Parser::parseAttributeExpr() {
 
   loc.End = curToken.getEndLoc();
   if (failed(
           parseToken(Token::greater, "expected `>` after attribute literal")))
     return failure();
   return ast::AttributeExpr::create(ctx, loc, attrExpr);
 }
 
-FailureOr<ast::Expr *> Parser::parseCallExpr(ast::Expr *parentExpr) {
+FailureOr<ast::Expr *> Parser::parseCallExpr(ast::Expr *parentExpr,
+                                             bool isNegated) {
   consumeToken(Token::l_paren);
 
   // Parse the arguments of the call.
   SmallVector<ast::Expr *> arguments;
   if (curToken.isNot(Token::r_paren)) {
     do {
       // Handle code completion for the call arguments.
       if (curToken.is(Token::code_complete)) {
         codeCompleteCallSignature(parentExpr, arguments.size());
         return failure();
       }
 
       FailureOr<ast::Expr *> argument = parseExpr();
       if (failed(argument))
         return failure();
       arguments.push_back(*argument);
     } while (consumeIf(Token::comma));
   }
 
   SMRange loc(parentExpr->getLoc().Start, curToken.getEndLoc());
   if (failed(parseToken(Token::r_paren, "expected `)` after argument list")))
     return failure();
 
-  return createCallExpr(loc, parentExpr, arguments);
+  return createCallExpr(loc, parentExpr, arguments, isNegated);
 }
 
 FailureOr<ast::Expr *> Parser::parseDeclRefExpr(StringRef name, SMRange loc) {
   ast::Decl *decl = curDeclScope->lookup(name);
   if (!decl)
     return emitError(loc, "undefined reference to `" + name + "`");
 
   return createDeclRefExpr(loc, decl);
if (memberNameTok.isNot(Token::identifier, Token::integer) &&
     return emitError(dotLoc, "expected identifier or numeric member name");
   StringRef memberName = memberNameTok.getSpelling();
   SMRange loc(parentExpr->getLoc().Start, curToken.getEndLoc());
   consumeToken();
 
   return createMemberAccessExpr(parentExpr, memberName, loc);
 }
 
+FailureOr<ast::Expr *> Parser::parseNegatedExpr() {
+  consumeToken(Token::kw_not);
+  // Only native constraints are supported after negation
+  if (!curToken.is(Token::identifier))
+    return emitError("expected native constraint");
+  FailureOr<ast::Expr *> identifierExpr = parseIdentifierExpr();
+  if (failed(identifierExpr))
+    return failure();
+  return parseCallExpr(*identifierExpr, /*isNegated = */ true);
+}
+
 FailureOr<ast::OpNameDecl *> Parser::parseOperationName(bool allowEmptyName) {
   SMRange loc = curToken.getLoc();
 
   // Check for code completion for the dialect name.
   if (curToken.is(Token::code_complete))
     return codeCompleteDialectName();
 
   // Handle the case of an no operation name.
Parser::validateTypeRangeConstraintExpr(const ast::Expr *typeExpr) {
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // Exprs
 
 FailureOr<ast::CallExpr *>
 Parser::createCallExpr(SMRange loc, ast::Expr *parentExpr,
-                       MutableArrayRef<ast::Expr *> arguments) {
+                       MutableArrayRef<ast::Expr *> arguments, bool isNegated) {
   ast::Type parentType = parentExpr->getType();
 
   ast::CallableDecl *callableDecl = tryExtractCallableDecl(parentExpr);
   if (!callableDecl) {
     return emitError(loc,
                      llvm::formatv("expected a reference to a callable "
                                    "`Constraint` or `Rewrite`, but got: `{0}`",
                                    parentType));
   }
   if (parserContext == ParserContext::Rewrite) {
     if (isa<ast::UserConstraintDecl>(callableDecl))
       return emitError(
           loc, "unable to invoke `Constraint` within a rewrite section");
-  } else if (isa<ast::UserRewriteDecl>(callableDecl)) {
-    return emitError(loc, "unable to invoke `Rewrite` within a match section");
+    if (isNegated)
+      return emitError(loc, "unable to negate a Rewrite");
+  } else {
+    if (isa<ast::UserRewriteDecl>(callableDecl))
+      return emitError(loc,
+                       "unable to invoke `Rewrite` within a match section");
+    if (isNegated && cast<ast::UserConstraintDecl>(callableDecl)->getBody())
+      return emitError(loc, "unable to negate non native constraints");
   }
 
   // Verify the arguments of the call.
   /// Handle size mismatch.
   ArrayRef<ast::VariableDecl *> callArgs = callableDecl->getInputs();
   if (callArgs.size() != arguments.size()) {
     return emitErrorAndNote(
         loc,
Parser::createCallExpr(SMRange loc, ast::Expr *parentExpr,
   };
   for (auto it : llvm::zip(callArgs, arguments)) {
     if (failed(convertExpressionTo(std::get<1>(it), std::get<0>(it)->getType(),
                                    attachDiagFn)))
       return failure();
   }
 
   return ast::CallExpr::create(ctx, loc, parentExpr, arguments,
-                               callableDecl->getResultType());
+                               callableDecl->getResultType(), isNegated);
 }
 
 FailureOr<ast::DeclRefExpr *> Parser::createDeclRefExpr(SMRange loc,
                                                         ast::Decl *decl) {
   // Check the type of decl being referenced.
   ast::Type declType;
   if (isa<ast::ConstraintDecl>(decl))
     declType = ast::ConstraintType::get(ctx);

mlir/test/mlir-pdll/CodeGen/MLIR/expr.pdll

 // CHECK: rewrite %[[ROOT]]
 // CHECK: %[[RESULTS:.*]] = apply_native_rewrite "TestRewrite"(%[[ROOT]] : !pdl.operation) : !pdl.range<value>
 // CHECK: replace %[[ROOT]] with(%[[RESULTS]] : !pdl.range<value>)
 Rewrite TestRewrite(op: Op) -> ValueRange;
 Pattern TestExternalCall => replace root: Op with TestRewrite(root);
 
 // -----
 
+// CHECK: pdl.pattern @TestExternalNegatedCall
+// CHECK: %[[ROOT:.*]] = operation
+// CHECK: apply_native_constraint  "TestConstraint"(%[[ROOT]] : !pdl.operation) {isNegated = true}
+// CHECK: rewrite %[[ROOT]]
+// CHECK: erase %[[ROOT]]
+Constraint TestConstraint(op: Op);
+Pattern TestExternalNegatedCall {
+    let root = op : Op;
+    not TestConstraint(root);
+    erase root;
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // MemberAccessExpr
 //===----------------------------------------------------------------------===//
 
 // Handle implicit "all" operation results access.
 // CHECK: pdl.pattern @OpAllResultMemberAccess
 // CHECK: %[[OP0:.*]] = operation
 // CHECK: %[[OP0_RES:.*]] = result 0 of %[[OP0]]

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

 Pattern {
   // CHECK: unable to build a tuple with `Constraint` element
   let tuple = (Foo);
   erase op<>;
 }
 
 // -----
 
+Constraint Foo(op: Op) {}
+
+Pattern {
+  // CHECK: unable to negate non native constraints
+  let root = op<>;
+  not Foo(root);
+}
+
+// -----
+
 Rewrite Foo();
 
 Pattern {
   // CHECK: unable to build a tuple with `Rewrite` element
   let tuple = (Foo);
   erase op<>;
 }
 
 // -----
 
+Rewrite Foo(op: Op);
+
+Pattern {
+  // CHECK: unable to negate a Rewrite
+  let root = op<>;
+  rewrite root with {
+     not Foo(root);
+  }
+}
+
+// -----
+
 Pattern {
   // CHECK: expected expression
   let tuple = (10 = _: Value);
   erase op<>;
 }
 
 // -----
 

mlir/test/mlir-pdll/Parser/expr.pdll

 
 Pattern {
   let inputOp = op<my_dialect.bar>;
   replace op<my_dialect.bar>(inputOp) with CreateNewOp(inputOp);
 }
 
 // -----
 
+// CHECK: Module {{.*}}
+// CHECK:  -UserConstraintDecl {{.*}} Name<TestConstraint> ResultType<Tuple<>>
+// CHECK: `-PatternDecl {{.*}}
+// CHECK:   -CallExpr {{.*}} Type<Tuple<>> Negated
+// CHECK:     `-DeclRefExpr {{.*}} Type<Constraint>
+// CHECK:       `-UserConstraintDecl {{.*}} Name<TestConstraint> ResultType<Tuple<>>
+Constraint TestConstraint(op: Op);
+
+Pattern {
+  let inputOp = op<my_dialect.bar>;
+  not TestConstraint(inputOp);
+  erase inputOp;
+}
+
+// -----
+
 //===----------------------------------------------------------------------===//
 // MemberAccessExpr
 //===----------------------------------------------------------------------===//
 
 // CHECK: Module
 // CHECK: `-VariableDecl {{.*}} Name<firstEltIndex> Type<Op>
 // CHECK:   `-MemberAccessExpr {{.*}} Member<0> Type<Op>
 // CHECK:     `-DeclRefExpr {{.*}} Type<Tuple<firstElt: Op>>