[MLIR][Presburger] Factor out various Space equality checks to PresburgerSpace::isEqual

[Groverkss]

Summary

This patch factors out various checks for dimension compatibility to
PresburgerSpace::isEqual and PresburgerLocalSpace::isEqual (for local
identifiers).

[arjunp]

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

100

inline keyword is redundant here

100

If I pass a PresburgerLocalSpace to this, it will return true even if their locals aren't equal. that seems misleading

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

161

Just use PresburgerSpace:: isEqual? That also makes it immediately clear we're checking if they're equal as PresburgerSpaces here, i.e., ignoring the locals.

[Groverkss]

• Address Arjun's comments

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

161

MultiAffineFunction is protected inheritance from IntegerPolyhedron so I cannot use that.

[arjunp]

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

161

Good point.

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

public:
 
   unsigned getNumInputs() const { return getNumDimAndSymbolIds(); }
   unsigned getNumOutputs() const { return output.getNumRows(); }
   bool isConsistent() const {
     return output.getNumColumns() == getNumIds() + 1;
   }
   const IntegerPolyhedron &getDomain() const { return *this; }
 
-  bool hasCompatibleDimensions(const MultiAffineFunction &f) const;
-
   /// Insert `num` identifiers of the specified kind at position `pos`.
   /// Positions are relative to the kind of identifier. The coefficient columns
   /// corresponding to the added identifiers are initialized to zero. 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;
 
   /// 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;
 
   /// 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
   /// 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.
   ///

 /// this represents a partial function and so could be undefined at some points.
 ///
 /// As in PresburgerSets, the input ids are partitioned into dimension ids and
 /// 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 : PresburgerSpace {
+class PWMAFunction : public PresburgerSpace {
 public:
   PWMAFunction(unsigned numDims, unsigned numSymbols, unsigned numOutputs)
       : PresburgerSpace(numDims, numSymbols), 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]; }
   unsigned getNumPieces() const { return pieces.size(); }
   unsigned getNumOutputs() const { return numOutputs; }
   unsigned getNumInputs() const { return getNumIds(); }
   MultiAffineFunction &getPiece(unsigned i) { return pieces[i]; }
 
   /// Return the domain of this piece-wise MultiAffineFunction. This is the
   /// union of the domains of all the pieces.
   PresburgerSet getDomain() const;
 
-  /// Check whether the `this` and the given function have compatible
-  /// dimensions, i.e., the same number of dimension inputs, symbol inputs, and
-  /// outputs.
-  bool hasCompatibleDimensions(const MultiAffineFunction &f) const;
-  bool hasCompatibleDimensions(const PWMAFunction &f) const;
-
   /// Return the value at the specified point and an empty optional if the
   /// point does not lie in the domain.
   ///
   /// Note: domains with local ids are not yet supported, and will assert-fail.
   Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
 
   /// Return whether `this` and `other` are equal as PWMAFunctions, i.e. whether
   /// they have the same dimensions, the same domain and they take the same

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

public:
   /// 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);
 
+  /// 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
   /// split become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);
 
   void print(llvm::raw_ostream &os) const;
   void dump() const;
protected:
   PresburgerLocalSpace(unsigned numDomain, unsigned numRange,
                        unsigned numSymbols, unsigned numLocals)
       : PresburgerSpace(Relation, numDomain, numRange, numSymbols, numLocals) {}
 
   /// Local Space constructor for Set space type.
   PresburgerLocalSpace(unsigned numDims, unsigned numSymbols,
                        unsigned numLocals)
       : PresburgerSpace(Set, /*numDomain=*/0, numDims, numSymbols, numLocals) {}
+
+  /// 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;
 };
 
 } // namespace presburger
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_PRESBURGERSPACE_H

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

 
 void IntegerPolyhedron::reset(unsigned newNumDims, unsigned newNumSymbols,
                               unsigned newNumLocals) {
   reset(0, 0, newNumDims + newNumSymbols + newNumLocals + 1, newNumDims,
         newNumSymbols, newNumLocals);
 }
 
 void IntegerPolyhedron::append(const IntegerPolyhedron &other) {
-  assert(other.getNumCols() == getNumCols());
-  assert(other.getNumDimIds() == getNumDimIds());
-  assert(other.getNumSymbolIds() == getNumSymbolIds());
+  assert(PresburgerLocalSpace::isEqual(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));
   }

 ///
 /// To detect local ids that always take the same in both sets, each local id is
 /// represented as a floordiv with constant denominator in terms of other ids.
 /// After extracting these divisions, local ids with the same division
 /// representation are considered duplicate and are merged. It is possible that
 /// division representation for some local id cannot be obtained, and thus these
 /// local ids are not considered for detecting duplicates.
 void IntegerPolyhedron::mergeLocalIds(IntegerPolyhedron &other) {
-  assert(getNumDimIds() == other.getNumDimIds() &&
-         "Number of dimension ids should match");
-  assert(getNumSymbolIds() == other.getNumSymbolIds() &&
-         "Number of symbol ids should match");
+  assert(PresburgerSpace::isEqual(other) && "Spaces should match.");
 
   IntegerPolyhedron &polyA = *this;
   IntegerPolyhedron &polyB = other;
 
   // Merge local ids of polyA and polyB without using division information,
   // i.e. append local ids of `polyB` to `polyA` and insert local ids of `polyA`
   // to `polyB` at start of its local ids.
   unsigned initLocals = polyA.getNumLocalIds();
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
 IntegerPolyhedron::unionBoundingBox(const IntegerPolyhedron &otherCst) {
-  assert(otherCst.getNumDimIds() == getNumDimIds() && "dims mismatch");
-  assert(otherCst.getNumLocalIds() == 0 && "local ids not supported here");
+  assert(PresburgerLocalSpace::isEqual(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.
   IntegerPolyhedron commonCst;
   getCommonConstraints(*this, otherCst, commonCst);
 
   std::vector<SmallVector<int64_t, 8>> boundingLbs;

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

void MultiAffineFunction::print(raw_ostream &os) const {
   os << "Output:\n";
   output.print(os);
   os << "\n";
 }
 
 void MultiAffineFunction::dump() const { print(llvm::errs()); }
 
 bool MultiAffineFunction::isEqual(const MultiAffineFunction &other) const {
-  return hasCompatibleDimensions(other) &&
+  return PresburgerSpace::isEqual(other) &&
          getDomain().isEqual(other.getDomain()) &&
          isEqualWhereDomainsOverlap(other);
 }
 
 unsigned MultiAffineFunction::insertId(IdKind kind, unsigned pos,
                                        unsigned num) {
   unsigned absolutePos = getIdKindOffset(kind) + pos;
   output.insertColumns(absolutePos, num);

 void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA,
                                                     unsigned posB) {
   output.addToColumn(posB, posA, /*scale=*/1);
   IntegerPolyhedron::eliminateRedundantLocalId(posA, posB);
 }
 
 bool MultiAffineFunction::isEqualWhereDomainsOverlap(
     MultiAffineFunction other) const {
-  if (!hasCompatibleDimensions(other))
+  if (!PresburgerSpace::isEqual(other))
     return false;
 
   // `commonFunc` has the same output as `this`.
   MultiAffineFunction commonFunc = *this;
   // After this merge, `commonFunc` and `other` have the same local ids; they
   // are merged.
   commonFunc.mergeLocalIds(other);
   // After this, the domain of `commonFunc` will be the intersection of the
bool MultiAffineFunction::isEqualWhereDomainsOverlap(
   // If the whole common domain is a subset of commonDomainMatching, then they
   // are equal and the two functions match on the whole common domain.
   return commonFunc.getDomain().isSubsetOf(commonDomainMatching);
 }
 
 /// Two PWMAFunctions are equal if they have the same dimensionalities,
 /// the same domain, and take the same value at every point in the domain.
 bool PWMAFunction::isEqual(const PWMAFunction &other) const {
-  if (!hasCompatibleDimensions(other))
+  if (!PresburgerSpace::isEqual(other))
     return false;
 
   if (!this->getDomain().isEqual(other.getDomain()))
     return false;
 
   // Check if, whenever the domains of a piece of `this` and a piece of `other`
   // overlap, they take the same output value. If `this` and `other` have the
   // same domain (checked above), then this check passes iff the two functions
   // have the same output at every point in the domain.
   for (const MultiAffineFunction &aPiece : this->pieces)
     for (const MultiAffineFunction &bPiece : other.pieces)
       if (!aPiece.isEqualWhereDomainsOverlap(bPiece))
         return false;
   return true;
 }
 
 void PWMAFunction::addPiece(const MultiAffineFunction &piece) {
-  assert(hasCompatibleDimensions(piece) &&
+  assert(piece.isSpaceEqual(*this) &&
          "Piece to be added is not compatible with this PWMAFunction!");
   assert(piece.isConsistent() && "Piece is internally inconsistent!");
   assert(this->getDomain()
              .intersect(PresburgerSet(piece.getDomain()))
              .isIntegerEmpty() &&
          "New piece's domain overlaps with that of existing pieces!");
   pieces.push_back(piece);
 }
 
 void PWMAFunction::addPiece(const IntegerPolyhedron &domain,
                             const Matrix &output) {
   addPiece(MultiAffineFunction(domain, output));
 }
 
 void PWMAFunction::print(raw_ostream &os) const {
   os << pieces.size() << " pieces:\n";
   for (const MultiAffineFunction &piece : pieces)
     piece.print(os);
 }
-
-/// The hasCompatibleDimensions functions don't check the number of local ids;
-/// functions are still compatible if they have differing number of locals.
-bool MultiAffineFunction::hasCompatibleDimensions(
-    const MultiAffineFunction &f) const {
-  return getNumDimIds() == f.getNumDimIds() &&
-         getNumSymbolIds() == f.getNumSymbolIds() &&
-         getNumOutputs() == f.getNumOutputs();
-}
-bool PWMAFunction::hasCompatibleDimensions(const MultiAffineFunction &f) const {
-  return getNumDimIds() == f.getNumDimIds() &&
-         getNumSymbolIds() == f.getNumSymbolIds() &&
-         getNumOutputs() == f.getNumOutputs();
-}
-bool PWMAFunction::hasCompatibleDimensions(const PWMAFunction &f) const {
-  return getNumDimIds() == f.getNumDimIds() &&
-         getNumSymbolIds() == f.getNumSymbolIds() &&
-         getNumOutputs() == f.getNumOutputs();
-}

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

 }
 
 const IntegerPolyhedron &
 PresburgerSet::getIntegerPolyhedron(unsigned index) const {
   assert(index < integerPolyhedrons.size() && "index out of bounds!");
   return integerPolyhedrons[index];
 }
 
-/// Assert that the IntegerPolyhedron and PresburgerSet live in
-/// compatible spaces.
-static void assertDimensionsCompatible(const IntegerPolyhedron &poly,
-                                       const PresburgerSet &set) {
-  assert(poly.getNumDimIds() == set.getNumDimIds() &&
-         "Number of dimensions of the IntegerPolyhedron and PresburgerSet"
-         "do not match!");
-  assert(poly.getNumSymbolIds() == set.getNumSymbolIds() &&
-         "Number of symbols of the IntegerPolyhedron and PresburgerSet"
-         "do not match!");
-}
-
-/// Assert that the two PresburgerSets live in compatible spaces.
-static void assertDimensionsCompatible(const PresburgerSet &setA,
-                                       const PresburgerSet &setB) {
-  assert(setA.getNumDimIds() == setB.getNumDimIds() &&
-         "Number of dimensions of the PresburgerSets do not match!");
-  assert(setA.getNumSymbolIds() == setB.getNumSymbolIds() &&
-         "Number of symbols of the PresburgerSets do not match!");
-}
-
 /// Mutate this set, turning it into the union of this set and the given
 /// IntegerPolyhedron.
 void PresburgerSet::unionPolyInPlace(const IntegerPolyhedron &poly) {
-  assertDimensionsCompatible(poly, *this);
+  assert(PresburgerSpace::isEqual(poly) && "Spaces should match");
   integerPolyhedrons.push_back(poly);
 }
 
 /// Mutate this set, turning it into the union of this set and the given set.
 ///
 /// This is accomplished by simply adding all the Poly of the given set to this
 /// set.
 void PresburgerSet::unionSetInPlace(const PresburgerSet &set) {
-  assertDimensionsCompatible(set, *this);
+  assert(PresburgerSpace::isEqual(set) && "Spaces should match");
   for (const IntegerPolyhedron &poly : set.integerPolyhedrons)
     unionPolyInPlace(poly);
 }
 
 /// Return the union of this set and the given set.
 PresburgerSet PresburgerSet::unionSet(const PresburgerSet &set) const {
-  assertDimensionsCompatible(set, *this);
+  assert(PresburgerSpace::isEqual(set) && "Spaces should match");
   PresburgerSet result = *this;
   result.unionSetInPlace(set);
   return result;
 }
 
 /// A point is contained in the union iff any of the parts contain the point.
 bool PresburgerSet::containsPoint(ArrayRef<int64_t> point) const {
   return llvm::any_of(integerPolyhedrons, [&](const IntegerPolyhedron &poly) {

 // Return the intersection of this set with the given set.
 //
 // We directly compute (S_1 or S_2 ...) and (T_1 or T_2 ...)
 // as (S_1 and T_1) or (S_1 and T_2) or ...
 //
 // If S_i or T_j have local variables, then S_i and T_j contains the local
 // variables of both.
 PresburgerSet PresburgerSet::intersect(const PresburgerSet &set) const {
-  assertDimensionsCompatible(set, *this);
+  assert(PresburgerSpace::isEqual(set) && "Spaces should match");
 
   PresburgerSet result(getNumDimIds(), getNumSymbolIds());
   for (const IntegerPolyhedron &csA : integerPolyhedrons) {
     for (const IntegerPolyhedron &csB : set.integerPolyhedrons) {
       IntegerPolyhedron csACopy = csA, csBCopy = csB;
       csACopy.mergeLocalIds(csBCopy);
       csACopy.append(csBCopy);
       if (!csACopy.isEmpty())

 
 /// Return the set difference poly \ set.
 ///
 /// The Poly here is modified in subtractRecursively, so it cannot be a const
 /// reference even though it is restored to its original state before returning
 /// from that function.
 PresburgerSet PresburgerSet::getSetDifference(IntegerPolyhedron poly,
                                               const PresburgerSet &set) {
-  assertDimensionsCompatible(poly, set);
+  assert(poly.PresburgerSpace::isEqual(set) && "Spaces should match");
   if (poly.isEmptyByGCDTest())
     return PresburgerSet::getEmptySet(poly.getNumDimIds(),
                                       poly.getNumSymbolIds());
 
   PresburgerSet result(poly.getNumDimIds(), poly.getNumSymbolIds());
   Simplex simplex(poly);
   subtractRecursively(poly, simplex, set, 0, result);
   return result;
 }
 
 /// Return the complement of this set.
 PresburgerSet PresburgerSet::complement() const {
   return getSetDifference(
       IntegerPolyhedron::getUniverse(getNumDimIds(), getNumSymbolIds()), *this);
 }
 
 /// Return the result of subtract the given set from this set, i.e.,
 /// return `this \ set`.
 PresburgerSet PresburgerSet::subtract(const PresburgerSet &set) const {
-  assertDimensionsCompatible(set, *this);
+  assert(PresburgerSpace::isEqual(set) && "Spaces should match");
   PresburgerSet result(getNumDimIds(), getNumSymbolIds());
   // We compute (U_i t_i) \ (U_i set_i) as U_i (t_i \ V_i set_i).
   for (const IntegerPolyhedron &poly : integerPolyhedrons)
     result.unionSetInPlace(getSetDifference(poly, set));
   return result;
 }
 
 /// T is a subset of S iff T \ S is empty, since if T \ S contains a
 /// point then this is a point that is contained in T but not S, and
 /// if T contains a point that is not in S, this also lies in T \ S.
 bool PresburgerSet::isSubsetOf(const PresburgerSet &set) const {
   return this->subtract(set).isIntegerEmpty();
 }
 
 /// Two sets are equal iff they are subsets of each other.
 bool PresburgerSet::isEqual(const PresburgerSet &set) const {
-  assertDimensionsCompatible(set, *this);
+  assert(PresburgerSpace::isEqual(set) && "Spaces should match");
   return this->isSubsetOf(set) && set.isSubsetOf(*this);
 }
 
 /// Return true if all the sets in the union are known to be integer empty,
 /// false otherwise.
 bool PresburgerSet::isIntegerEmpty() const {
   // The set is empty iff all of the disjuncts are empty.
   return std::all_of(

mlir/lib/Analysis/Presburger/PresburgerSpace.cpp

void PresburgerLocalSpace::removeIdRange(unsigned idStart, unsigned idLimit) {
 
   // Update space parameters.
   PresburgerSpace::removeIdRange(idStart, idLimit);
 
   // Update local ids.
   numLocals -= numLocalsEliminated;
 }
 
+bool PresburgerSpace::isEqual(const PresburgerSpace &other) const {
+  return getNumDomainIds() == other.getNumDomainIds() &&
+         getNumRangeIds() == other.getNumRangeIds() &&
+         getNumSymbolIds() == other.getNumSymbolIds();
+}
+
+bool PresburgerLocalSpace::isEqual(const PresburgerLocalSpace &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 {