[PDLL] Add support for literal Attribute and Type expressions

[rriddle]

Summary

This allows for using literal attributes and types within PDLL,
which simplifies building both constraints and rewriters. For
example, checking if an attribute is true is as simple as
attr<"true">.

Depends On D115294

[nicolasvasilache]

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

99

nit:nl

[rriddle]

update

[rriddle]

rebase

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

Expr(TypeID typeID, llvm::SMRange loc, Type type)
       : Stmt(typeID, loc), type(type) {}
 
 private:
   /// The type of this expression.
   Type type;
 };
 
 //===----------------------------------------------------------------------===//
+// AttributeExpr
+//===----------------------------------------------------------------------===//
+
+/// This expression represents a literal MLIR Attribute, and contains the
+/// textual assembly format of that attribute.
+class AttributeExpr : public Node::NodeBase<AttributeExpr, Expr> {
+public:
+  static AttributeExpr *create(Context &ctx, llvm::SMRange loc,
+                               StringRef value);
+
+  /// Get the raw value of this expression. This is the textual assembly format
+  /// of the MLIR Attribute.
+  StringRef getValue() const { return value; }
+
+private:
+  AttributeExpr(Context &ctx, llvm::SMRange loc, StringRef value)
+      : Base(loc, AttributeType::get(ctx)), value(value) {}
+
+  /// The value referenced by this expression.
+  StringRef value;
+};
+
+//===----------------------------------------------------------------------===//
 // DeclRefExpr
 //===----------------------------------------------------------------------===//
 
 /// This expression represents a reference to a Decl node.
 class DeclRefExpr : public Node::NodeBase<DeclRefExpr, Expr> {
 public:
   static DeclRefExpr *create(Context &ctx, llvm::SMRange loc, Decl *decl,
                              Type type);
private:
   /// The parent expression of this access.
   const Expr *parentExpr;
 
   /// The name of the member being accessed from the parent.
   StringRef memberName;
 };
 
 //===----------------------------------------------------------------------===//
+// TypeExpr
+//===----------------------------------------------------------------------===//
+
+/// This expression represents a literal MLIR Type, and contains the textual
+/// assembly format of that type.
+class TypeExpr : public Node::NodeBase<TypeExpr, Expr> {
+public:
+  static TypeExpr *create(Context &ctx, llvm::SMRange loc, StringRef value);
+
+  /// Get the raw value of this expression. This is the textual assembly format
+  /// of the MLIR Type.
+  StringRef getValue() const { return value; }
+
+private:
+  TypeExpr(Context &ctx, llvm::SMRange loc, StringRef value)
+      : Base(loc, TypeType::get(ctx)), value(value) {}
+
+  /// The value referenced by this expression.
+  StringRef value;
+};
+
+//===----------------------------------------------------------------------===//
 // Decl
 //===----------------------------------------------------------------------===//
 
 /// This class represents the base Decl node.
 class Decl : public Node {
 public:
   /// Return the name of the decl, or nullptr if it doesn't have one.
   const Name *getName() const { return name; }

 
 inline bool CoreConstraintDecl::classof(const Node *node) {
   return isa<AttrConstraintDecl, OpConstraintDecl, TypeConstraintDecl,
              TypeRangeConstraintDecl, ValueConstraintDecl,
              ValueRangeConstraintDecl>(node);
 }
 
 inline bool Expr::classof(const Node *node) {
-  return isa<DeclRefExpr, MemberAccessExpr>(node);
+  return isa<AttributeExpr, DeclRefExpr, MemberAccessExpr, TypeExpr>(node);
 }
 
 inline bool OpRewriteStmt::classof(const Node *node) {
   return isa<EraseStmt>(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 AttributeExpr *expr);
   void printImpl(const DeclRefExpr *expr);
   void printImpl(const MemberAccessExpr *expr);
+  void printImpl(const TypeExpr *expr);
 
   void printImpl(const AttrConstraintDecl *decl);
   void printImpl(const OpConstraintDecl *decl);
   void printImpl(const TypeConstraintDecl *decl);
   void printImpl(const TypeRangeConstraintDecl *decl);
   void printImpl(const ValueConstraintDecl *decl);
   void printImpl(const ValueRangeConstraintDecl *decl);
   void printImpl(const OpNameDecl *decl);
void NodePrinter::print(const Node *node) {
 
   elementIndentStack.push_back(/*isLastElt*/ false);
   TypeSwitch<const Node *>(node)
       .Case<
           // Statements.
           const CompoundStmt, const EraseStmt, const LetStmt,
 
           // Expressions.
-          const DeclRefExpr, const MemberAccessExpr,
+          const AttributeExpr, const DeclRefExpr, const MemberAccessExpr,
+          const TypeExpr,
 
           // Decls.
           const AttrConstraintDecl, const OpConstraintDecl,
           const TypeConstraintDecl, const TypeRangeConstraintDecl,
           const ValueConstraintDecl, const ValueRangeConstraintDecl,
           const OpNameDecl, const PatternDecl, const VariableDecl,
 
           const Module>([&](auto derivedNode) { this->printImpl(derivedNode); })
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 AttributeExpr *expr) {
+  os << "AttributeExpr " << expr << " Value<\"" << expr->getValue() << "\">\n";
+}
+
 void NodePrinter::printImpl(const DeclRefExpr *expr) {
   os << "DeclRefExpr " << expr << " Type<";
   print(expr->getType());
   os << ">\n";
   printChildren(expr->getDecl());
 }
 
 void NodePrinter::printImpl(const MemberAccessExpr *expr) {
   os << "MemberAccessExpr " << expr << " Member<" << expr->getMemberName()
      << "> Type<";
   print(expr->getType());
   os << ">\n";
   printChildren(expr->getParentExpr());
 }
 
+void NodePrinter::printImpl(const TypeExpr *expr) {
+  os << "TypeExpr " << expr << " Value<\"" << expr->getValue() << "\">\n";
+}
+
 void NodePrinter::printImpl(const AttrConstraintDecl *decl) {
   os << "AttrConstraintDecl " << decl << "\n";
   if (const auto *typeExpr = decl->getTypeExpr())
     printChildren(typeExpr);
 }
 
 void NodePrinter::printImpl(const OpConstraintDecl *decl) {
   os << "OpConstraintDecl " << decl << "\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);
 }
 
 //===----------------------------------------------------------------------===//
+// AttributeExpr
+//===----------------------------------------------------------------------===//
+
+AttributeExpr *AttributeExpr::create(Context &ctx, llvm::SMRange loc,
+                                     StringRef value) {
+  return new (ctx.getAllocator().Allocate<AttributeExpr>())
+      AttributeExpr(ctx, loc, copyStringWithNull(ctx, value));
+}
+
+//===----------------------------------------------------------------------===//
 // DeclRefExpr
 //===----------------------------------------------------------------------===//
 
 DeclRefExpr *DeclRefExpr::create(Context &ctx, llvm::SMRange loc, Decl *decl,
                                  Type type) {
   return new (ctx.getAllocator().Allocate<DeclRefExpr>())
       DeclRefExpr(loc, decl, type);
 }
 
 //===----------------------------------------------------------------------===//
 // MemberAccessExpr
 //===----------------------------------------------------------------------===//
 
 MemberAccessExpr *MemberAccessExpr::create(Context &ctx, llvm::SMRange loc,
                                            const Expr *parentExpr,
                                            StringRef memberName, Type type) {
   return new (ctx.getAllocator().Allocate<MemberAccessExpr>()) MemberAccessExpr(
       loc, parentExpr, memberName.copy(ctx.getAllocator()), type);
 }
 
 //===----------------------------------------------------------------------===//
+// TypeExpr
+//===----------------------------------------------------------------------===//
+
+TypeExpr *TypeExpr::create(Context &ctx, llvm::SMRange loc, StringRef value) {
+  return new (ctx.getAllocator().Allocate<TypeExpr>())
+      TypeExpr(ctx, loc, copyStringWithNull(ctx, value));
+}
+
+//===----------------------------------------------------------------------===//
 // AttrConstraintDecl
 //===----------------------------------------------------------------------===//
 
 AttrConstraintDecl *AttrConstraintDecl::create(Context &ctx, llvm::SMRange loc,
                                                Expr *typeExpr) {
   return new (ctx.getAllocator().Allocate<AttrConstraintDecl>())
       AttrConstraintDecl(loc, typeExpr);
 }

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

parseConstraint(Optional<llvm::SMRange> &typeConstraint,
                   ArrayRef<ast::ConstraintRef> existingConstraints);
 
   //===--------------------------------------------------------------------===//
   // Exprs
 
   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::Expr *> parseTypeExpr();
   FailureOr<ast::Expr *> parseUnderscoreExpr();
 
   //===--------------------------------------------------------------------===//
   // Stmts
 
   FailureOr<ast::Stmt *> parseStmt(bool expectTerminalSemicolon = true);
   FailureOr<ast::CompoundStmt *> parseCompoundStmt();
   FailureOr<ast::EraseStmt *> parseEraseStmt();
private:
   /// Advance the current lexer onto the next token, asserting what the expected
   /// current token is. This is preferred to the above method because it leads
   /// to more self-documenting code with better checking.
   void consumeToken(Token::Kind kind) {
     assert(curToken.is(kind) && "consumed an unexpected token");
     consumeToken();
   }
 
+  /// Reset the lexer to the location at the given position.
+  void resetToken(llvm::SMRange tokLoc) {
+    lexer.resetPointer(tokLoc.Start.getPointer());
+    curToken = lexer.lexToken();
+  }
+
   /// Consume the specified token if present and return success. On failure,
   /// output a diagnostic and return failure.
   LogicalResult parseToken(Token::Kind kind, const Twine &msg) {
     if (curToken.getKind() != kind)
       return emitError(curToken.getLoc(), msg);
     consumeToken();
     return success();
   }

 
 FailureOr<ast::Expr *> Parser::parseExpr() {
   if (curToken.is(Token::underscore))
     return parseUnderscoreExpr();
 
   // Parse the LHS expression.
   FailureOr<ast::Expr *> lhsExpr;
   switch (curToken.getKind()) {
+  case Token::kw_attr:
+    lhsExpr = parseAttributeExpr();
+    break;
   case Token::identifier:
     lhsExpr = parseIdentifierExpr();
     break;
+  case Token::kw_type:
+    lhsExpr = parseTypeExpr();
+    break;
   default:
     return emitError("expected expression");
   }
   if (failed(lhsExpr))
     return failure();
 
   // Check for an operator expression.
   while (true) {
     switch (curToken.getKind()) {
     case Token::dot:
       lhsExpr = parseMemberAccessExpr(*lhsExpr);
       break;
     default:
       return lhsExpr;
     }
     if (failed(lhsExpr))
       return failure();
   }
 }
 
+FailureOr<ast::Expr *> Parser::parseAttributeExpr() {
+  llvm::SMRange loc = curToken.getLoc();
+  consumeToken(Token::kw_attr);
+
+  // If we aren't followed by a `<`, the `attr` keyword is treated as a normal
+  // identifier.
+  if (!consumeIf(Token::less)) {
+    resetToken(loc);
+    return parseIdentifierExpr();
+  }
+
+  if (!curToken.isString())
+    return emitError("expected string literal containing MLIR attribute");
+  std::string attrExpr = curToken.getStringValue();
+  consumeToken();
+
+  if (failed(
+          parseToken(Token::greater, "expected `>` after attribute literal")))
+    return failure();
+  return ast::AttributeExpr::create(ctx, loc, attrExpr);
+}
+
 FailureOr<ast::Expr *> Parser::parseDeclRefExpr(StringRef name,
                                                 llvm::SMRange loc) {
   ast::Decl *decl = curDeclScope->lookup(name);
   if (!decl)
     return emitError(loc, "undefined reference to `" + name + "`");
 
   return createDeclRefExpr(loc, decl);
 }
FailureOr<ast::OpNameDecl *> Parser::parseWrappedOperationName() {
   if (failed(opNameDecl))
     return failure();
 
   if (failed(parseToken(Token::greater, "expected `>` after operation name")))
     return failure();
   return opNameDecl;
 }
 
+FailureOr<ast::Expr *> Parser::parseTypeExpr() {
+  llvm::SMRange loc = curToken.getLoc();
+  consumeToken(Token::kw_type);
+
+  // If we aren't followed by a `<`, the `type` keyword is treated as a normal
+  // identifier.
+  if (!consumeIf(Token::less)) {
+    resetToken(loc);
+    return parseIdentifierExpr();
+  }
+
+  if (!curToken.isString())
+    return emitError("expected string literal containing MLIR type");
+  std::string attrExpr = curToken.getStringValue();
+  consumeToken();
+
+  if (failed(parseToken(Token::greater, "expected `>` after type literal")))
+    return failure();
+  return ast::TypeExpr::create(ctx, loc, attrExpr);
+}
+
 FailureOr<ast::Expr *> Parser::parseUnderscoreExpr() {
   StringRef name = curToken.getSpelling();
   llvm::SMRange nameLoc = curToken.getLoc();
   consumeToken(Token::underscore);
 
   // Underscore expressions require a constraint list.
   if (failed(parseToken(Token::colon, "expected `:` after `_` variable")))
     return failure();

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

 
 // -----
 
 Pattern {
   // CHECK: invalid member access `unknown_result` on expression of type `Op`
   let root: Op;
   erase root.unknown_result;
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// `attr` Expr
+//===----------------------------------------------------------------------===//
+
+Pattern {
+  // CHECK: expected string literal containing MLIR attribute
+  let foo = attr<foo>;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `>` after attribute literal
+  let foo = attr<""<>;
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// `type` Expr
+//===----------------------------------------------------------------------===//
+
+Pattern {
+  // CHECK: expected string literal containing MLIR type
+  let foo = type<foo;
+}
+
+// -----
+
+Pattern {
+  // CHECK: expected `>` after type literal
+  let foo = type<"";
+}

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

+// RUN: mlir-pdll %s -I %S -split-input-file | FileCheck %s
+
+//===----------------------------------------------------------------------===//
+// AttrExpr
+//===----------------------------------------------------------------------===//
+
+// CHECK: Module
+// CHECK: `-AttributeExpr {{.*}} Value<"10: i32">
+Pattern {
+  let attr = attr<"10: i32">;
+
+  erase _: Op;
+}
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// TypeExpr
+//===----------------------------------------------------------------------===//
+
+// CHECK: Module
+// CHECK: `-TypeExpr {{.*}} Value<"i64">
+Pattern {
+  let type = type<"i64">;
+
+  erase _: Op;
+}