[mlir][affinexpr] add parseAffineExpr to parser API

[aartbik]

Summary

Similar to AffineMap and IntegerSet parsing, this change makes the more fine-grained AffineExpr available for general parsing, using a preset symbol set to recognize variables.

Motivation:
The AffineExpr parser will be used by the new sparse tensor encoding surface syntax. Originally, we planned to duplicate the affine parser completely, but that would be a terrible waste of a good thing. With this minor API change, we prepare the way for the sparse tensor dialect (and others) to reuse the AffineExpr parser outside the context of a more restricted AffineMap parser.

[Peiming]

mlir/lib/AsmParser/AffineParser.cpp

60–61

Was it dead code previously?

539

typo, symbols

[aartbik]

fix typos

[aartbik]

mlir/lib/AsmParser/AffineParser.cpp

60–61

yeah, not defined anywhere, bit surprising this is not linted somehow

[mehdi_amini]

Can you add a motivation in the description? Right now it seems like dead code in the repo

mlir/lib/AsmParser/AffineParser.cpp

60–61

Can you just push this ahead of time in a NFC commit? (as well as other NFC things)

[aartbik]

Will split out the NFC stuff

mlir/lib/AsmParser/AffineParser.cpp

60–61

sure, also the renaming

[aartbik]

I split out the NFC stuff into https://reviews.llvm.org/D154179
and I added motivation to this revision

[aartbik]

rebased with main (this removes the NFC part of the original revision)

mlir/include/mlir/IR/OpImplementation.h

public:
   /// Parse a named dictionary into 'result' if the `attributes` keyword is
   /// present.
   virtual ParseResult
   parseOptionalAttrDictWithKeyword(NamedAttrList &result) = 0;
 
   /// Parse an affine map instance into 'map'.
   virtual ParseResult parseAffineMap(AffineMap &map) = 0;
 
+  /// Parse an affine expr instance into 'expr' using the already computed
+  /// mapping from symbols to affine expressions in 'symbolSet'.
+  virtual ParseResult
+  parseAffineExpr(SmallVectorImpl<std::pair<StringRef, AffineExpr>> &symbolSet,
+                  AffineExpr &expr) = 0;
+
   /// Parse an integer set instance into 'set'.
   virtual ParseResult parseIntegerSet(IntegerSet &set) = 0;
 
   //===--------------------------------------------------------------------===//
   // Identifier Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse an @-identifier and store it (without the '@' symbol) in a string

mlir/lib/AsmParser/AffineParser.cpp

public:
   AffineParser(ParserState &state, bool allowParsingSSAIds = false,
                function_ref<ParseResult(bool)> parseElement = nullptr)
       : Parser(state), allowParsingSSAIds(allowParsingSSAIds),
         parseElement(parseElement) {}
 
   ParseResult parseAffineMapRange(unsigned numDims, unsigned numSymbols,
                                   AffineMap &result);
   ParseResult parseAffineMapOrIntegerSetInline(AffineMap &map, IntegerSet &set);
+  ParseResult parseAffineExprInline(
+      SmallVectorImpl<std::pair<StringRef, AffineExpr>> &symbolSet,
+      AffineExpr &expr);
   ParseResult parseIntegerSetConstraints(unsigned numDims, unsigned numSymbols,
                                          IntegerSet &result);
   ParseResult parseAffineMapOfSSAIds(AffineMap &map,
                                      OpAsmParser::Delimiter delimiter);
   ParseResult parseAffineExprOfSSAIds(AffineExpr &expr);
 
 private:
   // Binary affine op parsing.
   AffineLowPrecOp consumeIfLowPrecOp();
   AffineHighPrecOp consumeIfHighPrecOp();
 
   // Identifier lists for polyhedral structures.
   ParseResult parseDimIdList(unsigned &numDims);
   ParseResult parseSymbolIdList(unsigned &numSymbols);
ParseResult AffineParser::parseAffineMapOrIntegerSetInline(AffineMap &map,
   if (consumeIf(Token::arrow))
     return parseAffineMapRange(numDims, numSymbols, map);
 
   if (parseToken(Token::colon, "expected '->' or ':'"))
     return failure();
   return parseIntegerSetConstraints(numDims, numSymbols, set);
 }
 
+/// Parse an affine expresion definition inline, with given symbols.
+ParseResult AffineParser::parseAffineExprInline(
+    SmallVectorImpl<std::pair<StringRef, AffineExpr>> &symbolSet,
+    AffineExpr &expr) {
+  dimsAndSymbols.assign(symbolSet);
+  expr = parseAffineExpr();
+  return success(expr != nullptr);
+}
+
 /// Parse an AffineMap where the dim and symbol identifiers are SSA ids.
 ParseResult
 AffineParser::parseAffineMapOfSSAIds(AffineMap &map,
                                      OpAsmParser::Delimiter delimiter) {
 
   SmallVector<AffineExpr, 4> exprs;
   auto parseElt = [&]() -> ParseResult {
     auto elt = parseAffineExpr();
ParseResult Parser::parseAffineMapReference(AffineMap &map) {
   SMLoc curLoc = getToken().getLoc();
   IntegerSet set;
   if (parseAffineMapOrIntegerSetReference(map, set))
     return failure();
   if (set)
     return emitError(curLoc, "expected AffineMap, but got IntegerSet");
   return success();
 }
+ParseResult Parser::parseAffineExprReference(
+    SmallVectorImpl<std::pair<StringRef, AffineExpr>> &symbolSet,
+    AffineExpr &expr) {
+  return AffineParser(state).parseAffineExprInline(symbolSet, expr);
+}
 ParseResult Parser::parseIntegerSetReference(IntegerSet &set) {
   SMLoc curLoc = getToken().getLoc();
   AffineMap map;
   if (parseAffineMapOrIntegerSetReference(map, set))
     return failure();
   if (map)
     return emitError(curLoc, "expected IntegerSet, but got AffineMap");
   return success();

mlir/lib/AsmParser/AsmParserImpl.h

ParseResult parseOptionalAttrDictWithKeyword(NamedAttrList &result) override {
     return parser.parseAttributeDict(result);
   }
 
   /// Parse an affine map instance into 'map'.
   ParseResult parseAffineMap(AffineMap &map) override {
     return parser.parseAffineMapReference(map);
   }
 
+  /// Parse an affine expr instance into 'expr' using the already computed
+  /// mapping from symbols to affine expressions in 'symbolSet'.
+  ParseResult
+  parseAffineExpr(SmallVectorImpl<std::pair<StringRef, AffineExpr>> &symbolSet,
+                  AffineExpr &expr) override {
+    return parser.parseAffineExprReference(symbolSet, expr);
+  }
+
   /// Parse an integer set instance into 'set'.
   ParseResult parseIntegerSet(IntegerSet &set) override {
     return parser.parseIntegerSetReference(set);
   }
 
   //===--------------------------------------------------------------------===//
   // Identifier Parsing
   //===--------------------------------------------------------------------===//

mlir/lib/AsmParser/Parser.h

public:
 
   /// Parse a name or FileLineCol location instance.
   ParseResult parseNameOrFileLineColLocation(LocationAttr &loc);
 
   //===--------------------------------------------------------------------===//
   // Affine Parsing
   //===--------------------------------------------------------------------===//
 
-  /// Parse a reference to either an affine map, or an integer set.
+  /// Parse a reference to either an affine map, expr, or an integer set.
   ParseResult parseAffineMapOrIntegerSetReference(AffineMap &map,
                                                   IntegerSet &set);
   ParseResult parseAffineMapReference(AffineMap &map);
+  ParseResult parseAffineExprReference(
+      SmallVectorImpl<std::pair<StringRef, AffineExpr>> &symbolSet,
+      AffineExpr &expr);
   ParseResult parseIntegerSetReference(IntegerSet &set);
 
   /// Parse an AffineMap where the dim and symbol identifiers are SSA ids.
   ParseResult
   parseAffineMapOfSSAIds(AffineMap &map,
                          function_ref<ParseResult(bool)> parseElement,
                          Delimiter delimiter);