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