[MLIR][FlatAffineConstraints] Remove duplicate divisions while merging local ids

[Groverkss]

Summary

This patch implements detecting duplicate local identifiers by extracting their
division representation while merging local identifiers.

For example, given the FACs A, B:

A: (x, y)[s0] : (exists d0 = [x / 4], d1 = [y / 4]: d0 <= s0, d1 <= s0, x + y >= 2)
B: (x, y)[s0] : (exists d0 = [x / 4], d1 = [y / 4]: d0 <= s0, d1 <= s0, x + y >= 5)

The intersection of A and B without this patch would lead to the following FAC:

(x, y)[s0] : (exists d0 = [x / 4], d1 = [y / 4], d2 = [x / 4], d3 = [x / 4]: d0 <= s0, d1 <= s0, d2 <= s0, d3 <= s0, x + y >= 2, x + y >= 5)

after this patch, merging of local ids will detect that d0 = d2 and d1 = d3,
and the intersection of these two FACs will be (after removing duplicate constraints):

(x, y)[s0] : (exists d0 = [x / 4], d1 = [y / 4] : d0 <= s0, d1 <= s0, x + y >= 2, x + y >= 5)

This reduces the number of constraints by 2 (constraints) + 4 (2 constraints for each extra division) for this case.

This is used to reduce the output size representation of operations like
PresburgerSet::subtract, PresburgerSet::intersect which require merging local
variables.

[Groverkss]

Rebase

[Groverkss]

• Fix formatting errors

[Groverkss]

• Update docs for mergeLocalIds

[arjunp]

Hi, thanks for this patch. Maybe you can add a line at the end of the patch summary about why this is important, e.g. mentioning that when operating on sets with many divisions, it is common for blowups in the number of variables to occur, which this heuristic prevents.

mlir/include/mlir/Analysis/AffineStructures.h

445

This looks like an implementation detail that doesn't need to be described here (extracting div representations). I would just say "local ids with identical division representations are merged".

447

Please write complete sentences: "The number of dimensions and symbol ids should..."

539

nit: representationS

539

What is a local representation? Do you mean a division representation?

540

nit: tildes for i^th but not denominator[i]?

542

Why std::vector here?

mlir/lib/Analysis/AffineStructures.cpp

1379

I think factoring out this function should be done in a separate commit.

1921–1950

There seem to be a multitude of different functions related to obtaining division representations. Are they all really needed? Can we factor out common functionality from them?

[Groverkss]

• Move division representation to a common function
• Update docs

[Groverkss]

Addressed most comments.

mlir/include/mlir/Analysis/AffineStructures.h

542

I remember @bondhugula recommending not to use SmallVector<SmallVector> although I may be wrong. If you think I should use SmallVector instead, let me know.

[bondhugula]

mlir/include/mlir/Analysis/AffineStructures.h

542

SmallVector isn't really optimized to hold another SmallVector I had heard. Just using std::vector here would be fine.

[arjunp]

mlir/include/mlir/Analysis/AffineStructures.h

160–172

The dividend -> The dividends

162

For this overload, maybe mention that the denominator will be set to zero if no representation could be found.

163–164

nit: dividends is plural but denominator is singular?

445–446

nit: I think "The number of dimensions and symbol ids in this and other should match" reads more naturally

542

Maybe you could use Matrix here, if you want.

mlir/lib/Analysis/AffineStructures.cpp

1476–1495

I think this refactor can also be moved into a separate patch

1494–1499

Maybe it would be nicer for the users if we just resize it internally and don't require it to be initialized to that size.

1921–1950

Please explain what this is doing in more detail. At the moment, someone looking at this function has no idea what it has to do with divisions? Maybe add a note saying it's intended to be used to remove redundancy when we have local variables that always take the same values. Or you can just call it mergeLocals. BTW, we are now using "merge" with a different meaning in mergeLocalIds and here, which could get confusing. You could call this one, e.g., eliminateRedundantLocalId.

1963

nit: why 8 above and 4 here?

1988

1990

1992

1995

1998

2001

I don't understand why this is needed. If they both depend on the same local variables in the same way, they should still always take the same values and should be able to be merged right? (the ith local variable here, is the same in both now after merging, right?)

mlir/unittests/Analysis/AffineStructuresTest.cpp

814

nit: z, y is counterintuitive when trying to read the coefficients. Writing the constraint in a comment next to each line would help improve readability.

[Groverkss]

• Address arjun's comments
• Update tests for mergeLocalIds

[Groverkss]

Addressed @arjunp's comments.

mlir/lib/Analysis/AffineStructures.cpp

1963

My reasoning was that denoms only has columns for local ids while a division has columns for all ids. Usually, I would not expect local variables to exceed 3-4.

2001

I don't remember the reason why I did not consider such cases. But I don't find any now so I think I'll remove it.

[Groverkss]

• Update docs

[Groverkss]

Rebase

[Groverkss]

• Fix clang-format errors

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

1928–1929

Less than or equal to? Shouldn't this be strict?

2014–2015

Use range-based for here?

2014–2020

I see you added some test cases. Did you add one that was failing before this change?

[Groverkss]

• Fix asserts as suggested by Arjun

[Groverkss]

Addressed @arjunp's comments

mlir/lib/Analysis/AffineStructures.cpp

1928–1929

Yes, thanks for catching this.

2014–2015

I don't think I can use range-based for loops here since I need to check if the division representation actually exists using denoms.

2014–2020

No, I think they worked before. I had previously assumed that local ids wouldn't be checked when checking for the same division. To allow local ids too, I had to add change some functionality and required more tests.

[arjunp]

mlir/lib/Analysis/AffineStructures.cpp

1925

1951

A top-level comment to explain what is being done would be useful.

1966–1969

It only makes sense to use division information if the division inequalities actually get added to both sets here, right?

1970–1985

I think it's better to first insert locals in the FACs and then get the local representation, then you can skip the extra insertions into divs1 and divs2 here.

Also, if we are actually adding the division constraints, then you can skip the appending at the end and just get the div repr from any on of the sets. In fact, then you can factor out removing redundant divs into a separate function. Such a function would anyway be useful in other contexts as well. There might be a little bit of duplicated work in detecting redundancies but thinking about that right now is probably premature optimization.

1985

If you want, you can use SmallVector::append here

1996

2001

2008

Just use operator !=?

2024

Might be useful to put a comment saying this is fine because j is never zero here.

mlir/unittests/Analysis/AffineStructuresTest.cpp

885

889

914

Is the denominator supposed to be 2 or 3?

920

Same here, 2 or 3?

[Groverkss]

Addressed @arjunp's comments.

mlir/lib/Analysis/AffineStructures.cpp

1951

The doc comment for this function is in the include file. I checked with others on discord and was recommended keeping doc comments in the header file only (for functions that are documented in the header file).

Are you asking for something else?

1966–1969

No, I think we can use the division information without actually adding the division inequalities. mergeLocalIds is only used to align local ids. Its only use is to allow comparing inequalities between two sets. Adding division constraints has no real use there.

1970–1985

This cannot be done unless division constraints are added to each set from the other, which I personally think is unnecessary :/

1985

It's not really appending. It's inserting before the coefficient for the constant.

[Groverkss]

• Addressed more comments

[arjunp]

mlir/include/mlir/Analysis/AffineStructures.h

444–454

I think the meanings of both "merge" and "align" are a little bit mysterious here. How about this:

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 actually lie in `this` and `other`. The orderings of the local ids in the sets may also be changed.

The number of dimensions and symbol ids in `this` and `other` should match.

Also, please update the documentation of mergeAndAlignIds, it currently says all local ids of A will appear before all local ids of B, which is not true anymore.

mlir/lib/Analysis/AffineStructures.cpp

1951

Yes, I meant a comment giving a high-level explanation of the implementation of the function. (Whereas the header has an explanation of when/why/how to use the function.)

1970–1985

You can still do the first part I guess, but then you would first have to erase the suffix of divs1 and denoms1 corresponding to the newly added locals which have no representation in the first set, and then append their representations from the second set. I leave it to you to decide if that's better or not

1985

It's not really appending. It's inserting before the coefficient for the constant.

I don't understand. I'm saying you could do

denoms1.append(denoms2) instead of
denoms1.insert(denoms1.end(), denoms2.begin(), denoms2.end());

[Groverkss]

Rebase and address @arjunp's comments

mlir/include/mlir/Analysis/AffineStructures.h

444–454

Added this doc comment with some addition which I think is useful:

The ordering of the local ids in the sets may also be changed but is the same in both sets.

[arjunp]

mlir/include/mlir/Analysis/AffineStructures.h

444–454

I guess the question is, what is meant by "same in both sets". In what sense are two local ids same?

Maybe you can rather replace it with something like this:

The ordering of the local ids in the sets may be changed. After merging, if the ith local variable in one set has a known division representation, then the ith local variable in the other set either has the same division representation or no known division representation.

[Groverkss]

• Fix doc comment for mergeLocalIds.

[arjunp]

LGTM!

mlir/lib/Analysis/AffineStructures.cpp

1973

1986

[Groverkss]

• Addressed comments

[Groverkss]

This comment was removed by Groverkss.

[Groverkss]

Addressed suggested changes.

[Groverkss]

Thank you @arjunp for the review and helpful comments!

I'll wait for @bondhugula before landing in case he has any comments.

[Groverkss]

@bondhugula gentle ping. could you take a look at this patch? :)

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

496–498

The last sentence may be a bit unclear - rephrase along the lines of the other doc comment above.

1957–1958

Nit: Consider using facA and facB. Numbered suffixes aren't great in that it isn't clear if they are 0-based or 1-based.

1977–1980

Likewise.

[Groverkss]

• Address bondhugula's comments

[Groverkss]

This comment was removed by Groverkss.

[Mogball]

@Groverkss Please squash your commits next time :S

[Groverkss]

In D112867#3168270, @Mogball wrote:

@Groverkss Please squash your commits next time :S

Sorry for that! I accidentally sent the commits directly instead of sending them from arcanist which caused that. I reverted them and send them correctly. Will not let this happen again :(

[MaskRay]

There are many fix-up commits which seem to should be in one change.

You may find llvm/utils/git/pre-push.py useful.

[mehdi_amini]

In D112867#3168617, @MaskRay wrote:

There are many fix-up commits which seem to should be in one change.

You may find llvm/utils/git/pre-push.py useful.

To be clear, please do this: https://llvm.org/docs/GettingStarted.html#git-pre-push-hook

[bondhugula]

In D112867#3168366, @Groverkss wrote:

In D112867#3168270, @Mogball wrote:

@Groverkss Please squash your commits next time :S

Sorry for that! I accidentally sent the commits directly instead of sending them from arcanist which caused that. I reverted them and send them correctly. Will not let this happen again :(

I use the workflow here to commit (instead of using arc to land) and always keep my revision commit as a single one on which arc diff is repeatedly used to update to phab:
https://llvm.org/docs/GettingStarted.html#for-developers-to-commit-changes-from-git

mlir/include/mlir/Analysis/AffineStructures.h

public:
   /// otherwise.
   Optional<SmallVector<int64_t, 8>> findIntegerSample() const;
 
   /// Returns true if the given point satisfies the constraints, or false
   /// otherwise.
   bool containsPoint(ArrayRef<int64_t> point) const;
 
   /// Find pairs of inequalities identified by their position indices, using
   /// which an explicit representation for each local variable can be computed.
   /// The pairs are stored as indices of upperbound, lowerbound inequalities. If
   /// no such pair can be found, it is stored as llvm::None.
   ///
   /// The dividends of the explicit representations are stored in `dividends`
   /// and the denominators in `denominators`. If no explicit representation
   /// could be found for the `i^th` local identifier, `denominators[i]` is set
   /// to 0.
   void getLocalReprs(
       std::vector<SmallVector<int64_t, 8>> &dividends,
       SmallVector<unsigned, 4> &denominators,
       std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
   void getLocalReprs(
       std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
   void getLocalReprs(std::vector<SmallVector<int64_t, 8>> &dividends,
                      SmallVector<unsigned, 4> &denominators) const;
 
   // Clones this object.
   std::unique_ptr<FlatAffineConstraints> clone() const;
 
   /// Returns the value at the specified equality row and column.
public:
   /// in the polytope and C is the number of constraints, the algorithm takes
   /// O(VC) time.
   void removeRedundantConstraints();
 
   /// Converts identifiers in the column range [idStart, idLimit) to local
   /// variables.
   void convertDimToLocal(unsigned dimStart, unsigned dimLimit);
 
-  /// Merge local ids of `this` and `other`. This is done by appending local ids
-  /// of `other` to `this` and inserting local ids of `this` to `other` at start
-  /// of its local ids. Number of dimension and symbol ids should match in
-  /// `this` and `other`.
+  /// 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(FlatAffineConstraints &other);
 
   /// Removes all equalities and inequalities.
   void clearConstraints();
 
   void print(raw_ostream &os) const;
   void dump() const;
 
 protected:
protected:
   /// i.e.,
   ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
   /// fast method (linear in the number of coefficients).
   void gcdTightenInequalities();
 
   /// Normalized each constraints by the GCD of its coefficients.
   void normalizeConstraintsByGCD();
 
   /// Removes identifiers in the column range [idStart, idLimit), and copies any
   /// remaining valid data into place, updates member variables, and resizes
   /// arrays as needed.
   virtual void removeIdRange(unsigned idStart, unsigned idLimit);
 
   /// Total number of identifiers.
   unsigned numIds;
 
   /// Number of identifiers corresponding to real dimensions.
   unsigned numDims;
 
   /// Number of identifiers corresponding to symbols (unknown but constant for
public:
   using FlatAffineConstraints::unionBoundingBox;
 
   /// Merge and align the identifiers of `this` and `other` starting at
   /// `offset`, so that both constraint systems get the union of the contained
   /// identifiers that is dimension-wise and symbol-wise unique; both
   /// constraint systems are updated so that they have the union of all
   /// identifiers, with `this`'s original identifiers appearing first followed
   /// by any of `other`'s identifiers that didn't appear in `this`. Local
-  /// identifiers of each system are by design separate/local and are placed
-  /// one after other (`this`'s followed by `other`'s).
+  /// identifiers in `other` that have the same division representation as local
+  /// identifiers in `this` are merged into one.
   //  E.g.: Input: `this`  has (%i, %j) [%M, %N]
   //               `other` has (%k, %j) [%P, %N, %M]
   //        Output: both `this`, `other` have (%i, %j, %k) [%M, %N, %P]
   //
   void mergeAndAlignIdsWithOther(unsigned offset,
                                  FlatAffineValueConstraints *other);
 
   /// Returns true if this constraint system and `other` are in the same

mlir/lib/Analysis/AffineStructures.cpp

static bool LLVM_ATTRIBUTE_UNUSED areIdsUnique(
   llvm_unreachable("Unexpected IdKind");
 }
 
 /// Merge and align the identifiers of A and B starting at 'offset', so that
 /// both constraint systems get the union of the contained identifiers that is
 /// dimension-wise and symbol-wise unique; both constraint systems are updated
 /// so that they have the union of all identifiers, with A's original
 /// identifiers appearing first followed by any of B's identifiers that didn't
-/// appear in A. Local identifiers of each system are by design separate/local
-/// and are placed one after other (A's followed by B's).
+/// appear in A. Local identifiers in B that have the same division
+/// representation as local identifiers in A are merged into one.
 //  E.g.: Input: A has ((%i, %j) [%M, %N]) and B has (%k, %j) [%P, %N, %M])
 //        Output: both A, B have (%i, %j, %k) [%M, %N, %P]
 static void mergeAndAlignIds(unsigned offset, FlatAffineValueConstraints *a,
                              FlatAffineValueConstraints *b) {
   assert(offset <= a->getNumDimIds() && offset <= b->getNumDimIds());
   // A merge/align isn't meaningful if a cst's ids aren't distinct.
   assert(areIdsUnique(*a) && "A's values aren't unique");
   assert(areIdsUnique(*b) && "B's values aren't unique");
 

 ///
 ///       4q >= i + j - 2                       <-- Lower bound for 'q'
 ///       4q <= i + j + 1                       <-- Upper bound for 'q'
 ///       expr = i + j + 1, divisor = 4
 ///       q = (i + j + 1) floordiv 4
 ///
 /// If successful, `expr` is set to dividend of the division and `divisor` is
 /// set to the denominator of the division.
 static LogicalResult getDivRepr(const FlatAffineConstraints &cst, unsigned pos,
                                 unsigned ubIneq, unsigned lbIneq,
                                 SmallVector<int64_t, 8> &expr,
                                 unsigned &divisor) {
 
   assert(pos <= cst.getNumIds() && "Invalid identifier position");
   assert(ubIneq <= cst.getNumInequalities() &&
          "Invalid upper bound inequality position");
   assert(lbIneq <= cst.getNumInequalities() &&
for (unsigned lbPos : lbIndices) {
 
       return std::make_pair(ubPos, lbPos);
     }
   }
 
   return llvm::None;
 }
 
 void FlatAffineConstraints::getLocalReprs(
     std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
   std::vector<SmallVector<int64_t, 8>> dividends(getNumLocalIds());
   SmallVector<unsigned, 4> denominators(getNumLocalIds());
   getLocalReprs(dividends, denominators, repr);
 }
 
 void FlatAffineConstraints::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,
     SmallVector<unsigned, 4> &denominators) const {
   std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
       getNumLocalIds());
   getLocalReprs(dividends, denominators, repr);
 }
 
 void FlatAffineConstraints::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,
     SmallVector<unsigned, 4> &denominators,
     std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
 
   repr.resize(getNumLocalIds());
   dividends.resize(getNumLocalIds());
   denominators.resize(getNumLocalIds());
 
   SmallVector<bool, 8> foundRepr(getNumIds(), false);
   for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i)
     foundRepr[i] = true;
 
   unsigned divOffset = getNumDimAndSymbolIds();
   bool changed;
   do {
     // Each time changed is true, at end of this iteration, one or more local
void FlatAffineConstraints::removeRedundantConstraints() {
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
     if (!(simplex.isMarkedRedundant(numIneqs + 2 * r) &&
           simplex.isMarkedRedundant(numIneqs + 2 * r + 1)))
       copyEquality(r, pos++);
   }
   equalities.resizeVertically(pos);
 }
 
-/// Merge local ids of `this` and `other`. This is done by appending local ids
-/// of `other` to `this` and inserting local ids of `this` to `other` at start
-/// of its local ids. Number of dimension and symbol ids should match in
-/// `this` and `other`.
+/// Eliminate `pos2^th` local identifier, replacing its every instance with
+/// `pos1^th` local identifier. This function is intended to be used to remove
+/// redundancy when local variables at position `pos1` and `pos2` are restricted
+/// to have the same value.
+static void eliminateRedundantLocalId(FlatAffineConstraints &fac, unsigned pos1,
+                                      unsigned pos2) {
+
+  assert(pos1 < fac.getNumLocalIds() && "Invalid local id position");
+  assert(pos2 < fac.getNumLocalIds() && "Invalid local id position");
+
+  unsigned localOffset = fac.getNumDimAndSymbolIds();
+  pos1 += localOffset;
+  pos2 += localOffset;
+  for (unsigned i = 0, e = fac.getNumInequalities(); i < e; ++i)
+    fac.atIneq(i, pos1) += fac.atIneq(i, pos2);
+  for (unsigned i = 0, e = fac.getNumEqualities(); i < e; ++i)
+    fac.atEq(i, pos1) += fac.atEq(i, pos2);
+  fac.removeId(pos2);
+}
+
+/// 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 FlatAffineConstraints::mergeLocalIds(FlatAffineConstraints &other) {
   assert(getNumDimIds() == other.getNumDimIds() &&
          "Number of dimension ids should match");
   assert(getNumSymbolIds() == other.getNumSymbolIds() &&
          "Number of symbol ids should match");
-  unsigned initLocals = getNumLocalIds();
-  insertLocalId(getNumLocalIds(), other.getNumLocalIds());
-  other.insertLocalId(0, initLocals);
+
+  FlatAffineConstraints &facA = *this;
+  FlatAffineConstraints &facB = other;
+
+  // Merge local ids of facA and facB without using division information,
+  // i.e. append local ids of `facB` to `facA` and insert local ids of `facA`
+  // to `facB` at start of its local ids.
+  unsigned initLocals = facA.getNumLocalIds();
+  insertLocalId(facA.getNumLocalIds(), facB.getNumLocalIds());
+  facB.insertLocalId(0, initLocals);
+
+  // Get division representations from each FAC.
+  std::vector<SmallVector<int64_t, 8>> divsA, divsB;
+  SmallVector<unsigned, 4> denomsA, denomsB;
+  facA.getLocalReprs(divsA, denomsA);
+  facB.getLocalReprs(divsB, denomsB);
+
+  // Copy division information for facB into `divsA` and `denomsA`, so that
+  // these have the combined division information of both FACs. Since newly
+  // added local variables in facA and facB 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);
+
+  // Find and merge duplicate divisions.
+  // TODO: Add division normalization to support divisions that differ by
+  // a constant.
+  // TODO: Add division ordering such that a division representation for local
+  // identifier at position `i` only depends on local identifiers at position <
+  // `i`. This would make sure that all divisions depending on other local
+  // variables that can be merged, are merged.
+  unsigned localOffset = getIdKindOffset(IdKind::Local);
+  for (unsigned i = 0; i < divsA.size(); ++i) {
+    // Check if a division representation exists for the `i^th` local id.
+    if (denomsA[i] == 0)
+      continue;
+    // Check if a division exists which is a duplicate of the division at `i`.
+    for (unsigned j = i + 1; j < divsA.size(); ++j) {
+      // Check if a division representation exists for the `j^th` local id.
+      if (denomsA[j] == 0)
+        continue;
+      // Check if the denominators match.
+      if (denomsA[i] != denomsA[j])
+        continue;
+      // Check if the representations are equal.
+      if (divsA[i] != divsA[j])
+        continue;
+
+      // Merge divisions at position `j` into division at position `i`.
+      eliminateRedundantLocalId(facA, i, j);
+      eliminateRedundantLocalId(facB, i, j);
+      for (unsigned k = 0, g = divsA.size(); k < g; ++k) {
+        SmallVector<int64_t, 8> &div = divsA[k];
+        if (denomsA[k] != 0) {
+          div[localOffset + i] += div[localOffset + j];
+          div.erase(div.begin() + localOffset + j);
+        }
+      }
+
+      divsA.erase(divsA.begin() + j);
+      denomsA.erase(denomsA.begin() + j);
+      // Since `j` can never be zero, we do not need to worry about overflows.
+      --j;
+    }
+  }
 }
 
 /// Removes local variables using equalities. Each equality is checked if it
 /// can be reduced to the form: `e = affine-expr`, where `e` is a local
 /// variable and `affine-expr` is an affine expression not containing `e`.
 /// If an equality satisfies this form, the local variable is replaced in
 /// each constraint and then removed. The equality used to replace this local
 /// variable is also removed.
 void FlatAffineConstraints::removeRedundantLocalVars() {
   // Normalize the equality constraints to reduce coefficients of local

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, simplifyLocalsTest) {
   FlatAffineConstraints fac3(1, 0, 1);
   fac3.addInequality({1, -1, -1});
   fac3.addInequality({0, 1, 1});
   fac3.addEquality({2, 1, 0});
 
   EXPECT_TRUE(fac3.isEmpty());
 }
 
+TEST(FlatAffineConstraintsTest, mergeDivisionsSimple) {
+  {
+    // (x) : (exists z, y  = [x / 2] : x = 3y and x + z + 1 >= 0).
+    FlatAffineConstraints fac1(1, 0, 1);
+    fac1.addLocalFloorDiv({1, 0, 0}, 2); // y = [x / 2].
+    fac1.addEquality({1, 0, -3, 0});     // x = 3y.
+    fac1.addInequality({1, 1, 0, 1});    // x + z + 1 >= 0.
+
+    // (x) : (exists y = [x / 2], z : x = 5y).
+    FlatAffineConstraints fac2(1);
+    fac2.addLocalFloorDiv({1, 0}, 2); // y = [x / 2].
+    fac2.addEquality({1, -5, 0});     // x = 5y.
+    fac2.appendLocalId();             // Add local id z.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 1 division should be matched + 2 unmatched local ids.
+    EXPECT_EQ(fac1.getNumLocalIds(), 3u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 3u);
+  }
+
+  {
+    // (x) : (exists z = [x / 5], y = [x / 2] : x = 3y).
+    FlatAffineConstraints fac1(1);
+    fac1.addLocalFloorDiv({1, 0}, 5);    // z = [x / 5].
+    fac1.addLocalFloorDiv({1, 0, 0}, 2); // y = [x / 2].
+    fac1.addEquality({1, 0, -3, 0});     // x = 3y.
+
+    // (x) : (exists y = [x / 2], z = [x / 5]: x = 5z).
+    FlatAffineConstraints fac2(1);
+    fac2.addLocalFloorDiv({1, 0}, 2);    // y = [x / 2].
+    fac2.addLocalFloorDiv({1, 0, 0}, 5); // z = [x / 5].
+    fac2.addEquality({1, 0, -5, 0});     // x = 5z.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 2 divisions should be matched.
+    EXPECT_EQ(fac1.getNumLocalIds(), 2u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 2u);
+  }
+}
+
+TEST(FlatAffineConstraintsTest, mergeDivisionsNestedDivsions) {
+  {
+    // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z >= x).
+    FlatAffineConstraints fac1(1);
+    fac1.addLocalFloorDiv({1, 0}, 2);    // y = [x / 2].
+    fac1.addLocalFloorDiv({1, 1, 0}, 3); // z = [x + y / 3].
+    fac1.addInequality({-1, 1, 1, 0});   // y + z >= x.
+
+    // (x) : (exists y = [x / 2], z = [x + y / 3]: y + z <= x).
+    FlatAffineConstraints fac2(1);
+    fac2.addLocalFloorDiv({1, 0}, 2);    // y = [x / 2].
+    fac2.addLocalFloorDiv({1, 1, 0}, 3); // z = [x + y / 3].
+    fac2.addInequality({1, -1, -1, 0});  // y + z <= x.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 2 divisions should be matched.
+    EXPECT_EQ(fac1.getNumLocalIds(), 2u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 2u);
+  }
+
+  {
+    // (x) : (exists y = [x / 2], z = [x + y / 3], w = [z + 1 / 5]: y + z >= x).
+    FlatAffineConstraints fac1(1);
+    fac1.addLocalFloorDiv({1, 0}, 2);       // y = [x / 2].
+    fac1.addLocalFloorDiv({1, 1, 0}, 3);    // z = [x + y / 3].
+    fac1.addLocalFloorDiv({0, 0, 1, 1}, 5); // w = [z + 1 / 5].
+    fac1.addInequality({-1, 1, 1, 0, 0});   // y + z >= x.
+
+    // (x) : (exists y = [x / 2], z = [x + y / 3], w = [z + 1 / 5]: y + z <= x).
+    FlatAffineConstraints fac2(1);
+    fac2.addLocalFloorDiv({1, 0}, 2);       // y = [x / 2].
+    fac2.addLocalFloorDiv({1, 1, 0}, 3);    // z = [x + y / 3].
+    fac2.addLocalFloorDiv({0, 0, 1, 1}, 5); // w = [z + 1 / 5].
+    fac2.addInequality({1, -1, -1, 0, 0});  // y + z <= x.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 3 divisions should be matched.
+    EXPECT_EQ(fac1.getNumLocalIds(), 3u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 3u);
+  }
+}
+
+TEST(FlatAffineConstraintsTest, mergeDivisionsConstants) {
+  {
+    // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z >= x).
+    FlatAffineConstraints fac1(1);
+    fac1.addLocalFloorDiv({1, 1}, 2);    // y = [x + 1 / 2].
+    fac1.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
+    fac1.addInequality({-1, 1, 1, 0});   // y + z >= x.
+
+    // (x) : (exists y = [x + 1 / 3], z = [x + 2 / 3]: y + z <= x).
+    FlatAffineConstraints fac2(1);
+    fac2.addLocalFloorDiv({1, 1}, 2);    // y = [x + 1 / 2].
+    fac2.addLocalFloorDiv({1, 0, 2}, 3); // z = [x + 2 / 3].
+    fac2.addInequality({1, -1, -1, 0});  // y + z <= x.
+
+    fac1.mergeLocalIds(fac2);
+
+    // Local space should be same.
+    EXPECT_EQ(fac1.getNumLocalIds(), fac2.getNumLocalIds());
+
+    // 2 divisions should be matched.
+    EXPECT_EQ(fac1.getNumLocalIds(), 2u);
+    EXPECT_EQ(fac2.getNumLocalIds(), 2u);
+  }
+}
+
 } // namespace mlir