[Mlir-commits] [mlir] [MLIR][Presburger] optimize bound computation by pruning orthogonal constraints (PR #164199)

[donald chen]

================
@@ -1742,12 +1751,65 @@ std::optional<DynamicAPInt> IntegerRelation::getConstantBoundOnDimSize(
   return minDiff;
 }
 
+void IntegerRelation::pruneOrthogonalConstraints(unsigned pos) {
+  llvm::DenseSet<unsigned> relatedCols({pos}), relatedRows;
+
+  // Early exit if constraints is empty.
+  unsigned numConstraints = getNumConstraints();
+  if (numConstraints == 0)
+    return;
+
+  llvm::SmallVector<unsigned> rowStack, colStack({pos});
+  // The following code performs a graph traversal, starting from the target
+  // variable, to identify all variables(recorded in relatedCols) and
+  // constraints (recorded in relatedRows) belonging to the same connected
+  // component.
+  while (!rowStack.empty() || !colStack.empty()) {
+    if (!rowStack.empty()) {
+      unsigned currentRow = rowStack.pop_back_val();
+      // Push all variable that accociated to this constraints to relatedCols
+      // and colStack.
+      for (unsigned colIndex = 0; colIndex < getNumVars(); ++colIndex) {
+        if (atConstraint(currentRow, colIndex) != 0 &&
+            relatedCols.insert(colIndex).second) {
+          colStack.push_back(colIndex);
+        }
+      }
+    } else {
+      unsigned currentCol = colStack.pop_back_val();
+      // Push all constraints that are associated with this variable to related
+      // rows and the row stack.
+      for (unsigned rowIndex = 0; rowIndex < numConstraints; ++rowIndex) {
+        if (atConstraint(rowIndex, currentCol) != 0 &&
+            relatedRows.insert(rowIndex).second) {
+          rowStack.push_back(rowIndex);
+        }
+      }
+    }
+  }
+
+  // Prune all constraints not related to target variable.
+  for (int constraintId = numConstraints - 1; constraintId >= 0;
+       --constraintId) {
+    if (!relatedRows.contains(constraintId)) {
+      removeConstraint((unsigned)constraintId);
+    }
----------------
cxy-1993 wrote:

Done

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