[MLIR][Presburger] Remove inheritence in MultiAffineFunction

[Groverkss]

Summary

This patch removes inheritence of MultiAffineFunction from IntegerPolyhedron
and instead makes IntegerPolyhedron as a member.

This patch removes virtualization in MultiAffineFunction and also removes
unnecessary functions inherited from IntegerPolyhedron.

[ftynse]

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

98

Documentation, please.

[Groverkss]

• Address Alex's comment

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

 /// Since IntegerRelation makes a distinction between dimensions, IdKind::Range
 /// and IdKind::Domain should be used to refer to dimension identifiers.
 class IntegerRelation {
 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,

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

 /// In this way every point in the polyhedron has a tuple of integers associated
 /// with it. If the integer polyhedron has local ids, then the output
 /// expressions can use them as well. The output expressions are represented as
 /// a matrix with one row for every element in the output vector one column for
 /// each id, and an extra column at the end for the constant term.
 ///
 /// Checking equality of two such functions is supported, as well as finding the
 /// value of the function at a specified point.
-class MultiAffineFunction : protected IntegerPolyhedron {
+class MultiAffineFunction {
 public:
-  /// We use protected inheritance to avoid inheriting the whole public
-  /// interface of IntegerPolyhedron. These using declarations explicitly make
-  /// only the relevant functions part of the public interface.
-  using IntegerPolyhedron::getNumDimAndSymbolIds;
-  using IntegerPolyhedron::getNumDimIds;
-  using IntegerPolyhedron::getNumIds;
-  using IntegerPolyhedron::getNumLocalIds;
-  using IntegerPolyhedron::getNumSymbolIds;
-  using IntegerPolyhedron::getSpace;
-
   MultiAffineFunction(const IntegerPolyhedron &domain, const Matrix &output)
-      : IntegerPolyhedron(domain), output(output) {}
+      : domainSet(domain), output(output) {}
   MultiAffineFunction(const Matrix &output, const PresburgerSpace &space)
-      : IntegerPolyhedron(space), output(output) {}
-
-  ~MultiAffineFunction() override = default;
-  Kind getKind() const override { return Kind::MultiAffineFunction; }
-  bool classof(const IntegerRelation *rel) const {
-    return rel->getKind() == Kind::MultiAffineFunction;
-  }
+      : domainSet(space), output(output) {}
 
-  unsigned getNumInputs() const { return getNumDimAndSymbolIds(); }
+  unsigned getNumInputs() const { return domainSet.getNumDimAndSymbolIds(); }
   unsigned getNumOutputs() const { return output.getNumRows(); }
   bool isConsistent() const {
-    return output.getNumColumns() == getNumIds() + 1;
+    return output.getNumColumns() == domainSet.getNumIds() + 1;
   }
-  const IntegerPolyhedron &getDomain() const { return *this; }
+  const IntegerPolyhedron &getDomain() const { return domainSet; }
+  const PresburgerSpace &getDomainSpace() const { return domainSet.getSpace(); }
 
   /// 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;
+  unsigned insertId(IdKind kind, unsigned pos, unsigned num = 1);
 
   /// Remove the specified range of ids.
-  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) override;
-  using IntegerRelation::removeIdRange;
+  void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
 
-  /// 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;
+  /// Given a MAF `other`, merges local identifiers such that both funcitons
+  /// have union of local ids, without changing the set of points in domain or
+  /// the output.
+  void mergeLocalIds(MultiAffineFunction &other);
 
   /// 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;
 
   /// 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.
   Optional<SmallVector<int64_t, 8>> valueAt(ArrayRef<int64_t> point) const;
 
   /// Truncate the output dimensions to the first `count` dimensions.
   ///
   /// TODO: refactor so that this can be accomplished through removeIdRange.
   void truncateOutput(unsigned count);
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 private:
+  /// The IntegerPolyhedron representing the domain over which the function is
+  /// defined.
+  IntegerPolyhedron domainSet;
+
   /// The function's output is a tuple of integers, with the ith element of the
   /// tuple defined by the affine expression given by the ith row of this output
   /// matrix.
   Matrix output;
 };
 
 /// This class represents a piece-wise MultiAffineFunction. This can be thought
 /// of as a list of MultiAffineFunction with disjoint domains, with each having

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

 /// it's each instance is replaced by `i^th` division. If it returns `false`,
 /// the divisions are not merged. `merge` can also do side effects, For example
 /// it can merge the local identifiers in IntegerRelation.
 void removeDuplicateDivs(
     std::vector<SmallVector<int64_t, 8>> &divs,
     SmallVectorImpl<unsigned> &denoms, unsigned localOffset,
     llvm::function_ref<bool(unsigned i, unsigned j)> merge);
 
+/// Given two relations, A and B, add additional local ids to the sets such
+/// that both have the union of the local ids in each set, without changing
+/// the set of points that lie in A and B.
+///
+/// While taking union, if a local id in any set has a division representation
+/// which is a duplicate of division representation, of another local id in any
+/// set, it is not added to the final union of local ids and is instead merged.
+///
+/// On every possible merge, `merge(i, j)` is called. `i`, `j` are position
+/// of local identifiers in both sets which are being merged. If `merge(i, j)`
+/// returns true, the divisions are merged, otherwise the divisions are not
+/// merged.
+void mergeLocalIds(IntegerRelation &relA, IntegerRelation &relB,
+                   llvm::function_ref<bool(unsigned i, unsigned j)> merge);
+
 /// Compute the gcd of the range.
 int64_t gcdRange(ArrayRef<int64_t> range);
 
 /// Divide the range by its gcd and return the gcd.
 int64_t normalizeRange(MutableArrayRef<int64_t> range);
 
 /// Normalize the given (numerator, denominator) pair by dividing out the
 /// common factors between them. The numerator here is an affine expression

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

 /// After extracting these divisions, local ids in `other` with the same
 /// division representation as some other local id in any set 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.
 unsigned IntegerRelation::mergeLocalIds(IntegerRelation &other) {
-  assert(space.isCompatible(other.getSpace()) &&
-         "Spaces should be compatible.");
-
   IntegerRelation &relA = *this;
   IntegerRelation &relB = other;
 
   unsigned oldALocals = relA.getNumLocalIds();
 
-  // Merge local ids of relA and relB without using division information,
-  // i.e. append local ids of `relB` to `relA` and insert local ids of `relA`
-  // to `relB` at start of its local ids.
-  unsigned initLocals = relA.getNumLocalIds();
-  insertId(IdKind::Local, relA.getNumLocalIds(), relB.getNumLocalIds());
-  relB.insertId(IdKind::Local, 0, initLocals);
-
-  // Get division representations from each rel.
-  std::vector<SmallVector<int64_t, 8>> divsA, divsB;
-  SmallVector<unsigned, 4> denomsA, denomsB;
-  relA.getLocalReprs(divsA, denomsA);
-  relB.getLocalReprs(divsB, denomsB);
-
-  // Copy division information for relB into `divsA` and `denomsA`, so that
-  // these have the combined division information of both rels. Since newly
-  // added local variables in relA and relB have no constraints, they will not
-  // have any division representation.
-  std::copy(divsB.begin() + initLocals, divsB.end(),
-            divsA.begin() + initLocals);
-  std::copy(denomsB.begin() + initLocals, denomsB.end(),
-            denomsA.begin() + initLocals);
-
   // Merge function that merges the local variables in both sets by treating
   // them as the same identifier.
   auto merge = [&relA, &relB, oldALocals](unsigned i, unsigned j) -> bool {
     // We only merge from local at pos j to local at pos i, where j > i.
     if (i >= j)
       return false;
 
     // If i < oldALocals, we are trying to merge duplicate divs. Since we do not
     // want to merge duplicates in A, we ignore this call.
     if (j < oldALocals)
       return false;
 
     // Merge local at pos j into local at position i.
     relA.eliminateRedundantLocalId(i, j);
     relB.eliminateRedundantLocalId(i, j);
     return true;
   };
 
-  // Merge all divisions by removing duplicate divisions.
-  unsigned localOffset = getIdKindOffset(IdKind::Local);
-  presburger::removeDuplicateDivs(divsA, denomsA, localOffset, merge);
+  presburger::mergeLocalIds(*this, other, merge);
 
   // Since we do not remove duplicate divisions in relA, this is guranteed to be
   // non-negative.
   return relA.getNumLocalIds() - oldALocals;
 }
 
 void IntegerRelation::removeDuplicateDivs() {
   std::vector<SmallVector<int64_t, 8>> divs;

mlir/lib/Analysis/Presburger/PWMAFunction.cpp

PresburgerSet PWMAFunction::getDomain() const {
   PresburgerSet domain = PresburgerSet::getEmpty(getSpace());
   for (const MultiAffineFunction &piece : pieces)
     domain.unionInPlace(piece.getDomain());
   return domain;
 }
 
 Optional<SmallVector<int64_t, 8>>
 MultiAffineFunction::valueAt(ArrayRef<int64_t> point) const {
-  assert(point.size() == getNumDimAndSymbolIds() &&
+  assert(point.size() == domainSet.getNumDimAndSymbolIds() &&
          "Point has incorrect dimensionality!");
 
   Optional<SmallVector<int64_t, 8>> maybeLocalValues =
       getDomain().containsPointNoLocal(point);
   if (!maybeLocalValues)
     return {};
 
   // The point lies in the domain, so we need to compute the output value.
PWMAFunction::valueAt(ArrayRef<int64_t> point) const {
   for (const MultiAffineFunction &piece : pieces)
     if (Optional<SmallVector<int64_t, 8>> output = piece.valueAt(point))
       return output;
   return {};
 }
 
 void MultiAffineFunction::print(raw_ostream &os) const {
   os << "Domain:";
-  IntegerPolyhedron::print(os);
+  domainSet.print(os);
   os << "Output:\n";
   output.print(os);
   os << "\n";
 }
 
 void MultiAffineFunction::dump() const { print(llvm::errs()); }
 
 bool MultiAffineFunction::isEqual(const MultiAffineFunction &other) const {
-  return space.isCompatible(other.getSpace()) &&
+  return getDomainSpace().isCompatible(other.getDomainSpace()) &&
          getDomain().isEqual(other.getDomain()) &&
          isEqualWhereDomainsOverlap(other);
 }
 
 unsigned MultiAffineFunction::insertId(IdKind kind, unsigned pos,
                                        unsigned num) {
-  assert((kind != IdKind::Domain || num == 0) &&
-         "Domain has to be zero in a set");
-  unsigned absolutePos = getIdKindOffset(kind) + pos;
+  assert(kind != IdKind::Domain && "Domain has to be zero in a set");
+  unsigned absolutePos = domainSet.getIdKindOffset(kind) + pos;
   output.insertColumns(absolutePos, num);
-  return IntegerPolyhedron::insertId(kind, pos, num);
-}
-
-void MultiAffineFunction::swapId(unsigned posA, unsigned posB) {
-  output.swapColumns(posA, posB);
-  IntegerPolyhedron::swapId(posA, posB);
+  return domainSet.insertId(kind, pos, num);
 }
 
 void MultiAffineFunction::removeIdRange(IdKind kind, unsigned idStart,
                                         unsigned idLimit) {
-  output.removeColumns(idStart + getIdKindOffset(kind), idLimit - idStart);
-  IntegerPolyhedron::removeIdRange(kind, idStart, idLimit);
-}
-
-void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA,
-                                                    unsigned posB) {
-  unsigned localOffset = getIdKindOffset(IdKind::Local);
-  output.addToColumn(localOffset + posB, localOffset + posA, /*scale=*/1);
-  IntegerPolyhedron::eliminateRedundantLocalId(posA, posB);
+  output.removeColumns(idStart + domainSet.getIdKindOffset(kind),
+                       idLimit - idStart);
+  domainSet.removeIdRange(kind, idStart, idLimit);
 }
 
 void MultiAffineFunction::truncateOutput(unsigned count) {
   assert(count <= output.getNumRows());
   output.resizeVertically(count);
 }
 
 void PWMAFunction::truncateOutput(unsigned count) {
   assert(count <= numOutputs);
   for (MultiAffineFunction &piece : pieces)
     piece.truncateOutput(count);
   numOutputs = count;
 }
 
+void MultiAffineFunction::mergeLocalIds(MultiAffineFunction &other) {
+  // Merge output local ids of both functions without using division
+  // information i.e. append local ids of `other` to `this` and insert
+  // local ids of `this` to `other` at the start of it's local ids.
+  output.insertColumns(domainSet.getIdKindEnd(IdKind::Local),
+                       other.domainSet.getNumLocalIds());
+  other.output.insertColumns(other.domainSet.getIdKindOffset(IdKind::Local),
+                             domainSet.getNumLocalIds());
+
+  auto merge = [this, &other](unsigned i, unsigned j) -> bool {
+    // Merge local at position j into local at position i in function domain.
+    domainSet.eliminateRedundantLocalId(i, j);
+    other.domainSet.eliminateRedundantLocalId(i, j);
+
+    unsigned localOffset = domainSet.getIdKindOffset(IdKind::Local);
+
+    // Merge local at position j into local at position i in output domain.
+    output.addToColumn(localOffset + j, localOffset + i, 1);
+    output.removeColumn(localOffset + j);
+    other.output.addToColumn(localOffset + j, localOffset + i, 1);
+    other.output.removeColumn(localOffset + j);
+
+    return true;
+  };
+
+  presburger::mergeLocalIds(domainSet, other.domainSet, merge);
+}
+
 bool MultiAffineFunction::isEqualWhereDomainsOverlap(
     MultiAffineFunction other) const {
-  if (!space.isCompatible(other.getSpace()))
+  if (!getDomainSpace().isCompatible(other.getDomainSpace()))
     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
   // domains of `this` and `other`.
-  commonFunc.IntegerPolyhedron::append(other);
+  commonFunc.domainSet.append(other.domainSet);
 
   // `commonDomainMatching` contains the subset of the common domain
   // where the outputs of `this` and `other` match.
   //
   // We want to add constraints equating the outputs of `this` and `other`.
   // However, `this` may have difference local ids from `other`, whereas we
   // need both to have the same locals. Accordingly, we use `commonFunc.output`
   // in place of `this->output`, since `commonFunc` has the same output but also
bool PWMAFunction::isEqual(const PWMAFunction &other) const {
   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(space.isCompatible(piece.getSpace()) &&
+  assert(space.isCompatible(piece.getDomainSpace()) &&
          "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);
 }

mlir/lib/Analysis/Presburger/Utils.cpp

for (unsigned j = i + 1; j < divs.size(); ++j) {
       divs.erase(divs.begin() + j);
       denoms.erase(denoms.begin() + j);
       // Since `j` can never be zero, we do not need to worry about overflows.
       --j;
     }
   }
 }
 
+void presburger::mergeLocalIds(
+    IntegerRelation &relA, IntegerRelation &relB,
+    llvm::function_ref<bool(unsigned i, unsigned j)> merge) {
+  assert(relA.getSpace().isCompatible(relB.getSpace()) &&
+         "Spaces should be compatible.");
+
+  // Merge local ids of relA and relB without using division information,
+  // i.e. append local ids of `relB` to `relA` and insert local ids of `relA`
+  // to `relB` at start of its local ids.
+  unsigned initLocals = relA.getNumLocalIds();
+  relA.insertId(IdKind::Local, relA.getNumLocalIds(), relB.getNumLocalIds());
+  relB.insertId(IdKind::Local, 0, initLocals);
+
+  // Get division representations from each rel.
+  std::vector<SmallVector<int64_t, 8>> divsA, divsB;
+  SmallVector<unsigned, 4> denomsA, denomsB;
+  relA.getLocalReprs(divsA, denomsA);
+  relB.getLocalReprs(divsB, denomsB);
+
+  // Copy division information for relB into `divsA` and `denomsA`, so that
+  // these have the combined division information of both rels. Since newly
+  // added local variables in relA and relB have no constraints, they will not
+  // have any division representation.
+  std::copy(divsB.begin() + initLocals, divsB.end(),
+            divsA.begin() + initLocals);
+  std::copy(denomsB.begin() + initLocals, denomsB.end(),
+            denomsA.begin() + initLocals);
+
+  // Merge all divisions by removing duplicate divisions.
+  unsigned localOffset = relA.getIdKindOffset(IdKind::Local);
+  presburger::removeDuplicateDivs(divsA, denomsA, localOffset, merge);
+}
+
 int64_t presburger::gcdRange(ArrayRef<int64_t> range) {
   int64_t gcd = 0;
   for (int64_t elem : range) {
     gcd = llvm::GreatestCommonDivisor64(gcd, std::abs(elem));
     if (gcd == 1)
       return gcd;
   }
   return gcd;