diff --git a/mlir/lib/Parser/AffineParser.cpp b/mlir/lib/Parser/AffineParser.cpp --- a/mlir/lib/Parser/AffineParser.cpp +++ b/mlir/lib/Parser/AffineParser.cpp @@ -265,17 +265,25 @@ 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; } } @@ -342,8 +350,6 @@ // -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: @@ -357,6 +363,8 @@ 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) @@ -365,6 +373,10 @@ 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 @@ -458,7 +470,7 @@ /// 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(); @@ -466,7 +478,7 @@ if (entry.first == name) return emitError("redefinition of identifier '" + name + "'"); } - consumeToken(Token::bare_identifier); + consumeToken(); dimsAndSymbols.push_back({name, idExpr}); return success(); diff --git a/mlir/test/IR/affine-map.mlir b/mlir/test/IR/affine-map.mlir --- a/mlir/test/IR/affine-map.mlir +++ b/mlir/test/IR/affine-map.mlir @@ -192,6 +192,25 @@ // 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>) @@ -379,3 +398,18 @@ // 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} : () -> () diff --git a/mlir/test/IR/invalid-affinemap.mlir b/mlir/test/IR/invalid-affinemap.mlir --- a/mlir/test/IR/invalid-affinemap.mlir +++ b/mlir/test/IR/invalid-affinemap.mlir @@ -54,13 +54,13 @@ #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)> @@ -71,13 +71,13 @@ #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}}