[Mlir-commits] [mlir] [mlir][Presburger] Add IntegerRelation::reduceUsingEqualities helper (PR #235956)

[Jeremy Kun]

https://github.com/j2kun updated https://github.com/llvm/llvm-project/pull/235956

>From 7f0e7a326e8ce15ed62ad0f9ce11815dcd147576 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <jkun at google.com>
Date: Mon, 2 Nov 2026 19:10:32 -0800
Subject: [PATCH] [mlir][Presburger] Add IntegerRelation::reduceUsingEqualities
 helper

Add a protected helper that rewrites all constraints of a relation using
the equalities of another relation in the same space, without changing
the intersection of the two. This is a port of ISL's
isl_basic_set_reduce_using_equalities together with
project_onto_involved: the equalities are first projected onto the
variables that occur in the relation being reduced, so that no new
variables (in particular, no locals of the other relation) are
introduced, and each remaining equality is then used to eliminate its
last variable with a non-zero coefficient from every constraint, so that
later variables (locals) are eliminated in favor of earlier ones.

This is a building block for context-based simplification (gist) and is
exercised by the gist tests added in a follow-up commit, since there is
no existing pattern for unit-testing protected members of
IntegerRelation.
---
 .../Analysis/Presburger/IntegerRelation.h     |  17 +++
 .../Analysis/Presburger/IntegerRelation.cpp   | 108 ++++++++++++++++++
 2 files changed, 125 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index 2d995e24b8afec..3d1ae4c8358cbf 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -938,6 +938,23 @@ class IntegerRelation {
   void removeRedundantConstraintsFromSimplex(const Simplex &simplex,
                                              unsigned offset);
 
+  /// Rewrites every constraint of `this` using the equalities of `eqs`, which
+  /// must have the same space as `this` (including local variables). This does
+  /// not change the intersection of `this` with `eqs`, but may change `this`
+  /// outside of `eqs`. Inequalities of `eqs` are ignored.
+  ///
+  /// Only those equalities of `eqs` are used that can be expressed in terms of
+  /// the variables that already occur in some constraint of `this`; i.e., the
+  /// equalities of `eqs` are first projected onto the variables involved in
+  /// `this`, so that no new variables are introduced into `this`. Each
+  /// remaining equality is used to eliminate its last variable with a non-zero
+  /// coefficient from every constraint of `this`, so later variables (in
+  /// particular, locals) are eliminated in favor of earlier ones.
+  ///
+  /// Constraints that become trivially true are removed; if `eqs` is found to
+  /// be inconsistent, `this` becomes empty.
+  void reduceUsingEqualities(const IntegerRelation &eqs);
+
   /// Returns false if the fields corresponding to various variable counts, or
   /// equality/inequality buffer sizes aren't consistent; true otherwise. This
   /// is meant to be used within an assert internally.
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 88012b846133c2..e68cb0c5a42657 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -1263,6 +1263,114 @@ void IntegerRelation::removeRedundantConstraintsFromSimplex(
   equalities.resizeVertically(pos);
 }
 
+/// Eliminates the variable at column `col` from `row` using the equality `eq`,
+/// i.e., replaces `row` by `a * row + b * eq` with `a > 0` chosen such that the
+/// coefficient of `col` becomes zero. Since `a` is positive, this preserves
+/// the meaning of `row` as an inequality (or equality) wherever `eq` holds.
+static void eliminateColumnUsingEquality(MutableArrayRef<DynamicAPInt> row,
+                                         ArrayRef<DynamicAPInt> eq,
+                                         unsigned col) {
+  assert(row.size() == eq.size() && "row and equality must have same size");
+  assert(eq[col] != 0 && "equality must involve the eliminated column");
+  if (row[col] == 0)
+    return;
+  DynamicAPInt lcm = llvm::lcm(abs(row[col]), abs(eq[col]));
+  DynamicAPInt rowMultiplier = lcm / abs(row[col]);
+  DynamicAPInt eqMultiplier = lcm / abs(eq[col]);
+  if ((row[col] > 0) == (eq[col] > 0))
+    eqMultiplier = -eqMultiplier;
+  for (unsigned j = 0, e = row.size(); j < e; ++j)
+    row[j] = rowMultiplier * row[j] + eqMultiplier * eq[j];
+  assert(row[col] == 0 && "column should have been eliminated");
+}
+
+/// Brings the equality rows of `eqs` that are not marked in `used` into a
+/// reduced form with respect to the columns in `cols`, processed in order: for
+/// each column, the first unused row with a non-zero coefficient in it becomes
+/// the pivot row of that column, the column is eliminated from all other rows,
+/// and the row is marked used. The pivot rows and columns are appended to
+/// `pivotRows` and `pivotCols`. Afterwards, no row other than its pivot row has
+/// a non-zero coefficient in any pivot column.
+static void reduceEqualityRows(IntMatrix &eqs, ArrayRef<unsigned> cols,
+                               llvm::SmallBitVector &used,
+                               SmallVectorImpl<unsigned> &pivotRows,
+                               SmallVectorImpl<unsigned> &pivotCols) {
+  for (unsigned col : cols) {
+    std::optional<unsigned> pivotRow;
+    for (unsigned r = 0, e = eqs.getNumRows(); r < e; ++r) {
+      if (!used[r] && eqs(r, col) != 0) {
+        pivotRow = r;
+        break;
+      }
+    }
+    if (!pivotRow)
+      continue;
+    // Copy the pivot row, since the matrix rows are modified below.
+    SmallVector<DynamicAPInt, 8> pivot(eqs.getRow(*pivotRow));
+    for (unsigned r = 0, e = eqs.getNumRows(); r < e; ++r) {
+      if (r == *pivotRow)
+        continue;
+      eliminateColumnUsingEquality(eqs.getRow(r), pivot, col);
+      eqs.normalizeRow(r);
+    }
+    used[*pivotRow] = true;
+    pivotRows.emplace_back(*pivotRow);
+    pivotCols.emplace_back(col);
+  }
+}
+
+void IntegerRelation::reduceUsingEqualities(const IntegerRelation &eqs) {
+  assert(space.isCompatible(eqs.getSpace()) &&
+         getNumLocalVars() == eqs.getNumLocalVars() &&
+         "spaces (including locals) must match");
+  if (eqs.getNumEqualities() == 0)
+    return;
+
+  unsigned numVars = getNumVars();
+  unsigned numCols = getNumCols();
+
+  // Work on a copy of the equalities of `eqs`. We bring them into a reduced
+  // row echelon form where the pivot of each row is its last variable with a
+  // non-zero coefficient, and no other row has a non-zero coefficient in that
+  // column. Variables that do not occur in `this` are eliminated first and
+  // their pivot rows are dropped (projecting the equalities onto the variables
+  // involved in `this`), so that reducing `this` never introduces them.
+  IntMatrix ctxEqs = eqs.equalities;
+  llvm::SmallBitVector rowUsed(ctxEqs.getNumRows());
+  SmallVector<unsigned, 8> uninvolvedCols, involvedCols;
+  for (unsigned col : llvm::reverse(llvm::seq<unsigned>(0, numVars)))
+    (isColZero(col) ? uninvolvedCols : involvedCols).emplace_back(col);
+
+  SmallVector<unsigned, 8> pivotRows, pivotCols;
+  reduceEqualityRows(ctxEqs, uninvolvedCols, rowUsed, pivotRows, pivotCols);
+  // Drop the pivot rows of the uninvolved variables.
+  pivotRows.clear();
+  pivotCols.clear();
+  reduceEqualityRows(ctxEqs, involvedCols, rowUsed, pivotRows, pivotCols);
+
+  // Any unused row now has all-zero variable coefficients; a non-zero constant
+  // means the equalities of `eqs` are inconsistent.
+  for (unsigned r = 0, e = ctxEqs.getNumRows(); r < e; ++r) {
+    if (!rowUsed[r] && ctxEqs(r, numCols - 1) != 0) {
+      *this = getEmpty(getSpace());
+      return;
+    }
+  }
+
+  // Eliminate each pivot column from every constraint of `this`. Since the
+  // pivot rows are in reduced form, no pivot column is reintroduced.
+  for (auto [pivotRow, pivotCol] : llvm::zip_equal(pivotRows, pivotCols)) {
+    ArrayRef<DynamicAPInt> pivot = ctxEqs.getRow(pivotRow);
+    for (unsigned i = 0, e = getNumEqualities(); i < e; ++i)
+      eliminateColumnUsingEquality(equalities.getRow(i), pivot, pivotCol);
+    for (unsigned i = 0, e = getNumInequalities(); i < e; ++i)
+      eliminateColumnUsingEquality(inequalities.getRow(i), pivot, pivotCol);
+  }
+
+  removeTrivialRedundancy();
+  removeTrivialEqualities();
+}
+
 std::optional<DynamicAPInt> IntegerRelation::computeVolume() const {
   assert(getNumSymbolVars() == 0 && "Symbols are not yet supported!");