[mlir] FlatAffineConstraint parsing for unit tests

This patch adds functionality to parse FlatAffineConstraints from a
StringRef with the intention to be used for unit tests. This should
make the construction of FlatAffineConstraints easier for testing
purposes.

The patch contains an example usage of the functionality in a unit test that
uses FlatAffineConstraints.

Reviewed By: bondhugula, grosser

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

mlir/include/mlir/Parser.h

@@ -256,6 +256,16 @@ Type parseType(llvm::StringRef typeStr, MLIRContext *context);
 /// `typeStr`. The number of characters of `typeStr` parsed in the process is
 /// 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);
+
 } // end namespace mlir
 
 #endif // MLIR_PARSER_H

mlir/lib/Parser/AffineParser.cpp

@@ -13,10 +13,13 @@
 #include "Parser.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/IntegerSet.h"
+#include "llvm/Support/SourceMgr.h"
 
 using namespace mlir;
 using namespace mlir::detail;
+using llvm::MemoryBuffer;
 using llvm::SMLoc;
+using llvm::SourceMgr;
 
 namespace {
 
@@ -717,3 +720,29 @@ Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
       .parseAffineExprOfSSAIds(expr);
 }
+
+IntegerSet mlir::parseIntegerSet(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);
+  IntegerSet set;
+  if (parser.parseIntegerSetReference(set))
+    return IntegerSet();
+
+  Token endTok = parser.getToken();
+  if (endTok.isNot(Token::eof)) {
+    parser.emitError(endTok.getLoc(), "encountered unexpected token");
+    return IntegerSet();
+  }
+
+  return set;
+}

mlir/unittests/Analysis/AffineStructuresParser.cpp

@@ -0,0 +1,24 @@
+//===- AffineStructuresParser.cpp - Parser for AffineStructures -*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "./AffineStructuresParser.h"
+#include "mlir/IR/IntegerSet.h"
+#include "mlir/Parser.h"
+
+using namespace mlir;
+
+FailureOr<FlatAffineConstraints>
+mlir::parseIntegerSetToFAC(llvm::StringRef str, MLIRContext *context,
+                           bool printDiagnosticInfo) {
+  IntegerSet set = parseIntegerSet(str, context, printDiagnosticInfo);
+
+  if (!set)
+    return failure();
+
+  return FlatAffineConstraints(set);
+}

mlir/unittests/Analysis/AffineStructuresParser.h

@@ -0,0 +1,33 @@
+//===- AffineStructuresParser.h - Parser for AffineStructures ---*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file defines helper functions to parse AffineStructures from
+// StringRefs.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_UNITTEST_ANALYSIS_AFFINESTRUCTURESPARSER_H
+#define MLIR_UNITTEST_ANALYSIS_AFFINESTRUCTURESPARSER_H
+
+#include "mlir/Analysis/AffineStructures.h"
+#include "mlir/Support/LogicalResult.h"
+
+namespace mlir {
+/// This parses a single IntegerSet to an MLIR context and transforms it to
+/// FlatAffineConstraints if it was valid. If not, a failure is returned. 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.
+
+FailureOr<FlatAffineConstraints>
+parseIntegerSetToFAC(llvm::StringRef, MLIRContext *context,
+                     bool printDiagnosticInfo = true);
+
+} // namespace mlir
+
+#endif // MLIR_UNITTEST_ANALYSIS_AFFINESTRUCTURESPARSER_H

mlir/unittests/Analysis/AffineStructuresParserTest.cpp

@@ -0,0 +1,137 @@
+//===- AffineStructuresParserTest.cpp - FAC parsing unit tests --*- C++ -*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file contains tests for parsing IntegerSets to FlatAffineConstraints.
+// The tests with invalid input check that the parser only accepts well-formed
+// IntegerSets. The tests with well-formed input compare the returned FACs to
+// manually constructed FACs with a PresburgerSet equality check.
+//
+//===----------------------------------------------------------------------===//
+
+#include "./AffineStructuresParser.h"
+#include "mlir/Analysis/PresburgerSet.h"
+
+#include <gtest/gtest.h>
+
+namespace mlir {
+
+/// Construct a FlatAffineConstraints from a set of inequality, equality, and
+/// division onstraints.
+static FlatAffineConstraints makeFACFromConstraints(
+    unsigned dims, unsigned syms, ArrayRef<SmallVector<int64_t, 4>> ineqs,
+    ArrayRef<SmallVector<int64_t, 4>> eqs = {},
+    ArrayRef<std::pair<SmallVector<int64_t, 4>, int64_t>> divs = {}) {
+  FlatAffineConstraints fac(ineqs.size(), eqs.size(), dims + syms + 1, dims,
+                            syms, 0);
+  for (const auto &div : divs)
+    fac.addLocalFloorDiv(div.first, div.second);
+  for (const auto &eq : eqs)
+    fac.addEquality(eq);
+  for (const auto &ineq : ineqs)
+    fac.addInequality(ineq);
+  return fac;
+}
+
+TEST(ParseFACTest, InvalidInputTest) {
+  MLIRContext context;
+  FailureOr<FlatAffineConstraints> fac;
+
+  fac = parseIntegerSetToFAC("(x)", &context, false);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings with no constraint list";
+
+  fac = parseIntegerSetToFAC("(x)[] : ())", &context, false);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain remaining characters";
+
+  fac = parseIntegerSetToFAC("(x)[] : (x - >= 0)", &context, false);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain incomplete constraints";
+
+  fac = parseIntegerSetToFAC("(x)[] : (y == 0)", &context, false);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain unkown identifiers";
+
+  fac = parseIntegerSetToFAC("(x, x) : (2 * x >= 0)", &context, false);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain repeated identifier names";
+
+  fac = parseIntegerSetToFAC("(x)[x] : (2 * x >= 0)", &context, false);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain repeated identifier names";
+
+  fac = parseIntegerSetToFAC("(x) : (2 * x + 9223372036854775808 >= 0)",
+                             &context, false);
+  EXPECT_TRUE(failed(fac)) << "should not accept strings with integer literals "
+                              "that do not fit into int64_t";
+}
+
+/// Parses and compares the `str` to the `ex`. The equality check is performed
+/// by using PresburgerSet::isEqual
+static bool parseAndCompare(StringRef str, FlatAffineConstraints ex,
+                            MLIRContext *context) {
+  FailureOr<FlatAffineConstraints> fac = parseIntegerSetToFAC(str, context);
+
+  EXPECT_TRUE(succeeded(fac));
+
+  return PresburgerSet(*fac).isEqual(PresburgerSet(ex));
+}
+
+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)",
+                              makeFACFromConstraints(1, 0, {}, {{1, 0}}),
+                              &context));
+
+  // multiple constraints
+  EXPECT_TRUE(parseAndCompare("(x)[] : (7 * x >= 0, -7 * x + 5 >= 0)",
+                              makeFACFromConstraints(1, 0, {{7, 0}, {-7, 5}}),
+                              &context));
+
+  // multiple dimensions
+  EXPECT_TRUE(parseAndCompare("(x,y,z)[] : (x + y - z >= 0)",
+                              makeFACFromConstraints(3, 0, {{1, 1, -1, 0}}),
+                              &context));
+
+  // dimensions and symbols
+  EXPECT_TRUE(parseAndCompare(
+      "(x,y,z)[a,b] : (x + y - z + 2 * a - 15 * b >= 0)",
+      makeFACFromConstraints(3, 2, {{1, 1, -1, 2, -15, 0}}), &context));
+
+  // only symbols
+  EXPECT_TRUE(parseAndCompare("()[a] : (2 * a - 4 == 0)",
+                              makeFACFromConstraints(0, 1, {}, {{2, -4}}),
+                              &context));
+
+  // simple floordiv
+  EXPECT_TRUE(parseAndCompare(
+      "(x, y) : (y - 3 * ((x + y - 13) floordiv 3) - 42 == 0)",
+      makeFACFromConstraints(2, 0, {}, {{0, 1, -3, -42}}, {{{1, 1, -13}, 3}}),
+      &context));
+
+  // multiple floordiv
+  EXPECT_TRUE(parseAndCompare(
+      "(x, y) : (y - x floordiv 3 - y floordiv 2 == 0)",
+      makeFACFromConstraints(2, 0, {}, {{0, 1, -1, -1, 0}},
+                             {{{1, 0, 0}, 3}, {{0, 1, 0, 0}, 2}}),
+      &context));
+
+  // 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));
+}
+
+} // namespace mlir

mlir/unittests/Analysis/AffineStructuresTest.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineStructures.h"
+#include "./AffineStructuresParser.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/MLIRContext.h"
 
@@ -98,11 +99,24 @@ static void checkPermutationsSample(bool hasSample, unsigned nDim,
   } while (std::next_permutation(perm.begin(), perm.end()));
 }
 
+/// Parses a FlatAffineConstraints from a StringRef. It is expected that the
+/// string represents a valid IntegerSet, otherwise it will violate a gtest
+/// assertion.
+static FlatAffineConstraints parseFAC(StringRef str, MLIRContext *context) {
+  FailureOr<FlatAffineConstraints> fac = parseIntegerSetToFAC(str, context);
+
+  EXPECT_TRUE(succeeded(fac));
+
+  return *fac;
+}
+
 TEST(FlatAffineConstraintsTest, FindSampleTest) {
   // Bounded sets with only inequalities.
 
+  MLIRContext context;
+
   // 0 <= 7x <= 5
-  checkSample(true, makeFACFromConstraints(1, {{7, 0}, {-7, 5}}, {}));
+  checkSample(true, parseFAC("(x) : (7 * x >= 0, -7 * x + 5 >= 0)", &context));
 
   // 1 <= 5x and 5x <= 4 (no solution).
   checkSample(false, makeFACFromConstraints(1, {{5, -1}, {-5, 4}}, {}));

mlir/unittests/Analysis/CMakeLists.txt

@@ -1,10 +1,15 @@
 add_mlir_unittest(MLIRAnalysisTests
+  AffineStructuresParser.cpp
+  AffineStructuresParserTest.cpp
   AffineStructuresTest.cpp
   LinearTransformTest.cpp
   PresburgerSetTest.cpp
 )
 
 target_link_libraries(MLIRAnalysisTests
-  PRIVATE MLIRLoopAnalysis)
+  PRIVATE 
+  MLIRLoopAnalysis
+  MLIRParser
+  )
 
 add_subdirectory(Presburger)