[MLIR][Presburger] Make IntegerRelation::mergeLocalIds not delete duplicates

[Groverkss]

Summary

This patch modifies mergeLocalIds to not delete duplicate local ids in
this relation. This allows the ordering of the final local ids for this
to be determined more easily, which is generally required when other objects
refer to these local ids.

[arjunp]

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

477–478

Should be this and other

486

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1085

1088

I can't understand this. What are local id sets?

1132–1133

"duplicates in this"*

[Groverkss]

• Address Arjun's comments

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

1085

I don't think the "both sets" is true. Local ids which are duplicated in other are also merged.

[arjunp]

We should probably consider not this and other and also A and B in this

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

public:
   void convertIdKind(IdKind srcKind, unsigned idStart, unsigned idLimit,
                      IdKind dstKind);
   void convertToLocal(IdKind kind, unsigned idStart, unsigned idLimit) {
     convertIdKind(kind, idStart, idLimit, IdKind::Local);
   }
 
   /// Adds additional local ids to the sets such that they both have the union
   /// of the local ids in each set, without changing the set of points that
-  /// lie in `this` and `other`. The ordering of the local ids in the
-  /// sets may also be changed. After merging, if the `i^th` local variable in
-  /// one set has a known division representation, then the `i^th` local
-  /// variable in the other set either has the same division representation or
-  /// no known division representation.
-  ///
-  /// The number of dimensions and symbol ids in `this` and `other` should
-  /// match.
-  void mergeLocalIds(IntegerRelation &other);
+  /// lie in `this` and `other`.
+  ///
+  /// While taking union, if a local id in `other` has a division representation
+  /// which is a duplicate of division representation, of another local id, it
+  /// is not added to the final union of local ids and is instead merged. The
+  /// new ordering of local ids is:
+  ///
+  /// [Local ids of A] [Non-merged local ids of B with same relative ordering]
+  ///
+  /// After merging, if the `i^th` local variable in one set has a known
+  /// division representation, then the `i^th` local variable in the other set
+  /// either has the same division representation or no known division
+  /// representation.
+  ///
+  /// The spaces of both relations should be compatible.
+  ///
+  /// Returns the number of non-merged local ids of B.
+  unsigned mergeLocalIds(IntegerRelation &other);
 
   /// 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) {
     space.setDimSymbolSeparation(newSymbolCount);
   }

mlir/lib/Analysis/Presburger/IntegerRelation.cpp

void IntegerRelation::eliminateRedundantLocalId(unsigned posA, unsigned posB) {
   equalities.addToColumn(posB, posA, 1);
   removeId(posB);
 }
 
 /// Adds additional local ids to the sets such that they both have the union
 /// of the local ids in each set, without changing the set of points that
 /// lie in `this` and `other`.
 ///
-/// 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 IntegerRelation::mergeLocalIds(IntegerRelation &other) {
+/// To detect local ids that always take the value, each local id is represented
+/// as a floordiv with constant denominator in terms of other ids. After
+/// extracting these divisions, local ids in `other` with the same division
+/// representation with some other local id sets 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](unsigned i, unsigned j) -> bool {
+  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, 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);
+
+  // 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;
   SmallVector<unsigned, 4> denoms;
 
   getLocalReprs(divs, denoms);
   auto merge = [this](unsigned i, unsigned j) -> bool {

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

TEST(IntegerPolyhedronTest, mergeDivisionsConstants) {
     EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
 
     // 2 divisions should be matched.
     EXPECT_EQ(poly1.getNumLocalIds(), 2u);
     EXPECT_EQ(poly2.getNumLocalIds(), 2u);
   }
 }
 
+TEST(IntegerPolyhedronTest, mergeDivisionsDuplicateInSameSet) {
+  // (x) : (exists y = [x + 1 / 3], z = [x + 1 / 3]: y + z >= x).
+  IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
+  poly1.addLocalFloorDiv({1, 1}, 3);    // y = [x + 1 / 2].
+  poly1.addLocalFloorDiv({1, 0, 1}, 3); // z = [x + 1 / 3].
+  poly1.addInequality({-1, 1, 1, 0});   // y + z >= x.
+
+  // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z <= x).
+  IntegerPolyhedron poly2(PresburgerSpace::getSetSpace(1));
+  poly2.addLocalFloorDiv({1, 1}, 3);    // y = [x + 1 / 3].
+  poly2.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
+  poly2.addInequality({1, -1, -1, 0});  // y + z <= x.
+
+  poly1.mergeLocalIds(poly2);
+
+  // Local space should be same.
+  EXPECT_EQ(poly1.getNumLocalIds(), poly2.getNumLocalIds());
+
+  // 1 divisions should be matched.
+  EXPECT_EQ(poly1.getNumLocalIds(), 3u);
+  EXPECT_EQ(poly2.getNumLocalIds(), 3u);
+}
+
 TEST(IntegerPolyhedronTest, negativeDividends) {
   // (x) : (exists y = [-x + 1 / 2], z = [-x - 2 / 3]: y + z >= x).
   IntegerPolyhedron poly1(PresburgerSpace::getSetSpace(1));
   poly1.addLocalFloorDiv({-1, 1}, 2); // y = [x + 1 / 2].
   // Normalization test with negative dividends
   poly1.addLocalFloorDiv({-3, 0, -6}, 9); // z = [3x + 6 / 9] -> [x + 2 / 3].
   poly1.addInequality({-1, 1, 1, 0});     // y + z >= x.