[PDLL] Add a `rewrite` statement to enable complex rewrites

[rriddle]

Summary

The rewrite statement allows for rewriting a given root
operation with a block of nested rewriters. The root operation is
not implicitly erased or replaced, and any transformations to it
must be expressed within the nested rewrite block. The inner body
may contain any number of other rewrite statements, variables, or
expressions.

Depends On D115298

[jpienaar]

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

198

What happens in this case? This is not as direct given the comment about no replacement happening.

If I were to have read this it would be

rewrite any op with (undefined op?)

but rewrite would make me thing a replacement is happening.

[rriddle]

Update

[rriddle]

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

198

Thanks for the catch. op<> should be an error in a rewrite. Also made the test more representative of something real (to prevent confusion).

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

private:
   /// The number of replacement values within this statement.
   unsigned numReplExprs;
 
   /// TrailingObject utilities.
   friend class llvm::TrailingObjects<ReplaceStmt, Expr *>;
 };
 
 //===----------------------------------------------------------------------===//
+// RewriteStmt
+
+/// This statement represents an operation rewrite that contains a block of
+/// nested rewrite commands. This allows for building more complex operation
+/// rewrites that span across multiple statements, which may be unconnected.
+class RewriteStmt final : public Node::NodeBase<RewriteStmt, OpRewriteStmt> {
+public:
+  static RewriteStmt *create(Context &ctx, llvm::SMRange loc, Expr *rootOp,
+                             CompoundStmt *rewriteBody);
+
+  /// Return the compound rewrite body.
+  CompoundStmt *getRewriteBody() const { return rewriteBody; }
+
+private:
+  RewriteStmt(llvm::SMRange loc, Expr *rootOp, CompoundStmt *rewriteBody)
+      : Base(loc, rootOp), rewriteBody(rewriteBody) {}
+
+  /// The body of nested rewriters within this statement.
+  CompoundStmt *rewriteBody;
+};
+
+//===----------------------------------------------------------------------===//
 // 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, ReplaceStmt>(node);
+  return isa<EraseStmt, ReplaceStmt, RewriteStmt>(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) {
     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 RewriteStmt *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 ReplaceStmt,
+          const RewriteStmt,
 
           // 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 ReplaceStmt *stmt) {
   os << "ReplaceStmt " << stmt << "\n";
   printChildren(stmt->getRootOpExpr());
   printChildren("ReplValues", stmt->getReplExprs());
 }
 
+void NodePrinter::printImpl(const RewriteStmt *stmt) {
+  os << "RewriteStmt " << stmt << "\n";
+  printChildren(stmt->getRootOpExpr(), stmt->getRewriteBody());
+}
+
 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

ReplaceStmt *ReplaceStmt::create(Context &ctx, llvm::SMRange loc, Expr *rootOp,
 
   ReplaceStmt *stmt = new (rawData) ReplaceStmt(loc, rootOp, replExprs.size());
   std::uninitialized_copy(replExprs.begin(), replExprs.end(),
                           stmt->getReplExprs().begin());
   return stmt;
 }
 
 //===----------------------------------------------------------------------===//
+// RewriteStmt
+
+RewriteStmt *RewriteStmt::create(Context &ctx, llvm::SMRange loc, Expr *rootOp,
+                                 CompoundStmt *rewriteBody) {
+  return new (ctx.getAllocator().Allocate<RewriteStmt>())
+      RewriteStmt(loc, rootOp, rewriteBody);
+}
+
+//===----------------------------------------------------------------------===//
 // 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:
   //===--------------------------------------------------------------------===//
   // Stmts
 
   FailureOr<ast::Stmt *> parseStmt(bool expectTerminalSemicolon = true);
   FailureOr<ast::CompoundStmt *> parseCompoundStmt();
   FailureOr<ast::EraseStmt *> parseEraseStmt();
   FailureOr<ast::LetStmt *> parseLetStmt();
   FailureOr<ast::ReplaceStmt *> parseReplaceStmt();
+  FailureOr<ast::RewriteStmt *> parseRewriteStmt();
 
   //===--------------------------------------------------------------------===//
   // Creation+Analysis
   //===--------------------------------------------------------------------===//
 
   //===--------------------------------------------------------------------===//
   // Decls
 
private:
   //===--------------------------------------------------------------------===//
   // 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);
+  FailureOr<ast::RewriteStmt *>
+  createRewriteStmt(llvm::SMRange loc, ast::Expr *rootOp,
+                    ast::CompoundStmt *rewriteBody);
 
   //===--------------------------------------------------------------------===//
   // 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_erase:
     stmt = parseEraseStmt();
     break;
   case Token::kw_let:
     stmt = parseLetStmt();
     break;
   case Token::kw_replace:
     stmt = parseReplaceStmt();
     break;
+  case Token::kw_rewrite:
+    stmt = parseRewriteStmt();
+    break;
   default:
     stmt = parseExpr();
     break;
   }
   if (failed(stmt) ||
       (expectTerminalSemicolon &&
        failed(parseToken(Token::semicolon, "expected `;` after statement"))))
     return failure();
if (consumeIf(Token::l_paren)) {
     if (failed(replExpr))
       return failure();
     replValues.emplace_back(*replExpr);
   }
 
   return createReplaceStmt(loc, *rootOp, replValues);
 }
 
+FailureOr<ast::RewriteStmt *> Parser::parseRewriteStmt() {
+  llvm::SMRange loc = curToken.getLoc();
+  consumeToken(Token::kw_rewrite);
+
+  // Parse the root operation.
+  FailureOr<ast::Expr *> rootOp = parseExpr();
+  if (failed(rootOp))
+    return failure();
+
+  if (failed(parseToken(Token::kw_with, "expected `with` before rewrite body")))
+    return failure();
+
+  if (curToken.isNot(Token::l_brace))
+    return emitError("expected `{` to start rewrite body");
+
+  // The rewrite body of this statement is within a rewrite context.
+  llvm::SaveAndRestore<ParserContext> saveCtx(parserContext,
+                                              ParserContext::Rewrite);
+
+  FailureOr<ast::CompoundStmt *> rewriteBody = parseCompoundStmt();
+  if (failed(rewriteBody))
+    return failure();
+
+  return createRewriteStmt(loc, *rootOp, *rewriteBody);
+}
+
 //===----------------------------------------------------------------------===//
 // Creation+Analysis
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Decls
 
 FailureOr<ast::PatternDecl *>
if (replType != valueTy && replType != valueRangeTy) {
                                      "expression, but got `{0}`",
                                      replType));
     }
   }
 
   return ast::ReplaceStmt::create(ctx, loc, rootOp, replValues);
 }
 
+FailureOr<ast::RewriteStmt *>
+Parser::createRewriteStmt(llvm::SMRange loc, ast::Expr *rootOp,
+                          ast::CompoundStmt *rewriteBody) {
+  // 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));
+  }
+
+  return ast::RewriteStmt::create(ctx, loc, rootOp, rewriteBody);
+}
+
 //===----------------------------------------------------------------------===//
 // 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: expected dialect namespace
   replace op<>(input: Value) with op<>;
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// `rewrite`
+//===----------------------------------------------------------------------===//
+
+Pattern {
+  // CHECK: expected `Op` expression
+  rewrite attr<""> with { op<toy.reshape>; };
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `with` before rewrite body
+  rewrite op<>;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `{` to start rewrite body
+  rewrite op<> with;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected dialect namespace
+  rewrite root: Op with {
+      op<>;
+  };
+}

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

 // 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>);
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// RewriteStmt
+//===----------------------------------------------------------------------===//
+
+// CHECK: Module
+// CHECK: `-RewriteStmt
+// CHECK:   |-DeclRefExpr {{.*}} Type<Op>
+// CHECK:   `-CompoundStmt
+// CHECK:     |-OperationExpr {{.*}} Type<Op<my_dialect.some_op>>
+// CHECK:     `-ReplaceStmt {{.*}}
+// CHECK:       `-DeclRefExpr {{.*}} Type<Op>
+// CHECK:       `ReplValues`
+// CHECK:          `-OperationExpr {{.*}} Type<Op<my_dialect.foo>>
+Pattern {
+  rewrite root: Op with {
+    op<my_dialect.some_op>;
+    replace root with op<my_dialect.foo>;
+  };
+}