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