[mlir][presburger] add functionality to compute local mod in IntegerRelation

[asraa]

Similar to IntegerRelation::addLocalFloorDiv, this adds a utility IntegerRelation::addLocalModulo that adds and returns a local variable that is the modulus of an affine function of the variables modulo some constant modulus. The function returns the absolute index of the new var in the relation.

This is computed by first finding the floordiv of exprs // modulus = q and then computing the remainder result = exprs - q * modulus.

[llvmbot]

@llvm/pr-subscribers-mlir-presburger

@llvm/pr-subscribers-mlir

Author: None (asraa)

Changes

Similar to IntegerRelation::addLocalFloorDiv, this adds a utility IntegerRelation::addLocalModulo that adds and returns a local variable that is the modulus of an affine function of the variables modulo some constant modulus. The function returns the absolute index of the new var in the relation.

This is computed by first finding the floordiv of exprs // modulus = q and then computing the remainder result = exprs - q * modulus.

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Analysis/Presburger/IntegerRelation.h (+16)
• (modified) mlir/lib/Analysis/Presburger/IntegerRelation.cpp (+20)
• (modified) mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp (+12)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index 4b1802413f75f..32dc629260277 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -485,6 +485,22 @@ class IntegerRelation {
     addLocalFloorDiv(getDynamicAPIntVec(dividend), DynamicAPInt(divisor));
   }
 
+  /// Adds a new local variable as the modulus of an affine function of other
+  /// variables, the coefficients of which are provided in `exprs`. The modulus is
+  /// with respect to a positive constant `modulus`. Two new local variables are
+  /// added to the system, one representing the floor div with respect to the
+  /// modulus and one representing the mod. Three constraints are added to the
+  /// system to capture the equivalance. The first two are required to compute the
+  /// result of the floor division `q`, and the third computes the equality
+  /// relation:
+  /// result = exprs - modulus * q
+  /// The function returns the absolute index of the new local variable
+  /// representing the result of the modulus operation.
+  unsigned addLocalModulo(ArrayRef<DynamicAPInt> exprs, const DynamicAPInt &modulus);
+  unsigned addLocalModulo(ArrayRef<int64_t> exprs, int64_t modulus) {
+    return addLocalModulo(getDynamicAPIntVec(exprs), DynamicAPInt(modulus));
+  }
+
   /// Projects out (aka eliminates) `num` variables starting at position
   /// `pos`. The resulting constraint system is the shadow along the dimensions
   /// that still exist. This method may not always be integer exact.
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 5c4d4d13580a0..72a164623d619 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -1515,6 +1515,26 @@ void IntegerRelation::addLocalFloorDiv(ArrayRef<DynamicAPInt> dividend,
       getDivUpperBound(dividendCopy, divisor, dividendCopy.size() - 2));
 }
 
+unsigned IntegerRelation::addLocalModulo(ArrayRef<DynamicAPInt> exprs,
+                                           const DynamicAPInt &modulus) {
+  assert(exprs.size() == getNumCols() && "incorrect exprs size");
+  assert(modulus > 0 && "positive modulus expected");
+
+  /// Add a local variable for q = expr floordiv modulus
+  addLocalFloorDiv(exprs, modulus);
+
+  /// Add a local var to represent the result
+  auto resultIndex = appendVar(VarKind::Local);
+
+  SmallVector<DynamicAPInt, 8> exprsCopy(exprs);
+  /// Insert the two new locals before the constant
+  /// Add locals that correspond to `q` and `result` to compute
+  /// 0 = (expr - modulus * q) - result
+  exprsCopy.insert(exprsCopy.end() - 1, {DynamicAPInt(-modulus), DynamicAPInt(-1)});
+  addEquality(exprsCopy);
+  return resultIndex;
+}
+
 int IntegerRelation::findEqualityToConstant(unsigned pos, bool symbolic) const {
   assert(pos < getNumVars() && "invalid position");
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index a6ed5c5b21e79..a1375ed05e5b0 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -714,3 +714,15 @@ TEST(IntegerRelationTest, getVarKindRange) {
   }
   EXPECT_THAT(actual, ElementsAre(2, 3, 4));
 }
+
+TEST(IntegerRelationTest, addLocalModulo) {
+  IntegerRelation rel =
+      parseRelationFromSet("(x) : (x >= 0, 100 - x >= 0)", 1);
+  unsigned result = rel.addLocalModulo({1, 0}, 32);  // x % 32
+  rel.convertVarKind(VarKind::Local,
+                     result - rel.getVarKindOffset(VarKind::Local),
+                     rel.getNumVarKind(VarKind::Local), VarKind::Range);
+  for (unsigned x = 0; x <= 100; ++x) {
+    EXPECT_TRUE(rel.containsPointNoLocal({x, x % 32}));
+  }
+}

[github-actions]

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

[j2kun]

⚠️ C/C++ code formatter, clang-format found issues in your code. ⚠️

You can test this locally with the following command:

git-clang-format --diff HEAD~1 HEAD --extensions cpp,h -- mlir/include/mlir/Analysis/Presburger/IntegerRelation.h mlir/lib/Analysis/Presburger/IntegerRelation.cpp mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

View the diff from clang-format here.

Please apply the formatter (I can't get my formatter to match the CI so I always just pipe the diff in the error message to git apply >_<)

[asraa]

Done!