[MLIR][Presburger] Simplex: refactor (symbolic)lex to support specifying multiple varKinds as symbols

[arjunp]

Summary

This is also required to support lexmin for relations.

[arjunp]

Fix some stuff

[Groverkss]

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

30–44

I do not like this. I think for PIP we can just use a bitvector. We can later add these flags if we need, but currently, I do not see any good uses.

mlir/include/mlir/Analysis/Presburger/Simplex.h

483–484

Why not assert here then?

[arjunp]

Address kunwar's comments

[arjunp]

mlir/include/mlir/Analysis/Presburger/Simplex.h

483–484

Clarified the doc

[Groverkss]

mlir/include/mlir/Analysis/Presburger/Simplex.h

225

nit: drop "llvm::" here

[arjunp]

mlir/include/mlir/Analysis/Presburger/Simplex.h

225

Not yet possible. Will need to add it to LLVM.h. Let's do it in another patch.

[Groverkss]

LGTM.

[arjunp]

• format

mlir/include/mlir/Analysis/Presburger/PresburgerSpace.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 #define MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H
 
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/Support/ErrorHandling.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace presburger {
 
 /// Kind of variable. Implementation wise SetDims are treated as Range
 /// vars, and spaces with no distinction between dimension vars are treated
 /// as relations with zero domain vars.
-enum class VarKind { Symbol, Local, Domain, Range, SetDim = Range };
+enum class VarKind : unsigned {
+  Symbol = 1 << 0,
+  Domain = 1 << 1,
+  Range = 1 << 2,
+  Local = 1 << 3,
+  SetDim = Range
+};
+/// Note: this reflects the number of distinct VarKinds. SetDim doesn't count.
+const unsigned numVarKinds = 4;
+
+/// These are useful when a mask of VarKinds needs to be passed.
+const unsigned varKindFlagSymbol = unsigned(VarKind::Symbol);
+const unsigned varKindFlagDomain = unsigned(VarKind::Domain);
+const unsigned varKindFlagRange = unsigned(VarKind::Range);
+const unsigned varKindFlagLocal = unsigned(VarKind::Local);
 
 /// PresburgerSpace is the space of all possible values of a tuple of integer
 /// valued variables/variables. Each variable has one of the three types:
 ///
 /// Dimension: Ordinary variables over which the space is represented.
 ///
 /// Symbol: Symbol variables correspond to fixed but unknown values.
 /// Mathematically, a space with symbolic variables is like a
public:
   bool isEqual(const PresburgerSpace &other) const;
 
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of dimensional variables immediately before
   /// the split become symbols, or some of the symbols immediately after the
   /// split become dimensions.
   void setVarSymbolSeperation(unsigned newSymbolCount);
 
+  /// Given a mask representing a subset of VarKinds, return a mask representing
+  /// the subset of this space's variables that are of one of the specified
+  /// VarKinds.
+  llvm::SmallBitVector getVarMaskFromVarKindMask(unsigned varKindMask) const;
+
   void print(llvm::raw_ostream &os) const;
   void dump() const;
 
   //===--------------------------------------------------------------------===//
   //     Attachment Interactions
   //===--------------------------------------------------------------------===//
 
   /// Set the attachment for `i^th` variable to `attachment`. `T` here should

mlir/include/mlir/Analysis/Presburger/Simplex.h

 #include "mlir/Analysis/Presburger/Fraction.h"
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/Analysis/Presburger/PWMAFunction.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
+#include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/StringSaver.h"
 #include "llvm/Support/raw_ostream.h"
 
 namespace mlir {
 namespace presburger {
 
 class GBRSimplex;
public:
   void intersectIntegerRelation(const IntegerRelation &rel);
 
   /// Print the tableau's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// Construct a SimplexBase with the specified number of variables and fixed
-  /// columns.
+  /// columns. The first overload should be used when there are nosymbols.
+  /// With the second overload, the specified range of vars will be marked
+  /// as symbols. With the third overload, `isSymbol` is a bitmask denoting
+  /// which vars are symbols. The size of `isSymbol` must be `nVar`.
   ///
   /// For example, Simplex uses two fixed columns: the denominator and the
   /// constant term, whereas LexSimplex has an extra fixed column for the
   /// so-called big M parameter. For more information see the documentation for
   /// LexSimplex.
-  SimplexBase(unsigned nVar, bool mustUseBigM, unsigned symbolOffset,
-              unsigned nSymbol);
+  SimplexBase(unsigned nVar, bool mustUseBigM);
+  SimplexBase(unsigned nVar, bool mustUseBigM,
+              const llvm::SmallBitVector &isSymbol);
 
   enum class Orientation { Row, Column };
 
   /// An Unknown is either a variable or a constraint. It is always associated
   /// with either a row or column. Whether it's a row or a column is specified
   /// by the orientation and pos identifies the specific row or column it is
   /// associated with. If the unknown is restricted, then it has a
   /// non-negativity constraint associated with it, i.e., its sample value must
public:
   /// This just adds the inequality to the tableau and does not try to create a
   /// consistent tableau configuration.
   void addInequality(ArrayRef<int64_t> coeffs) final;
 
   /// Get a snapshot of the current state. This is used for rolling back.
   unsigned getSnapshot() { return SimplexBase::getSnapshotBasis(); }
 
 protected:
-  LexSimplexBase(unsigned nVar, unsigned symbolOffset, unsigned nSymbol)
-      : SimplexBase(nVar, /*mustUseBigM=*/true, symbolOffset, nSymbol) {}
+  LexSimplexBase(unsigned nVar) : SimplexBase(nVar, /*mustUseBigM=*/true) {}
+  LexSimplexBase(unsigned nVar, const llvm::SmallBitVector &isSymbol)
+      : SimplexBase(nVar, /*mustUseBigM=*/true, isSymbol) {}
   explicit LexSimplexBase(const IntegerRelation &constraints)
-      : LexSimplexBase(constraints.getNumVars(),
-                       constraints.getVarKindOffset(VarKind::Symbol),
-                       constraints.getNumSymbolVars()) {
+      : LexSimplexBase(constraints.getNumVars()) {
+    intersectIntegerRelation(constraints);
+  }
+  explicit LexSimplexBase(const IntegerRelation &constraints,
+                          const llvm::SmallBitVector &isSymbol)
+      : LexSimplexBase(constraints.getNumVars(), isSymbol) {
     intersectIntegerRelation(constraints);
   }
 
   /// Add new symbolic variables to the end of the list of variables.
   void appendSymbol();
 
   /// Try to move the specified row to column orientation while preserving the
   /// lexicopositivity of the basis transform. The row must have a non-positive
protected:
   unsigned getLexMinPivotColumn(unsigned row, unsigned colA,
                                 unsigned colB) const;
 };
 
 /// A class for lexicographic optimization without any symbols. This also
 /// provides support for integer-exact redundancy and separateness checks.
 class LexSimplex : public LexSimplexBase {
 public:
-  explicit LexSimplex(unsigned nVar)
-      : LexSimplexBase(nVar, /*symbolOffset=*/0, /*nSymbol=*/0) {}
+  explicit LexSimplex(unsigned nVar) : LexSimplexBase(nVar) {}
+  // Note that LexSimplex does NOT support symbols.
+  // Symbols in `constraints` will be IGNORED.
   explicit LexSimplex(const IntegerRelation &constraints)
-      : LexSimplexBase(constraints) {
-    assert(constraints.getNumSymbolVars() == 0 &&
-           "LexSimplex does not support symbols!");
-  }
+      : LexSimplexBase(constraints) {}
 
   /// Return the lexicographically minimum rational solution to the constraints.
   MaybeOptimum<SmallVector<Fraction, 8>> findRationalLexMin();
 
   /// Return the lexicographically minimum integer solution to the constraints.
   ///
   /// Note: this should be used only when the lexmin is really needed. To obtain
   /// any integer sample, use Simplex::findIntegerSample as that is more robust.
private:
 
   /// Get a row corresponding to a var that has a non-integral sample value, if
   /// one exists. Otherwise, return an empty optional.
   Optional<unsigned> maybeGetNonIntegralVarRow() const;
 };
 
 /// Represents the result of a symbolic lexicographic minimization computation.
 struct SymbolicLexMin {
-  SymbolicLexMin(unsigned nSymbols, unsigned nNonSymbols)
-      : lexmin(PresburgerSpace::getSetSpace(nSymbols), nNonSymbols),
-        unboundedDomain(
-            PresburgerSet::getEmpty(PresburgerSpace::getSetSpace(nSymbols))) {}
+  SymbolicLexMin(const PresburgerSpace &domainSpace, unsigned numOutputs)
+      : lexmin(domainSpace, numOutputs),
+        unboundedDomain(PresburgerSet::getEmpty(domainSpace)) {}
 
   /// This maps assignments of symbols to the corresponding lexmin.
   /// Takes no value when no integer sample exists for the assignment or if the
   /// lexmin is unbounded.
   PWMAFunction lexmin;
   /// Contains all assignments to the symbols that made the lexmin unbounded.
   /// Note that the symbols of the input set to the symbolic lexmin are dims
   /// of this PrebsurgerSet.

 /// handle this by splitting the symbol domain. We run the algorithm for the
 /// case where the row isn't violated, and then come back and run the case
 /// where it is.
 class SymbolicLexSimplex : public LexSimplexBase {
 public:
   /// `constraints` is the set for which the symbolic lexmin will be computed.
   /// `symbolDomain` is the set of values of the symbols for which the lexmin
   /// will be computed. `symbolDomain` should have a dim var for every symbol in
-  /// `constraints`, and no other vars.
+  /// `constraints`, and no other vars. `isSymbol` specifies which vars of
+  /// `constraints` should be considered as symbols.
+  ///
+  /// The resulting SymbolicLexMin's space will be compatible with that of
+  /// symbolDomain.
   SymbolicLexSimplex(const IntegerRelation &constraints,
-                     const IntegerPolyhedron &symbolDomain)
-      : SymbolicLexSimplex(constraints,
-                           constraints.getVarKindOffset(VarKind::Symbol),
-                           symbolDomain) {
-    assert(constraints.getNumSymbolVars() == symbolDomain.getNumVars());
+                     const IntegerPolyhedron &symbolDomain,
+                     const llvm::SmallBitVector &isSymbol)
+      : LexSimplexBase(constraints, isSymbol), domainPoly(symbolDomain),
+        domainSimplex(symbolDomain) {
+    // TODO consider supporting this case. It amounts
+    // to just returning the input constraints.
+    assert(domainPoly.getNumVars() > 0 &&
+           "there must be some non-symbols to optimize!");
   }
 
+  /// Same as above, but varKindMask specifies which VarKinds in `constraints`
+  /// should be considered as symbols; by default the symbols are considered
+  /// as symbols.
+  SymbolicLexSimplex(const IntegerRelation &constraints,
+                     const IntegerPolyhedron &symbolDomain,
+                     unsigned isSymbolKindMask = varKindFlagSymbol)
+      : SymbolicLexSimplex(constraints, symbolDomain,
+                           constraints.getSpace().getVarMaskFromVarKindMask(
+                               isSymbolKindMask)) {}
+
   /// An overload to select some other subrange of ids as symbols for lexmin.
   /// The symbol ids are the range of ids with absolute index
   /// [symbolOffset, symbolOffset + symbolDomain.getNumVars())
   /// symbolDomain should only have dim ids.
   SymbolicLexSimplex(const IntegerRelation &constraints, unsigned symbolOffset,
                      const IntegerPolyhedron &symbolDomain)
-      : LexSimplexBase(/*nVar=*/constraints.getNumVars(), symbolOffset,
-                       symbolDomain.getNumVars()),
-        domainPoly(symbolDomain), domainSimplex(symbolDomain) {
-    // TODO consider supporting this case. It amounts
-    // to just returning the input constraints.
-    assert(domainPoly.getNumVars() > 0 &&
-           "there must be some non-symbols to optimize!");
-    assert(domainPoly.getNumVars() == domainPoly.getNumDimVars());
-    intersectIntegerRelation(constraints);
-  }
+      : SymbolicLexSimplex(constraints, symbolDomain,
+                           getSubrangeBitVector(constraints.getNumVars(),
+                                                symbolOffset,
+                                                symbolDomain.getNumVars())) {}
 
   /// The lexmin will be stored as a function `lexmin` from symbols to
   /// non-symbols in the result.
   ///
   /// For some values of the symbols, the lexmin may be unbounded.
   /// These parts of the symbol domain will be stored in `unboundedDomain`.
   ///
   /// The spaces of the sets in the result are compatible with the symbolDomain

 ///
 /// Finding an integer sample is done with the Generalized Basis Reduction
 /// algorithm. See the documentation for findIntegerSample and reduceBasis.
 class Simplex : public SimplexBase {
 public:
   enum class Direction { Up, Down };
 
   Simplex() = delete;
-  explicit Simplex(unsigned nVar)
-      : SimplexBase(nVar, /*mustUseBigM=*/false, /*symbolOffset=*/0,
-                    /*nSymbol=*/0) {}
+  explicit Simplex(unsigned nVar) : SimplexBase(nVar, /*mustUseBigM=*/false) {}
   explicit Simplex(const IntegerRelation &constraints)
       : Simplex(constraints.getNumVars()) {
     intersectIntegerRelation(constraints);
   }
   ~Simplex() override = default;
 
   /// Add an inequality to the tableau. If coeffs is c_0, c_1, ... c_n, where n
   /// is the current number of variables, then the corresponding inequality is

mlir/include/mlir/Analysis/Presburger/Utils.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_UTILS_H
 #define MLIR_ANALYSIS_PRESBURGER_UTILS_H
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SmallBitVector.h"
 
 namespace mlir {
 namespace presburger {
 
 class IntegerRelation;
 
 /// This class represents the result of operations optimizing something subject
 /// to some constraints. If the constraints were not satisfiable the, kind will

 ///
 /// The coefficient of `q` in `dividend` must be zero, as it is not allowed for
 /// local variable to be a floor division of an expression involving itself.
 SmallVector<int64_t, 8> getDivUpperBound(ArrayRef<int64_t> dividend,
                                          int64_t divisor, unsigned localVarIdx);
 SmallVector<int64_t, 8> getDivLowerBound(ArrayRef<int64_t> dividend,
                                          int64_t divisor, unsigned localVarIdx);
 
+llvm::SmallBitVector getSubrangeBitVector(unsigned len, unsigned setOffset,
+                                          unsigned numSet);
+
 /// Check if the pos^th variable can be expressed as a floordiv of an affine
 /// function of other variables (where the divisor is a positive constant).
 /// `foundRepr` contains a boolean for each variable indicating if the
 /// explicit representation for that variable has already been computed.
 /// Returns the `MaybeLocalRepr` struct which contains the indices of the
 /// constraints that can be expressed as a floordiv of an affine function. If
 /// the representation could be computed, `dividend` and `denominator` are set.
 /// If the representation could not be computed, the kind attribute in

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

if (kind == VarKind::Range)
     return getNumDomainVars();
   if (kind == VarKind::Symbol)
     return getNumDimVars();
   if (kind == VarKind::Local)
     return getNumDimAndSymbolVars();
   llvm_unreachable("VarKind does not exist!");
 }
 
+llvm::SmallBitVector
+PresburgerSpace::getVarMaskFromVarKindMask(unsigned varKindMask) const {
+  llvm::SmallBitVector result(getNumVars(), false);
+  for (unsigned varKind = 1; varKind < (1 << numVarKinds); varKind <<= 1)
+    if (varKindMask & varKind)
+      result.set(getVarKindOffset(VarKind(varKind)),
+                 getVarKindEnd(VarKind(varKind)));
+  return result;
+}
+
 unsigned PresburgerSpace::getVarKindEnd(VarKind kind) const {
   return getVarKindOffset(kind) + getNumVarKind(kind);
 }
 
 unsigned PresburgerSpace::getVarKindOverlap(VarKind kind, unsigned varStart,
                                             unsigned varLimit) const {
   unsigned varRangeStart = getVarKindOffset(kind);
   unsigned varRangeEnd = getVarKindEnd(kind);

mlir/lib/Analysis/Presburger/Simplex.cpp

scaleAndAddForAssert(ArrayRef<int64_t> a, int64_t scale, ArrayRef<int64_t> b) {
   assert(a.size() == b.size());
   SmallVector<int64_t, 8> res;
   res.reserve(a.size());
   for (unsigned i = 0, e = a.size(); i < e; ++i)
     res.push_back(a[i] + scale * b[i]);
   return res;
 }
 
-SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM, unsigned symbolOffset,
-                         unsigned nSymbol)
-    : usingBigM(mustUseBigM), nRedundant(0), nSymbol(nSymbol),
+SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM)
+    : usingBigM(mustUseBigM), nRedundant(0), nSymbol(0),
       tableau(0, getNumFixedCols() + nVar), empty(false) {
-  assert(symbolOffset + nSymbol <= nVar);
-
   colUnknown.insert(colUnknown.begin(), getNumFixedCols(), nullIndex);
   for (unsigned i = 0; i < nVar; ++i) {
     var.emplace_back(Orientation::Column, /*restricted=*/false,
                      /*pos=*/getNumFixedCols() + i);
     colUnknown.push_back(i);
   }
+}
 
-  // Move the symbols to be in columns [3, 3 + nSymbol).
-  for (unsigned i = 0; i < nSymbol; ++i) {
-    var[symbolOffset + i].isSymbol = true;
-    swapColumns(var[symbolOffset + i].pos, getNumFixedCols() + i);
+SimplexBase::SimplexBase(unsigned nVar, bool mustUseBigM,
+                         const llvm::SmallBitVector &isSymbol)
+    : SimplexBase(nVar, mustUseBigM) {
+  assert(isSymbol.size() == nVar && "invalid bitmask!");
+  // Invariant: nSymbol is the number of symbols that have been marked
+  // already and these occupy the columns
+  // [getNumFixedCols(), getNumFixedCols() + nSymbol).
+  for (unsigned symbolIdx : isSymbol.set_bits()) {
+    var[symbolIdx].isSymbol = true;
+    swapColumns(var[symbolIdx].pos, getNumFixedCols() + nSymbol);
+    ++nSymbol;
   }
 }
 
 const Simplex::Unknown &SimplexBase::unknownFromIndex(int index) const {
   assert(index != nullIndex && "nullIndex passed to unknownFromIndex");
   return index >= 0 ? var[index] : con[~index];
 }
 

 LogicalResult SymbolicLexSimplex::doNonBranchingPivots() {
   while (Optional<unsigned> row = maybeGetAlwaysViolatedRow())
     if (moveRowUnknownToColumn(*row).failed())
       return failure();
   return success();
 }
 
 SymbolicLexMin SymbolicLexSimplex::computeSymbolicIntegerLexMin() {
-  SymbolicLexMin result(nSymbol, var.size() - nSymbol);
+  SymbolicLexMin result(domainPoly.getSpace(), var.size() - nSymbol);
 
   /// The algorithm is more naturally expressed recursively, but we implement
   /// it iteratively here to avoid potential issues with stack overflows in the
   /// compiler. We explicitly maintain the stack frames in a vector.
   ///
   /// To "recurse", we store the current "stack frame", i.e., state variables
   /// that we will need when we "return", into `stack`, increment `level`, and
   /// `continue`. To "tail recurse", we just `continue`.

mlir/lib/Analysis/Presburger/Utils.cpp

for (unsigned eqPos : eqIndices) {
 
     repr.kind = ReprKind::Equality;
     repr.repr.equalityIdx = eqPos;
     return repr;
   }
   return repr;
 }
 
+llvm::SmallBitVector presburger::getSubrangeBitVector(unsigned len,
+                                                      unsigned setOffset,
+                                                      unsigned numSet) {
+  llvm::SmallBitVector vec(len, false);
+  vec.set(setOffset, setOffset + numSet);
+  return vec;
+}
+
 void presburger::removeDuplicateDivs(
     std::vector<SmallVector<int64_t, 8>> &divs,
     SmallVectorImpl<unsigned> &denoms, unsigned localOffset,
     llvm::function_ref<bool(unsigned i, unsigned j)> merge) {
 
   // Find and merge duplicate divisions.
   // TODO: Add division normalization to support divisions that differ by
   // a constant.