mlir/lib/AsmParser/AffineParser.cpp

   1 //===- AffineParser.cpp - MLIR Affine Parser ------------------------------===//
   2 //
   3 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
   4 // See https://llvm.org/LICENSE.txt for license information.
   5 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
   6 //
   7 //===----------------------------------------------------------------------===//
   8 //
   9 // This file implements a parser for Affine structures.
  10 //
  11 //===----------------------------------------------------------------------===//
  12
  13 #include "Parser.h"
  14 #include "ParserState.h"
  15 #include "mlir/IR/AffineExpr.h"
  16 #include "mlir/IR/AffineMap.h"
  17 #include "mlir/IR/AsmState.h"
  18 #include "mlir/IR/Diagnostics.h"
  19 #include "mlir/IR/IntegerSet.h"
  20 #include "mlir/IR/OpImplementation.h"
  21 #include "mlir/Support/LLVM.h"
  22 #include "llvm/Support/ErrorHandling.h"
  23 #include "llvm/Support/MemoryBuffer.h"
  24 #include "llvm/Support/SourceMgr.h"
  25 #include "llvm/Support/raw_ostream.h"
  26 #include <cassert>
  27 #include <cstdint>
  28 #include <utility>
  29
  30 using namespace mlir;
  31 using namespace mlir::detail;
  32
  33 namespace {
  34
  35 /// Lower precedence ops (all at the same precedence level). LNoOp is false in
  36 /// the boolean sense.
  37 enum AffineLowPrecOp {
  38   /// Null value.
  39   LNoOp,
  40   Add,
  41   Sub
  42 };
  43
  44 /// Higher precedence ops - all at the same precedence level. HNoOp is false
  45 /// in the boolean sense.
  46 enum AffineHighPrecOp {
  47   /// Null value.
  48   HNoOp,
  49   Mul,
  50   FloorDiv,
  51   CeilDiv,
  52   Mod
  53 };
  54
  55 /// This is a specialized parser for affine structures (affine maps, affine
  56 /// expressions, and integer sets), maintaining the state transient to their
  57 /// bodies.
  58 class AffineParser : public Parser {
  59 public:
  60   AffineParser(ParserState &state, bool allowParsingSSAIds = false,
  61                function_ref<ParseResult(bool)> parseElement = nullptr)
  62       : Parser(state), allowParsingSSAIds(allowParsingSSAIds),
  63         parseElement(parseElement) {}
  64
  65   ParseResult parseAffineMapRange(unsigned numDims, unsigned numSymbols,
  66                                   AffineMap &result);
  67   ParseResult parseAffineMapOrIntegerSetInline(AffineMap &map, IntegerSet &set);
  68   ParseResult
  69   parseAffineExprInline(ArrayRef<std::pair<StringRef, AffineExpr>> symbolSet,
  70                         AffineExpr &expr);
  71   ParseResult parseIntegerSetConstraints(unsigned numDims, unsigned numSymbols,
  72                                          IntegerSet &result);
  73   ParseResult parseAffineMapOfSSAIds(AffineMap &map,
  74                                      OpAsmParser::Delimiter delimiter);
  75   ParseResult parseAffineExprOfSSAIds(AffineExpr &expr);
  76
  77 private:
  78   // Binary affine op parsing.
  79   AffineLowPrecOp consumeIfLowPrecOp();
  80   AffineHighPrecOp consumeIfHighPrecOp();
  81
  82   // Identifier lists for polyhedral structures.
  83   ParseResult parseDimIdList(unsigned &numDims);
  84   ParseResult parseSymbolIdList(unsigned &numSymbols);
  85   ParseResult parseDimAndOptionalSymbolIdList(unsigned &numDims,
  86                                               unsigned &numSymbols);
  87   ParseResult parseIdentifierDefinition(AffineExpr idExpr);
  88
  89   AffineExpr parseAffineExpr();
  90   AffineExpr parseParentheticalExpr();
  91   AffineExpr parseNegateExpression(AffineExpr lhs);
  92   AffineExpr parseIntegerExpr();
  93   AffineExpr parseBareIdExpr();
  94   AffineExpr parseSSAIdExpr(bool isSymbol);
  95   AffineExpr parseSymbolSSAIdExpr();
  96
  97   AffineExpr getAffineBinaryOpExpr(AffineHighPrecOp op, AffineExpr lhs,
  98                                    AffineExpr rhs, SMLoc opLoc);
  99   AffineExpr getAffineBinaryOpExpr(AffineLowPrecOp op, AffineExpr lhs,
 100                                    AffineExpr rhs);
 101   AffineExpr parseAffineOperandExpr(AffineExpr lhs);
 102   AffineExpr parseAffineLowPrecOpExpr(AffineExpr llhs, AffineLowPrecOp llhsOp);
 103   AffineExpr parseAffineHighPrecOpExpr(AffineExpr llhs, AffineHighPrecOp llhsOp,
 104                                        SMLoc llhsOpLoc);
 105   AffineExpr parseAffineConstraint(bool *isEq);
 106
 107 private:
 108   bool allowParsingSSAIds;
 109   function_ref<ParseResult(bool)> parseElement;
 110   unsigned numDimOperands = 0;
 111   unsigned numSymbolOperands = 0;
 112   SmallVector<std::pair<StringRef, AffineExpr>, 4> dimsAndSymbols;
 113 };
 114 } // namespace
 115
 116 /// Create an affine binary high precedence op expression (mul's, div's, mod).
 117 /// opLoc is the location of the op token to be used to report errors
 118 /// for non-conforming expressions.
 119 AffineExpr AffineParser::getAffineBinaryOpExpr(AffineHighPrecOp op,
 120                                                AffineExpr lhs, AffineExpr rhs,
 121                                                SMLoc opLoc) {
 122   // TODO: make the error location info accurate.
 123   switch (op) {
 124   case Mul:
 125     if (!lhs.isSymbolicOrConstant() && !rhs.isSymbolicOrConstant()) {
 126       emitError(opLoc, "non-affine expression: at least one of the multiply "
 127                        "operands has to be either a constant or symbolic");
 128       return nullptr;
 129     }
 130     return lhs * rhs;
 131   case FloorDiv:
 132     if (!rhs.isSymbolicOrConstant()) {
 133       emitError(opLoc, "non-affine expression: right operand of floordiv "
 134                        "has to be either a constant or symbolic");
 135       return nullptr;
 136     }
 137     return lhs.floorDiv(rhs);
 138   case CeilDiv:
 139     if (!rhs.isSymbolicOrConstant()) {
 140       emitError(opLoc, "non-affine expression: right operand of ceildiv "
 141                        "has to be either a constant or symbolic");
 142       return nullptr;
 143     }
 144     return lhs.ceilDiv(rhs);
 145   case Mod:
 146     if (!rhs.isSymbolicOrConstant()) {
 147       emitError(opLoc, "non-affine expression: right operand of mod "
 148                        "has to be either a constant or symbolic");
 149       return nullptr;
 150     }
 151     return lhs % rhs;
 152   case HNoOp:
 153     llvm_unreachable("can't create affine expression for null high prec op");
 154     return nullptr;
 155   }
 156   llvm_unreachable("Unknown AffineHighPrecOp");
 157 }
 158
 159 /// Create an affine binary low precedence op expression (add, sub).
 160 AffineExpr AffineParser::getAffineBinaryOpExpr(AffineLowPrecOp op,
 161                                                AffineExpr lhs, AffineExpr rhs) {
 162   switch (op) {
 163   case AffineLowPrecOp::Add:
 164     return lhs + rhs;
 165   case AffineLowPrecOp::Sub:
 166     return lhs - rhs;
 167   case AffineLowPrecOp::LNoOp:
 168     llvm_unreachable("can't create affine expression for null low prec op");
 169     return nullptr;
 170   }
 171   llvm_unreachable("Unknown AffineLowPrecOp");
 172 }
 173
 174 /// Consume this token if it is a lower precedence affine op (there are only
 175 /// two precedence levels).
 176 AffineLowPrecOp AffineParser::consumeIfLowPrecOp() {
 177   switch (getToken().getKind()) {
 178   case Token::plus:
 179     consumeToken(Token::plus);
 180     return AffineLowPrecOp::Add;
 181   case Token::minus:
 182     consumeToken(Token::minus);
 183     return AffineLowPrecOp::Sub;
 184   default:
 185     return AffineLowPrecOp::LNoOp;
 186   }
 187 }
 188
 189 /// Consume this token if it is a higher precedence affine op (there are only
 190 /// two precedence levels)
 191 AffineHighPrecOp AffineParser::consumeIfHighPrecOp() {
 192   switch (getToken().getKind()) {
 193   case Token::star:
 194     consumeToken(Token::star);
 195     return Mul;
 196   case Token::kw_floordiv:
 197     consumeToken(Token::kw_floordiv);
 198     return FloorDiv;
 199   case Token::kw_ceildiv:
 200     consumeToken(Token::kw_ceildiv);
 201     return CeilDiv;
 202   case Token::kw_mod:
 203     consumeToken(Token::kw_mod);
 204     return Mod;
 205   default:
 206     return HNoOp;
 207   }
 208 }
 209
 210 /// Parse a high precedence op expression list: mul, div, and mod are high
 211 /// precedence binary ops, i.e., parse a
 212 ///   expr_1 op_1 expr_2 op_2 ... expr_n
 213 /// where op_1, op_2 are all a AffineHighPrecOp (mul, div, mod).
 214 /// All affine binary ops are left associative.
 215 /// Given llhs, returns (llhs llhsOp lhs) op rhs, or (lhs op rhs) if llhs is
 216 /// null. If no rhs can be found, returns (llhs llhsOp lhs) or lhs if llhs is
 217 /// null. llhsOpLoc is the location of the llhsOp token that will be used to
 218 /// report an error for non-conforming expressions.
 219 AffineExpr AffineParser::parseAffineHighPrecOpExpr(AffineExpr llhs,
 220                                                    AffineHighPrecOp llhsOp,
 221                                                    SMLoc llhsOpLoc) {
 222   AffineExpr lhs = parseAffineOperandExpr(llhs);
 223   if (!lhs)
 224     return nullptr;
 225
 226   // Found an LHS. Parse the remaining expression.
 227   auto opLoc = getToken().getLoc();
 228   if (AffineHighPrecOp op = consumeIfHighPrecOp()) {
 229     if (llhs) {
 230       AffineExpr expr = getAffineBinaryOpExpr(llhsOp, llhs, lhs, opLoc);
 231       if (!expr)
 232         return nullptr;
 233       return parseAffineHighPrecOpExpr(expr, op, opLoc);
 234     }
 235     // No LLHS, get RHS
 236     return parseAffineHighPrecOpExpr(lhs, op, opLoc);
 237   }
 238
 239   // This is the last operand in this expression.
 240   if (llhs)
 241     return getAffineBinaryOpExpr(llhsOp, llhs, lhs, llhsOpLoc);
 242
 243   // No llhs, 'lhs' itself is the expression.
 244   return lhs;
 245 }
 246
 247 /// Parse an affine expression inside parentheses.
 248 ///
 249 ///   affine-expr ::= `(` affine-expr `)`
 250 AffineExpr AffineParser::parseParentheticalExpr() {
 251   if (parseToken(Token::l_paren, "expected '('"))
 252     return nullptr;
 253   if (getToken().is(Token::r_paren))
 254     return emitError("no expression inside parentheses"), nullptr;
 255
 256   auto expr = parseAffineExpr();
 257   if (!expr || parseToken(Token::r_paren, "expected ')'"))
 258     return nullptr;
 259
 260   return expr;
 261 }
 262
 263 /// Parse the negation expression.
 264 ///
 265 ///   affine-expr ::= `-` affine-expr
 266 AffineExpr AffineParser::parseNegateExpression(AffineExpr lhs) {
 267   if (parseToken(Token::minus, "expected '-'"))
 268     return nullptr;
 269
 270   AffineExpr operand = parseAffineOperandExpr(lhs);
 271   // Since negation has the highest precedence of all ops (including high
 272   // precedence ops) but lower than parentheses, we are only going to use
 273   // parseAffineOperandExpr instead of parseAffineExpr here.
 274   if (!operand)
 275     // Extra error message although parseAffineOperandExpr would have
 276     // complained. Leads to a better diagnostic.
 277     return emitError("missing operand of negation"), nullptr;
 278   return (-1) * operand;
 279 }
 280
 281 /// Returns true if the given token can be represented as an identifier.
 282 static bool isIdentifier(const Token &token) {
 283   // We include only `inttype` and `bare_identifier` here since they are the
 284   // only non-keyword tokens that can be used to represent an identifier.
 285   return token.isAny(Token::bare_identifier, Token::inttype) ||
 286          token.isKeyword();
 287 }
 288
 289 /// Parse a bare id that may appear in an affine expression.
 290 ///
 291 ///   affine-expr ::= bare-id
 292 AffineExpr AffineParser::parseBareIdExpr() {
 293   if (!isIdentifier(getToken()))
 294     return emitWrongTokenError("expected bare identifier"), nullptr;
 295
 296   StringRef sRef = getTokenSpelling();
 297   for (auto entry : dimsAndSymbols) {
 298     if (entry.first == sRef) {
 299       consumeToken();
 300       return entry.second;
 301     }
 302   }
 303
 304   return emitWrongTokenError("use of undeclared identifier"), nullptr;
 305 }
 306
 307 /// Parse an SSA id which may appear in an affine expression.
 308 AffineExpr AffineParser::parseSSAIdExpr(bool isSymbol) {
 309   if (!allowParsingSSAIds)
 310     return emitWrongTokenError("unexpected ssa identifier"), nullptr;
 311   if (getToken().isNot(Token::percent_identifier))
 312     return emitWrongTokenError("expected ssa identifier"), nullptr;
 313   auto name = getTokenSpelling();
 314   // Check if we already parsed this SSA id.
 315   for (auto entry : dimsAndSymbols) {
 316     if (entry.first == name) {
 317       consumeToken(Token::percent_identifier);
 318       return entry.second;
 319     }
 320   }
 321   // Parse the SSA id and add an AffineDim/SymbolExpr to represent it.
 322   if (parseElement(isSymbol))
 323     return nullptr;
 324   auto idExpr = isSymbol
 325                     ? getAffineSymbolExpr(numSymbolOperands++, getContext())
 326                     : getAffineDimExpr(numDimOperands++, getContext());
 327   dimsAndSymbols.push_back({name, idExpr});
 328   return idExpr;
 329 }
 330
 331 AffineExpr AffineParser::parseSymbolSSAIdExpr() {
 332   if (parseToken(Token::kw_symbol, "expected symbol keyword") ||
 333       parseToken(Token::l_paren, "expected '(' at start of SSA symbol"))
 334     return nullptr;
 335   AffineExpr symbolExpr = parseSSAIdExpr(/*isSymbol=*/true);
 336   if (!symbolExpr)
 337     return nullptr;
 338   if (parseToken(Token::r_paren, "expected ')' at end of SSA symbol"))
 339     return nullptr;
 340   return symbolExpr;
 341 }
 342
 343 /// Parse a positive integral constant appearing in an affine expression.
 344 ///
 345 ///   affine-expr ::= integer-literal
 346 AffineExpr AffineParser::parseIntegerExpr() {
 347   auto val = getToken().getUInt64IntegerValue();
 348   if (!val.has_value() || (int64_t)*val < 0)
 349     return emitError("constant too large for index"), nullptr;
 350
 351   consumeToken(Token::integer);
 352   return builder.getAffineConstantExpr((int64_t)*val);
 353 }
 354
 355 /// Parses an expression that can be a valid operand of an affine expression.
 356 /// lhs: if non-null, lhs is an affine expression that is the lhs of a binary
 357 /// operator, the rhs of which is being parsed. This is used to determine
 358 /// whether an error should be emitted for a missing right operand.
 359 //  Eg: for an expression without parentheses (like i + j + k + l), each
 360 //  of the four identifiers is an operand. For i + j*k + l, j*k is not an
 361 //  operand expression, it's an op expression and will be parsed via
 362 //  parseAffineHighPrecOpExpression(). However, for i + (j*k) + -l, (j*k) and
 363 //  -l are valid operands that will be parsed by this function.
 364 AffineExpr AffineParser::parseAffineOperandExpr(AffineExpr lhs) {
 365   switch (getToken().getKind()) {
 366   case Token::kw_symbol:
 367     return parseSymbolSSAIdExpr();
 368   case Token::percent_identifier:
 369     return parseSSAIdExpr(/*isSymbol=*/false);
 370   case Token::integer:
 371     return parseIntegerExpr();
 372   case Token::l_paren:
 373     return parseParentheticalExpr();
 374   case Token::minus:
 375     return parseNegateExpression(lhs);
 376   case Token::kw_ceildiv:
 377   case Token::kw_floordiv:
 378   case Token::kw_mod:
 379     // Try to treat these tokens as identifiers.
 380     return parseBareIdExpr();
 381   case Token::plus:
 382   case Token::star:
 383     if (lhs)
 384       emitError("missing right operand of binary operator");
 385     else
 386       emitError("missing left operand of binary operator");
 387     return nullptr;
 388   default:
 389     // If nothing matches, we try to treat this token as an identifier.
 390     if (isIdentifier(getToken()))
 391       return parseBareIdExpr();
 392
 393     if (lhs)
 394       emitError("missing right operand of binary operator");
 395     else
 396       emitError("expected affine expression");
 397     return nullptr;
 398   }
 399 }
 400
 401 /// Parse affine expressions that are bare-id's, integer constants,
 402 /// parenthetical affine expressions, and affine op expressions that are a
 403 /// composition of those.
 404 ///
 405 /// All binary op's associate from left to right.
 406 ///
 407 /// {add, sub} have lower precedence than {mul, div, and mod}.
 408 ///
 409 /// Add, sub'are themselves at the same precedence level. Mul, floordiv,
 410 /// ceildiv, and mod are at the same higher precedence level. Negation has
 411 /// higher precedence than any binary op.
 412 ///
 413 /// llhs: the affine expression appearing on the left of the one being parsed.
 414 /// This function will return ((llhs llhsOp lhs) op rhs) if llhs is non null,
 415 /// and lhs op rhs otherwise; if there is no rhs, llhs llhsOp lhs is returned
 416 /// if llhs is non-null; otherwise lhs is returned. This is to deal with left
 417 /// associativity.
 418 ///
 419 /// Eg: when the expression is e1 + e2*e3 + e4, with e1 as llhs, this function
 420 /// will return the affine expr equivalent of (e1 + (e2*e3)) + e4, where
 421 /// (e2*e3) will be parsed using parseAffineHighPrecOpExpr().
 422 AffineExpr AffineParser::parseAffineLowPrecOpExpr(AffineExpr llhs,
 423                                                   AffineLowPrecOp llhsOp) {
 424   AffineExpr lhs;
 425   if (!(lhs = parseAffineOperandExpr(llhs)))
 426     return nullptr;
 427
 428   // Found an LHS. Deal with the ops.
 429   if (AffineLowPrecOp lOp = consumeIfLowPrecOp()) {
 430     if (llhs) {
 431       AffineExpr sum = getAffineBinaryOpExpr(llhsOp, llhs, lhs);
 432       return parseAffineLowPrecOpExpr(sum, lOp);
 433     }
 434     // No LLHS, get RHS and form the expression.
 435     return parseAffineLowPrecOpExpr(lhs, lOp);
 436   }
 437   auto opLoc = getToken().getLoc();
 438   if (AffineHighPrecOp hOp = consumeIfHighPrecOp()) {
 439     // We have a higher precedence op here. Get the rhs operand for the llhs
 440     // through parseAffineHighPrecOpExpr.
 441     AffineExpr highRes = parseAffineHighPrecOpExpr(lhs, hOp, opLoc);
 442     if (!highRes)
 443       return nullptr;
 444
 445     // If llhs is null, the product forms the first operand of the yet to be
 446     // found expression. If non-null, the op to associate with llhs is llhsOp.
 447     AffineExpr expr =
 448         llhs ? getAffineBinaryOpExpr(llhsOp, llhs, highRes) : highRes;
 449
 450     // Recurse for subsequent low prec op's after the affine high prec op
 451     // expression.
 452     if (AffineLowPrecOp nextOp = consumeIfLowPrecOp())
 453       return parseAffineLowPrecOpExpr(expr, nextOp);
 454     return expr;
 455   }
 456   // Last operand in the expression list.
 457   if (llhs)
 458     return getAffineBinaryOpExpr(llhsOp, llhs, lhs);
 459   // No llhs, 'lhs' itself is the expression.
 460   return lhs;
 461 }
 462
 463 /// Parse an affine expression.
 464 ///  affine-expr ::= `(` affine-expr `)`
 465 ///                | `-` affine-expr
 466 ///                | affine-expr `+` affine-expr
 467 ///                | affine-expr `-` affine-expr
 468 ///                | affine-expr `*` affine-expr
 469 ///                | affine-expr `floordiv` affine-expr
 470 ///                | affine-expr `ceildiv` affine-expr
 471 ///                | affine-expr `mod` affine-expr
 472 ///                | bare-id
 473 ///                | integer-literal
 474 ///
 475 /// Additional conditions are checked depending on the production. For eg.,
 476 /// one of the operands for `*` has to be either constant/symbolic; the second
 477 /// operand for floordiv, ceildiv, and mod has to be a positive integer.
 478 AffineExpr AffineParser::parseAffineExpr() {
 479   return parseAffineLowPrecOpExpr(nullptr, AffineLowPrecOp::LNoOp);
 480 }
 481
 482 /// Parse a dim or symbol from the lists appearing before the actual
 483 /// expressions of the affine map. Update our state to store the
 484 /// dimensional/symbolic identifier.
 485 ParseResult AffineParser::parseIdentifierDefinition(AffineExpr idExpr) {
 486   if (!isIdentifier(getToken()))
 487     return emitWrongTokenError("expected bare identifier");
 488
 489   auto name = getTokenSpelling();
 490   for (auto entry : dimsAndSymbols) {
 491     if (entry.first == name)
 492       return emitError("redefinition of identifier '" + name + "'");
 493   }
 494   consumeToken();
 495
 496   dimsAndSymbols.push_back({name, idExpr});
 497   return success();
 498 }
 499
 500 /// Parse the list of dimensional identifiers to an affine map.
 501 ParseResult AffineParser::parseDimIdList(unsigned &numDims) {
 502   auto parseElt = [&]() -> ParseResult {
 503     auto dimension = getAffineDimExpr(numDims++, getContext());
 504     return parseIdentifierDefinition(dimension);
 505   };
 506   return parseCommaSeparatedList(Delimiter::Paren, parseElt,
 507                                  " in dimensional identifier list");
 508 }
 509
 510 /// Parse the list of symbolic identifiers to an affine map.
 511 ParseResult AffineParser::parseSymbolIdList(unsigned &numSymbols) {
 512   auto parseElt = [&]() -> ParseResult {
 513     auto symbol = getAffineSymbolExpr(numSymbols++, getContext());
 514     return parseIdentifierDefinition(symbol);
 515   };
 516   return parseCommaSeparatedList(Delimiter::Square, parseElt,
 517                                  " in symbol list");
 518 }
 519
 520 /// Parse the list of symbolic identifiers to an affine map.
 521 ParseResult
 522 AffineParser::parseDimAndOptionalSymbolIdList(unsigned &numDims,
 523                                               unsigned &numSymbols) {
 524   if (parseDimIdList(numDims)) {
 525     return failure();
 526   }
 527   if (!getToken().is(Token::l_square)) {
 528     numSymbols = 0;
 529     return success();
 530   }
 531   return parseSymbolIdList(numSymbols);
 532 }
 533
 534 /// Parses an ambiguous affine map or integer set definition inline.
 535 ParseResult AffineParser::parseAffineMapOrIntegerSetInline(AffineMap &map,
 536                                                            IntegerSet &set) {
 537   unsigned numDims = 0, numSymbols = 0;
 538
 539   // List of dimensional and optional symbol identifiers.
 540   if (parseDimAndOptionalSymbolIdList(numDims, numSymbols))
 541     return failure();
 542
 543   if (consumeIf(Token::arrow))
 544     return parseAffineMapRange(numDims, numSymbols, map);
 545
 546   if (parseToken(Token::colon, "expected '->' or ':'"))
 547     return failure();
 548   return parseIntegerSetConstraints(numDims, numSymbols, set);
 549 }
 550
 551 /// Parse an affine expresion definition inline, with given symbols.
 552 ParseResult AffineParser::parseAffineExprInline(
 553     ArrayRef<std::pair<StringRef, AffineExpr>> symbolSet, AffineExpr &expr) {
 554   dimsAndSymbols.assign(symbolSet.begin(), symbolSet.end());
 555   expr = parseAffineExpr();
 556   return success(expr != nullptr);
 557 }
 558
 559 /// Parse an AffineMap where the dim and symbol identifiers are SSA ids.
 560 ParseResult
 561 AffineParser::parseAffineMapOfSSAIds(AffineMap &map,
 562                                      OpAsmParser::Delimiter delimiter) {
 563
 564   SmallVector<AffineExpr, 4> exprs;
 565   auto parseElt = [&]() -> ParseResult {
 566     auto elt = parseAffineExpr();
 567     exprs.push_back(elt);
 568     return elt ? success() : failure();
 569   };
 570
 571   // Parse a multi-dimensional affine expression (a comma-separated list of
 572   // 1-d affine expressions); the list can be empty. Grammar:
 573   // multi-dim-affine-expr ::= `(` `)`
 574   //                         | `(` affine-expr (`,` affine-expr)* `)`
 575   if (parseCommaSeparatedList(delimiter, parseElt, " in affine map"))
 576     return failure();
 577
 578   // Parsed a valid affine map.
 579   map = AffineMap::get(numDimOperands, dimsAndSymbols.size() - numDimOperands,
 580                        exprs, getContext());
 581   return success();
 582 }
 583
 584 /// Parse an AffineExpr where the dim and symbol identifiers are SSA ids.
 585 ParseResult AffineParser::parseAffineExprOfSSAIds(AffineExpr &expr) {
 586   expr = parseAffineExpr();
 587   return success(expr != nullptr);
 588 }
 589
 590 /// Parse the range and sizes affine map definition inline.
 591 ///
 592 ///  affine-map ::= dim-and-symbol-id-lists `->` multi-dim-affine-expr
 593 ///
 594 ///  multi-dim-affine-expr ::= `(` `)`
 595 ///  multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `)`
 596 ParseResult AffineParser::parseAffineMapRange(unsigned numDims,
 597                                               unsigned numSymbols,
 598                                               AffineMap &result) {
 599   SmallVector<AffineExpr, 4> exprs;
 600   auto parseElt = [&]() -> ParseResult {
 601     auto elt = parseAffineExpr();
 602     ParseResult res = elt ? success() : failure();
 603     exprs.push_back(elt);
 604     return res;
 605   };
 606
 607   // Parse a multi-dimensional affine expression (a comma-separated list of
 608   // 1-d affine expressions). Grammar:
 609   // multi-dim-affine-expr ::= `(` `)`
 610   //                         | `(` affine-expr (`,` affine-expr)* `)`
 611   if (parseCommaSeparatedList(Delimiter::Paren, parseElt,
 612                               " in affine map range"))
 613     return failure();
 614
 615   // Parsed a valid affine map.
 616   result = AffineMap::get(numDims, numSymbols, exprs, getContext());
 617   return success();
 618 }
 619
 620 /// Parse an affine constraint.
 621 ///  affine-constraint ::= affine-expr `>=` `affine-expr`
 622 ///                      | affine-expr `<=` `affine-expr`
 623 ///                      | affine-expr `==` `affine-expr`
 624 ///
 625 /// The constraint is normalized to
 626 ///  affine-constraint ::= affine-expr `>=` `0`
 627 ///                      | affine-expr `==` `0`
 628 /// before returning.
 629 ///
 630 /// isEq is set to true if the parsed constraint is an equality, false if it
 631 /// is an inequality (greater than or equal).
 632 ///
 633 AffineExpr AffineParser::parseAffineConstraint(bool *isEq) {
 634   AffineExpr lhsExpr = parseAffineExpr();
 635   if (!lhsExpr)
 636     return nullptr;
 637
 638   // affine-constraint ::= `affine-expr` `>=` `affine-expr`
 639   if (consumeIf(Token::greater) && consumeIf(Token::equal)) {
 640     AffineExpr rhsExpr = parseAffineExpr();
 641     if (!rhsExpr)
 642       return nullptr;
 643     *isEq = false;
 644     return lhsExpr - rhsExpr;
 645   }
 646
 647   // affine-constraint ::= `affine-expr` `<=` `affine-expr`
 648   if (consumeIf(Token::less) && consumeIf(Token::equal)) {
 649     AffineExpr rhsExpr = parseAffineExpr();
 650     if (!rhsExpr)
 651       return nullptr;
 652     *isEq = false;
 653     return rhsExpr - lhsExpr;
 654   }
 655
 656   // affine-constraint ::= `affine-expr` `==` `affine-expr`
 657   if (consumeIf(Token::equal) && consumeIf(Token::equal)) {
 658     AffineExpr rhsExpr = parseAffineExpr();
 659     if (!rhsExpr)
 660       return nullptr;
 661     *isEq = true;
 662     return lhsExpr - rhsExpr;
 663   }
 664
 665   return emitError("expected '== affine-expr' or '>= affine-expr' at end of "
 666                    "affine constraint"),
 667          nullptr;
 668 }
 669
 670 /// Parse the constraints that are part of an integer set definition.
 671 ///  integer-set-inline
 672 ///                ::= dim-and-symbol-id-lists `:`
 673 ///                '(' affine-constraint-conjunction? ')'
 674 ///  affine-constraint-conjunction ::= affine-constraint (`,`
 675 ///                                       affine-constraint)*
 676 ///
 677 ParseResult AffineParser::parseIntegerSetConstraints(unsigned numDims,
 678                                                      unsigned numSymbols,
 679                                                      IntegerSet &result) {
 680   SmallVector<AffineExpr, 4> constraints;
 681   SmallVector<bool, 4> isEqs;
 682   auto parseElt = [&]() -> ParseResult {
 683     bool isEq;
 684     auto elt = parseAffineConstraint(&isEq);
 685     ParseResult res = elt ? success() : failure();
 686     if (elt) {
 687       constraints.push_back(elt);
 688       isEqs.push_back(isEq);
 689     }
 690     return res;
 691   };
 692
 693   // Parse a list of affine constraints (comma-separated).
 694   if (parseCommaSeparatedList(Delimiter::Paren, parseElt,
 695                               " in integer set constraint list"))
 696     return failure();
 697
 698   // If no constraints were parsed, then treat this as a degenerate 'true' case.
 699   if (constraints.empty()) {
 700     /* 0 == 0 */
 701     auto zero = getAffineConstantExpr(0, getContext());
 702     result = IntegerSet::get(numDims, numSymbols, zero, true);
 703     return success();
 704   }
 705
 706   // Parsed a valid integer set.
 707   result = IntegerSet::get(numDims, numSymbols, constraints, isEqs);
 708   return success();
 709 }
 710
 711 //===----------------------------------------------------------------------===//
 712 // Parser
 713 //===----------------------------------------------------------------------===//
 714
 715 /// Parse an ambiguous reference to either and affine map or an integer set.
 716 ParseResult Parser::parseAffineMapOrIntegerSetReference(AffineMap &map,
 717                                                         IntegerSet &set) {
 718   return AffineParser(state).parseAffineMapOrIntegerSetInline(map, set);
 719 }
 720 ParseResult Parser::parseAffineMapReference(AffineMap &map) {
 721   SMLoc curLoc = getToken().getLoc();
 722   IntegerSet set;
 723   if (parseAffineMapOrIntegerSetReference(map, set))
 724     return failure();
 725   if (set)
 726     return emitError(curLoc, "expected AffineMap, but got IntegerSet");
 727   return success();
 728 }
 729 ParseResult Parser::parseAffineExprReference(
 730     ArrayRef<std::pair<StringRef, AffineExpr>> symbolSet, AffineExpr &expr) {
 731   return AffineParser(state).parseAffineExprInline(symbolSet, expr);
 732 }
 733 ParseResult Parser::parseIntegerSetReference(IntegerSet &set) {
 734   SMLoc curLoc = getToken().getLoc();
 735   AffineMap map;
 736   if (parseAffineMapOrIntegerSetReference(map, set))
 737     return failure();
 738   if (map)
 739     return emitError(curLoc, "expected IntegerSet, but got AffineMap");
 740   return success();
 741 }
 742
 743 /// Parse an AffineMap of SSA ids. The callback 'parseElement' is used to
 744 /// parse SSA value uses encountered while parsing affine expressions.
 745 ParseResult
 746 Parser::parseAffineMapOfSSAIds(AffineMap &map,
 747                                function_ref<ParseResult(bool)> parseElement,
 748                                OpAsmParser::Delimiter delimiter) {
 749   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
 750       .parseAffineMapOfSSAIds(map, delimiter);
 751 }
 752
 753 /// Parse an AffineExpr of SSA ids. The callback `parseElement` is used to parse
 754 /// SSA value uses encountered while parsing.
 755 ParseResult
 756 Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
 757                                 function_ref<ParseResult(bool)> parseElement) {
 758   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
 759       .parseAffineExprOfSSAIds(expr);
 760 }
 761
 762 static void parseAffineMapOrIntegerSet(StringRef inputStr, MLIRContext *context,
 763                                        AffineMap &map, IntegerSet &set) {
 764   llvm::SourceMgr sourceMgr;
 765   auto memBuffer = llvm::MemoryBuffer::getMemBuffer(
 766       inputStr, /*BufferName=*/"<mlir_parser_buffer>",
 767       /*RequiresNullTerminator=*/false);
 768   sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
 769   SymbolState symbolState;
 770   ParserConfig config(context);
 771   ParserState state(sourceMgr, config, symbolState, /*asmState=*/nullptr,
 772                     /*codeCompleteContext=*/nullptr);
 773   Parser parser(state);
 774
 775   SourceMgrDiagnosticHandler handler(sourceMgr, context, llvm::errs());
 776   if (parser.parseAffineMapOrIntegerSetReference(map, set))
 777     return;
 778
 779   Token endTok = parser.getToken();
 780   if (endTok.isNot(Token::eof)) {
 781     parser.emitError(endTok.getLoc(), "encountered unexpected token");
 782     return;
 783   }
 784 }
 785
 786 AffineMap mlir::parseAffineMap(StringRef inputStr, MLIRContext *context) {
 787   AffineMap map;
 788   IntegerSet set;
 789   parseAffineMapOrIntegerSet(inputStr, context, map, set);
 790   assert(!set &&
 791          "expected string to represent AffineMap, but got IntegerSet instead");
 792   return map;
 793 }
 794
 795 IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context) {
 796   AffineMap map;
 797   IntegerSet set;
 798   parseAffineMapOrIntegerSet(inputStr, context, map, set);
 799   assert(!map &&
 800          "expected string to represent IntegerSet, but got AffineMap instead");
 801   return set;
 802 }