[mlir][presburger] Optimize the compilation time for calculating bounds of an Integer Relation

[cxy-1993]

IntegerRelation uses Fourier-Motzkin elimination and Gaussian elimination to simplify constraints. These methods may repeatedly perform calculations and elimination on irrelevant variables. Preemptively eliminating irrelevant variables and their associated constraints can speed up up the calculation process.

[llvmbot]

@llvm/pr-subscribers-mlir-presburger

@llvm/pr-subscribers-mlir

Author: donald chen (cxy-1993)

Changes

IntegerRelation uses Fourier-Motzkin elimination and Gaussian elimination to simplify constraints. These methods may repeatedly perform calculations and elimination on irrelevant variables. Preemptively eliminating irrelevant variables and their associated constraints can speed up up the calculation process.

---

Full diff: https://github.com/llvm/llvm-project/pull/164199.diff

2 Files Affected:

• (modified) mlir/include/mlir/Analysis/Presburger/IntegerRelation.h (+3)
• (modified) mlir/lib/Analysis/Presburger/IntegerRelation.cpp (+55)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index f86535740fec9..026d84529edfb 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -511,6 +511,9 @@ class IntegerRelation {
   void projectOut(unsigned pos, unsigned num);
   inline void projectOut(unsigned pos) { return projectOut(pos, 1); }
 
+  /// Prune constraints that are irrelevant to the target variable.
+  void pruneConstraints(unsigned pos);
+
   /// Tries to fold the specified variable to a constant using a trivial
   /// equality detection; if successful, the constant is substituted for the
   /// variable everywhere in the constraint system and then removed from the
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 0dcdd5bb97bc8..b05e872323aff 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -21,6 +21,7 @@
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/Sequence.h"
 #include "llvm/ADT/SmallBitVector.h"
@@ -1723,12 +1724,66 @@ std::optional<DynamicAPInt> IntegerRelation::getConstantBoundOnDimSize(
   return minDiff;
 }
 
+void IntegerRelation::pruneConstraints(unsigned pos) {
+  llvm::DenseSet<unsigned> relatedCols({pos}), relatedRows;
+
+  llvm::SmallVector<unsigned> rowStack, colStack({pos});
+  unsigned numConstraints = getNumConstraints();
+  if (numConstraints == 0) return;
+  while (!rowStack.empty() || !colStack.empty()) {
+    if (!rowStack.empty()) {
+      unsigned currentRow = rowStack.pop_back_val();
+      for (uint64_t colIndex = 0; colIndex < getNumVars(); ++colIndex) {
+        if (currentRow < getNumInequalities()) {
+          if (atIneq(currentRow, colIndex) != 0 &&
+              relatedCols.insert(colIndex).second) {
+            colStack.push_back(colIndex);
+          }
+        } else {
+          if (atEq(currentRow - getNumInequalities(), colIndex) != 0 &&
+              relatedCols.insert(colIndex).second) {
+            colStack.push_back(colIndex);
+          }
+        }
+      }
+    } else {
+      unsigned currentCol = colStack.pop_back_val();
+      for (uint64_t rowIndex = 0; rowIndex < numConstraints; ++rowIndex) {
+        if (rowIndex < getNumInequalities()) {
+          if (atIneq(rowIndex, currentCol) != 0 &&
+              relatedRows.insert(rowIndex).second) {
+            rowStack.push_back(rowIndex);
+          }
+        } else {
+          if (atEq(rowIndex - getNumInequalities(), currentCol) != 0 &&
+              relatedRows.insert(rowIndex).second) {
+            rowStack.push_back(rowIndex);
+          }
+        }
+      }
+    }
+  }
+
+  for (int64_t constraintId = numConstraints - 1; constraintId >= 0;
+       --constraintId) {
+    if (!relatedRows.contains(constraintId)) {
+      if (constraintId >= getNumInequalities()) {
+        removeEquality(constraintId - getNumInequalities());
+      } else {
+        removeInequality(constraintId);
+      }
+    }
+  }
+}
+
 template <bool isLower>
 std::optional<DynamicAPInt>
 IntegerRelation::computeConstantLowerOrUpperBound(unsigned pos) {
   assert(pos < getNumVars() && "invalid position");
   // Project to 'pos'.
+  pruneConstraints(pos);
   projectOut(0, pos);
+  pruneConstraints(0);
   projectOut(1, getNumVars() - 1);
   // Check if there's an equality equating the '0'^th variable to a constant.
   int eqRowIdx = findEqualityToConstant(/*pos=*/0, /*symbolic=*/false);

[github-actions]

✅ With the latest revision this PR passed the C/C++ code formatter.

[cxy-1993]

ping

[Groverkss]

@Superty This seems like something we already have implemented somewhere? Can you review this PR?

[Superty]

Off the top of my head I don't remember such a thing

[Superty]

Hi @cxy-1993, sorry about the delay here. In "Prune constraints that are irrelevant to the target variable.", how is irrelevant defined?

[cxy-1993]

Hi @cxy-1993, sorry about the delay here. In "Prune constraints that are irrelevant to the target variable.", how is irrelevant defined?

In my opinion, the definition of variable irrelevance is as follows: If variables are considered as nodes and equations/inequalities as edges between nodes, a set of constraints forms an undirected graph. Variables in different connected components are irrelevant. This patch is about finding all nodes within the same connected component and then removing all nodes and related edges (equations and inequalities) that are not in the same connected component.

[cxy-1993]

ping @Superty

[cxy-1993]

ping @Superty

[Superty]

I probably won't be able to look at this for a few more weeks, sorry. I would suggest adding some more documentation explaining what irrelevant means though. Is there a geometric interpretation of the definition? Maybe it makes sense to explain it with reference to Fourier motzkin?

The implementation could also use some more documentation. Thanks!