[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

480

variables -> variable

mlir/lib/Analysis/AffineStructures.cpp

1794

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.

1798

unsigned

1798

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

1798

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.

1798

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

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1791

variables -> variable

1821–1822

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

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1799–1821

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

1799–1821

Oh, Sorry. I misunderstood something. Thanks.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

1803

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

1808

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

1808

Code comment here.

[Groverkss]

• Addressed bondhugula's comments

[Groverkss]

Addressed comments.

[bondhugula]

mlir/lib/Analysis/AffineStructures.cpp

945

Missing period at the end.

950

Reflow: wse the width.

1821

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.
+  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.
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);
 }
 
+/// 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.
+void FlatAffineConstraints::removeRedundantLocalVars() {
+  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;
+
+      // Break if local variable found
+      if (j < f)
+        break;
+    }
+
+    if (i >= e)
+      break;
+
+    // Use the ith equality to simplify other constraints
+    for (unsigned k = 0, t = getNumEqualities(); k < t; ++k)
+      eliminateFromConstraint(this, k, i, j, j, /*isEq=*/true);
+    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);
+    normalizeConstraintsByGCD();
+  }
+}
+
 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());
+}
+
 } // namespace mlir