[PDLL] Add support for parsing pattern metadata

[rriddle]

Summary

This allows for overriding the metadata of a pattern and
providing information such as the benefit, bounded recursion,
and more in the future.

Depends On D115093

[nicolasvasilache]

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

454–455

Enclose these in a Metadata helper struct?

552

Side note, it feels somewhat strange to me that we still have to explicitly manipulate loc information so intrusively everywhere (either here or other places in MLIR).
I would have expected that by now we'd have a bunch of Loc-implicit methods and types that just do the right thing and we'd put more efforts in tests to really ensure this is the case.

[jpienaar]

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

454–455

In DRR currently the benefit is a heuristic based on number of nodes matched (so that preference is given to larger matches in case of ties) and one can specify additional benefit to refine the matching, but it isn't an absolute but a diff on the heuristic. Is all explicit and fixed here?

552

That would seem to work for 1-N and N-1 patterns (well and this one you'd want to potentially record the pass name too), not sure in general how one could avoid being explicit in ensuring good tracing.

578

Are you also checking signed or not?

585

Line exceeding 80?

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

73

nl ?

mlir/test/mlir-pdll/Parser/pattern.pdll

23

Would it be possible to have the example show recursion? (in isolation this revision doesn't show it)

[rriddle]

update

[rriddle]

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

454–455

Yeah, for now benefit is the explicit final benefit. We will likely want to add another construct for added benefit (addBenefit or something), but I want to ensure that we can ditch the heuristics in cases where we want explicit control.

552

Explicitly managing locations is just part of writing a frontend and getting good diagnostics. There are a lot of cases that we need to explicitly manage where a specific token or construct came from, and moving any of that to implicit will lead to less ideal error messages. There isn't any way that I can think of that would remove the need of explicit location handling in a frontend, it's just one of the things you have to deal with.

578

The integer token doesn't support a leading -, so it can't be signed.

mlir/test/mlir-pdll/Parser/pattern.pdll

23

Yeah, but I should note that these are not and don't intend to be end-user examples. These are tests that ensure that we structurally build the expected AST nodes, which may or may not make sense as real patterns.

The revision with the end user doc will have actual examples that showcase the more end to end things, just getting the structural constructs down here.

[rriddle]

rebase

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

private:
   // Directives
 
   LogicalResult parseDirective(SmallVector<ast::Decl *> &decls);
   LogicalResult parseInclude(SmallVector<ast::Decl *> &decls);
 
   //===--------------------------------------------------------------------===//
   // Decls
 
+  /// This structure contains the set of pattern metadata that may be parsed.
+  struct ParsedPatternMetadata {
+    Optional<uint16_t> benefit;
+    bool hasBoundedRecursion = false;
+  };
+
   FailureOr<ast::Decl *> parseTopLevelDecl();
   FailureOr<ast::Decl *> parsePatternDecl();
+  LogicalResult parsePatternDeclMetadata(ParsedPatternMetadata &metadata);
 
   /// Check to see if a decl has already been defined with the given name, if
   /// one has emit and error and return failure. Returns success otherwise.
   LogicalResult checkDefineNamedDecl(const ast::Name &name);
 
   /// Try to define a variable decl with the given components, returns the
   /// variable on success.
   FailureOr<ast::VariableDecl *>
private:
   // Creation+Analysis
   //===--------------------------------------------------------------------===//
 
   //===--------------------------------------------------------------------===//
   // Decls
 
   /// Try to create a pattern decl with the given components, returning the
   /// Pattern on success.
-  FailureOr<ast::PatternDecl *> createPatternDecl(llvm::SMRange loc,
-                                                  const ast::Name *name,
-                                                  Optional<uint16_t> benefit,
-                                                  bool hasBoundedRecursion,
-                                                  ast::CompoundStmt *body);
+  FailureOr<ast::PatternDecl *>
+  createPatternDecl(llvm::SMRange loc, const ast::Name *name,
+                    const ParsedPatternMetadata &metadata,
 
   /// Try to create a variable decl with the given components, returning the
   /// Variable on success.
   FailureOr<ast::VariableDecl *>
   createVariableDecl(StringRef name, llvm::SMRange loc, ast::Expr *initializer,
                      ArrayRef<ast::ConstraintRef> constraints);
 
   /// Validate the constraints used to constraint a variable decl.
llvm::SaveAndRestore<ParserContext> saveCtx(parserContext,
                                               ParserContext::PatternMatch);
 
   // Check for an optional identifier for the pattern name.
   const ast::Name *name = nullptr;
   if (curToken.is(Token::identifier)) {
     name = &ast::Name::create(ctx, curToken.getSpelling(), curToken.getLoc());
     consumeToken(Token::identifier);
   }
 
-  // TODO: Parse any pattern metadata.
-  Optional<uint16_t> benefit;
-  bool hasBoundedRecursion = false;
+  // Parse any pattern metadata.
+  ParsedPatternMetadata metadata;
+  if (consumeIf(Token::kw_with) && failed(parsePatternDeclMetadata(metadata)))
+    return failure();
 
   // Parse the pattern body.
   ast::CompoundStmt *body;
 
   if (curToken.isNot(Token::l_brace))
     return emitError("expected `{` to start pattern body");
   FailureOr<ast::CompoundStmt *> bodyResult = parseCompoundStmt();
   if (failed(bodyResult))
return emitError(loc,
                      "rewrite statement, such as `erase`");
   }
   if (std::next(bodyIt) != bodyE) {
     return emitError((*std::next(bodyIt))->getLoc(),
                      "Pattern body was terminated by an operation "
                      "rewrite statement, but found trailing statements");
   }
 
-  return createPatternDecl(loc, name, benefit, hasBoundedRecursion, body);
+  return createPatternDecl(loc, name, metadata, body);
+}
+
+LogicalResult
+Parser::parsePatternDeclMetadata(ParsedPatternMetadata &metadata) {
+  Optional<llvm::SMRange> benefitLoc;
+  Optional<llvm::SMRange> hasBoundedRecursionLoc;
+
+  do {
+    if (curToken.isNot(Token::identifier))
+      return emitError("expected pattern metadata identifier");
+    StringRef metadataStr = curToken.getSpelling();
+    llvm::SMRange metadataLoc = curToken.getLoc();
+    consumeToken(Token::identifier);
+
+    // Parse the benefit metadata: benefit(<integer-value>)
+    if (metadataStr == "benefit") {
+      if (benefitLoc) {
+        return emitErrorAndNote(metadataLoc,
+                                "pattern benefit has already been specified",
+                                *benefitLoc, "see previous definition here");
+      }
+      if (failed(parseToken(Token::l_paren,
+                            "expected `(` before pattern benefit")))
+        return failure();
+
+      uint16_t benefitValue = 0;
+      if (curToken.isNot(Token::integer))
+        return emitError("expected integral pattern benefit");
+      if (curToken.getSpelling().getAsInteger(/*Radix=*/10, benefitValue))
+        return emitError(
+            "expected pattern benefit to fit within a 16-bit integer");
+      consumeToken(Token::integer);
+
+      metadata.benefit = benefitValue;
+      benefitLoc = metadataLoc;
+
+      if (failed(
+              parseToken(Token::r_paren, "expected `)` after pattern benefit")))
+        return failure();
+      continue;
+    }
+
+    // Parse the bounded recursion metadata: recursion
+    if (metadataStr == "recursion") {
+      if (hasBoundedRecursionLoc) {
+        return emitErrorAndNote(
+            metadataLoc,
+            "pattern recursion metadata has already been specified",
+            *hasBoundedRecursionLoc, "see previous definition here");
+      }
+      metadata.hasBoundedRecursion = true;
+      hasBoundedRecursionLoc = metadataLoc;
+      continue;
+    }
+
+    return emitError(metadataLoc, "unknown pattern metadata");
+  } while (consumeIf(Token::comma));
+
+  return success();
 }
 
 FailureOr<ast::Expr *> Parser::parseTypeConstraintExpr() {
   consumeToken(Token::less);
 
   FailureOr<ast::Expr *> typeExpr = parseExpr();
   if (failed(typeExpr) ||
       failed(parseToken(Token::greater,
                         "expected `>` after variable type constraint")))
     return failure();
   return typeExpr;
 }
 

 
 FailureOr<ast::VariableDecl *>
 Parser::defineVariableDecl(StringRef name, llvm::SMRange nameLoc,
                            ast::Type type, ast::Expr *initExpr,
                            ArrayRef<ast::ConstraintRef> constraints) {
   assert(curDeclScope && "defining variable outside of decl scope");
   const ast::Name &nameDecl = ast::Name::create(ctx, name, nameLoc);
 
   // If the name of the variable indicates a special variable, we don't add it
   // to the scope. This variable is local to the definition point.
   if (name.empty() || name == "_") {
     return ast::VariableDecl::create(ctx, nameDecl, type, initExpr,
                                      constraints);
   }
   if (failed(checkDefineNamedDecl(nameDecl)))
     return failure();
 
   auto *varDecl =
       ast::VariableDecl::create(ctx, nameDecl, type, initExpr, constraints);
   curDeclScope->add(varDecl);
   return varDecl;
 }
 
 FailureOr<ast::VariableDecl *>

 // Creation+Analysis
 //===----------------------------------------------------------------------===//
 
 //===----------------------------------------------------------------------===//
 // Decls
 
 FailureOr<ast::PatternDecl *>
 Parser::createPatternDecl(llvm::SMRange loc, const ast::Name *name,
-                          Optional<uint16_t> benefit, bool hasBoundedRecursion,
+                          const ParsedPatternMetadata &metadata,
                           ast::CompoundStmt *body) {
-  return ast::PatternDecl::create(ctx, loc, name, benefit, hasBoundedRecursion,
-                                  body);
+  return ast::PatternDecl::create(ctx, loc, name, metadata.benefit,
+                                  metadata.hasBoundedRecursion, body);
 }
 
 FailureOr<ast::VariableDecl *>
 Parser::createVariableDecl(StringRef name, llvm::SMRange loc,
                            ast::Expr *initializer,
                            ArrayRef<ast::ConstraintRef> constraints) {
   // The type of the variable, which is expected to be inferred by either a
   // constraint or an initializer expression.

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

 
 // -----
 
 // CHECK: Pattern body was terminated by an operation rewrite statement, but found trailing statements
 Pattern {
   erase root: Op;
   let value: Value;
 }
+
+// -----
+
+//===----------------------------------------------------------------------===//
+// Metadata
+//===----------------------------------------------------------------------===//
+
+// CHECK: expected pattern metadata identifier
+Pattern with {}
+
+// -----
+
+// CHECK: unknown pattern metadata
+Pattern with unknown {}
+
+// -----
+
+// CHECK: expected `(` before pattern benefit
+Pattern with benefit) {}
+
+// -----
+
+// CHECK: expected integral pattern benefit
+Pattern with benefit(foo) {}
+
+// -----
+
+// CHECK: expected pattern benefit to fit within a 16-bit integer
+Pattern with benefit(65536) {}
+
+// -----
+
+// CHECK: expected `)` after pattern benefit
+Pattern with benefit(1( {}
+
+// -----
+
+// CHECK: pattern benefit has already been specified
+// CHECK: see previous definition here
+Pattern with benefit(1), benefit(1) {}
+
+// -----
+
+// CHECK: pattern recursion metadata has already been specified
+// CHECK: see previous definition here
+Pattern with recursion, recursion {}

mlir/test/mlir-pdll/Parser/pattern.pdll

 
 // -----
 
 // CHECK: Module
 // CHECK: `-PatternDecl {{.*}} Name<NamedPattern>
 Pattern NamedPattern {
   erase _: Op;
 }
+
+// -----
+
+// CHECK: Module
+// CHECK: `-PatternDecl {{.*}} Name<NamedPattern> Benefit<10> Recursion
+Pattern NamedPattern with benefit(10), recursion {
+  erase _: Op;
+}