[MLIR][Presburger] Improve unittest parsing

This patch adds better functions for parsing MultiAffineFunctions and
PWMAFunctions in Presburger unittests.

A PWMAFunction can now be parsed as:

```
PWMAFunction result = parsePWMAF({
    {"(x, y) : (x >= 10, x <= 20, y >= 1)", "(x, y) -> (x + y)"},
    {"(x, y) : (x >= 21)", "(x, y) -> (x + y)"},
    {"(x, y) : (x <= 9)", "(x, y) -> (x - y)"},
    {"(x, y) : (x >= 10, x <= 20, y <= 0)", "(x, y) -> (x - y)"},
});
```

which is much more readable than the old format since the output can be
described as an AffineMap, instead of coefficients.

This patch also adds support for parsing divisions in MultiAffineFunctions
and PWMAFunctions which was previously not possible.

Reviewed By: arjunp

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

Author: Groverkss

mlir/include/mlir/AsmParser/AsmParser.h

@@ -76,14 +76,13 @@ Type parseType(llvm::StringRef typeStr, MLIRContext *context);
 /// returned in `numRead`.
 Type parseType(llvm::StringRef typeStr, MLIRContext *context, size_t &numRead);
 
-/// This parses a single IntegerSet to an MLIR context if it was valid. If not,
-/// an error message is emitted through a new SourceMgrDiagnosticHandler
-/// constructed from a new SourceMgr with a single MemoryBuffer wrapping
-/// `str`. If the passed `str` has additional tokens that were not part of the
-/// IntegerSet, a failure is returned. Diagnostics are printed on failure if
-/// `printDiagnosticInfo` is true.
-IntegerSet parseIntegerSet(llvm::StringRef str, MLIRContext *context,
-                           bool printDiagnosticInfo = true);
+/// This parses a single IntegerSet/AffineMap to an MLIR context if it was
+/// valid. If not, an error message is emitted through a new
+/// SourceMgrDiagnosticHandler constructed from a new SourceMgr with a single
+/// MemoryBuffer wrapping `str`. If the passed `str` has additional tokens that
+/// were not part of the IntegerSet/AffineMap, a failure is returned.
+AffineMap parseAffineMap(llvm::StringRef str, MLIRContext *context);
+IntegerSet parseIntegerSet(llvm::StringRef str, MLIRContext *context);
 
 } // namespace mlir
 

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

@@ -32,6 +32,10 @@ class Value;
 class MemRefType;
 struct MutableAffineMap;
 
+namespace presburger {
+class MultiAffineFunction;
+} // namespace presburger
+
 /// FlatAffineValueConstraints represents an extension of IntegerPolyhedron
 /// where each non-local variable can have an SSA Value attached to it.
 class FlatAffineValueConstraints : public presburger::IntegerPolyhedron {
@@ -615,6 +619,10 @@ getFlattenedAffineExprs(IntegerSet set,
                         std::vector<SmallVector<int64_t, 8>> *flattenedExprs,
                         FlatAffineValueConstraints *cst = nullptr);
 
+LogicalResult
+getMultiAffineFunctionFromMap(AffineMap map,
+                              presburger::MultiAffineFunction &multiAff);
+
 /// Re-indexes the dimensions and symbols of an affine map with given `operands`
 /// values to align with `dims` and `syms` values.
 ///

mlir/lib/AsmParser/AffineParser.cpp

@@ -734,8 +734,8 @@ Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
       .parseAffineExprOfSSAIds(expr);
 }
 
-IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context,
-                                 bool printDiagnosticInfo) {
+static void parseAffineMapOrIntegerSet(StringRef inputStr, MLIRContext *context,
+                                       AffineMap &map, IntegerSet &set) {
   llvm::SourceMgr sourceMgr;
   auto memBuffer = llvm::MemoryBuffer::getMemBuffer(
       inputStr, /*BufferName=*/"<mlir_parser_buffer>",
@@ -747,17 +747,31 @@ IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context,
                     /*codeCompleteContext=*/nullptr);
   Parser parser(state);
 
-  raw_ostream &os = printDiagnosticInfo ? llvm::errs() : llvm::nulls();
-  SourceMgrDiagnosticHandler handler(sourceMgr, context, os);
-  IntegerSet set;
-  if (parser.parseIntegerSetReference(set))
-    return IntegerSet();
+  SourceMgrDiagnosticHandler handler(sourceMgr, context, llvm::errs());
+  if (parser.parseAffineMapOrIntegerSetReference(map, set))
+    return;
 
   Token endTok = parser.getToken();
   if (endTok.isNot(Token::eof)) {
     parser.emitError(endTok.getLoc(), "encountered unexpected token");
-    return IntegerSet();
+    return;
   }
+}
+
+AffineMap mlir::parseAffineMap(StringRef inputStr, MLIRContext *context) {
+  AffineMap map;
+  IntegerSet set;
+  parseAffineMapOrIntegerSet(inputStr, context, map, set);
+  assert(!set &&
+         "expected string to represent AffineMap, but got IntegerSet instead");
+  return map;
+}
 
+IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context) {
+  AffineMap map;
+  IntegerSet set;
+  parseAffineMapOrIntegerSet(inputStr, context, map, set);
+  assert(!map &&
+         "expected string to represent IntegerSet, but got AffineMap instead");
   return set;
 }

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

@@ -1801,3 +1801,31 @@ LogicalResult mlir::getRelationFromMap(const AffineValueMap &map,
 
   return success();
 }
+
+LogicalResult
+mlir::getMultiAffineFunctionFromMap(AffineMap map,
+                                    MultiAffineFunction &multiAff) {
+  FlatAffineValueConstraints cst;
+  std::vector<SmallVector<int64_t, 8>> flattenedExprs;
+  LogicalResult result = getFlattenedAffineExprs(map, &flattenedExprs, &cst);
+
+  if (result.failed())
+    return failure();
+
+  DivisionRepr divs = cst.getLocalReprs();
+  assert(divs.hasAllReprs() &&
+         "AffineMap cannot produce divs without local representation");
+
+  // TODO: We shouldn't have to do this conversion.
+  Matrix mat(map.getNumResults(), map.getNumInputs() + divs.getNumDivs() + 1);
+  for (unsigned i = 0, e = flattenedExprs.size(); i < e; ++i)
+    for (unsigned j = 0, f = flattenedExprs[i].size(); j < f; ++j)
+      mat(i, j) = flattenedExprs[i][j];
+
+  multiAff = MultiAffineFunction(
+      PresburgerSpace::getRelationSpace(map.getNumDims(), map.getNumResults(),
+                                        map.getNumSymbols(), divs.getNumDivs()),
+      mat, divs);
+
+  return success();
+}

mlir/unittests/Analysis/Presburger/CMakeLists.txt

@@ -4,11 +4,12 @@ add_mlir_unittest(MLIRPresburgerTests
   LinearTransformTest.cpp
   MatrixTest.cpp
   MPIntTest.cpp
+  Parser.h
+  ParserTest.cpp
   PresburgerSetTest.cpp
   PresburgerSpaceTest.cpp
   PWMAFunctionTest.cpp
   SimplexTest.cpp
-  ../../Dialect/Affine/Analysis/AffineStructuresParser.cpp
 )
 
 target_link_libraries(MLIRPresburgerTests