[mlir] FlatAffineConstraint parsing for unit tests

[Dinistro]

Summary

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.

[Dinistro]

The usage of the parser in mlir/unittests/Analysis/AffineStructuresTest.cpp serves as an example. If this change to the parser and its usage for FAC unit tests is approved, we can replace the remaining creations of FACs with calls to parseFAC.

[bondhugula]

mlir/include/mlir/Parser.h

16

This is going to be a layering issue. The parser library can't depend on the Affine Analysis library, which is where FlatAffineConstraints resides. We'll have to define the parser for FlatAffineConstraints elsewhere. The MLIR parser is meant for just parsing IR AFAIK, and in any case, it can't depend on such analysis structures.

[bondhugula]

This is useful to have but please define this in the same library as FlatAffineConstraints.

[bondhugula]

mlir/unittests/Analysis/AffineStructuresTest.cpp

101

Missing doc comment.

mlir/unittests/Analysis/CMakeLists.txt

6

Sorted order here.

[Dinistro]

Address bondhugula's review comments

[Dinistro]

Thanks for your review. I addressed the comments by only exposing parseIntegerSet from the parser library and using it in AffineStructuresParser which is now part of the Affine Analysis library.

[bondhugula]

The diff is being shown on the wrong base. Please squash the multiple commits that you have on this revision. phab itself is able to show your revisions as you change them.

mlir/lib/Analysis/CMakeLists.txt

5

You don't need a separate file for a single function. It can be created when more are to be added. Could you just add the parsing function to AffineStructures.h/.cpp?

[Dinistro]

Squash commits and moving parsing functionality in AffineStructures.(h|cpp)

[Dinistro]

Add missing MLIRParser dependency to MLIRLoopAnalysis

[Groverkss]

mlir/unittests/Analysis/AffineStructuresParserTest.cpp

33

"stings" should be "strings". Same in other places.

[Dinistro]

Fix typos in error messages of AffineStructuresParserTest.

[bondhugula]

mlir/include/mlir/Analysis/AffineStructures.h

1079

This isn't parsing a FlatAffineConstraints -- instead, parseIntegerSetToFAC?

mlir/unittests/Analysis/AffineStructuresParserTest.cpp

2

Missing header block along with a file-level summary comment.

17–19

Drop trivial braces.

64–88

Creating an MLIRContext for each parseAndCompare doesn't look great. Instead, create it outside and pass it to parseAndCompare.

mlir/unittests/Analysis/AffineStructuresTest.cpp

104

This should be taking an MLIRContext as an argument instead of creating one each time -- also to be consistent with other similar methods.

mlir/unittests/Analysis/CMakeLists.txt

4

Sorted order here.

[Dinistro]

Addressed additional comments from bondhugula.

[Dinistro]

Thanks for your comments. All of them should be addressed now.

[bondhugula]

LGTM.

[rriddle]

mlir/lib/Analysis/AffineStructures.cpp

3832

Analysis should not be depending on the parser, please revert this.

[rriddle]

Commented inline, but we should not be depending on the parser from Analysis. Can you revert this and expose the functionality via the parser lib? We shouldn't be exposing parsing functionality in Analysis.

[mehdi_amini]

mlir/lib/Analysis/AffineStructures.cpp

3832

I reverted in d5730647accf for now.

[Dinistro]

Exposing this from the Parser library would make it dependent on Analysis, see bondhugula's earlier comments. Would this be an acceptable dependency or do you have another suggestion on how to solve this cleanly?

[rriddle]

In D113275#3130553, @Dinistro wrote:

Exposing this from the Parser library would make it dependent on Analysis, see bondhugula's earlier comments. Would this be an acceptable dependency or do you have another suggestion on how to solve this cleanly?

Thanks for pointing that out (I didn't read the full comments). I don't think we want the dependency in either case. Do we really need the combined method? Seems fairly simple (2-3 lines), and removing it also removes the need for any kind of dependency management.

[Dinistro]

The whole point of this patch is to provide a helper function for unit tests involving FlatAffineConstraints as construction them manually is error prone. Parsing them from a string is a lot easier and simplifies test writing. Maybe one could make the tests themselves depend on both libraries and do the transformation from IntegerSets to FACs there, but I'm not sure if that's clean either.

[rriddle]

In D113275#3130559, @Dinistro wrote:

The whole point of this patch is to provide a helper function for unit tests involving FlatAffineConstraints as construction them manually is error prone. Parsing them from a string is a lot easier and simplifies test writing. Maybe one could make the tests themselves depend on both libraries and do the transformation from IntegerSets to FACs there, but I'm not sure if that's clean either.

If it's just for unit tests, we shouldn't be adding it in the public API at all. This should be limited, and likely just local to the unit test files/directory.

[Dinistro]

Moved the parseIntegerSetToFAC function into the unittests directory to avoid coupling the Analysis and Parser libraries.

[Dinistro]

I implemented River's suggestions and reopened this revision to ensure that all the comments and previous versions stay in the same place.

[Dinistro]

Fix linting error and revert unrelated changes.

[Groverkss]

mlir/lib/Parser/AffineParser.cpp

733

For unit tests, we probably do not want diagnostic information to be printed. Can you somehow pass llvm::nulls() to SourceMgrDiagnosticHandler when it is used for unittests?

[bondhugula]

In D113275#3131183, @Dinistro wrote:

Moved the parseIntegerSetToFAC function into the unittests directory to avoid coupling the Analysis and Parser libraries.

This looks fine to me with the unit tests library depending on Parser. A dependency on the entire parser library is a bit weird, but if needed, the affine parser library could be split out into a smaller library in the future.

[Dinistro]

Add flag for printing diagnostics. This ensures that parser tests do not flood llvm:errs() with unimportant infos.

[grosser]

This looks good from my side as well. Let's see what @rriddle says.

[Dinistro]

What do you think @rriddle ?

[rriddle]

Having a dependency on the parser from a unit test looks fine, LGTM.

mlir/include/mlir/Analysis/AffineStructures.h

 ///
 /// For AffineValueMap, the domain and symbols have Value set corresponding to
 /// the Value in `map`. Returns failure if the AffineMap could not be flattened
 /// (i.e., semi-affine is not yet handled).
 LogicalResult getRelationFromMap(AffineMap &map, FlatAffineRelation &rel);
 LogicalResult getRelationFromMap(const AffineValueMap &map,
                                  FlatAffineRelation &rel);
 
+/// 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.
+FailureOr<FlatAffineConstraints> parseIntegerSetToFAC(llvm::StringRef,
+                                                      MLIRContext *context);
+
 } // end namespace mlir.
 
 #endif // MLIR_ANALYSIS_AFFINESTRUCTURES_H

mlir/include/mlir/Parser.h

 //===- Parser.h - MLIR Parser Library Interface -----------------*- 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 is contains the interface to the MLIR parser library.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_PARSER_H
 #define MLIR_PARSER_H
 
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include <cstddef>
 
 namespace llvm {
 class SourceMgr;
 class SMDiagnostic;
 class StringRef;
 } // end namespace llvm

 Type parseType(llvm::StringRef typeStr, MLIRContext *context);
 
 /// This parses a single MLIR type 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 a MemoryBuffer wrapping
 /// `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.
+IntegerSet parseIntegerSet(llvm::StringRef str, MLIRContext *context);
+
 } // end namespace mlir
 
 #endif // MLIR_PARSER_H

mlir/lib/Analysis/AffineStructures.cpp

 #include "mlir/Analysis/LinearTransform.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/IntegerSet.h"
+#include "mlir/Parser.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
LogicalResult mlir::getRelationFromMap(const AffineValueMap &map,
   for (unsigned i = 0, e = rel.getNumDomainDims(); i < e; ++i)
     rel.setValue(i, map.getOperand(i));
   for (unsigned i = rel.getNumDimIds(), e = rel.getNumDimAndSymbolIds(); i < e;
        ++i)
     rel.setValue(i, map.getOperand(i - rel.getNumRangeDims()));
 
   return success();
 }
+
+FailureOr<FlatAffineConstraints>
+mlir::parseIntegerSetToFAC(llvm::StringRef str, MLIRContext *context) {
+  IntegerSet set = parseIntegerSet(str, context);
+
+  if (!set)
+    return failure();
+
+  return FlatAffineConstraints(set);
+}

mlir/lib/Analysis/CMakeLists.txt

 set(LLVM_OPTIONAL_SOURCES
   AliasAnalysis.cpp
   AffineAnalysis.cpp
   AffineStructures.cpp
   BufferViewFlowAnalysis.cpp
   CallGraph.cpp
   DataFlowAnalysis.cpp
   DataLayoutAnalysis.cpp
   LinearTransform.cpp
   Liveness.cpp
   LoopAnalysis.cpp
   NestedMatcher.cpp
   NumberOfExecutions.cpp
add_mlir_library(MLIRLoopAnalysis
   mlir-headers
 
   LINK_LIBS PUBLIC
   MLIRAffine
   MLIRAnalysis
   MLIRCallInterfaces
   MLIRControlFlowInterfaces
   MLIRInferTypeOpInterface
+  MLIRParser
   MLIRPresburger
   MLIRSCF
   )
 
 add_subdirectory(Presburger)

mlir/lib/Parser/AffineParser.cpp

 //===- AffineParser.cpp - MLIR Affine Parser ------------------------------===//
 //
 // 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 implements a parser for Affine structures.
 //
 //===----------------------------------------------------------------------===//
 
 #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 {
 
 /// Lower precedence ops (all at the same precedence level). LNoOp is false in
 /// the boolean sense.
 enum AffineLowPrecOp {
   /// Null value.
   LNoOp,

 /// Parse an AffineExpr of SSA ids. The callback `parseElement` is used to parse
 /// SSA value uses encountered while parsing.
 ParseResult
 Parser::parseAffineExprOfSSAIds(AffineExpr &expr,
                                 function_ref<ParseResult(bool)> parseElement) {
   return AffineParser(state, /*allowParsingSSAIds=*/true, parseElement)
       .parseAffineExprOfSSAIds(expr);
 }
+
+IntegerSet mlir::parseIntegerSet(StringRef inputStr, MLIRContext *context) {
+  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);
+  SourceMgrDiagnosticHandler handler(sourceMgr, context);
+  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/AffineStructuresParserTest.cpp

+//===- 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 "mlir/Analysis/AffineStructures.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);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings with no constraint list";
+
+  fac = parseIntegerSetToFAC("(x)[] : ())", &context);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain remaining characters";
+
+  fac = parseIntegerSetToFAC("(x)[] : (x - >= 0)", &context);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain incomplete constraints";
+
+  fac = parseIntegerSetToFAC("(x)[] : (y == 0)", &context);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain unkown identifiers";
+
+  fac = parseIntegerSetToFAC("(x, x) : (2 * x >= 0)", &context);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain repeated identifier names";
+
+  fac = parseIntegerSetToFAC("(x)[x] : (2 * x >= 0)", &context);
+  EXPECT_TRUE(failed(fac))
+      << "should not accept strings that contain repeated identifier names";
+
+  fac = parseIntegerSetToFAC("(x) : (2 * x + 9223372036854775808 >= 0)",
+                             &context);
+  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

do {
     for (const auto &eq : eqs)
       permutedEqs.push_back(permute(eq));
 
     checkSample(hasSample,
                 makeFACFromConstraints(nDim, permutedIneqs, permutedEqs), fn);
   } 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}}, {}));
 
   // 1 <= 5x and 5x <= 9 (solution: x = 1).
   checkSample(true, makeFACFromConstraints(1, {{5, -1}, {-5, 9}}, {}));
 
   // Bounded sets with equalities.

mlir/unittests/Analysis/CMakeLists.txt

 add_mlir_unittest(MLIRAnalysisTests
+  AffineStructuresParserTest.cpp
   AffineStructuresTest.cpp
   LinearTransformTest.cpp
   PresburgerSetTest.cpp
 )
 
 target_link_libraries(MLIRAnalysisTests
   PRIVATE MLIRLoopAnalysis)
 
 add_subdirectory(Presburger)