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
30using namespace mlir;
31using namespace mlir::detail;
32
33namespace {
34
35/// Lower precedence ops (all at the same precedence level). LNoOp is false in
36/// the boolean sense.
37enum 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.
46enum 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.
58class AffineParser : public Parser {
59public:
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
77private:
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
107private:
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.
119AffineExpr 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(loc: opLoc, message: "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(loc: opLoc, message: "non-affine expression: right operand of floordiv "
134 "has to be either a constant or symbolic");
135 return nullptr;
136 }
137 return lhs.floorDiv(other: rhs);
138 case CeilDiv:
139 if (!rhs.isSymbolicOrConstant()) {
140 emitError(loc: opLoc, message: "non-affine expression: right operand of ceildiv "
141 "has to be either a constant or symbolic");
142 return nullptr;
143 }
144 return lhs.ceilDiv(other: rhs);
145 case Mod:
146 if (!rhs.isSymbolicOrConstant()) {
147 emitError(loc: opLoc, message: "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).
160AffineExpr 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).
176AffineLowPrecOp AffineParser::consumeIfLowPrecOp() {
177 switch (getToken().getKind()) {
178 case Token::plus:
179 consumeToken(kind: Token::plus);
180 return AffineLowPrecOp::Add;
181 case Token::minus:
182 consumeToken(kind: 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)
191AffineHighPrecOp AffineParser::consumeIfHighPrecOp() {
192 switch (getToken().getKind()) {
193 case Token::star:
194 consumeToken(kind: Token::star);
195 return Mul;
196 case Token::kw_floordiv:
197 consumeToken(kind: Token::kw_floordiv);
198 return FloorDiv;
199 case Token::kw_ceildiv:
200 consumeToken(kind: Token::kw_ceildiv);
201 return CeilDiv;
202 case Token::kw_mod:
203 consumeToken(kind: 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.
219AffineExpr AffineParser::parseAffineHighPrecOpExpr(AffineExpr llhs,
220 AffineHighPrecOp llhsOp,
221 SMLoc llhsOpLoc) {
222 AffineExpr lhs = parseAffineOperandExpr(lhs: 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(op: llhsOp, lhs: llhs, rhs: lhs, opLoc);
231 if (!expr)
232 return nullptr;
233 return parseAffineHighPrecOpExpr(llhs: expr, llhsOp: op, llhsOpLoc: opLoc);
234 }
235 // No LLHS, get RHS
236 return parseAffineHighPrecOpExpr(llhs: lhs, llhsOp: op, llhsOpLoc: opLoc);
237 }
238
239 // This is the last operand in this expression.
240 if (llhs)
241 return getAffineBinaryOpExpr(op: llhsOp, lhs: llhs, rhs: lhs, opLoc: 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 `)`
250AffineExpr AffineParser::parseParentheticalExpr() {
251 if (parseToken(expectedToken: Token::l_paren, message: "expected '('"))
252 return nullptr;
253 if (getToken().is(k: Token::r_paren))
254 return emitError(message: "no expression inside parentheses"), nullptr;
255
256 auto expr = parseAffineExpr();
257 if (!expr || parseToken(expectedToken: Token::r_paren, message: "expected ')'"))
258 return nullptr;
259
260 return expr;
261}
262
263/// Parse the negation expression.
264///
265/// affine-expr ::= `-` affine-expr
266AffineExpr AffineParser::parseNegateExpression(AffineExpr lhs) {
267 if (parseToken(expectedToken: Token::minus, message: "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(message: "missing operand of negation"), nullptr;
278 return (-1) * operand;
279}
280
281/// Returns true if the given token can be represented as an identifier.
282static 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(k1: Token::bare_identifier, k2: 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
292AffineExpr AffineParser::parseBareIdExpr() {
293 if (!isIdentifier(token: getToken()))
294 return emitWrongTokenError(message: "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(message: "use of undeclared identifier"), nullptr;
305}
306
307/// Parse an SSA id which may appear in an affine expression.
308AffineExpr AffineParser::parseSSAIdExpr(bool isSymbol) {
309 if (!allowParsingSSAIds)
310 return emitWrongTokenError(message: "unexpected ssa identifier"), nullptr;
311 if (getToken().isNot(k: Token::percent_identifier))
312 return emitWrongTokenError(message: "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(kind: 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(position: numSymbolOperands++, context: getContext())
326 : getAffineDimExpr(position: numDimOperands++, context: getContext());
327 dimsAndSymbols.push_back(Elt: {name, idExpr});
328 return idExpr;
329}
330
331AffineExpr AffineParser::parseSymbolSSAIdExpr() {
332 if (parseToken(expectedToken: Token::kw_symbol, message: "expected symbol keyword") ||
333 parseToken(expectedToken: Token::l_paren, message: "expected '(' at start of SSA symbol"))
334 return nullptr;
335 AffineExpr symbolExpr = parseSSAIdExpr(/*isSymbol=*/true);
336 if (!symbolExpr)
337 return nullptr;
338 if (parseToken(expectedToken: Token::r_paren, message: "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
346AffineExpr AffineParser::parseIntegerExpr() {
347 auto val = getToken().getUInt64IntegerValue();
348 if (!val.has_value() || (int64_t)*val < 0)
349 return emitError(message: "constant too large for index"), nullptr;
350
351 consumeToken(kind: Token::integer);
352 return builder.getAffineConstantExpr(constant: (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.
364AffineExpr 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(message: "missing right operand of binary operator");
385 else
386 emitError(message: "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(token: getToken()))
391 return parseBareIdExpr();
392
393 if (lhs)
394 emitError(message: "missing right operand of binary operator");
395 else
396 emitError(message: "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().
422AffineExpr AffineParser::parseAffineLowPrecOpExpr(AffineExpr llhs,
423 AffineLowPrecOp llhsOp) {
424 AffineExpr lhs;
425 if (!(lhs = parseAffineOperandExpr(lhs: 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(op: llhsOp, lhs: llhs, rhs: lhs);
432 return parseAffineLowPrecOpExpr(llhs: sum, llhsOp: lOp);
433 }
434 // No LLHS, get RHS and form the expression.
435 return parseAffineLowPrecOpExpr(llhs: lhs, llhsOp: 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(llhs: lhs, llhsOp: hOp, llhsOpLoc: 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(op: llhsOp, lhs: llhs, rhs: 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(llhs: expr, llhsOp: nextOp);
454 return expr;
455 }
456 // Last operand in the expression list.
457 if (llhs)
458 return getAffineBinaryOpExpr(op: llhsOp, lhs: llhs, rhs: 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.
478AffineExpr AffineParser::parseAffineExpr() {
479 return parseAffineLowPrecOpExpr(llhs: nullptr, llhsOp: 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.
485ParseResult AffineParser::parseIdentifierDefinition(AffineExpr idExpr) {
486 if (!isIdentifier(token: getToken()))
487 return emitWrongTokenError(message: "expected bare identifier");
488
489 auto name = getTokenSpelling();
490 for (auto entry : dimsAndSymbols) {
491 if (entry.first == name)
492 return emitError(message: "redefinition of identifier '" + name + "'");
493 }
494 consumeToken();
495
496 dimsAndSymbols.push_back(Elt: {name, idExpr});
497 return success();
498}
499
500/// Parse the list of dimensional identifiers to an affine map.
501ParseResult AffineParser::parseDimIdList(unsigned &numDims) {
502 auto parseElt = [&]() -> ParseResult {
503 auto dimension = getAffineDimExpr(position: numDims++, context: getContext());
504 return parseIdentifierDefinition(idExpr: dimension);
505 };
506 return parseCommaSeparatedList(delimiter: Delimiter::Paren, parseElementFn: parseElt,
507 contextMessage: " in dimensional identifier list");
508}
509
510/// Parse the list of symbolic identifiers to an affine map.
511ParseResult AffineParser::parseSymbolIdList(unsigned &numSymbols) {
512 auto parseElt = [&]() -> ParseResult {
513 auto symbol = getAffineSymbolExpr(position: numSymbols++, context: getContext());
514 return parseIdentifierDefinition(idExpr: symbol);
515 };
516 return parseCommaSeparatedList(delimiter: Delimiter::Square, parseElementFn: parseElt,
517 contextMessage: " in symbol list");
518}
519
520/// Parse the list of symbolic identifiers to an affine map.
521ParseResult
522AffineParser::parseDimAndOptionalSymbolIdList(unsigned &numDims,
523 unsigned &numSymbols) {
524 if (parseDimIdList(numDims)) {
525 return failure();
526 }
527 if (!getToken().is(k: 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.
535ParseResult 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(kind: Token::arrow))
544 return parseAffineMapRange(numDims, numSymbols, result&: map);
545
546 if (parseToken(expectedToken: Token::colon, message: "expected '->' or ':'"))
547 return failure();
548 return parseIntegerSetConstraints(numDims, numSymbols, result&: set);
549}
550
551/// Parse an affine expresion definition inline, with given symbols.
552ParseResult AffineParser::parseAffineExprInline(
553 ArrayRef<std::pair<StringRef, AffineExpr>> symbolSet, AffineExpr &expr) {
554 dimsAndSymbols.assign(in_start: symbolSet.begin(), in_end: symbolSet.end());
555 expr = parseAffineExpr();
556 return success(IsSuccess: expr != nullptr);
557}
558
559/// Parse an AffineMap where the dim and symbol identifiers are SSA ids.
560ParseResult
561AffineParser::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: 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, parseElementFn: parseElt, contextMessage: " in affine map"))
576 return failure();
577
578 // Parsed a valid affine map.
579 map = AffineMap::get(dimCount: numDimOperands, symbolCount: dimsAndSymbols.size() - numDimOperands,
580 results: exprs, context: getContext());
581 return success();
582}
583
584/// Parse an AffineExpr where the dim and symbol identifiers are SSA ids.
585ParseResult AffineParser::parseAffineExprOfSSAIds(AffineExpr &expr) {
586 expr = parseAffineExpr();
587 return success(IsSuccess: 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)* `)`
596ParseResult 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: 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: Delimiter::Paren, parseElementFn: parseElt,
612 contextMessage: " in affine map range"))
613 return failure();
614
615 // Parsed a valid affine map.
616 result = AffineMap::get(dimCount: numDims, symbolCount: numSymbols, results: exprs, context: 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///
633AffineExpr 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(kind: Token::greater) && consumeIf(kind: 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(kind: Token::less) && consumeIf(kind: 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(kind: Token::equal) && consumeIf(kind: Token::equal)) {
658 AffineExpr rhsExpr = parseAffineExpr();
659 if (!rhsExpr)
660 return nullptr;
661 *isEq = true;
662 return lhsExpr - rhsExpr;
663 }
664
665 return emitError(message: "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///
677ParseResult 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: &isEq);
685 ParseResult res = elt ? success() : failure();
686 if (elt) {
687 constraints.push_back(Elt: elt);
688 isEqs.push_back(Elt: isEq);
689 }
690 return res;
691 };
692
693 // Parse a list of affine constraints (comma-separated).
694 if (parseCommaSeparatedList(delimiter: Delimiter::Paren, parseElementFn: parseElt,
695 contextMessage: " 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(constant: 0, context: getContext());
702 result = IntegerSet::get(dimCount: numDims, symbolCount: numSymbols, constraints: zero, eqFlags: true);
703 return success();
704 }
705
706 // Parsed a valid integer set.
707 result = IntegerSet::get(dimCount: numDims, symbolCount: numSymbols, constraints, eqFlags: 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.
716ParseResult Parser::parseAffineMapOrIntegerSetReference(AffineMap &map,
717 IntegerSet &set) {
718 return AffineParser(state).parseAffineMapOrIntegerSetInline(map, set);
719}
720ParseResult 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(loc: curLoc, message: "expected AffineMap, but got IntegerSet");
727 return success();
728}
729ParseResult Parser::parseAffineExprReference(
730 ArrayRef<std::pair<StringRef, AffineExpr>> symbolSet, AffineExpr &expr) {
731 return AffineParser(state).parseAffineExprInline(symbolSet, expr);
732}
733ParseResult 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(loc: curLoc, message: "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.
745ParseResult
746Parser::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.
755ParseResult
756Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
757 function_ref<ParseResult(bool)> parseElement) {
758 return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
759 .parseAffineExprOfSSAIds(expr);
760}
761
762static void parseAffineMapOrIntegerSet(StringRef inputStr, MLIRContext *context,
763 AffineMap &map, IntegerSet &set) {
764 llvm::SourceMgr sourceMgr;
765 auto memBuffer = llvm::MemoryBuffer::getMemBuffer(
766 InputData: inputStr, /*BufferName=*/"<mlir_parser_buffer>",
767 /*RequiresNullTerminator=*/false);
768 sourceMgr.AddNewSourceBuffer(F: std::move(memBuffer), IncludeLoc: 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(k: Token::eof)) {
781 parser.emitError(loc: endTok.getLoc(), message: "encountered unexpected token");
782 return;
783 }
784}
785
786AffineMap 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
795IntegerSet 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}
803

source code of mlir/lib/AsmParser/AffineParser.cpp