[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

423

Enclose these in a Metadata helper struct?

462

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

423

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?

462

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.

488

Are you also checking signed or not?

495

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

423

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.

462

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.

488

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:
   LogicalResult parseDirective(SmallVector<ast::Decl *> &decls);
   LogicalResult parseInclude(SmallVector<ast::Decl *> &decls);
 
   //===--------------------------------------------------------------------===//
   // Decls
 
   FailureOr<ast::Decl *> parseTopLevelDecl();
   FailureOr<ast::Decl *> parsePatternDecl();
+  LogicalResult parsePatternDeclMetadata(Optional<uint16_t> &benefit,
+                                         bool &hasBoundedRecursion);
 
   LogicalResult checkDefineNamedDecl(const ast::Name &name);
   FailureOr<ast::VariableDecl *>
   defineVariableDecl(StringRef name, llvm::SMRange nameLoc, ast::Type type,
                      ast::Expr *initExpr,
                      ArrayRef<ast::ConstraintRef> constraints);
   FailureOr<ast::VariableDecl *>
   defineVariableDecl(StringRef name, llvm::SMRange nameLoc, ast::Type type,
FailureOr<ast::Decl *> Parser::parsePatternDecl() {
 
   // Check for an optional identifier for the pattern name.
   Optional<ast::Name> name;
   if (curToken.is(Token::identifier)) {
     name.emplace(curToken.getSpelling(), curToken.getLoc());
     consumeToken(Token::identifier);
   }
 
-  // TODO: Parse any pattern metadata.
+  // Parse any pattern metadata.
   Optional<uint16_t> benefit;
   bool hasBoundedRecursion = false;
+  if (consumeIf(Token::kw_with) &&
+      failed(parsePatternDeclMetadata(benefit, hasBoundedRecursion)))
+    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))
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);
 }
 
+LogicalResult Parser::parsePatternDeclMetadata(Optional<uint16_t> &benefit,
+                                               bool &hasBoundedRecursion) {
+  Optional<llvm::SMRange> benefitLoc;
+  Optional<llvm::SMRange> hasBoundedRecursionLoc;
+
+  do {
+    if (curToken.isNot(Token::identifier))
+      return emitError("expected pattern metadata identifier");
+    StringRef metadata = curToken.getSpelling();
+    llvm::SMRange metadataLoc = curToken.getLoc();
+    consumeToken(Token::identifier);
+
+    // Parse the benefit metadata: benefit(<integer-value>)
+    if (metadata == "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);
+
+      benefit = benefitValue;
+      benefitLoc = metadataLoc;
+
+      if (failed(
+              parseToken(Token::r_paren, "expected `)` after pattern benefit")))
+        return failure();
+      continue;
+    }
+
+    // Parse the bounded recursion metadata: recursion
+    if (metadata == "recursion") {
+      if (hasBoundedRecursionLoc) {
+        return emitErrorAndNote(
+            metadataLoc,
+            "pattern recursion metadata has already been specified",
+            *hasBoundedRecursionLoc, "see previous definition here");
+      }
+      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();

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 {}
 No newline at end of file

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;
+}