[MLIR] Add functionality to remove redundant local variables

[Groverkss]

Summary

This patch adds functionality to FlatAffineConstraints to remove local
variables using equalities. This helps in keeping output representation of
FlatAffineConstraints smaller.

This patch is part of a series of patches aimed at generalizing affine
dependence analysis.

[bondhugula]

mlir/include/mlir/Analysis/AffineStructures.h

485

variables -> variable

mlir/lib/Analysis/AffineStructures.cpp

1802

In the form this is being added, this would be dead code as it isn't being used anywhere. Please update a path to use it so that this gets tested if there is no separate test case just for this.

1806

unsigned

1806

Use i = 0, e = ... form.

[Groverkss]

• Addressed bondhugula's comments

Used removeRedundantLocalVars to improve isEmpty coverage. Added a test that
would not work without simplification.

[Groverkss]

Addressed comments

mlir/lib/Analysis/AffineStructures.cpp

1806

The reason for not using unsigned is that there is a --i inside the loop. If i = 0 and i is unsigned, it will overflow.

1806

The reason for not using this form is that an equality may be removed inside the loop.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1799

variables -> variable

1829–1830

Why do you need to decrement here if you are going to anyway break out below.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1807–1829

This whole loop's logic can be simplified/broken apart for better readability.

For eg., invert the check:

for ...
  if (foundOne)
    break;

Then, have a block below outside of the for loop.

if (i < getNumEqualities()) {
  ...
}

getNumEqualities() in the for loop header isn't changing at all. You are breaking out when you need to remove an equality.

[Groverkss]

• Simplified loop in removeRedundantLocalVars

[Groverkss]

Addressed comments.

mlir/lib/Analysis/AffineStructures.cpp

1807–1829

Oh, Sorry. I misunderstood something. Thanks.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1811

Please terminate all comments with a full stop - here and above. They should all be properly punctuated per LLVM/MLIR style.

1816

This is again confusing. i can't be greater than e. i == e here.

1816

Code comment here.

[Groverkss]

• Addressed bondhugula's comments

[Groverkss]

Addressed comments.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

954

Missing period at the end.

959

Reflow: wse the width.

1829

Do we need this normalization here? If yes, can you check the call sites to make sure it's not being redundantly called? This is doing some more basic simplification beyond just removing local vars -- it's fine to do as long as it's trivial and inexpensive to be considered more like a cleanup. Best to add a comment here or on the function.

mlir/unittests/Analysis/AffineStructuresTest.cpp

772–778

It'd be good to a couple more test cases to exercise it better.

773

Full stop at the end.

[Groverkss]

• Addressed comments

[Groverkss]

Addressed @bondhugula's comments.

[Groverkss]

• Fix incomplete comment

[Groverkss]

• Fix comment typos

mlir/include/mlir/Analysis/AffineStructures.h

protected:
 
   /// Eliminates a single identifier at `position` from equality and inequality
   /// constraints. Returns `success` if the identifier was eliminated, and
   /// `failure` otherwise.
   inline LogicalResult gaussianEliminateId(unsigned position) {
     return success(gaussianEliminateIds(position, position + 1) == 1);
   }
 
+  /// 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 removeRedundantLocalVars();
+
   /// Eliminates identifiers from equality and inequality constraints
   /// in column range [posStart, posLimit).
   /// Returns the number of variables eliminated.
   unsigned gaussianEliminateIds(unsigned posStart, unsigned posLimit);
 
   /// Eliminates the identifier at the specified position using Fourier-Motzkin
   /// variable elimination, but uses Gaussian elimination if there is an
   /// equality involving that identifier. If the result of the elimination is

mlir/lib/Analysis/AffineStructures.cpp

 // Checks for emptiness of the set by eliminating identifiers successively and
 // using the GCD test (on all equality constraints) and checking for trivially
 // invalid constraints. Returns 'true' if the constraint system is found to be
 // empty; false otherwise.
 bool FlatAffineConstraints::isEmpty() const {
   if (isEmptyByGCDTest() || hasInvalidConstraint())
     return true;
 
-  // First, eliminate as many identifiers as possible using Gaussian
-  // elimination.
   FlatAffineConstraints tmpCst(*this);
+
+  // First, eliminate as many local variables as possible using equalities.
+  tmpCst.removeRedundantLocalVars();
+  if (tmpCst.isEmptyByGCDTest() || tmpCst.hasInvalidConstraint())
+    return true;
+
+  // Eliminate as many identifiers as possible using Gaussian elimination.
   unsigned currentPos = 0;
   while (currentPos < tmpCst.getNumIds()) {
     tmpCst.gaussianEliminateIds(currentPos, tmpCst.getNumIds());
     ++currentPos;
     // We check emptiness through trivial checks after eliminating each ID to
     // detect emptiness early. Since the checks isEmptyByGCDTest() and
     // hasInvalidConstraint() are linear time and single sweep on the constraint
     // buffer, this appears reasonable - but can optimize in the future.
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.
 void FlatAffineConstraints::mergeLocalIds(FlatAffineConstraints &other) {
   unsigned initLocals = getNumLocalIds();
   insertLocalId(getNumLocalIds(), other.getNumLocalIds());
   other.insertLocalId(0, initLocals);
 }
 
+/// 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
+  // variables to 1 wherever possible.
+  for (unsigned i = 0, e = getNumEqualities(); i < e; ++i)
+    normalizeConstraintByGCD</*isEq=*/true>(this, i);
+
+  while (true) {
+    unsigned i, e, j, f;
+    for (i = 0, e = getNumEqualities(); i < e; ++i) {
+      // Find a local variable to eliminate using ith equality.
+      for (j = getNumDimAndSymbolIds(), f = getNumIds(); j < f; ++j)
+        if (std::abs(atEq(i, j)) == 1)
+          break;
+
+      // Local variable can be eliminated using ith equality.
+      if (j < f)
+        break;
+    }
+
+    // No equality can be used to eliminate a local variable.
+    if (i == e)
+      break;
+
+    // Use the ith equality to simplify other equalities. If any changes
+    // are made to an equality constraint, it is normalized by GCD.
+    for (unsigned k = 0, t = getNumEqualities(); k < t; ++k) {
+      if (atEq(k, j) != 0) {
+        eliminateFromConstraint(this, k, i, j, j, /*isEq=*/true);
+        normalizeConstraintByGCD</*isEq=*/true>(this, k);
+      }
+    }
+
+    // Use the ith equality to simplify inequalities.
+    for (unsigned k = 0, t = getNumInequalities(); k < t; ++k)
+      eliminateFromConstraint(this, k, i, j, j, /*isEq=*/false);
+
+    // Remove the ith equality and the found local variable.
+    removeId(j);
+    removeEquality(i);
+  }
+}
+
 std::pair<AffineMap, AffineMap> FlatAffineConstraints::getLowerAndUpperBound(
     unsigned pos, unsigned offset, unsigned num, unsigned symStartPos,
     ArrayRef<AffineExpr> localExprs, MLIRContext *context) const {
   assert(pos + offset < getNumDimIds() && "invalid dim start pos");
   assert(symStartPos >= (pos + offset) && "invalid sym start pos");
   assert(getNumLocalIds() == localExprs.size() &&
          "incorrect local exprs count");
 

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, removeIdRange) {
 
   fac.removeIdRange(FlatAffineConstraints::IdKind::Local, 1, 1);
   EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 30, 40));
 
   fac.removeIdRange(FlatAffineConstraints::IdKind::Local, 0, 1);
   EXPECT_THAT(fac.getInequality(0), testing::ElementsAre(12, 20, 40));
 }
 
+TEST(FlatAffineConstraintsTest, simplifyLocalsTest) {
+  // (x) : (exists y: 2x + y = 1 and y = 2).
+  FlatAffineConstraints fac(1, 0, 1);
+  fac.addEquality({2, 1, -1});
+  fac.addEquality({0, 1, -2});
+
+  EXPECT_TRUE(fac.isEmpty());
+
+  // (x) : (exists y, z, w: 3x + y = 1 and 2y = z and 3y = w and z = w).
+  FlatAffineConstraints fac2(1, 0, 3);
+  fac2.addEquality({3, 1, 0, 0, -1});
+  fac2.addEquality({0, 2, -1, 0, 0});
+  fac2.addEquality({0, 3, 0, -1, 0});
+  fac2.addEquality({0, 0, 1, -1, 0});
+
+  EXPECT_TRUE(fac2.isEmpty());
+
+  // (x) : (exists y: x >= y + 1 and 2x + y = 0 and y >= -1).
+  FlatAffineConstraints fac3(1, 0, 1);
+  fac3.addInequality({1, -1, -1});
+  fac3.addInequality({0, 1, 1});
+  fac3.addEquality({2, 1, 0});
+
+  EXPECT_TRUE(fac3.isEmpty());
+}
+
 } // namespace mlir