mlir/Presburger: contribute a free-standing parser

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 `elems` to the specified row.
+  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,37 @@
 //
 //===----------------------------------------------------------------------===//
 
-#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;
+
+/// Parse an IntegerPolyhedron from a StringRef.
+///
+///  integer-set ::= dim-and-symbol-id-lists `:`
+///                  '(' affine-constraint-conjunction? ')'
+///  affine-constraint-conjunction ::= affine-constraint (`,`
+///                                       affine-constraint)*
+///
+IntegerPolyhedron parseIntegerPolyhedron(StringRef str);
+
+/// Parse a MultiAffineFunction from a StringRef.
+///
+///  affine-map ::= dim-and-symbol-id-lists `->` multi-dim-affine-expr
+///
+///  multi-dim-affine-expr ::= `(` `)`
+///  multi-dim-affine-expr ::= `(` affine-expr (`,` affine-expr)* `)`
+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 +52,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 +90,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/include/mlir/Analysis/Presburger/PresburgerSpace.h

@@ -173,6 +173,7 @@ class PresburgerSpace {
     return PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols,
                            numLocals);
   }
+  PresburgerSpace() = default;
 
   /// Get the domain/range space of this space. The returned space is a set
   /// space.
@@ -182,18 +183,43 @@ class PresburgerSpace {
   /// Get the space without local variables.
   PresburgerSpace getSpaceWithoutLocals() const;
 
-  unsigned getNumDomainVars() const { return numDomain; }
-  unsigned getNumRangeVars() const { return numRange; }
-  unsigned getNumSetDimVars() const { return numRange; }
-  unsigned getNumSymbolVars() const { return numSymbols; }
-  unsigned getNumLocalVars() const { return numLocals; }
+  constexpr unsigned getNumDomainVars() const { return numDomain; }
+  constexpr unsigned &numDomainVars() { return numDomain; }
+  constexpr unsigned getNumRangeVars() const { return numRange; }
+  constexpr unsigned &numRangeVars() { return numRange; }
+  constexpr unsigned getNumSetDimVars() const { return numRange; }
+  constexpr unsigned &numSetDimVars() { return numRange; }
+  constexpr unsigned getNumSymbolVars() const { return numSymbols; }
+  constexpr unsigned &numSymbolVars() { return numSymbols; }
+  constexpr unsigned getNumLocalVars() const { return numLocals; }
+  constexpr unsigned &numLocalVars() { return numLocals; }
+  constexpr unsigned getNumDimVars() const { return numDomain + numRange; }
+  constexpr unsigned getNumDimAndSymbolVars() const {
+    return getNumDimVars() + getNumSymbolVars();
+  }
+  constexpr unsigned getNumVars() const {
+    return getNumDimAndSymbolVars() + getNumLocalVars();
+  }
+  constexpr unsigned getNumCols() const { return getNumVars() + 1; }
+
+  constexpr unsigned getSetDimStartIdx() const { return 0; }
+  constexpr unsigned getSymbolStartIdx() const { return getNumDimVars(); }
+  constexpr unsigned getLocalVarStartIdx() const {
+    return getNumDimAndSymbolVars();
+  }
+  constexpr unsigned getConstantIdx() const { return getNumCols() - 1; }
 
-  unsigned getNumDimVars() const { return numDomain + numRange; }
-  unsigned getNumDimAndSymbolVars() const {
-    return numDomain + numRange + numSymbols;
+  constexpr bool isSetDimIdx(unsigned i) const {
+    return i < getSymbolStartIdx();
+  }
+  constexpr bool isSymbolIdx(unsigned i) const {
+    return i >= getSymbolStartIdx() && i < getLocalVarStartIdx();
+  }
+  constexpr bool isLocalVarIdx(unsigned i) const {
+    return i >= getLocalVarStartIdx() && i < getConstantIdx();
   }
-  unsigned getNumVars() const {
-    return numDomain + numRange + numSymbols + numLocals;
+  constexpr bool isConstantIdx(unsigned i) const {
+    return i == getConstantIdx();
   }
 
   /// Get the number of vars of the specified kind.
@@ -321,18 +347,18 @@ class PresburgerSpace {
 
 private:
   // Number of variables corresponding to domain variables.
-  unsigned numDomain;
+  unsigned numDomain = 0;
 
   // Number of variables corresponding to range variables.
-  unsigned numRange;
+  unsigned numRange = 0;
 
   /// Number of variables corresponding to symbols (unknown but constant for
   /// analysis).
-  unsigned numSymbols;
+  unsigned numSymbols = 0;
 
   /// Number of variables corresponding to locals (variables corresponding
   /// to existentially quantified variables).
-  unsigned numLocals;
+  unsigned numLocals = 0;
 
   /// Stores whether or not identifiers are being used in this space.
   bool usingIds = false;

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(
+      space.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(space.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(exprs.size(), space.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(space.getNumCols())));
+  }
+  void setRow(unsigned row, const CoefficientVector &l) {
+    flatMatrix.setRow(row, getDynamicAPIntVec(l.getPadded(space.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 space.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