[MLIR][Parser] Fix AffineParser colliding bare identifiers with primitive types

[Groverkss]

Summary

The parser currently can't parse bare identifiers like 'i0' in affine
maps and sets, and similarly ids like f16/f32. But these bare ids are
part of the grammar - although they are primitive types.

error: expected bare identifier
set = affine_set<(i0, i1) : ()>
                   ^

This patch allows the parser for AffineMap/IntegerSet to parse bare
identifiers as defined by the grammer.

[Groverkss]

• Remove debug header

[rriddle]

mlir/lib/Parser/Token.cpp

193–203 ↗

(On Diff #434354)

We should really avoid this, this function is effectively re-lexing the token. Take a look at what the AsmParserImpl does, we should do the same thing (which shouldn't require touch the Token class)

[rriddle]

mlir/lib/Parser/Token.cpp

193–203 ↗

(On Diff #434354)

More specifically, the isCurrentTokenAKeyword function which checks for the tokens that could be treated as a "keyword".

[Groverkss]

Use implementation of isCurrentTokenAKeyword from AsmParserImpl to check
for bare identifier token and make it a static function instead of modifying the
Token class.

[bondhugula]

mlir/lib/Parser/AffineParser.cpp

268–271

This can't be the right approach. There is nothing special about inttype; you could have the same issue with floattype. This is also missing doc and code comments.

A test case is missing as well.

[rriddle]

mlir/lib/Parser/AffineParser.cpp

268–271

There is something special in that inttype is the only non-keyword token (aside from bare_identifier). You would not run into this for floattype because there is no floattype token, there are various float type "keywords" which are covered by the isKeyword call on Token.

[bondhugula]

mlir/lib/Parser/AffineParser.cpp

268–271

I see, thanks.

It looks odd though here to define a bare identifier as either a bare identifier or an int type or a keyword.

[Groverkss]

• Added doc comments + tests
• Fixed some bugs found while adding tests

[Groverkss]

While adding tests, I noticed affine_map's like:

#map = affine_map<(d0, d1)[mod] -> (d0 + d1 + mod)>

wouldn't be parsed, even though they are correct according to the grammar. I added support to parse these affine_maps too.

[Groverkss]

• Fix small bug in previous implementation

[bondhugula]

In D127076#3561425, @Groverkss wrote:

While adding tests, I noticed affine_map's like:

#map = affine_map<(d0, d1)[mod] -> (d0 + d1 + mod)>

wouldn't be parsed, even though they are correct according to the grammar. I added support to parse these affine_maps too.

Interesting; that's good. I'll defer to @rriddle for the review here given his earlier comments. Dismissing my change requests.

mlir/lib/Parser/AffineParser.cpp

366

Nit: this token ... -> these tokens ... as identifiers

[Groverkss]

Address Uday's comments and rebase

[Groverkss]

@rriddle gentle ping :)

[Groverkss]

@rriddle ping

[rriddle]

LG (I've been OOO for the past few weeks)

mlir/lib/Parser/AffineParser.cpp

AffineExpr AffineParser::parseNegateExpression(AffineExpr lhs) {
   // parseAffineOperandExpr instead of parseAffineExpr here.
   if (!operand)
     // Extra error message although parseAffineOperandExpr would have
     // complained. Leads to a better diagnostic.
     return emitError("missing operand of negation"), nullptr;
   return (-1) * operand;
 }
 
+/// Returns true if the given token can be represented as an identifier.
+static bool isIdentifier(const Token &token) {
+  // We include only `inttype` and `bare_identifier` here since they are the
+  // only non-keyword tokens that can be used to represent an identifier.
+  return token.isAny(Token::bare_identifier, Token::inttype) ||
+         token.isKeyword();
+}
+
 /// Parse a bare id that may appear in an affine expression.
 ///
 ///   affine-expr ::= bare-id
 AffineExpr AffineParser::parseBareIdExpr() {
-  if (getToken().isNot(Token::bare_identifier))
+  if (!isIdentifier(getToken()))
     return emitWrongTokenError("expected bare identifier"), nullptr;
 
   StringRef sRef = getTokenSpelling();
   for (auto entry : dimsAndSymbols) {
     if (entry.first == sRef) {
-      consumeToken(Token::bare_identifier);
+      consumeToken();
       return entry.second;
     }
   }
 
   return emitWrongTokenError("use of undeclared identifier"), nullptr;
 }
 
 /// Parse an SSA id which may appear in an affine expression.

 /// whether an error should be emitted for a missing right operand.
 //  Eg: for an expression without parentheses (like i + j + k + l), each
 //  of the four identifiers is an operand. For i + j*k + l, j*k is not an
 //  operand expression, it's an op expression and will be parsed via
 //  parseAffineHighPrecOpExpression(). However, for i + (j*k) + -l, (j*k) and
 //  -l are valid operands that will be parsed by this function.
 AffineExpr AffineParser::parseAffineOperandExpr(AffineExpr lhs) {
   switch (getToken().getKind()) {
-  case Token::bare_identifier:
-    return parseBareIdExpr();
   case Token::kw_symbol:
     return parseSymbolSSAIdExpr();
   case Token::percent_identifier:
     return parseSSAIdExpr(/*isSymbol=*/false);
   case Token::integer:
     return parseIntegerExpr();
   case Token::l_paren:
     return parseParentheticalExpr();
   case Token::minus:
     return parseNegateExpression(lhs);
   case Token::kw_ceildiv:
   case Token::kw_floordiv:
   case Token::kw_mod:
+    // Try to treat these tokens as identifiers.
+    return parseBareIdExpr();
   case Token::plus:
   case Token::star:
     if (lhs)
       emitError("missing right operand of binary operator");
     else
       emitError("missing left operand of binary operator");
     return nullptr;
   default:
+    // If nothing matches, we try to treat this token as an identifier.
+    if (isIdentifier(getToken()))
+      return parseBareIdExpr();
+
     if (lhs)
       emitError("missing right operand of binary operator");
     else
       emitError("expected affine expression");
     return nullptr;
   }
 }
 

 AffineExpr AffineParser::parseAffineExpr() {
   return parseAffineLowPrecOpExpr(nullptr, AffineLowPrecOp::LNoOp);
 }
 
 /// Parse a dim or symbol from the lists appearing before the actual
 /// expressions of the affine map. Update our state to store the
 /// dimensional/symbolic identifier.
 ParseResult AffineParser::parseIdentifierDefinition(AffineExpr idExpr) {
-  if (getToken().isNot(Token::bare_identifier))
+  if (!isIdentifier(getToken()))
     return emitWrongTokenError("expected bare identifier");
 
   auto name = getTokenSpelling();
   for (auto entry : dimsAndSymbols) {
     if (entry.first == name)
       return emitError("redefinition of identifier '" + name + "'");
   }
-  consumeToken(Token::bare_identifier);
+  consumeToken();
 
   dimsAndSymbols.push_back({name, idExpr});
   return success();
 }
 
 /// Parse the list of dimensional identifiers to an affine map.
 ParseResult AffineParser::parseDimIdList(unsigned &numDims) {
   auto parseElt = [&]() -> ParseResult {

mlir/test/IR/affine-map.mlir

 #map57 = affine_map<(d0, d1) -> (d0 - d0 + d1, -d0 + d0 + d0, (1 + d0 + d1 floordiv 4) - (d0 + d1 floordiv 4 + 1))>
 
 // CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d0 * 3, (d0 + d1) * 2, d0 mod 2)>
 #map58 = affine_map<(d0, d1) -> (4*d0 - 2*d0 + d0, (d0 + d1) + (d0 + d1), 2 * (d0 mod 2) - d0 mod 2)>
 
 // CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d0 mod 5, (d1 mod 35) mod 4)>
 #map59 = affine_map<(d0, d1) -> ((d0 mod 35) mod 5, (d1 mod 35) mod 4)>
 
+// Check if parser can parse affine_map with identifiers that collide with
+// integer types.
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1) -> (d0, d1)>
+#map60 = affine_map<(i0, i1) -> (i0, i1)>
+
+// Check if parser can parse affine_map with identifiers that collide with
+// reserved keywords.
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1)[s0, s1] -> (-d0 + s0, -d1 + s1)>
+#map61 = affine_map<(d0, d1)[step, loc] -> (step - d0, loc - d1)>
+
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1)[s0, s1] -> (-d0 + s0 floordiv 2, -d1 + s1 mod 3)>
+#map62 = affine_map<(d0, d1)[mod, floordiv] -> (mod floordiv 2 - d0, floordiv mod 3 - d1)>
+
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1)[s0, s1] -> (-d0 + s1 floordiv 2, -d1 + s0 mod 3)>
+#map63 = affine_map<(d0, d1)[mod, floordiv] -> (floordiv floordiv 2 - d0, mod mod 3 - d1)>
+
+// CHECK: #map{{[0-9]+}} = affine_map<(d0, d1)[s0] -> (d0 + d1 + s0)>
+#map64 = affine_map<(i0, i1)[mod] -> (i0 + i1 + mod)>
+
 // Single identity maps are removed.
 // CHECK: @f0(memref<2x4xi8, 1>)
 func.func private @f0(memref<2x4xi8, #map0, 1>)
 
 // Single identity maps are removed.
 // CHECK: @f1(memref<2x4xi8, 1>)
 func.func private @f1(memref<2x4xi8, #map1, 1>)
 

 // CHECK: "f57"() {map = #map{{[0-9]+}}} : () -> ()
 "f57"() {map = #map57} : () -> ()
 
 // CHECK: "f58"() {map = #map{{[0-9]+}}} : () -> ()
 "f58"() {map = #map58} : () -> ()
 
 // CHECK: "f59"() {map = #map{{[0-9]+}}} : () -> ()
 "f59"() {map = #map59} : () -> ()
+
+// CHECK: "f60"() {map = #map{{[0-9]+}}} : () -> ()
+"f60"() {map = #map60} : () -> ()
+
+// CHECK: "f61"() {map = #map{{[0-9]+}}} : () -> ()
+"f61"() {map = #map61} : () -> ()
+
+// CHECK: "f62"() {map = #map{{[0-9]+}}} : () -> ()
+"f62"() {map = #map62} : () -> ()
+
+// CHECK: "f63"() {map = #map{{[0-9]+}}} : () -> ()
+"f63"() {map = #map63} : () -> ()
+
+// CHECK: "f64"() {map = #map{{[0-9]+}}} : () -> ()
+"f64"() {map = #map64} : () -> ()

mlir/test/IR/invalid-affinemap.mlir

 #hello_world = affine_map<(i, j) [s0, s1] -> (+i, j)> // expected-error {{missing left operand of binary op}}
 
 // -----
 
 
 #hello_world = affine_map<(i, j) [s0, s1] -> (i, *j)> // expected-error {{missing left operand of binary op}}
 
 // -----
-#hello_world = affine_map<(i, j) [s0, s1] -> (floordiv i 2, j)> // expected-error {{missing left operand of binary op}}
+#hello_world = affine_map<(i, j) [s0, s1] -> (floordiv i 2, j)> // expected-error {{use of undeclared identifier}}
 
 // -----
-#hello_world = affine_map<(i, j) [s0, s1] -> (ceildiv i 2, j)> // expected-error {{missing left operand of binary op}}
+#hello_world = affine_map<(i, j) [s0, s1] -> (ceildiv i 2, j)> // expected-error {{use of undeclared identifier}}
 
 // -----
-#hello_world = affine_map<(i, j) [s0, s1] -> (mod i 2, j)> // expected-error {{missing left operand of binary op}}
+#hello_world = affine_map<(i, j) [s0, s1] -> (mod i 2, j)> // expected-error {{use of undeclared identifier}}
 
 // -----
 #hello_world = affine_map<(i, j) [s0, s1] -> (-(), j)>
 // expected-error@-1 {{no expression inside parentheses}}
 // expected-error@-2 {{missing operand of negation}}
 
 // -----
 #hello_world = affine_map<(i, j) [s0, s1] -> (i, *j+5)> // expected-error {{missing left operand of binary op}}
 
 // -----
-#hello_world = affine_map<(i, j) [s0, s1] -> (i, floordiv j+5)> // expected-error {{missing left operand of binary op}}
+#hello_world = affine_map<(i, j) [s0, s1] -> (i, floordiv j+5)> // expected-error {{use of undeclared identifier}}
 
 // -----
-#hello_world = affine_map<(i, j) [s0, s1] -> (i, ceildiv j+5)> // expected-error {{missing left operand of binary op}}
+#hello_world = affine_map<(i, j) [s0, s1] -> (i, ceildiv j+5)> // expected-error {{use of undeclared identifier}}
 
 // -----
-#hello_world = affine_map<(i, j) [s0, s1] -> (i, mod j+5)> // expected-error {{missing left operand of binary op}}
+#hello_world = affine_map<(i, j) [s0, s1] -> (i, mod j+5)> // expected-error {{use of undeclared identifier}}
 
 // -----
 #hello_world = affine_map<(i, j) [s0, s1] -> (i*j, j)> // expected-error {{non-affine expression: at least one of the multiply operands has to be either a constant or symbolic}}
 
 // -----
 #hello_world = affine_map<(i, j) [s0, s1] -> (i, j + j ceildiv 128 mod 16 * i - 4)> // expected-error {{non-affine expression: at least one of the multiply operands has to be either a constant or symbolic}}
 
 // -----