[MLIR][Presburger] Use Identifiers outside Presburger library

Author: iambrj

mlir/include/mlir/Analysis/FlatLinearValueConstraints.h

@@ -205,6 +205,8 @@ class FlatLinearConstraints : public presburger::IntegerPolyhedron {
 /// where each non-local variable can have an SSA Value attached to it.
 class FlatLinearValueConstraints : public FlatLinearConstraints {
 public:
+  using Identifier = presburger::Identifier;
+
   /// Constructs a constraint system reserving memory for the specified number
   /// of constraints and variables. `valArgs` are the optional SSA values
   /// associated with each dimension/symbol. These must either be empty or match
@@ -217,11 +219,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       : FlatLinearConstraints(numReservedInequalities, numReservedEqualities,
                               numReservedCols, numDims, numSymbols, numLocals) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    values.reserve(numReservedCols);
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    // Use values in space for FlatLinearValueConstraints.
+    space.resetIds();
+    // Set the values for the non-local variables.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, *valArgs[i]);
   }
 
   /// Constructs a constraint system reserving memory for the specified number
@@ -236,11 +239,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       : FlatLinearConstraints(numReservedInequalities, numReservedEqualities,
                               numReservedCols, numDims, numSymbols, numLocals) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    values.reserve(numReservedCols);
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    // Use values in space for FlatLinearValueConstraints.
+    space.resetIds();
+    // Set the values for the non-local variables.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, valArgs[i]);
   }
 
   /// Constructs a constraint system with the specified number of dimensions
@@ -273,10 +277,12 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
                              ArrayRef<std::optional<Value>> valArgs = {})
       : FlatLinearConstraints(fac) {
     assert(valArgs.empty() || valArgs.size() == getNumDimAndSymbolVars());
-    if (valArgs.empty())
-      values.resize(getNumDimAndSymbolVars(), std::nullopt);
-    else
-      values.append(valArgs.begin(), valArgs.end());
+    // Use values in space for FlatLinearValueConstraints.
+    space.resetIds();
+    // Set the values for the non-local variables.
+    for (unsigned i = 0, e = valArgs.size(); i < e; ++i)
+      if (valArgs[i])
+        setValue(i, *valArgs[i]);
   }
 
   /// Creates an affine constraint system from an IntegerSet.
@@ -302,7 +308,9 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   inline Value getValue(unsigned pos) const {
     assert(pos < getNumDimAndSymbolVars() && "Invalid position");
     assert(hasValue(pos) && "variable's Value not set");
-    return *values[pos];
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    return space.getId(kind, relativePos).getValue<Value>();
   }
 
   /// Returns the Values associated with variables in range [start, end).
@@ -317,21 +325,44 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
       values->push_back(getValue(i));
   }
 
-  inline ArrayRef<std::optional<Value>> getMaybeValues() const {
-    return {values.data(), values.size()};
+  inline SmallVector<std::optional<Value>> getMaybeValues() const {
+    SmallVector<std::optional<Value>> maybeValues;
+    maybeValues.reserve(getNumDimAndSymbolVars());
+    for (unsigned i = 0, e = getNumDimAndSymbolVars(); i < e; i++) {
+      if (hasValue(i))
+        maybeValues.push_back(getValue(i));
+      else
+        maybeValues.push_back(std::nullopt);
+    }
+    return maybeValues;
   }
 
-  inline ArrayRef<std::optional<Value>>
-  getMaybeValues(presburger::VarKind kind) const {
-    assert(kind != VarKind::Local &&
-           "Local variables do not have any value attached to them.");
-    return {values.data() + getVarKindOffset(kind), getNumVarKind(kind)};
-  }
+  inline SmallVector<std::optional<Value>>
+    getMaybeValues(presburger::VarKind kind) const {
+      assert(kind != VarKind::Local &&
+          "Local variables do not have any value attached to them.");
+      SmallVector<std::optional<Value>> maybeValues;
+      maybeValues.reserve(getNumVarKind(kind));
+      for (unsigned i = 0, e = getNumVarKind(kind); i < e; i++) {
+        Identifier id = space.getId(kind, i);
+        if (id.hasValue())
+          maybeValues.push_back(space.getId(kind, i).getValue<Value>());
+        else
+          maybeValues.push_back(std::nullopt);
+      }
+      return maybeValues;
+    }
 
   /// Returns true if the pos^th variable has an associated Value.
   inline bool hasValue(unsigned pos) const {
     assert(pos < getNumDimAndSymbolVars() && "Invalid position");
-    return values[pos].has_value();
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    return space.getId(kind, relativePos).hasValue();
+  }
+
+  void resetValues() {
+    space.resetIds();
   }
 
   unsigned appendDimVar(ValueRange vals);
@@ -360,7 +391,9 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   /// Sets the Value associated with the pos^th variable.
   inline void setValue(unsigned pos, Value val) {
     assert(pos < getNumDimAndSymbolVars() && "invalid var position");
-    values[pos] = val;
+    VarKind kind = getVarKindAt(pos);
+    unsigned relativePos = pos - getVarKindOffset(kind);
+    space.getId(kind, relativePos) = presburger::Identifier(val);
   }
 
   /// Sets the Values associated with the variables in the range [start, end).
@@ -455,17 +488,6 @@ class FlatLinearValueConstraints : public FlatLinearConstraints {
   // See implementation comments for more details.
   void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
                                bool *isResultIntegerExact = nullptr) override;
-
-  /// Returns false if the fields corresponding to various variable counts, or
-  /// equality/inequality buffer sizes aren't consistent; true otherwise. This
-  /// is meant to be used within an assert internally.
-  bool hasConsistentState() const override;
-
-  /// Values corresponding to the (column) non-local variables of this
-  /// constraint system appearing in the order the variables correspond to
-  /// columns. Variables that aren't associated with any Value are set to
-  /// std::nullopt.
-  SmallVector<std::optional<Value>, 8> values;
 };
 
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the

mlir/include/mlir/Dialect/Affine/Analysis/AffineAnalysis.h

@@ -15,6 +15,7 @@
 #ifndef MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H
 #define MLIR_DIALECT_AFFINE_ANALYSIS_AFFINEANALYSIS_H
 
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/IR/Value.h"
 #include "llvm/ADT/SmallVector.h"
@@ -115,7 +116,7 @@ struct MemRefAccess {
   ///
   /// Returns failure for yet unimplemented/unsupported cases (see docs of
   /// mlir::getIndexSet and mlir::getRelationFromMap for these cases).
-  LogicalResult getAccessRelation(FlatAffineRelation &accessRel) const;
+  LogicalResult getAccessRelation(presburger::IntegerRelation &accessRel) const;
 
   /// Populates 'accessMap' with composition of AffineApplyOps reachable from
   /// 'indices'.

mlir/include/mlir/Dialect/Affine/Analysis/AffineStructures.h

@@ -251,9 +251,9 @@ class FlatAffineRelation : public FlatAffineValueConstraints {
 /// 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(AffineMap &map, presburger::IntegerRelation &rel);
 LogicalResult getRelationFromMap(const AffineValueMap &map,
-                                 FlatAffineRelation &rel);
+                                 presburger::IntegerRelation &rel);
 
 } // namespace affine
 } // namespace mlir

mlir/lib/Analysis/FlatLinearValueConstraints.cpp

@@ -817,13 +817,12 @@ FlatLinearValueConstraints::FlatLinearValueConstraints(IntegerSet set,
                             set.getNumDims() + set.getNumSymbols() + 1,
                             set.getNumDims(), set.getNumSymbols(),
                             /*numLocals=*/0) {
-  // Populate values.
-  if (operands.empty()) {
-    values.resize(getNumDimAndSymbolVars(), std::nullopt);
-  } else {
-    assert(set.getNumInputs() == operands.size() && "operand count mismatch");
-    values.assign(operands.begin(), operands.end());
-  }
+  // Use values in space for FlatLinearValueConstraints.
+  space.resetIds();
+  // Set the values for the non-local variables.
+  for (unsigned i = 0, e = operands.size(); i < e; ++i)
+    setValue(i, operands[i]);
+
 
   // Flatten expressions and add them to the constraint system.
   std::vector<SmallVector<int64_t, 8>> flatExprs;
@@ -873,11 +872,6 @@ unsigned FlatLinearValueConstraints::insertVar(VarKind kind, unsigned pos,
                                                unsigned num) {
   unsigned absolutePos = IntegerPolyhedron::insertVar(kind, pos, num);
 
-  if (kind != VarKind::Local) {
-    values.insert(values.begin() + absolutePos, num, std::nullopt);
-    assert(values.size() == getNumDimAndSymbolVars());
-  }
-
   return absolutePos;
 }
 
@@ -890,22 +884,17 @@ unsigned FlatLinearValueConstraints::insertVar(VarKind kind, unsigned pos,
   unsigned absolutePos = IntegerPolyhedron::insertVar(kind, pos, num);
 
   // If a Value is provided, insert it; otherwise use std::nullopt.
-  for (unsigned i = 0; i < num; ++i)
-    values.insert(values.begin() + absolutePos + i,
-                  vals[i] ? std::optional<Value>(vals[i]) : std::nullopt);
+  for (unsigned i = 0, e = vals.size(); i < e; ++i)
+    setValue(absolutePos + i, vals[i]);
 
-  assert(values.size() == getNumDimAndSymbolVars());
   return absolutePos;
 }
 
 /// Checks if two constraint systems are in the same space, i.e., if they are
 /// associated with the same set of variables, appearing in the same order.
 static bool areVarsAligned(const FlatLinearValueConstraints &a,
                            const FlatLinearValueConstraints &b) {
-  return a.getNumDimVars() == b.getNumDimVars() &&
-         a.getNumSymbolVars() == b.getNumSymbolVars() &&
-         a.getNumVars() == b.getNumVars() &&
-         a.getMaybeValues().equals(b.getMaybeValues());
+  return a.getSpace().isAligned(b.getSpace());
 }
 
 /// Calls areVarsAligned to check if two constraint systems have the same set
@@ -928,12 +917,14 @@ static bool LLVM_ATTRIBUTE_UNUSED areVarsUnique(
     return true;
 
   SmallPtrSet<Value, 8> uniqueVars;
-  ArrayRef<std::optional<Value>> maybeValues =
-      cst.getMaybeValues().slice(start, end - start);
-  for (std::optional<Value> val : maybeValues) {
+  SmallVector<std::optional<Value>, 8> maybeValuesAll = cst.getMaybeValues();
+  ArrayRef<std::optional<Value>> maybeValues = {maybeValuesAll.data() + start,
+                                                maybeValuesAll.data() + end};
+
+  for (std::optional<Value> val : maybeValues)
     if (val && !uniqueVars.insert(*val).second)
       return false;
-  }
+
   return true;
 }
 
@@ -1058,20 +1049,9 @@ void FlatLinearValueConstraints::mergeSymbolVars(
          "expected same number of symbols");
 }
 
-bool FlatLinearValueConstraints::hasConsistentState() const {
-  return IntegerPolyhedron::hasConsistentState() &&
-         values.size() == getNumDimAndSymbolVars();
-}
-
 void FlatLinearValueConstraints::removeVarRange(VarKind kind, unsigned varStart,
                                                 unsigned varLimit) {
   IntegerPolyhedron::removeVarRange(kind, varStart, varLimit);
-  unsigned offset = getVarKindOffset(kind);
-
-  if (kind != VarKind::Local) {
-    values.erase(values.begin() + varStart + offset,
-                 values.begin() + varLimit + offset);
-  }
 }
 
 AffineMap
@@ -1089,14 +1069,15 @@ FlatLinearValueConstraints::computeAlignedMap(AffineMap map,
 
   dims.reserve(getNumDimVars());
   syms.reserve(getNumSymbolVars());
-  for (unsigned i = getVarKindOffset(VarKind::SetDim),
-                e = getVarKindEnd(VarKind::SetDim);
-       i < e; ++i)
-    dims.push_back(values[i] ? *values[i] : Value());
-  for (unsigned i = getVarKindOffset(VarKind::Symbol),
-                e = getVarKindEnd(VarKind::Symbol);
-       i < e; ++i)
-    syms.push_back(values[i] ? *values[i] : Value());
+  for (unsigned i = 0, e = getNumVarKind(VarKind::SetDim); i < e; ++i) {
+    Identifier id = space.getId(VarKind::SetDim, i);
+    dims.push_back(id.hasValue() ? Value(id.getValue<Value>()) : Value());
+  }
+  for (unsigned i = 0, e = getNumVarKind(VarKind::Symbol); i < e; ++i) {
+    Identifier id = space.getId(VarKind::Symbol, i);
+    syms.push_back(id.hasValue() ? Value(id.getValue<Value>()) : Value());
+  }
+
 
   AffineMap alignedMap =
       alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr);
@@ -1110,8 +1091,7 @@ FlatLinearValueConstraints::computeAlignedMap(AffineMap map,
 bool FlatLinearValueConstraints::findVar(Value val, unsigned *pos,
                                          unsigned offset) const {
   unsigned i = offset;
-  for (const auto &mayBeVar :
-       ArrayRef<std::optional<Value>>(values).drop_front(offset)) {
+  for (const auto &mayBeVar : getMaybeValues()) {
     if (mayBeVar && *mayBeVar == val) {
       *pos = i;
       return true;
@@ -1122,25 +1102,12 @@ bool FlatLinearValueConstraints::findVar(Value val, unsigned *pos,
 }
 
 bool FlatLinearValueConstraints::containsVar(Value val) const {
-  return llvm::any_of(values, [&](const std::optional<Value> &mayBeVar) {
-    return mayBeVar && *mayBeVar == val;
-  });
+  unsigned pos;
+  return findVar(val, &pos, 0);
 }
 
 void FlatLinearValueConstraints::swapVar(unsigned posA, unsigned posB) {
   IntegerPolyhedron::swapVar(posA, posB);
-
-  if (getVarKindAt(posA) == VarKind::Local &&
-      getVarKindAt(posB) == VarKind::Local)
-    return;
-
-  // Treat value of a local variable as std::nullopt.
-  if (getVarKindAt(posA) == VarKind::Local)
-    values[posB] = std::nullopt;
-  else if (getVarKindAt(posB) == VarKind::Local)
-    values[posA] = std::nullopt;
-  else
-    std::swap(values[posA], values[posB]);
 }
 
 void FlatLinearValueConstraints::addBound(BoundType type, Value val,
@@ -1182,27 +1149,13 @@ void FlatLinearValueConstraints::printSpace(raw_ostream &os) const {
 
 void FlatLinearValueConstraints::clearAndCopyFrom(
     const IntegerRelation &other) {
-
-  if (auto *otherValueSet =
-          dyn_cast<const FlatLinearValueConstraints>(&other)) {
-    *this = *otherValueSet;
-  } else {
-    *static_cast<IntegerRelation *>(this) = other;
-    values.clear();
-    values.resize(getNumDimAndSymbolVars(), std::nullopt);
-  }
+  IntegerPolyhedron::clearAndCopyFrom(other);
 }
 
 void FlatLinearValueConstraints::fourierMotzkinEliminate(
     unsigned pos, bool darkShadow, bool *isResultIntegerExact) {
-  SmallVector<std::optional<Value>, 8> newVals = values;
-  if (getVarKindAt(pos) != VarKind::Local)
-    newVals.erase(newVals.begin() + pos);
-  // Note: Base implementation discards all associated Values.
   IntegerPolyhedron::fourierMotzkinEliminate(pos, darkShadow,
                                              isResultIntegerExact);
-  values = newVals;
-  assert(values.size() == getNumDimAndSymbolVars());
 }
 
 void FlatLinearValueConstraints::projectOut(Value val) {
@@ -1215,11 +1168,7 @@ void FlatLinearValueConstraints::projectOut(Value val) {
 
 LogicalResult FlatLinearValueConstraints::unionBoundingBox(
     const FlatLinearValueConstraints &otherCst) {
-  assert(otherCst.getNumDimVars() == getNumDimVars() && "dims mismatch");
-  assert(otherCst.getMaybeValues()
-             .slice(0, getNumDimVars())
-             .equals(getMaybeValues().slice(0, getNumDimVars())) &&
-         "dim values mismatch");
+  assert(otherCst.getSpace().isAligned(getSpace(), VarKind::SetDim) && "dims mismatch");
   assert(otherCst.getNumLocalVars() == 0 && "local vars not supported here");
   assert(getNumLocalVars() == 0 && "local vars not supported yet here");