[PDLL] Add a `replace` rewrite statement for replacing operations

[rriddle]

Summary

This statement acts as a companion to the existing erase
statement, and is the corresponding PDLL construct for the
PatternRewriter::replaceOp C++ API. This statement replaces a
given operation with a set of values.

Depends On D115297

[nicolasvasilache]

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

924

nl

[jpienaar]

It would be good to add a multi-op replacement example utilizing single op replacement too

[rriddle]

Update

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

 public:
   static EraseStmt *create(Context &ctx, llvm::SMRange loc, Expr *rootOp);
 
 private:
   EraseStmt(llvm::SMRange loc, Expr *rootOp) : Base(loc, rootOp) {}
 };
 
 //===----------------------------------------------------------------------===//
+// ReplaceStmt
+
+/// This statement represents the `replace` statement in PDLL. This statement
+/// replace the given root operation with a set of values, corresponding roughly
+/// to the PatternRewriter::replaceOp API.
+class ReplaceStmt final : public Node::NodeBase<ReplaceStmt, OpRewriteStmt>,
+                          private llvm::TrailingObjects<ReplaceStmt, Expr *> {
+public:
+  static ReplaceStmt *create(Context &ctx, llvm::SMRange loc, Expr *rootOp,
+                             ArrayRef<Expr *> replExprs);
+
+  /// Return the replacement values of this statement.
+  MutableArrayRef<Expr *> getReplExprs() {
+    return {getTrailingObjects<Expr *>(), numReplExprs};
+  }
+  ArrayRef<Expr *> getReplExprs() const {
+    return const_cast<ReplaceStmt *>(this)->getReplExprs();
+  }
+
+private:
+  ReplaceStmt(llvm::SMRange loc, Expr *rootOp, unsigned numReplExprs)
+      : Base(loc, rootOp), numReplExprs(numReplExprs) {}
+
+  /// The number of replacement values within this statement.
+  unsigned numReplExprs;
+
+  /// TrailingObject utilities.
+  friend class llvm::TrailingObjects<ReplaceStmt, Expr *>;
+};
+
+//===----------------------------------------------------------------------===//
 // Expr
 //===----------------------------------------------------------------------===//
 
 /// This class represents a base AST Expression node.
 class Expr : public Stmt {
 public:
   /// Return the type of this expression.
   Type getType() const { return type; }

 }
 
 inline bool Expr::classof(const Node *node) {
   return isa<AttributeExpr, DeclRefExpr, MemberAccessExpr, OperationExpr,
              TupleExpr, TypeExpr>(node);
 }
 
 inline bool OpRewriteStmt::classof(const Node *node) {
-  return isa<EraseStmt>(node);
+  return isa<EraseStmt, ReplaceStmt>(node);
 }
 
 inline bool Stmt::classof(const Node *node) {
   return isa<CompoundStmt, LetStmt, OpRewriteStmt, Expr>(node);
 }
 
 } // namespace ast
 } // namespace pdll
 } // namespace mlir
 
 #endif // MLIR_TOOLS_PDLL_AST_NODES_H_

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

void printChildren(StringRef label, RangeT &&range) {
     printChildren(std::forward<RangeT>(range));
     elementIndentStack.pop_back();
   }
 
   /// Print the given derived node to the stream.
   void printImpl(const CompoundStmt *stmt);
   void printImpl(const EraseStmt *stmt);
   void printImpl(const LetStmt *stmt);
+  void printImpl(const ReplaceStmt *stmt);
 
   void printImpl(const AttributeExpr *expr);
   void printImpl(const DeclRefExpr *expr);
   void printImpl(const MemberAccessExpr *expr);
   void printImpl(const OperationExpr *expr);
   void printImpl(const TupleExpr *expr);
   void printImpl(const TypeExpr *expr);
 

 void NodePrinter::print(const Node *node) {
   printIndent();
   os << "-";
 
   elementIndentStack.push_back(/*isLastElt*/ false);
   TypeSwitch<const Node *>(node)
       .Case<
           // Statements.
-          const CompoundStmt, const EraseStmt, const LetStmt,
+          const CompoundStmt, const EraseStmt, const LetStmt, const ReplaceStmt,
 
           // Expressions.
           const AttributeExpr, const DeclRefExpr, const MemberAccessExpr,
           const OperationExpr, const TupleExpr, const TypeExpr,
 
           // Decls.
           const AttrConstraintDecl, const OpConstraintDecl,
           const TypeConstraintDecl, const TypeRangeConstraintDecl,
void NodePrinter::printImpl(const EraseStmt *stmt) {
   printChildren(stmt->getRootOpExpr());
 }
 
 void NodePrinter::printImpl(const LetStmt *stmt) {
   os << "LetStmt " << stmt << "\n";
   printChildren(stmt->getVarDecl());
 }
 
+void NodePrinter::printImpl(const ReplaceStmt *stmt) {
+  os << "ReplaceStmt " << stmt << "\n";
+  printChildren(stmt->getRootOpExpr());
+  printChildren("ReplValues", stmt->getReplExprs());
+}
+
 void NodePrinter::printImpl(const AttributeExpr *expr) {
   os << "AttributeExpr " << expr << " Value<\"" << expr->getValue() << "\">\n";
 }
 
 void NodePrinter::printImpl(const DeclRefExpr *expr) {
   os << "DeclRefExpr " << expr << " Type<";
   print(expr->getType());
   os << ">\n";

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

 //===----------------------------------------------------------------------===//
 // EraseStmt
 
 EraseStmt *EraseStmt::create(Context &ctx, llvm::SMRange loc, Expr *rootOp) {
   return new (ctx.getAllocator().Allocate<EraseStmt>()) EraseStmt(loc, rootOp);
 }
 
 //===----------------------------------------------------------------------===//
+// ReplaceStmt
+
+ReplaceStmt *ReplaceStmt::create(Context &ctx, llvm::SMRange loc, Expr *rootOp,
+                                 ArrayRef<Expr *> replExprs) {
+  unsigned allocSize = ReplaceStmt::totalSizeToAlloc<Expr *>(replExprs.size());
+  void *rawData = ctx.getAllocator().Allocate(allocSize, alignof(ReplaceStmt));
+
+  ReplaceStmt *stmt = new (rawData) ReplaceStmt(loc, rootOp, replExprs.size());
+  std::uninitialized_copy(replExprs.begin(), replExprs.end(),
+                          stmt->getReplExprs().begin());
+  return stmt;
+}
+
+//===----------------------------------------------------------------------===//
 // AttributeExpr
 //===----------------------------------------------------------------------===//
 
 AttributeExpr *AttributeExpr::create(Context &ctx, llvm::SMRange loc,
                                      StringRef value) {
   return new (ctx.getAllocator().Allocate<AttributeExpr>())
       AttributeExpr(ctx, loc, copyStringWithNull(ctx, value));
 }

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

private:
   /// a rewrite section, etc.
   enum class ParserContext {
     /// The parser is in the global context.
     Global,
     /// The parser is currently within the matcher portion of a Pattern, which
     /// is allows a terminal operation rewrite statement but no other rewrite
     /// transformations.
     PatternMatch,
+    /// The parser is currently within a Rewrite, which disallows calls to
+    /// constraints, requires operation expressions to have names, etc.
+    Rewrite,
   };
 
   //===--------------------------------------------------------------------===//
   // Parsing
   //===--------------------------------------------------------------------===//
 
   /// Push a new decl scope onto the lexer.
   ast::DeclScope *pushDeclScope() {
private:
 
   FailureOr<ast::Expr *> parseExpr();
 
   /// Identifier expressions.
   FailureOr<ast::Expr *> parseAttributeExpr();
   FailureOr<ast::Expr *> parseDeclRefExpr(StringRef name, llvm::SMRange loc);
   FailureOr<ast::Expr *> parseIdentifierExpr();
   FailureOr<ast::Expr *> parseMemberAccessExpr(ast::Expr *parentExpr);
-  FailureOr<ast::OpNameDecl *> parseOperationName();
-  FailureOr<ast::OpNameDecl *> parseWrappedOperationName();
+  FailureOr<ast::OpNameDecl *> parseOperationName(bool allowEmptyName = false);
+  FailureOr<ast::OpNameDecl *> parseWrappedOperationName(bool allowEmptyName);
   FailureOr<ast::Expr *> parseOperationExpr();
   FailureOr<ast::Expr *> parseTupleExpr();
   FailureOr<ast::Expr *> parseTypeExpr();
   FailureOr<ast::Expr *> parseUnderscoreExpr();
 
   //===--------------------------------------------------------------------===//
   // Stmts
 
   FailureOr<ast::Stmt *> parseStmt(bool expectTerminalSemicolon = true);
   FailureOr<ast::CompoundStmt *> parseCompoundStmt();
   FailureOr<ast::EraseStmt *> parseEraseStmt();
   FailureOr<ast::LetStmt *> parseLetStmt();
+  FailureOr<ast::ReplaceStmt *> parseReplaceStmt();
 
   //===--------------------------------------------------------------------===//
   // Creation+Analysis
   //===--------------------------------------------------------------------===//
 
   //===--------------------------------------------------------------------===//
   // Decls
 
FailureOr<ast::TupleExpr *> createTupleExpr(llvm::SMRange loc,
                                               ArrayRef<ast::Expr *> elements,
                                               ArrayRef<StringRef> elementNames);
 
   //===--------------------------------------------------------------------===//
   // Stmts
 
   FailureOr<ast::EraseStmt *> createEraseStmt(llvm::SMRange loc,
                                               ast::Expr *rootOp);
+  FailureOr<ast::ReplaceStmt *>
+  createReplaceStmt(llvm::SMRange loc, ast::Expr *rootOp,
+                    MutableArrayRef<ast::Expr *> replValues);
 
   //===--------------------------------------------------------------------===//
   // Lexer Utilities
   //===--------------------------------------------------------------------===//
 
   /// If the current token has the specified kind, consume it and return true.
   /// If not, return false.
   bool consumeIf(Token::Kind kind) {
case Token::kw_Attr: {
     if (curToken.is(Token::less) && failed(parseTypeConstraint(typeExpr)))
       return failure();
     return ast::ConstraintRef(
         ast::AttrConstraintDecl::create(ctx, loc, typeExpr), loc);
   }
   case Token::kw_Op: {
     consumeToken(Token::kw_Op);
 
-    // Parse an optional operation name.
-    FailureOr<ast::OpNameDecl *> opName = parseWrappedOperationName();
+    // Parse an optional operation name. If the name isn't provided, this refers
+    // to "any" operation.
+    FailureOr<ast::OpNameDecl *> opName =
+        parseWrappedOperationName(/*allowEmptyName=*/true);
     if (failed(opName))
       return failure();
 
     return ast::ConstraintRef(ast::OpConstraintDecl::create(ctx, loc, *opName),
                               loc);
   }
   case Token::kw_Type:
     consumeToken(Token::kw_Type);
if (memberNameTok.isNot(Token::identifier, Token::integer) &&
       !memberNameTok.isKeyword())
     return emitError(loc, "expected identifier or numeric member name");
   StringRef memberName = memberNameTok.getSpelling();
   consumeToken();
 
   return createMemberAccessExpr(parentExpr, memberName, loc);
 }
 
-FailureOr<ast::OpNameDecl *> Parser::parseOperationName() {
+FailureOr<ast::OpNameDecl *> Parser::parseOperationName(bool allowEmptyName) {
   llvm::SMRange loc = curToken.getLoc();
 
   // Handle the case of an no operation name.
-  if (curToken.isNot(Token::identifier) && !curToken.isKeyword())
+  if (curToken.isNot(Token::identifier) && !curToken.isKeyword()) {
+    if (allowEmptyName)
       return ast::OpNameDecl::create(ctx, llvm::SMRange());
-
+    return emitError("expected dialect namespace");
+  }
   StringRef name = curToken.getSpelling();
   consumeToken();
 
   // Otherwise, this is a literal operation name.
   if (failed(parseToken(Token::dot, "expected `.` after dialect namespace")))
     return failure();
 
   if (curToken.isNot(Token::identifier) && !curToken.isKeyword())
     return emitError("expected operation name after dialect namespace");
 
   name = StringRef(name.data(), name.size() + 1);
   do {
     name = StringRef(name.data(), name.size() + curToken.getSpelling().size());
     loc.End = curToken.getEndLoc();
     consumeToken();
   } while (curToken.isAny(Token::identifier, Token::dot) ||
            curToken.isKeyword());
   return ast::OpNameDecl::create(ctx, ast::Name::create(ctx, name, loc));
 }
 
-FailureOr<ast::OpNameDecl *> Parser::parseWrappedOperationName() {
+FailureOr<ast::OpNameDecl *>
+Parser::parseWrappedOperationName(bool allowEmptyName) {
   if (!consumeIf(Token::less))
     return ast::OpNameDecl::create(ctx, llvm::SMRange());
 
-  FailureOr<ast::OpNameDecl *> opNameDecl = parseOperationName();
+  FailureOr<ast::OpNameDecl *> opNameDecl = parseOperationName(allowEmptyName);
   if (failed(opNameDecl))
     return failure();
 
   if (failed(parseToken(Token::greater, "expected `>` after operation name")))
     return failure();
   return opNameDecl;
 }
 
 FailureOr<ast::Expr *> Parser::parseOperationExpr() {
   llvm::SMRange loc = curToken.getLoc();
   consumeToken(Token::kw_op);
 
   // If it isn't followed by a `<`, the `op` keyword is treated as a normal
   // identifier.
   if (curToken.isNot(Token::less)) {
     resetToken(loc);
     return parseIdentifierExpr();
   }
 
   // Parse the operation name. The name may be elided, in which case the
   // operation refers to "any" operation(i.e. a difference between `MyOp` and
-  // `Operation*`).
-  FailureOr<ast::OpNameDecl *> opNameDecl = parseWrappedOperationName();
+  // `Operation*`). Operation names within a rewrite context must be named.
+  bool allowEmptyName = parserContext != ParserContext::Rewrite;
+  FailureOr<ast::OpNameDecl *> opNameDecl =
+      parseWrappedOperationName(allowEmptyName);
   if (failed(opNameDecl))
     return failure();
 
   // Check for the optional list of operands.
   SmallVector<ast::Expr *> operands;
   if (consumeIf(Token::l_paren)) {
     do {
       FailureOr<ast::Expr *> operand = parseExpr();
FailureOr<ast::Stmt *> Parser::parseStmt(bool expectTerminalSemicolon) {
   FailureOr<ast::Stmt *> stmt;
   switch (curToken.getKind()) {
   case Token::kw_erase:
     stmt = parseEraseStmt();
     break;
   case Token::kw_let:
     stmt = parseLetStmt();
     break;
+  case Token::kw_replace:
+    stmt = parseReplaceStmt();
+    break;
   default:
     stmt = parseExpr();
     break;
   }
   if (failed(stmt) ||
       (expectTerminalSemicolon &&
        failed(parseToken(Token::semicolon, "expected `;` after statement"))))
     return failure();
FailureOr<ast::LetStmt *> Parser::parseLetStmt() {
 
   FailureOr<ast::VariableDecl *> varDecl =
       createVariableDecl(varName, varLoc, initializer, constraints);
   if (failed(varDecl))
     return failure();
   return ast::LetStmt::create(ctx, loc, *varDecl);
 }
 
+FailureOr<ast::ReplaceStmt *> Parser::parseReplaceStmt() {
+  llvm::SMRange loc = curToken.getLoc();
+  consumeToken(Token::kw_replace);
+
+  // Parse the root operation expression.
+  FailureOr<ast::Expr *> rootOp = parseExpr();
+  if (failed(rootOp))
+    return failure();
+
+  if (failed(
+          parseToken(Token::kw_with, "expected `with` after root operation")))
+    return failure();
+
+  // The replacement portion of this statement is within a rewrite context.
+  llvm::SaveAndRestore<ParserContext> saveCtx(parserContext,
+                                              ParserContext::Rewrite);
+
+  // Parse the replacement values.
+  SmallVector<ast::Expr *> replValues;
+  if (consumeIf(Token::l_paren)) {
+    if (consumeIf(Token::r_paren)) {
+      return emitError(
+          loc, "expected at least one replacement value, consider using "
+               "`erase` if no replacement values are desired");
+    }
+
+    do {
+      FailureOr<ast::Expr *> replExpr = parseExpr();
+      if (failed(replExpr))
+        return failure();
+      replValues.emplace_back(*replExpr);
+    } while (consumeIf(Token::comma));
+
+    if (failed(parseToken(Token::r_paren,
+                          "expected `)` after replacement values")))
+      return failure();
+  } else {
+    FailureOr<ast::Expr *> replExpr = parseExpr();
+    if (failed(replExpr))
+      return failure();
+    replValues.emplace_back(*replExpr);
+  }
+
+  return createReplaceStmt(loc, *rootOp, replValues);
+}
+
 //===----------------------------------------------------------------------===//
 // Creation+Analysis
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Decls
 
 FailureOr<ast::PatternDecl *>
FailureOr<ast::EraseStmt *> Parser::createEraseStmt(llvm::SMRange loc,
   // Check that root is an Operation.
   ast::Type rootType = rootOp->getType();
   if (!rootType.isa<ast::OperationType>())
     return emitError(rootOp->getLoc(), "expected `Op` expression");
 
   return ast::EraseStmt::create(ctx, loc, rootOp);
 }
 
+FailureOr<ast::ReplaceStmt *>
+Parser::createReplaceStmt(llvm::SMRange loc, ast::Expr *rootOp,
+                          MutableArrayRef<ast::Expr *> replValues) {
+  // Check that root is an Operation.
+  ast::Type rootType = rootOp->getType();
+  if (!rootType.isa<ast::OperationType>()) {
+    return emitError(
+        rootOp->getLoc(),
+        llvm::formatv("expected `Op` expression, but got `{0}`", rootType));
+  }
+
+  // If there are multiple replacement values, we implicitly convert any Op
+  // expressions to the value form.
+  bool shouldConvertOpToValues = replValues.size() > 1;
+  for (ast::Expr *&replExpr : replValues) {
+    ast::Type replType = replExpr->getType();
+
+    // Check that replExpr is an Operation, Value, or ValueRange.
+    if (replType.isa<ast::OperationType>()) {
+      if (shouldConvertOpToValues)
+        replExpr = convertOpToValue(replExpr);
+      continue;
+    }
+
+    if (replType != valueTy && replType != valueRangeTy) {
+      return emitError(replExpr->getLoc(),
+                       llvm::formatv("expected `Op`, `Value` or `ValueRange` "
+                                     "expression, but got `{0}`",
+                                     replType));
+    }
+  }
+
+  return ast::ReplaceStmt::create(ctx, loc, rootOp, replValues);
+}
+
 //===----------------------------------------------------------------------===//
 // Parser
 //===----------------------------------------------------------------------===//
 
 FailureOr<ast::Module *> mlir::pdll::parsePDLAST(ast::Context &ctx,
                                                  llvm::SourceMgr &sourceMgr) {
   Parser parser(ctx, sourceMgr);
   return parser.parseModule();
 }

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

 // -----
 
 Pattern {
   // CHECK: :7:7: error: `foo` has already been defined
   // CHECK: :6:7: note: see previous definition here
   let foo: Attr;
   let foo: Attr;
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// `replace`
+//===----------------------------------------------------------------------===//
+
+Pattern {
+  // CHECK: expected `Op` expression
+  replace attr<""> with attr<"">;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `with` after root operation
+  replace op<>;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `Op`, `Value` or `ValueRange` expression
+  replace op<> with attr<"">;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `Op`, `Value` or `ValueRange` expression
+  replace op<> with (attr<"">);
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `)` after replacement values
+  replace op<>(input: Value) with (input;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected at least one replacement value, consider using `erase` if no replacement values are desired
+  replace op<>(input: Value) with ();
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected dialect namespace
+  replace op<>(input: Value) with op<>;
+}

mlir/test/mlir-pdll/Parser/stmt.pdll

 // CHECK:         `-DeclRefExpr {{.*}} Type<TypeRange>
 // CHECK:           `-VariableDecl {{.*}} Name<_> Type<TypeRange>
 // CHECK:             `Constraints`
 // CHECK:               `-TypeRangeConstraintDecl
 Pattern {
   let var: ValueRange<_: TypeRange>;
   erase _: Op;
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// ReplaceStmt
+//===----------------------------------------------------------------------===//
+
+// CHECK: Module
+// CHECK: `-ReplaceStmt
+// CHECK:   `-DeclRefExpr {{.*}} Type<Op>
+// CHECK:   ReplValues
+// CHECK:     `-OperationExpr {{.*}} Type<Op<my_dialect.foo>>
+Pattern {
+  replace _: Op with op<my_dialect.foo>;
+}
+
+// -----
+
+// CHECK: Module
+// CHECK: `-ReplaceStmt
+// CHECK:   `-DeclRefExpr {{.*}} Type<Op>
+// CHECK:   ReplValues
+// CHECK:     |-DeclRefExpr {{.*}} Type<Value>
+// CHECK:     |-DeclRefExpr {{.*}} Type<ValueRange>
+// CHECK:     `-MemberAccessExpr {{.*}} Member<$results> Type<ValueRange>
+// CHECK:       `-OperationExpr {{.*}} Type<Op<my_dialect.foo>>
+Pattern {
+  replace _: Op with (_: Value, _: ValueRange, op<my_dialect.foo>);
+}