[Mlir-commits] [mlir] 9b3a080 - [mlir][affinexpr] add parseAffineExpr to parser API

[Aart Bik]

Author: Aart Bik
Date: 2023-06-29T23:44:04-07:00
New Revision: 9b3a080328370f0257bde2528b9c0b1616c2d98b

URL: https://github.com/llvm/llvm-project/commit/9b3a080328370f0257bde2528b9c0b1616c2d98b
DIFF: https://github.com/llvm/llvm-project/commit/9b3a080328370f0257bde2528b9c0b1616c2d98b.diff

LOG: [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

Added: 
    

Modified: 
    mlir/include/mlir/IR/OpImplementation.h
    mlir/lib/AsmParser/AffineParser.cpp
    mlir/lib/AsmParser/AsmParserImpl.h
    mlir/lib/AsmParser/Parser.h

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/IR/OpImplementation.h b/mlir/include/mlir/IR/OpImplementation.h
index a05c5a5c5bfea..dda4f3782ea79 100644
--- a/mlir/include/mlir/IR/OpImplementation.h
+++ b/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;
 

diff  --git a/mlir/lib/AsmParser/AffineParser.cpp b/mlir/lib/AsmParser/AffineParser.cpp
index 9c6058fbaf53e..e648a1a9d00d3 100644
--- a/mlir/lib/AsmParser/AffineParser.cpp
+++ b/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;

diff  --git a/mlir/lib/AsmParser/AsmParserImpl.h b/mlir/lib/AsmParser/AsmParserImpl.h
index aff3cb9a0b135..ef70a7126cba8 100644
--- a/mlir/lib/AsmParser/AsmParserImpl.h
+++ b/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);

diff  --git a/mlir/lib/AsmParser/Parser.h b/mlir/lib/AsmParser/Parser.h
index 9704cea0e2b55..c856e1f1f3054 100644
--- a/mlir/lib/AsmParser/Parser.h
+++ b/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.