mlir/Presburger: contribute a free-standing parser

The Presburger library is already quite independent of MLIR, with the
exception of the MLIR Support library. There is, however, one major
exception: the test suite for the library depends on the core MLIR
parser. To free it of this dependency, extract the parts of the core
MLIR parser that are applicable to the Presburger test suite, author
custom parsing data structures, and adapt the new parser to parse into
these structures.

This patch is part of a project to move the Presburger library into
LLVM.

mlir/include/mlir/Analysis/Presburger/Matrix.h

@@ -103,6 +103,9 @@ class Matrix {
   /// Set the specified row to `elems`.
   void setRow(unsigned row, ArrayRef<T> elems);
 
+  /// Add the specified row to `elems`.
+  void addToRow(unsigned row, ArrayRef<T> elems);
+
   /// Insert columns having positions pos, pos + 1, ... pos + count - 1.
   /// Columns that were at positions 0 to pos - 1 will stay where they are;
   /// columns that were at positions pos to nColumns - 1 will be pushed to the

mlir/unittests/Analysis/Presburger/Parser.h → mlir/include/mlir/Analysis/Presburger/Parser.h

@@ -11,32 +11,26 @@
 //
 //===----------------------------------------------------------------------===//
 
-#ifndef MLIR_UNITTESTS_ANALYSIS_PRESBURGER_PARSER_H
-#define MLIR_UNITTESTS_ANALYSIS_PRESBURGER_PARSER_H
+#ifndef MLIR_ANALYSIS_PRESBURGER_PARSER_H
+#define MLIR_ANALYSIS_PRESBURGER_PARSER_H
 
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/PresburgerRelation.h"
-#include "mlir/AsmParser/AsmParser.h"
-#include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
-#include "mlir/IR/AffineExpr.h"
-#include "mlir/IR/AffineMap.h"
-#include "mlir/IR/IntegerSet.h"
-
-namespace mlir {
-namespace presburger {
-
-/// Parses an IntegerPolyhedron from a StringRef. It is expected that the string
-/// represents a valid IntegerSet.
-inline IntegerPolyhedron parseIntegerPolyhedron(StringRef str) {
-  MLIRContext context(MLIRContext::Threading::DISABLED);
-  return affine::FlatAffineValueConstraints(parseIntegerSet(str, &context));
-}
+
+namespace mlir::presburger {
+using llvm::StringRef;
+
+/// Parses an IntegerPolyhedron from a StringRef.
+IntegerPolyhedron parseIntegerPolyhedron(StringRef str);
+
+/// Parses a MultiAffineFunction from a StringRef.
+MultiAffineFunction parseMultiAffineFunction(StringRef str);
 
 /// Parse a list of StringRefs to IntegerRelation and combine them into a
-/// PresburgerSet by using the union operation. It is expected that the strings
-/// are all valid IntegerSet representation and that all of them have compatible
-/// spaces.
+/// PresburgerSet by using the union operation. It is expected that the
+/// strings are all valid IntegerSet representation and that all of them have
+/// compatible spaces.
 inline PresburgerSet parsePresburgerSet(ArrayRef<StringRef> strs) {
   assert(!strs.empty() && "strs should not be empty");
 
@@ -47,25 +41,10 @@ inline PresburgerSet parsePresburgerSet(ArrayRef<StringRef> strs) {
   return result;
 }
 
-inline MultiAffineFunction parseMultiAffineFunction(StringRef str) {
-  MLIRContext context(MLIRContext::Threading::DISABLED);
-
-  // TODO: Add default constructor for MultiAffineFunction.
-  MultiAffineFunction multiAff(PresburgerSpace::getRelationSpace(),
-                               IntMatrix(0, 1));
-  if (getMultiAffineFunctionFromMap(parseAffineMap(str, &context), multiAff)
-          .failed())
-    llvm_unreachable(
-        "Failed to parse MultiAffineFunction because of semi-affinity");
-  return multiAff;
-}
-
 inline PWMAFunction
 parsePWMAF(ArrayRef<std::pair<StringRef, StringRef>> pieces) {
   assert(!pieces.empty() && "At least one piece should be present.");
 
-  MLIRContext context(MLIRContext::Threading::DISABLED);
-
   IntegerPolyhedron initDomain = parseIntegerPolyhedron(pieces[0].first);
   MultiAffineFunction initMultiAff = parseMultiAffineFunction(pieces[0].second);
 
@@ -100,8 +79,6 @@ parsePresburgerRelationFromPresburgerSet(ArrayRef<StringRef> strs,
   result.convertVarKind(VarKind::SetDim, 0, numDomain, VarKind::Domain, 0);
   return result;
 }
+} // namespace mlir::presburger
 
-} // namespace presburger
-} // namespace mlir
-
-#endif // MLIR_UNITTESTS_ANALYSIS_PRESBURGER_PARSER_H
+#endif // MLIR_ANALYSIS_PRESBURGER_PARSER_H

mlir/lib/Analysis/Presburger/CMakeLists.txt

@@ -1,3 +1,5 @@
+add_subdirectory(Parser)
+
 add_mlir_library(MLIRPresburger
   Barvinok.cpp
   IntegerRelation.cpp

mlir/lib/Analysis/Presburger/Matrix.cpp

@@ -144,6 +144,14 @@ void Matrix<T>::setRow(unsigned row, ArrayRef<T> elems) {
     at(row, i) = elems[i];
 }
 
+template <typename T>
+void Matrix<T>::addToRow(unsigned row, ArrayRef<T> elems) {
+  assert(elems.size() == getNumColumns() &&
+         "elems size must match row length!");
+  for (unsigned i = 0, e = getNumColumns(); i < e; ++i)
+    at(row, i) += elems[i];
+}
+
 template <typename T>
 void Matrix<T>::insertColumn(unsigned pos) {
   insertColumns(pos, 1);

mlir/lib/Analysis/Presburger/Parser/CMakeLists.txt

@@ -0,0 +1,6 @@
+add_mlir_library(MLIRPresburgerParser
+  Flattener.cpp
+  Lexer.cpp
+  ParserImpl.cpp
+  ParseStructs.cpp
+  Token.cpp)

mlir/lib/Analysis/Presburger/Parser/Flattener.cpp

@@ -0,0 +1,88 @@
+//===- Flattener.cpp - Presburger ParseStruct Flattener ---------*- 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 implements the Flattener class for flattening the parse tree
+// produced by the parser for the Presburger library.
+//
+//===----------------------------------------------------------------------===//
+
+#include "Flattener.h"
+#include "ParseStructs.h"
+
+using namespace mlir;
+using namespace presburger;
+
+void Flattener::visitDiv(const PureAffineExprImpl &div) {
+  int64_t divisor = div.getDivisor();
+
+  // First construct the linear part of the divisor.
+  auto dividend = div.collectLinearTerms().getPadded(
+      info.getLocalVarStartIdx() + localExprs.size() + 1);
+
+  // Next, insert the non-linear coefficients.
+  for (const auto &[hash, adjustedMulFactor, adjustedLinearTerm] :
+       div.getNonLinearCoeffs()) {
+
+    // adjustedMulFactor will be mulFactor * -divisor in case of mod, and
+    // mulFactor in case of floordiv.
+    dividend[lookupLocal(hash)] = adjustedMulFactor;
+
+    // The expansion is either a modExpansion which we previously stored, or the
+    // adjustedLinearTerm, which is correct in the case when we're encountering
+    // the innermost mod for the first time.
+    CoefficientVector expansion =
+        lookupModExpansion(hash)
+            .value_or(adjustedLinearTerm)
+            .getPadded(info.getLocalVarStartIdx() + localExprs.size() + 1);
+    dividend += expansion;
+
+    // If this is a mod, insert the new computed expansion, which is the
+    // dividend * mulFactor.
+    if (div.isMod())
+      localModExpansion.insert({div.hash(), dividend * div.getMulFactor()});
+  }
+
+  cst.addLocalFloorDiv(dividend, divisor);
+  localExprs.insert(div.hash());
+}
+
+void Flattener::flatten(unsigned row, PureAffineExprImpl &div) {
+  // Visit divs inner to outer.
+  for (auto &nestedDiv : div.getNestedDivTerms())
+    flatten(row, *nestedDiv);
+
+  if (div.hasDivisor()) {
+    visitDiv(div);
+    return;
+  }
+
+  // Hit multiple times every time we have a linear sub-expression, but the
+  // row is overwritten to consider only the outermost div, which is hit
+  // last.
+
+  // Set the linear part of the row.
+  setRow(row, div.getLinearDividend());
+
+  // Set the non-linear coefficients.
+  for (const auto &[hash, adjustedMulFactor, adjustedLinearTerm] :
+       div.getNonLinearCoeffs()) {
+    flatMatrix(row, lookupLocal(hash)) = adjustedMulFactor;
+    CoefficientVector expansion =
+        lookupModExpansion(hash).value_or(adjustedLinearTerm);
+    addToRow(row, expansion);
+  }
+}
+
+std::pair<IntMatrix, IntegerPolyhedron> Flattener::flatten() {
+  // Use the same flattener to simplify each expression successively. This way
+  // local variables / expressions are shared.
+  for (const auto &[row, expr] : enumerate(exprs))
+    flatten(row, *expr);
+
+  return {flatMatrix, cst};
+}

mlir/lib/Analysis/Presburger/Parser/Flattener.h

@@ -0,0 +1,67 @@
+//===- Flattener.h - Presburger ParseStruct Flattener -----------*- 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
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_ANALYSIS_PRESBURGER_PARSER_FLATTENER_H
+#define MLIR_ANALYSIS_PRESBURGER_PARSER_FLATTENER_H
+
+#include "ParseStructs.h"
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
+#include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/SetVector.h"
+
+namespace mlir::presburger {
+using llvm::SmallDenseMap;
+using llvm::SmallSetVector;
+
+class Flattener : public FinalParseResult {
+public:
+  // The final flattened result is stored here.
+  IntMatrix flatMatrix;
+
+  // We maintain a set of divs that we have seen while flattening. The size of
+  // this set is at most info.numDivs, hitting info.numDivs at the end of the
+  // flattening, if that expression contains all the possible divs.
+  SmallSetVector<size_t, 4> localExprs;
+
+  // We maintain a mapping between local mods and their expansions. The vector
+  // is the dividend.
+  SmallDenseMap<size_t, CoefficientVector, 2> localModExpansion;
+
+  Flattener(FinalParseResult &&parseResult)
+      : FinalParseResult(std::move(parseResult)),
+        flatMatrix(info.numExprs, info.getNumCols()) {}
+
+  std::pair<IntMatrix, IntegerPolyhedron> flatten();
+
+private:
+  void flatten(unsigned row, PureAffineExprImpl &div);
+  void visitDiv(const PureAffineExprImpl &div);
+
+  void addToRow(unsigned row, const CoefficientVector &l) {
+    flatMatrix.addToRow(row,
+                        getDynamicAPIntVec(l.getPadded(info.getNumCols())));
+  }
+  void setRow(unsigned row, const CoefficientVector &l) {
+    flatMatrix.setRow(row, getDynamicAPIntVec(l.getPadded(info.getNumCols())));
+  }
+
+  unsigned lookupLocal(size_t hash) {
+    const auto *it = find(localExprs, hash);
+    assert(it != localExprs.end() &&
+           "Local expression not found; walking from inner to outer?");
+    return info.getLocalVarStartIdx() + it - localExprs.begin();
+  }
+  std::optional<CoefficientVector> lookupModExpansion(size_t hash) {
+    return localModExpansion.contains(hash)
+               ? std::make_optional(localModExpansion.at(hash))
+               : std::nullopt;
+  }
+};
+} // namespace mlir::presburger
+
+#endif // MLIR_ANALYSIS_PRESBURGER_PARSER_FLATTENER_H