[MLIR][Presburger] Merge PresburgerLocalSpace and PresburgerSpace

[Groverkss]

Summary

This patch is a cleanup patch that merges PresburgerLocalSpace and
PresburgerSpace. Asserting that there are no locals is shifted to the
users of PresburgerSpace themselves.

The reasoning for this patch is that PresburgerLocalSpace did not contribute
much and only introduced additional complexity as locals could still be present
in PresburgerSpace, just not writable. This could introduce problems if a
PresburgerSpace with locals was copied to PresburgerLocalSpace which expected
no locals in a PresburgerSpace.

[arjunp]

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

1–2

Please fix this.

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

113–121

I prefer to call isEqual -> isCompatible and isEqualWithLocals -> isEqual.

I find it weird to say things are equal when they have different locals. I would rather say they are compatible in that case.

[Groverkss]

• Address Arjun's comments.
• Move PresburgerSpace::isEqual -> PresburgerSpace::isSpaceCompatible.
• Move PresburgerSpace::isEqualWithLocal -> PresburgerSpace::isSpaceEqual.

[arjunp]

PresburgerSpace::isSpaceEqual/Compatible seems a little weird, to have Space twice. I understand that after inheriting we want IntegerRelation::isSpaceEqual, though. Is isSpaceEqual even used anywhere? Can we just drop it and just keep isCompatible? IntegerRelation::isCompatible seems like a reasonable name

[Groverkss]

In D122353#3406021, @arjunp wrote:

PresburgerSpace::isSpaceEqual/Compatible seems a little weird, to have Space twice. I understand that after inheriting we want IntegerRelation::isSpaceEqual, though. Is isSpaceEqual even used anywhere? Can we just drop it and just keep isCompatible? IntegerRelation::isCompatible seems like a reasonable name

IsSpaceEqual is required for functions that assume number of locals are equal, for example: IntegerRelation::append, IntegerRelation::isEqual, IntegerRelation::isSubsetOf, IntegerRelation::unionBoundingBox.

[arjunp]

Those probably need to be updated to support differing locals by merging.

Anyway, let's land this as is for now.

[Groverkss]

Rebase

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

 #include "mlir/Analysis/Presburger/Matrix.h"
 #include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 namespace presburger {
 
-/// An IntegerRelation is a PresburgerLocalSpace subject to affine constraints.
+/// An IntegerRelation is a PresburgerSpace subject to affine constraints.
 /// Affine constraints can be inequalities or equalities in the form:
 ///
 /// Inequality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n >= 0
 /// Equality  : c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n == 0
 ///
 /// where c_0, c_1, ..., c_n are integers and n is the total number of
 /// identifiers in the space.
 ///
 /// Such a relation corresponds to the set of integer points lying in a convex
 /// polyhedron. For example, consider the relation:
 ///         (x) -> (y) : (1 <= x <= 7, x = 2y)
 /// These can be thought of as points in the polyhedron:
 ///         (x, y) : (1 <= x <= 7, x = 2y)
 /// This relation contains the pairs (2, 1), (4, 2), and (6, 3).
 ///
 /// Since IntegerRelation makes a distinction between dimensions, IdKind::Range
 /// and IdKind::Domain should be used to refer to dimension identifiers.
-class IntegerRelation : public PresburgerLocalSpace {
+class IntegerRelation : public PresburgerSpace {
 public:
   /// All derived classes of IntegerRelation.
   enum class Kind {
     FlatAffineConstraints,
     FlatAffineValueConstraints,
     MultiAffineFunction,
     IntegerRelation,
     IntegerPolyhedron,
   };
 
   /// Constructs a relation reserving memory for the specified number
   /// of constraints and identifiers.
   IntegerRelation(unsigned numReservedInequalities,
                   unsigned numReservedEqualities, unsigned numReservedCols,
                   unsigned numDomain, unsigned numRange, unsigned numSymbols,
                   unsigned numLocals)
-      : PresburgerLocalSpace(numDomain, numRange, numSymbols, numLocals),
+      : PresburgerSpace(numDomain, numRange, numSymbols, numLocals),
         equalities(0, getNumIds() + 1, numReservedEqualities, numReservedCols),
         inequalities(0, getNumIds() + 1, numReservedInequalities,
                      numReservedCols) {
     assert(numReservedCols >= getNumIds() + 1);
   }
 
   /// Constructs a relation with the specified number of dimensions and symbols.
   IntegerRelation(unsigned numDomain = 0, unsigned numRange = 0,
protected:
 
   /// Coefficients of affine equalities (in == 0 form).
   Matrix equalities;
 
   /// Coefficients of affine inequalities (in >= 0 form).
   Matrix inequalities;
 };
 
-/// An IntegerPolyhedron is a PresburgerLocalSpace subject to affine
+/// An IntegerPolyhedron is a PresburgerSpace subject to affine
 /// constraints. Affine constraints can be inequalities or equalities in the
 /// form:
 ///
 /// Inequality: c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n >= 0
 /// Equality  : c_0*x_0 + c_1*x_1 + .... + c_{n-1}*x_{n-1} + c_n == 0
 ///
 /// where c_0, c_1, ..., c_n are integers and n is the total number of
 /// identifiers in the space.

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

public:
   /// the intersection of the domains.
   bool isEqualWhereDomainsOverlap(MultiAffineFunction other) const;
 
   /// Returns whether the underlying PresburgerSpace is equal to `other`.
   bool isSpaceEqual(const PresburgerSpace &other) const {
     return PresburgerSpace::isEqual(other);
   };
 
-  /// Returns whether the underlying PresburgerLocalSpace is equal to `other`.
-  bool isSpaceEqual(const PresburgerLocalSpace &other) const {
-    return PresburgerLocalSpace::isEqual(other);
-  };
-
   /// Return whether the `this` and `other` are equal. This is the case if
   /// they lie in the same space, i.e. have the same dimensions, and their
   /// domains are identical and their outputs are equal on their domain.
   bool isEqual(const MultiAffineFunction &other) const;
 
   /// Get the value of the function at the specified point. If the point lies
   /// outside the domain, an empty optional is returned.
   ///

 /// symbolic ids.
 ///
 /// Support is provided to compare equality of two such functions as well as
 /// finding the value of the function at a point. Note that local ids in the
 /// piece are not supported for the latter.
 class PWMAFunction : public PresburgerSpace {
 public:
   PWMAFunction(unsigned numDims, unsigned numSymbols, unsigned numOutputs)
-      : PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols),
+      : PresburgerSpace(/*numDomain=*/0, /*numRange=*/numDims, numSymbols,
+                        /*numLocals=*/0),
         numOutputs(numOutputs) {
     assert(numOutputs >= 1 && "The function must output something!");
   }
 
   void addPiece(const MultiAffineFunction &piece);
   void addPiece(const IntegerPolyhedron &domain, const Matrix &output);
 
   const MultiAffineFunction &getPiece(unsigned i) const { return pieces[i]; }

mlir/include/mlir/Analysis/Presburger/PresburgerRelation.h

-//===- PresburgerRelation.h - MLIR PresburgerRelation Class -----*- C++ -*-===//
+//===- PresburgerRelation.h - MLIR PresburgerRelation Class -----*- C++
+//-*-===//:%s
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // A class to represent unions of IntegerRelations.
public:
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
   /// Construct an empty PresburgerRelation with the specified number of
   /// dimension and symbols.
   PresburgerRelation(unsigned numDomain = 0, unsigned numRange = 0,
                      unsigned numSymbols = 0)
-      : PresburgerSpace(numDomain, numRange, numSymbols) {}
+      : PresburgerSpace(numDomain, numRange, numSymbols, /*numLocals=*/0) {}
 
   /// The list of disjuncts that this set is the union of.
   SmallVector<IntegerRelation, 2> integerRelations;
 
   friend class SetCoalescer;
 };
 
 class PresburgerSet : public PresburgerRelation {

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

 ///
 /// Identifiers are stored in the following order:
 ///       [Domain, Range, Symbols, Locals]
 ///
 /// A space with no distinction between types of dimension identifiers can
 /// be implemented as a space with zero domain. IdKind::SetDim should be used
 /// to refer to dimensions in such spaces.
 ///
-/// PresburgerSpace does not allow identifiers of kind Local. See
-/// PresburgerLocalSpace for an extension that does allow local identifiers.
+/// Since locals represent existential quantification, equality of two spaces
+/// implies that all kinds of identifers other than locals are equal.
 class PresburgerSpace {
-  friend PresburgerLocalSpace;
-
 public:
-  PresburgerSpace(unsigned numDomain, unsigned numRange, unsigned numSymbols)
-      : PresburgerSpace(numDomain, numRange, numSymbols, 0) {}
+  PresburgerSpace(unsigned numDomain = 0, unsigned numRange = 0,
+                  unsigned numSymbols = 0, unsigned numLocals = 0)
+      : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
+        numLocals(numLocals) {}
 
   virtual ~PresburgerSpace() = default;
 
   unsigned getNumDomainIds() const { return numDomain; }
   unsigned getNumRangeIds() const { return numRange; }
-  unsigned getNumSymbolIds() const { return numSymbols; }
   unsigned getNumSetDimIds() const { return numRange; }
+  unsigned getNumSymbolIds() const { return numSymbols; }
+  unsigned getNumLocalIds() const { return numLocals; }
 
   unsigned getNumDimIds() const { return numDomain + numRange; }
   unsigned getNumDimAndSymbolIds() const {
     return numDomain + numRange + numSymbols;
   }
   unsigned getNumIds() const {
     return numDomain + numRange + numSymbols + numLocals;
   }
public:
   /// 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 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).
+  /// same number of identifiers of each kind (excluding locals).
   bool isEqual(const PresburgerSpace &other) const;
 
+  /// Returns true if both the spaces are equal including local identifiers i.e.
+  /// if both spaces have the same number of identifiers of each kind (including
+  /// locals).
+  bool isEqualWithLocals(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
   /// split become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   void print(llvm::raw_ostream &os) const;
   void dump() const;
 
 private:
-  PresburgerSpace(unsigned numDomain, unsigned numRange, unsigned numSymbols,
-                  unsigned numLocals)
-      : numDomain(numDomain), numRange(numRange), numSymbols(numSymbols),
-        numLocals(numLocals) {}
-
   // Number of identifiers corresponding to domain identifiers.
   unsigned numDomain;
 
   // Number of identifiers corresponding to range identifiers.
   unsigned numRange;
 
   /// Number of identifiers corresponding to symbols (unknown but constant for
   /// analysis).
   unsigned numSymbols;
 
   /// Number of identifers corresponding to locals (identifiers corresponding
   /// to existentially quantified variables).
   unsigned numLocals;
 };
 
-/// Extension of PresburgerSpace supporting Local identifiers.
-class PresburgerLocalSpace : public PresburgerSpace {
-public:
-  PresburgerLocalSpace(unsigned numDomain, unsigned numRange,
-                       unsigned numSymbols, unsigned numLocals)
-      : PresburgerSpace(numDomain, numRange, numSymbols, numLocals) {}
-
-  unsigned getNumLocalIds() const { return numLocals; }
-
-  /// 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(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;
-};
-
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

std::unique_ptr<IntegerRelation> IntegerRelation::clone() const {
   return std::make_unique<IntegerRelation>(*this);
 }
 
 std::unique_ptr<IntegerPolyhedron> IntegerPolyhedron::clone() const {
   return std::make_unique<IntegerPolyhedron>(*this);
 }
 
 void IntegerRelation::append(const IntegerRelation &other) {
-  assert(PresburgerLocalSpace::isEqual(other) && "Spaces must be equal.");
+  assert(PresburgerSpace::isEqualWithLocals(other) && "Spaces must be equal.");
 
   inequalities.reserveRows(inequalities.getNumRows() +
                            other.getNumInequalities());
   equalities.reserveRows(equalities.getNumRows() + other.getNumEqualities());
 
   for (unsigned r = 0, e = other.getNumInequalities(); r < e; r++) {
     addInequality(other.getInequality(r));
   }
   for (unsigned r = 0, e = other.getNumEqualities(); r < e; r++) {
     addEquality(other.getEquality(r));
   }
 }
 
 IntegerRelation IntegerRelation::intersect(IntegerRelation other) const {
   IntegerRelation result = *this;
   result.mergeLocalIds(other);
   result.append(other);
   return result;
 }
 
 bool IntegerRelation::isEqual(const IntegerRelation &other) const {
-  assert(PresburgerLocalSpace::isEqual(other) && "Spaces must be equal.");
+  assert(PresburgerSpace::isEqualWithLocals(other) && "Spaces must be equal.");
   return PresburgerRelation(*this).isEqual(PresburgerRelation(other));
 }
 
 bool IntegerRelation::isSubsetOf(const IntegerRelation &other) const {
-  assert(PresburgerLocalSpace::isEqual(other) && "Spaces must be equal.");
+  assert(PresburgerSpace::isEqualWithLocals(other) && "Spaces must be equal.");
   return PresburgerRelation(*this).isSubsetOf(PresburgerRelation(other));
 }
 
 MaybeOptimum<SmallVector<Fraction, 8>>
 IntegerRelation::findRationalLexMin() const {
   assert(getNumSymbolIds() == 0 && "Symbols are not supported!");
   MaybeOptimum<SmallVector<Fraction, 8>> maybeLexMin =
       LexSimplex(*this).findRationalLexMin();
if (!rangeIsZero(poly.getEquality(i - 1).slice(begin, count)))
       poly.removeEquality(i - 1);
   for (unsigned i = poly.getNumInequalities(); i > 0; i--)
     if (!rangeIsZero(poly.getInequality(i - 1).slice(begin, count)))
       poly.removeInequality(i - 1);
 }
 unsigned IntegerRelation::insertId(IdKind kind, unsigned pos, unsigned num) {
   assert(pos <= getNumIdKind(kind));
 
-  unsigned insertPos = PresburgerLocalSpace::insertId(kind, pos, num);
+  unsigned insertPos = PresburgerSpace::insertId(kind, pos, num);
   inequalities.insertColumns(insertPos, num);
   equalities.insertColumns(insertPos, num);
   return insertPos;
 }
 
 unsigned IntegerRelation::appendId(IdKind kind, unsigned num) {
   unsigned pos = getNumIdKind(kind);
   return insertId(kind, pos, num);
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);
+  PresburgerSpace::removeIdRange(kind, idStart, idLimit);
 }
 
 void IntegerRelation::removeIdRange(unsigned idStart, unsigned idLimit) {
   assert(idLimit <= getNumIds());
 
   if (idStart >= idLimit)
     return;
 
for (unsigned r = 0, e = a.getNumEqualities(); r < e; ++r) {
     }
   }
 }
 
 // Computes the bounding box with respect to 'other' by finding the min of the
 // lower bounds and the max of the upper bounds along each of the dimensions.
 LogicalResult
 IntegerRelation::unionBoundingBox(const IntegerRelation &otherCst) {
-  assert(PresburgerLocalSpace::isEqual(otherCst) && "Spaces should match.");
+  assert(PresburgerSpace::isEqualWithLocals(otherCst) &&
+         "Spaces should match.");
   assert(getNumLocalIds() == 0 && "local ids not supported yet here");
 
   // Get the constraints common to both systems; these will be added as is to
   // the union.
   IntegerRelation commonCst;
   getCommonConstraints(*this, otherCst, commonCst);
 
   std::vector<SmallVector<int64_t, 8>> boundingLbs;
void IntegerRelation::removeIndependentConstraints(unsigned pos, unsigned num) {
   // inequality leads to much shifting/copying in the underlying buffer).
   for (auto nbIndex : llvm::reverse(nbIneqIndices))
     removeInequality(nbIndex);
   for (auto nbIndex : llvm::reverse(nbEqIndices))
     removeEquality(nbIndex);
 }
 
 void IntegerRelation::printSpace(raw_ostream &os) const {
-  PresburgerLocalSpace::print(os);
+  PresburgerSpace::print(os);
   os << getNumConstraints() << " constraints\n";
 }
 
 void IntegerRelation::print(raw_ostream &os) const {
   assert(hasConsistentState());
   printSpace(os);
   for (unsigned i = 0, e = getNumEqualities(); i < e; ++i) {
     for (unsigned j = 0, f = getNumCols(); j < f; ++j) {

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

unsigned PresburgerSpace::insertId(IdKind kind, unsigned pos, unsigned num) {
 
   if (kind == IdKind::Domain)
     numDomain += num;
   else if (kind == IdKind::Range)
     numRange += num;
   else if (kind == IdKind::Symbol)
     numSymbols += num;
   else
-    llvm_unreachable("PresburgerSpace does not support local identifiers!");
+    numLocals += num;
 
   return absolutePos;
 }
 
 void PresburgerSpace::removeIdRange(IdKind kind, unsigned idStart,
                                     unsigned idLimit) {
   assert(idLimit <= getNumIdKind(kind) && "invalid id limit");
 
   if (idStart >= idLimit)
     return;
 
   unsigned numIdsEliminated = idLimit - idStart;
   if (kind == IdKind::Domain)
     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(IdKind kind, unsigned idStart,
-                                         unsigned idLimit) {
-  assert(idLimit <= getNumIdKind(kind) && "invalid id limit");
-
-  if (idStart >= idLimit)
-    return;
-
-  if (kind == IdKind::Local)
-    numLocals -= idLimit - idStart;
-  else
-    PresburgerSpace::removeIdRange(kind, idStart, idLimit);
+    numLocals -= numIdsEliminated;
 }
 
 bool PresburgerSpace::isEqual(const PresburgerSpace &other) const {
   return getNumDomainIds() == other.getNumDomainIds() &&
          getNumRangeIds() == other.getNumRangeIds() &&
          getNumSymbolIds() == other.getNumSymbolIds();
 }
 
-bool PresburgerLocalSpace::isEqual(const PresburgerLocalSpace &other) const {
+bool PresburgerSpace::isEqualWithLocals(const PresburgerSpace &other) const {
   return PresburgerSpace::isEqual(other) &&
          getNumLocalIds() == other.getNumLocalIds();
 }
 
 void PresburgerSpace::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= getNumDimAndSymbolIds() &&
          "invalid separation position");
   numRange = numRange + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
 }
 
 void PresburgerSpace::print(llvm::raw_ostream &os) const {
   os << "Domain: " << getNumDomainIds() << ", "
      << "Range: " << getNumRangeIds() << ", "
-     << "Symbols: " << getNumSymbolIds() << "\n";
+     << "Symbols: " << getNumSymbolIds() << ", "
+     << "Locals: " << getNumLocalIds() << "\n";
 }
 
 void PresburgerSpace::dump() const { print(llvm::errs()); }
-
-void PresburgerLocalSpace::print(llvm::raw_ostream &os) const {
-  PresburgerSpace::print(os);
-  os << "Locals: " << getNumLocalIds() << "\n";
-}
-
-void PresburgerLocalSpace::dump() const { print(llvm::errs()); }