[MLIR][Presburger] Use IdKind for removeIdRange in PresburgerSpace

[Groverkss]

Summary

This patch moves PresburgerSpace::removeIdRange(idStart, idLimit) to
PresburgerSpace::removeIdRange(kind, idStart, idLimit), i.e. identifiers
can only be removed at once for a single kind.

This makes users of PresburgerSpace to not assume any inside ordering of
identifier kinds.

[Groverkss]

• Fix use of removeIdRange in convertDimToLocal

[arjunp]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

178

Please write the capture explicitly to indicate that idstart and idlimit are not being captured here

187

Drop parentheses for the condition

195

Write a short comment saying why idStart will never need to be updated

1152

Why is this SetDim and not Range? We are in Relation right?

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

139

Drop these parentheses?

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

1737

I don't understand why we would ever need to update other id kinds?

[Groverkss]

• Address Arjun's comments

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1152

This should ideally be shifted to IntegerPolyhedron or be given an IdKind to convert to Local. I can send this as a separate patch.

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

1737

We still need to keep it right? This is just a virtual function that just updates numDomain/numRange incase SetDim is updated.

[arjunp]

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

175

178

179–180

This is supposed to be start and limit I think.

1152

Okay.

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

1737

Ah, because you are considering SetDim to include both domain and range. This seems a little confusing, but I guess that's not for this patch.

[Groverkss]

• Adress Arjun's comments

[arjunp]

LGTM, please resolve the remaining comment below.

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

178

Please rename removeRange -> removeIdKindInRange

[Groverkss]

• rename removeRange -> removeIdKindInRange

mlir/include/mlir/Analysis/Presburger/IntegerRelation.h

public:
   /// with the `posA^th` local identifier. This should be used when the two
   /// local variables are known to always take the same values.
   virtual void eliminateRedundantLocalId(unsigned posA, unsigned posB);
 
   /// Removes identifiers of the specified kind with the specified pos (or
   /// within the specified range) from the system. The specified location is
   /// relative to the first identifier of the specified kind.
   void removeId(IdKind kind, unsigned pos);
-  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) override;
 
   /// Removes the specified identifier from the system.
   void removeId(unsigned pos);
 
   void removeEquality(unsigned pos);
   void removeInequality(unsigned pos);
 
   /// Remove the (in)equalities at positions [start, end).
protected:
 
   /// Prints the number of constraints, dimensions, symbols and locals in the
   /// IntegerRelation.
   virtual void printSpace(raw_ostream &os) const;
 
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes
   /// arrays as needed.
-  void removeIdRange(unsigned idStart, unsigned idLimit) override;
+  void removeIdRange(unsigned idStart, unsigned idLimit);
 
   /// A parameter that controls detection of an unrealistic number of
   /// constraints. If the number of constraints is this many times the number of
   /// variables, we consider such a system out of line with the intended use
   /// case of IntegerRelation.
   // The rationale for 32 is that in the typical simplest of cases, an
   // identifier is expected to have one lower bound and one upper bound
   // constraint. With a level of tiling or a connection to another identifier

mlir/include/mlir/Analysis/Presburger/PWMAFunction.h

public:
   /// absolute column position (i.e., not relative to the kind of identifier)
   /// of the first added identifier.
   unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
 
   /// Swap the posA^th identifier with the posB^th identifier.
   void swapId(unsigned posA, unsigned posB) override;
 
   /// Remove the specified range of ids.
-  void removeIdRange(unsigned idStart, unsigned idLimit) override;
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) override;
+  using IntegerRelation::removeIdRange;
 
   /// Eliminate the `posB^th` local identifier, replacing every instance of it
   /// with the `posA^th` local identifier. This should be used when the two
   /// local variables are known to always take the same values.
   void eliminateRedundantLocalId(unsigned posA, unsigned posB) override;
 
   /// Return whether the outputs of `this` and `other` agree wherever both
   /// functions are defined, i.e., the outputs should be equal for all points in

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

unsigned getIdKindOverlap(IdKind kind, unsigned idStart,
                             unsigned idLimit) const;
 
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. Return the absolute
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
   virtual unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1);
 
-  /// Removes identifiers in the column range [idStart, idLimit).
-  virtual void removeIdRange(unsigned idStart, unsigned idLimit);
+  /// Removes identifiers of the specified kind in the column range [idStart,
+  /// idLimit). The range is relative to the kind of identifier.
+  virtual void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
 
   /// Returns true if both the spaces are equal i.e. if both spaces have the
   /// same number of identifiers of each kind (excluding Local Identifiers).
   bool isEqual(const PresburgerSpace &other) const;
 
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of dimensional identifiers immediately before
   /// the split become symbols, or some of the symbols immediately after the
public:
 
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. Return the absolute
   /// column position (i.e., not relative to the kind of identifier) of the
   /// first added identifier.
   unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1) override;
 
   /// Removes identifiers in the column range [idStart, idLimit).
-  void removeIdRange(unsigned idStart, unsigned idLimit) override;
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) override;
 
   /// Returns true if both the spaces are equal i.e. if both spaces have the
   /// same number of identifiers of each kind.
   bool isEqual(const PresburgerLocalSpace &other) const;
 
   void print(llvm::raw_ostream &os) const;
   void dump() const;
 

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

public:
   /// identifiers. Return the position of the first added column.
   ///
   /// Note: Empty Values are allowed in `vals`.
   unsigned appendDimId(ValueRange vals);
   using FlatAffineConstraints::appendDimId;
   unsigned appendSymbolId(ValueRange vals);
   using FlatAffineConstraints::appendSymbolId;
 
+  /// Removes identifiers in the column range [idStart, idLimit), and copies any
+  /// remaining valid data into place, updates member variables, and resizes
+  /// arrays as needed.
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) override;
+  using IntegerRelation::removeIdRange;
+
   /// Add the specified values as a dim or symbol id depending on its nature, if
   /// it already doesn't exist in the system. `val` has to be either a terminal
   /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any
   /// symbols or loop IVs. The identifier is added to the end of the existing
   /// dims or symbols. Additional information on the identifier is extracted
   /// from the IR and added to the constraint system.
   void addInductionVarOrTerminalSymbol(Value val);
 
public:
   void mergeSymbolIds(FlatAffineValueConstraints &other);
 
 protected:
   /// Returns false if the fields corresponding to various identifier 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;
 
-  /// Removes identifiers in the column range [idStart, idLimit), and copies any
-  /// remaining valid data into place, updates member variables, and resizes
-  /// arrays as needed.
-  void removeIdRange(unsigned idStart, unsigned idLimit) override;
-
   /// Eliminates the identifier at the specified position using Fourier-Motzkin
   /// variable elimination, but uses Gaussian elimination if there is an
   /// equality involving that identifier. If the result of the elimination is
   /// integer exact, `*isResultIntegerExact` is set to true. If `darkShadow` is
   /// set to true, a potential under approximation (subset) of the rational
   /// shadow / exact integer shadow is computed.
   // See implementation comments for more details.
   void fourierMotzkinEliminate(unsigned pos, bool darkShadow = false,
public:
   void insertRangeId(unsigned pos, unsigned num = 1);
 
   /// Append `num` identifiers of the specified kind after the last identifier
   /// of that kind. The coefficient columns corresponding to the added
   /// identifiers are initialized to zero.
   void appendDomainId(unsigned num = 1);
   void appendRangeId(unsigned num = 1);
 
+  /// Removes identifiers in the column range [idStart, idLimit), and copies any
+  /// remaining valid data into place, updates member variables, and resizes
+  /// arrays as needed.
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) override;
+  using IntegerRelation::removeIdRange;
+
 protected:
   // Number of dimension identifers corresponding to domain identifers.
   unsigned numDomainDims;
 
   // Number of dimension identifers corresponding to range identifers.
   unsigned numRangeDims;
-
-  /// Removes identifiers in the column range [idStart, idLimit), and copies any
-  /// remaining valid data into place, updates member variables, and resizes
-  /// arrays as needed.
-  void removeIdRange(unsigned idStart, unsigned idLimit) override;
 };
 
 /// Flattens 'expr' into 'flattenedExpr', which contains the coefficients of the
 /// dimensions, symbols, and additional variables that represent floor divisions
 /// of dimensions, symbols, and in turn other floor divisions.  Returns failure
 /// if 'expr' could not be flattened (i.e., semi-affine is not yet handled).
 /// 'cst' contains constraints that connect newly introduced local identifiers
 /// to existing dimensional and symbolic identifiers. See documentation for

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

void IntegerRelation::removeId(IdKind kind, unsigned pos) {
   removeIdRange(kind, pos, pos + 1);
 }
 
 void IntegerRelation::removeId(unsigned pos) { removeIdRange(pos, pos + 1); }
 
 void IntegerRelation::removeIdRange(IdKind kind, unsigned idStart,
                                     unsigned idLimit) {
   assert(idLimit <= getNumIdKind(kind));
-  removeIdRange(getIdKindOffset(kind) + idStart,
-                getIdKindOffset(kind) + idLimit);
+
+  if (idStart >= idLimit)
+    return;
+
+  // Remove eliminated identifiers from the constraints.
+  unsigned offset = getIdKindOffset(kind);
+  equalities.removeColumns(offset + idStart, idLimit - idStart);
+  inequalities.removeColumns(offset + idStart, idLimit - idStart);
+
+  // Remove eliminated identifiers from the space.
+  PresburgerLocalSpace::removeIdRange(kind, idStart, idLimit);
 }
 
 void IntegerRelation::removeIdRange(unsigned idStart, unsigned idLimit) {
-  // Update space paramaters.
-  PresburgerLocalSpace::removeIdRange(idStart, idLimit);
+  assert(idLimit <= getNumIds());
+
+  if (idStart >= idLimit)
+    return;
+
+  // Helper function to remove ids of the specified kind in the given range
+  // [start, limit), The range is absolute (i.e. it is not relative to the kind
+  // of identifier). Also updates `limit` to reflect the deleted identifiers.
+  auto removeIdKindInRange = [this](IdKind kind, unsigned &start,
+                                    unsigned &limit) {
+    if (start >= limit)
+      return;
+
+    unsigned offset = getIdKindOffset(kind);
+    unsigned num = getNumIdKind(kind);
+
+    // Get `start`, `limit` relative to the specified kind.
+    unsigned relativeStart =
+        start <= offset ? 0 : std::min(num, start - offset);
+    unsigned relativeLimit =
+        limit <= offset ? 0 : std::min(num, limit - offset);
+
+    // Remove ids of the specified kind in the relative range.
+    removeIdRange(kind, relativeStart, relativeLimit);
+
+    // Update `limit` to reflect deleted identifiers.
+    // `start` does not need to be updated because any identifiers that are
+    // deleted are after position `start`.
+    limit -= relativeLimit - relativeStart;
+  };
 
-  // Remove eliminated identifiers from the constraints..
-  equalities.removeColumns(idStart, idLimit - idStart);
-  inequalities.removeColumns(idStart, idLimit - idStart);
+  removeIdKindInRange(IdKind::Domain, idStart, idLimit);
+  removeIdKindInRange(IdKind::Range, idStart, idLimit);
+  removeIdKindInRange(IdKind::Symbol, idStart, idLimit);
+  removeIdKindInRange(IdKind::Local, idStart, idLimit);
 }
 
 void IntegerRelation::removeEquality(unsigned pos) {
   equalities.removeRow(pos);
 }
 
 void IntegerRelation::removeInequality(unsigned pos) {
   inequalities.removeRow(pos);
void IntegerRelation::convertDimToLocal(unsigned dimStart, unsigned dimLimit) {
   unsigned newLocalIdStart = getNumIds();
   appendId(IdKind::Local, convertCount);
 
   // Swap the new local variables with dimensions.
   for (unsigned i = 0; i < convertCount; ++i)
     swapId(i + dimStart, i + newLocalIdStart);
 
   // Remove dimensions converted to local variables.
-  removeIdRange(dimStart, dimLimit);
+  removeIdRange(IdKind::SetDim, dimStart, dimLimit);
 }
 
 void IntegerRelation::addBound(BoundType type, unsigned pos, int64_t value) {
   assert(pos < getNumCols());
   if (type == BoundType::EQ) {
     unsigned row = equalities.appendExtraRow();
     equalities(row, pos) = 1;
     equalities(row, getNumCols() - 1) = -value;

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

unsigned MultiAffineFunction::insertId(IdKind kind, unsigned pos,
   return IntegerPolyhedron::insertId(kind, pos, num);
 }
 
 void MultiAffineFunction::swapId(unsigned posA, unsigned posB) {
   output.swapColumns(posA, posB);
   IntegerPolyhedron::swapId(posA, posB);
 }
 
-void MultiAffineFunction::removeIdRange(unsigned idStart, unsigned idLimit) {
-  output.removeColumns(idStart, idLimit - idStart);
+void MultiAffineFunction::removeIdRange(IdKind kind, unsigned idStart,
+                                        unsigned idLimit) {
+  output.removeColumns(idStart + getIdKindOffset(kind), idLimit - idStart);
   IntegerPolyhedron::removeIdRange(idStart, idLimit);
 }
 
 void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA,
                                                     unsigned posB) {
   output.addToColumn(posB, posA, /*scale=*/1);
   IntegerPolyhedron::eliminateRedundantLocalId(posA, posB);
 }

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

if (kind == IdKind::Domain) {
     numSymbols += num;
   } else {
     llvm_unreachable("PresburgerSpace does not support local identifiers!");
   }
 
   return absolutePos;
 }
 
-void PresburgerSpace::removeIdRange(unsigned idStart, unsigned idLimit) {
-  assert(idLimit <= getNumIds() && "invalid id limit");
+void PresburgerSpace::removeIdRange(IdKind kind, unsigned idStart,
+                                    unsigned idLimit) {
+  assert(idLimit <= getNumIdKind(kind) && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
-  // We are going to be removing one or more identifiers from the range.
-  assert(idStart < getNumIds() && "invalid idStart position");
-
-  // Update members numDomain, numRange, numSymbols and numIds.
-  unsigned numDomainEliminated = 0;
-  if (spaceKind == Relation)
-    numDomainEliminated = getIdKindOverlap(IdKind::Domain, idStart, idLimit);
-  unsigned numRangeEliminated =
-      getIdKindOverlap(IdKind::Range, idStart, idLimit);
-  unsigned numSymbolsEliminated =
-      getIdKindOverlap(IdKind::Symbol, idStart, idLimit);
-
-  numDomain -= numDomainEliminated;
-  numRange -= numRangeEliminated;
-  numSymbols -= numSymbolsEliminated;
+  unsigned numIdsEliminated = idLimit - idStart;
+  if (kind == IdKind::Domain) {
+    assert(spaceKind == Relation && "IdKind::Domain is not supported in Set.");
+    numDomain -= numIdsEliminated;
+  } else if (kind == IdKind::Range) {
+    numRange -= numIdsEliminated;
+  } else if (kind == IdKind::Symbol) {
+    numSymbols -= numIdsEliminated;
+  } else {
+    llvm_unreachable("PresburgerSpace does not support local identifiers!");
+  }
 }
 
 unsigned PresburgerLocalSpace::insertId(IdKind kind, unsigned pos,
                                         unsigned num) {
   if (kind == IdKind::Local) {
     numLocals += num;
     return getIdKindOffset(IdKind::Local) + pos;
   }
   return PresburgerSpace::insertId(kind, pos, num);
 }
 
-void PresburgerLocalSpace::removeIdRange(unsigned idStart, unsigned idLimit) {
-  assert(idLimit <= getNumIds() && "invalid id limit");
+void PresburgerLocalSpace::removeIdRange(IdKind kind, unsigned idStart,
+                                         unsigned idLimit) {
+  assert(idLimit <= getNumIdKind(kind) && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
-  // We are going to be removing one or more identifiers from the range.
-  assert(idStart < getNumIds() && "invalid idStart position");
-
-  unsigned numLocalsEliminated =
-      getIdKindOverlap(IdKind::Local, idStart, idLimit);
-
-  // Update space parameters.
-  PresburgerSpace::removeIdRange(idStart, idLimit);
-
-  // Update local ids.
-  numLocals -= numLocalsEliminated;
+  if (kind == IdKind::Local)
+    numLocals -= idLimit - idStart;
+  else
+    PresburgerSpace::removeIdRange(kind, idStart, idLimit);
 }
 
 bool PresburgerSpace::isEqual(const PresburgerSpace &other) const {
   return getNumDomainIds() == other.getNumDomainIds() &&
          getNumRangeIds() == other.getNumRangeIds() &&
          getNumSymbolIds() == other.getNumSymbolIds();
 }
 

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

void FlatAffineValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
   append(cst);
 }
 
 bool FlatAffineValueConstraints::hasConsistentState() const {
   return FlatAffineConstraints::hasConsistentState() &&
          values.size() == getNumIds();
 }
 
-void FlatAffineValueConstraints::removeIdRange(unsigned idStart,
+void FlatAffineValueConstraints::removeIdRange(IdKind kind, unsigned idStart,
                                                unsigned idLimit) {
-  FlatAffineConstraints::removeIdRange(idStart, idLimit);
-  values.erase(values.begin() + idStart, values.begin() + idLimit);
+  FlatAffineConstraints::removeIdRange(kind, idStart, idLimit);
+  unsigned offset = getIdKindOffset(kind);
+  values.erase(values.begin() + idStart + offset,
+               values.begin() + idLimit + offset);
 }
 
 // Determine whether the identifier at 'pos' (say id_r) can be expressed as
 // modulo of another known identifier (say id_n) w.r.t a constant. For example,
 // if the following constraints hold true:
 // ```
 // 0 <= id_r <= divisor - 1
 // id_n - (divisor * q_expr) = id_r
void FlatAffineRelation::appendDomainId(unsigned num) {
   numDomainDims += num;
 }
 
 void FlatAffineRelation::appendRangeId(unsigned num) {
   insertDimId(getNumDimIds(), num);
   numRangeDims += num;
 }
 
-void FlatAffineRelation::removeIdRange(unsigned idStart, unsigned idLimit) {
+void FlatAffineRelation::removeIdRange(IdKind kind, unsigned idStart,
+                                       unsigned idLimit) {
+  assert(idLimit <= getNumIdKind(kind));
   if (idStart >= idLimit)
     return;
 
+  FlatAffineValueConstraints::removeIdRange(kind, idStart, idLimit);
+
+  // If kind is not SetDim, domain and range don't need to be updated.
+  if (kind != IdKind::SetDim)
+    return;
+
   // Compute number of domain and range identifiers to remove. This is done by
   // intersecting the range of domain/range ids with range of ids to remove.
   unsigned intersectDomainLHS = std::min(idLimit, getNumDomainDims());
   unsigned intersectDomainRHS = idStart;
   unsigned intersectRangeLHS = std::min(idLimit, getNumDimIds());
   unsigned intersectRangeRHS = std::max(idStart, getNumDomainDims());
 
-  FlatAffineValueConstraints::removeIdRange(idStart, idLimit);
-
   if (intersectDomainLHS > intersectDomainRHS)
     numDomainDims -= intersectDomainLHS - intersectDomainRHS;
   if (intersectRangeLHS > intersectRangeRHS)
     numRangeDims -= intersectRangeLHS - intersectRangeRHS;
 }
 
 LogicalResult mlir::getRelationFromMap(AffineMap &map,
                                        FlatAffineRelation &rel) {

mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

TEST(PresburgerSpaceTest, insertIdSet) {
   space.insertId(IdKind::SetDim, 0, 2);
   EXPECT_EQ(space.getNumRangeIds(), 4u);
 }
 
 TEST(PresburgerSpaceTest, removeIdRange) {
   PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 3);
 
   // Remove 1 domain identifier.
-  space.removeIdRange(0, 1);
+  space.removeIdRange(IdKind::Domain, 0, 1);
   EXPECT_EQ(space.getNumDomainIds(), 1u);
 
   // Remove 1 symbol and 1 range identifier.
-  space.removeIdRange(1, 3);
+  space.removeIdRange(IdKind::Symbol, 0, 1);
+  space.removeIdRange(IdKind::Range, 0, 1);
   EXPECT_EQ(space.getNumDomainIds(), 1u);
   EXPECT_EQ(space.getNumRangeIds(), 0u);
   EXPECT_EQ(space.getNumSymbolIds(), 2u);
 }