[mlir][affinexpr] add parseAffineExpr to parser API

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.

Reviewed By: Peiming

Differential Revision: https://reviews.llvm.org/D154177

Author: aartbik

mlir/include/mlir/IR/OpImplementation.h

@@ -1032,6 +1032,12 @@ class AsmParser {
   /// 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;
 

mlir/lib/AsmParser/AffineParser.cpp

@@ -53,6 +53,9 @@ class AffineParser : public Parser {
   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,
@@ -533,6 +536,15 @@ ParseResult AffineParser::parseAffineMapOrIntegerSetInline(AffineMap &map,
   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,
@@ -703,6 +715,11 @@ ParseResult Parser::parseAffineMapReference(AffineMap &map) {
     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;

mlir/lib/AsmParser/AsmParserImpl.h

@@ -456,6 +456,14 @@ class AsmParserImpl : public BaseT {
     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);

mlir/lib/AsmParser/Parser.h

@@ -296,10 +296,13 @@ class Parser {
   // 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.