[MLIR][Presburger] optimize bound computation by pruning orthogonal constraints

[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

[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!

[cxy-1993]

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!

I have now added detailed information regarding the function's mechanism and usage. Please let me know if you feel anything further is needed. Thank you for your time.

[cxy-1993]

Could you please take some time to review this patch now? Thank you very much. @Superty

[Superty]

Thanks for your contribution and patience. I left a review!

(I removed the mention of "compilation time" from the title, I think we should reserve that phrase for the compilation time of the library)

[cxy-1993]

Thanks for your contribution and patience. I left a review!

(I removed the mention of "compilation time" from the title, I think we should reserve that phrase for the compilation time of the library)

Thanks for the valuable review. I've made the requested changes and left a comment about a duplicated section for discussion: While there is some duplication, the code block is quite small. Abstracting it into a lambda would involve too many parameters, which I think would make the code less readable. I prefer to leave it as is—what do you think? @Superty

[Superty]

Thanks for addressing the comments so quickly. Sorry, I should have explained what I meant a bit more clearly. I pushed a commit adding my proposed function (00f3410). Can you refactor your code to use this? Thanks!

[cxy-1993]

Thanks for addressing the comments so quickly. Sorry, I should have explained what I meant a bit more clearly. I pushed a commit adding my proposed function (00f3410). Can you refactor your code to use this? Thanks!

Excellent suggestion, I have made the modifications (including adding removeConstraint function).

[cxy-1993]

Thanks for addressing the comments so quickly. Sorry, I should have explained what I meant a bit more clearly. I pushed a commit adding my proposed function (00f3410). Can you refactor your code to use this? Thanks!

I rebased your commit and fixed the compilation errors. Could you review it for any further improvements? @Superty

[Superty]

PR will automerge once checks pass. Thanks again for the contribution!