[MLIR][PRESBURGER] Parse multiple integerSets.

[pashu123]

Summary

This patch adds functionality to parse multiple integer-sets of the form

multiple-integer-sets
    ::= dim-and-symbol-id-lists `:`
    '(' affine-constraint-conjunction ')' (`,` '(' affine-constraint-conjunction ')')*
    affine-constraint-conjunction ::= affine-constraint (`,`
                                       affine-constraint)*

For e.g., "(x,y) : (x + y - 4 >= 0), (x - 2 >= 0)".

[rriddle]

What is the intended use case of this? It isn't clear to me that we should expand the parser API in this way.

[pashu123]

@rriddle I guess the primary reason is for test cases and benchmarking. If you could suggest some other way to add this, that would be great.

mlir/include/mlir/Parser/Parser.h

 /// 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 comma separated IntegerSets 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
+/// IntegerSets, a failure is returned. Diagnostics are printed on failure if
+/// `printDiagnosticInfo` is true.
+SmallVector<IntegerSet, 4>
+parseMultipleIntegerSets(llvm::StringRef str, MLIRContext *context,
+                         bool printDiagnosticInfo = true);
+
 } // namespace mlir
 
 #endif // MLIR_PARSER_PARSER_H

mlir/lib/Parser/AffineParser.cpp

public:
   AffineParser(ParserState &state, bool allowParsingSSAIds = false,
                function_ref<ParseResult(bool)> parseElement = nullptr)
       : Parser(state), allowParsingSSAIds(allowParsingSSAIds),
         parseElement(parseElement) {}
 
   AffineMap parseAffineMapRange(unsigned numDims, unsigned numSymbols);
   ParseResult parseAffineMapOrIntegerSetInline(AffineMap &map, IntegerSet &set);
   IntegerSet parseIntegerSetConstraints(unsigned numDims, unsigned numSymbols);
+  ParseResult parseMultipleIntegerSets(SmallVectorImpl<IntegerSet> &unionSet);
   ParseResult parseAffineMapOfSSAIds(AffineMap &map,
                                      OpAsmParser::Delimiter delimiter);
   ParseResult parseAffineExprOfSSAIds(AffineExpr &expr);
   void getDimsAndSymbolSSAIds(SmallVectorImpl<StringRef> &dimAndSymbolSSAIds,
                               unsigned &numDims);
 
 private:
   // Binary affine op parsing.
if (parseToken(Token::colon, "expected ':' or '['"))
     return failure();
 
   if ((set = parseIntegerSetConstraints(numDims, numSymbols)))
     return success();
 
   return failure();
 }
 
+/// Parses the comma-separated integer sets.
+ParseResult
+AffineParser::parseMultipleIntegerSets(SmallVectorImpl<IntegerSet> &unionSet) {
+  unsigned numDims = 0, numSymbols = 0;
+
+  // List of dimensional and optional symbol identifiers.
+  if (parseDimAndOptionalSymbolIdList(numDims, numSymbols)) {
+    return failure();
+  }
+
+  if (parseToken(Token::colon, "expected ':'"))
+    return failure();
+
+  IntegerSet set;
+  do {
+    if ((set = parseIntegerSetConstraints(numDims, numSymbols)))
+      unionSet.push_back(set);
+    else
+      return failure();
+  } while (consumeIf(Token::comma));
+
+  return success();
+}
+
 /// 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::parseIntegerSetReference(IntegerSet &set) {
   AffineMap map;
   if (parseAffineMapOrIntegerSetReference(map, set))
     return failure();
   if (map)
     return emitError(curLoc, "expected IntegerSet, but got AffineMap");
   return success();
 }
 
+ParseResult
+Parser::parseMultipleIntegerSetsReference(SmallVectorImpl<IntegerSet> &set) {
+  if (AffineParser(state).parseMultipleIntegerSets(set))
+    return failure();
+  return success();
+}
+
 /// Parse an AffineMap of SSA ids. The callback 'parseElement' is used to
 /// parse SSA value uses encountered while parsing affine expressions.
 ParseResult
 Parser::parseAffineMapOfSSAIds(AffineMap &map,
                                function_ref<ParseResult(bool)> parseElement,
                                OpAsmParser::Delimiter delimiter) {
   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
       .parseAffineMapOfSSAIds(map, delimiter);
IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context,
   Token endTok = parser.getToken();
   if (endTok.isNot(Token::eof)) {
     parser.emitError(endTok.getLoc(), "encountered unexpected token");
     return IntegerSet();
   }
 
   return set;
 }
+
+SmallVector<IntegerSet, 4>
+mlir::parseMultipleIntegerSets(StringRef inputStr, MLIRContext *context,
+                               bool printDiagnosticInfo) {
+  llvm::SourceMgr sourceMgr;
+  auto memBuffer = llvm::MemoryBuffer::getMemBuffer(
+      inputStr, /*BufferName=*/"<mlir_parser_buffer>",
+      /*RequiresNullTerminator=*/false);
+  sourceMgr.AddNewSourceBuffer(std::move(memBuffer), SMLoc());
+  SymbolState symbolState;
+  ParserState state(sourceMgr, context, symbolState, /*asmState=*/nullptr);
+  Parser parser(state);
+
+  raw_ostream &os = printDiagnosticInfo ? llvm::errs() : llvm::nulls();
+  SourceMgrDiagnosticHandler handler(sourceMgr, context, os);
+  SmallVector<IntegerSet> set;
+  if (parser.parseMultipleIntegerSetsReference(set))
+    return SmallVector<IntegerSet>();
+
+  Token endTok = parser.getToken();
+  if (endTok.isNot(Token::eof)) {
+    parser.emitError(endTok.getLoc(), "encountered unexpected token");
+    return SmallVector<IntegerSet>();
+  }
+  return set;
+}

mlir/lib/Parser/Parser.h

public:
   // Affine Parsing
   //===--------------------------------------------------------------------===//
 
   /// Parse a reference to either an affine map, or an integer set.
   ParseResult parseAffineMapOrIntegerSetReference(AffineMap &map,
                                                   IntegerSet &set);
   ParseResult parseAffineMapReference(AffineMap &map);
   ParseResult parseIntegerSetReference(IntegerSet &set);
+  ParseResult
+  parseMultipleIntegerSetsReference(SmallVectorImpl<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);
 
   /// Parse an AffineExpr where dim and symbol identifiers are SSA ids.

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

static PresburgerSet makeSetFromPoly(unsigned numDims,
   PresburgerSet set =
       PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(numDims));
   for (const IntegerPolyhedron &poly : polys)
     set.unionInPlace(poly);
   return set;
 }
 
 TEST(SetTest, containsPoint) {
-  PresburgerSet setA = parsePresburgerSetFromPolyStrings(
-      1,
-      {"(x) : (x - 2 >= 0, -x + 8 >= 0)", "(x) : (x - 10 >= 0, -x + 20 >= 0)"});
+  PresburgerSet setA = parsePresburgerSet(
+      "(x) : (x - 2 >= 0, -x + 8 >= 0), (x - 10 >= 0, -x + 20 >= 0)");
   for (unsigned x = 0; x <= 21; ++x) {
     if ((2 <= x && x <= 8) || (10 <= x && x <= 20))
       EXPECT_TRUE(setA.containsPoint({x}));
     else
       EXPECT_FALSE(setA.containsPoint({x}));
   }
 
   // A parallelogram with vertices {(3, 1), (10, -6), (24, 8), (17, 15)} union
   // a square with opposite corners (2, 2) and (10, 10).
-  PresburgerSet setB = parsePresburgerSetFromPolyStrings(
-      2, {"(x,y) : (x + y - 4 >= 0, -x - y + 32 >= 0, "
-          "x - y - 2 >= 0, -x + y + 16 >= 0)",
-          "(x,y) : (x - 2 >= 0, y - 2 >= 0, -x + 10 >= 0, -y + 10 >= 0)"});
+  PresburgerSet setB = parsePresburgerSet(
+      "(x,y) : (x + y - 4 >= 0, -x - y + 32 >= 0, "
+      "x - y - 2 >= 0, -x + y + 16 >= 0),"
+      "(x - 2 >= 0, y - 2 >= 0, -x + 10 >= 0, -y + 10 >= 0)");
 
   for (unsigned x = 1; x <= 25; ++x) {
     for (unsigned y = -6; y <= 16; ++y) {
       if (4 <= x + y && x + y <= 32 && 2 <= x - y && x - y <= 16)
         EXPECT_TRUE(setB.containsPoint({x, y}));
       else if (2 <= x && x <= 10 && 2 <= y && y <= 10)
         EXPECT_TRUE(setB.containsPoint({x, y}));
       else
         EXPECT_FALSE(setB.containsPoint({x, y}));
     }
   }
 }
 
 TEST(SetTest, Union) {
-  PresburgerSet set = parsePresburgerSetFromPolyStrings(
-      1,
-      {"(x) : (x - 2 >= 0, -x + 8 >= 0)", "(x) : (x - 10 >= 0, -x + 20 >= 0)"});
+  PresburgerSet set = parsePresburgerSet(
+      "(x) : (x - 2 >= 0, -x + 8 >= 0), (x - 10 >= 0, -x + 20 >= 0)");
 
   // Universe union set.
   testUnionAtPoints(PresburgerSet::getUniverse(PresburgerSpace::getSetSpace(1)),
                     set, {{1}, {2}, {8}, {9}, {10}, {20}, {21}});
 
   // empty set union set.
   testUnionAtPoints(PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(1)),
                     set, {{1}, {2}, {8}, {9}, {10}, {20}, {21}});

mlir/unittests/Analysis/Presburger/Utils.h

 /// assertion.
 inline IntegerPolyhedron parsePoly(StringRef str) {
   MLIRContext context(MLIRContext::Threading::DISABLED);
   FailureOr<IntegerPolyhedron> poly = parseIntegerSetToFAC(str, &context);
   EXPECT_TRUE(succeeded(poly));
   return *poly;
 }
 
+/// Parses a list of comma separated IntegerSets to IntegerPolyhedron and
+/// combine them into a PresburgerSet by using the union operation. It is
+/// expected that the string has valid comma separated IntegerSet constraints
+/// and that all of them have the same number of dimensions as is specified by
+/// the numDims argument.
+inline PresburgerSet parsePresburgerSet(StringRef str) {
+  MLIRContext context(MLIRContext::Threading::DISABLED);
+  FailureOr<SmallVector<FlatAffineConstraints, 4>> facs =
+      parseMultipleIntegerSetsToFAC(str, &context);
+  SmallVector<IntegerPolyhedron, 4> ips;
+  EXPECT_TRUE(succeeded(facs));
+  for (auto fac : facs.getValue())
+    ips.push_back(IntegerPolyhedron(fac));
+
+  PresburgerSet set = PresburgerSet(ips.front());
+  for (int i = 1, m = facs.getValue().size(); i < m; i++)
+    set.unionInPlace(ips[i]);
+  return set;
+}
+
 /// Parse a list of StringRefs to IntegerRelation and combine them into a
 /// PresburgerSet be using the union operation. It is expected that the strings
 /// are all valid IntegerSet representation and that all of them have the same
 /// number of dimensions as is specified by the numDims argument.
 inline PresburgerSet
 parsePresburgerSetFromPolyStrings(unsigned numDims, ArrayRef<StringRef> strs) {
   PresburgerSet set =
       PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(numDims));

mlir/unittests/Dialect/Affine/Analysis/AffineStructuresParser.h

 /// 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.
 
 FailureOr<presburger::IntegerPolyhedron>
 parseIntegerSetToFAC(llvm::StringRef, MLIRContext *context,
                      bool printDiagnosticInfo = true);
 
+FailureOr<SmallVector<FlatAffineConstraints, 4>>
+parseMultipleIntegerSetsToFAC(llvm::StringRef str, MLIRContext *context,
+                              bool printDiagnosticInfo = true);
+
 } // namespace mlir
 
 #endif // MLIR_UNITTEST_ANALYSIS_AFFINESTRUCTURESPARSER_H

mlir/unittests/Dialect/Affine/Analysis/AffineStructuresParser.cpp

mlir::parseIntegerSetToFAC(llvm::StringRef str, MLIRContext *context,
                            bool printDiagnosticInfo) {
   IntegerSet set = parseIntegerSet(str, context, printDiagnosticInfo);
 
   if (!set)
     return failure();
 
   return FlatAffineValueConstraints(set);
 }
+
+FailureOr<SmallVector<FlatAffineConstraints, 4>>
+mlir::parseMultipleIntegerSetsToFAC(llvm::StringRef str, MLIRContext *context,
+                                    bool printDiagnosticInfo) {
+  SmallVector<IntegerSet, 4> set =
+      parseMultipleIntegerSets(str, context, printDiagnosticInfo);
+
+  if (set.empty())
+    return failure();
+
+  SmallVector<FlatAffineConstraints, 4> ret;
+  for (auto iSet : set)
+    ret.push_back(FlatAffineConstraints(iSet));
+
+  return ret;
+}

mlir/unittests/Dialect/Affine/Analysis/AffineStructuresParserTest.cpp

static bool parseAndCompare(StringRef str, const IntegerPolyhedron &ex,
                             MLIRContext *context) {
   FailureOr<IntegerPolyhedron> fac = parseIntegerSetToFAC(str, context);
 
   EXPECT_TRUE(succeeded(fac));
 
   return PresburgerSet(*fac).isEqual(PresburgerSet(ex));
 }
 
+static bool parseAndCompare(StringRef str,
+                            SmallVectorImpl<FlatAffineConstraints> &ex,
+                            MLIRContext *context) {
+  FailureOr<SmallVector<FlatAffineConstraints, 4>> fac =
+      parseMultipleIntegerSetsToFAC(str, context);
+
+  EXPECT_TRUE(succeeded(fac));
+  auto facVec = fac.getValue();
+  unsigned m = facVec.size();
+
+  EXPECT_TRUE(m == ex.size());
+
+  for (unsigned i = 0; i < m; i++)
+    if (!PresburgerSet(facVec[i]).isEqual(PresburgerSet(ex[i])))
+      return false;
+  return true;
+}
+
 TEST(ParseFACTest, ParseAndCompareTest) {
   MLIRContext context;
   // simple ineq
   EXPECT_TRUE(parseAndCompare(
       "(x)[] : (x >= 0)", makeFACFromConstraints(1, 0, {{1, 0}}), &context));
 
   // simple eq
   EXPECT_TRUE(parseAndCompare("(x)[] : (x == 0)",
TEST(ParseFACTest, ParseAndCompareTest) {
 
   // nested floordiv
   EXPECT_TRUE(parseAndCompare(
       "(x, y) : (y - (x + y floordiv 2) floordiv 3 == 0)",
       makeFACFromConstraints(2, 0, {}, {{0, 1, 0, -1, 0}},
                              {{{0, 1, 0}, 2}, {{1, 0, 1, 0}, 3}}),
       &context));
 }
+
+TEST(ParseMultipleFACTest, ParseAndCompareTest) {
+  MLIRContext context;
+
+  // Parses a simple ineq and eq respectively.
+  SmallVector<FlatAffineConstraints, 4> actualFac{
+      makeFACFromConstraints(1, 0, {{1, 0}}),
+      makeFACFromConstraints(1, 0, {}, {{1, 0}})};
+  EXPECT_TRUE(
+      parseAndCompare("(x)[] : (x >= 0), (x == 0)", actualFac, &context));
+
+  // simple floordiv, multiple floordiv and nested floordiv together.
+  SmallVector<FlatAffineConstraints, 4> divActualFac{
+      makeFACFromConstraints(2, 0, {}, {{0, 1, -3, -42}}, {{{1, 1, -13}, 3}}),
+      makeFACFromConstraints(2, 0, {}, {{0, 1, -1, -1, 0}},
+                             {{{1, 0, 0}, 3}, {{0, 1, 0, 0}, 2}}),
+      makeFACFromConstraints(2, 0, {}, {{0, 1, 0, -1, 0}},
+                             {{{0, 1, 0}, 2}, {{1, 0, 1, 0}, 3}})};
+
+  EXPECT_TRUE(
+      parseAndCompare("(x, y) : (y - 3 * ((x + y - 13) floordiv 3) - 42 == "
+                      "0), (y - x floordiv 3 - y floordiv 2 == 0), (y - (x + "
+                      "y floordiv 2) floordiv 3 == 0)",
+                      divActualFac, &context));
+}