[Mlir-commits] [mlir] [mlir] [presburger] Add IntegerRelation::rangeProduct (PR #148092)

Jeremy Kun llvmlistbot at llvm.org
Tue Jul 15 20:55:54 PDT 2025


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

>From 7e4bd749d57c70abced437a05e7d7d76f83a69b2 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun at users.noreply.github.com>
Date: Thu, 10 Jul 2025 17:24:50 -0700
Subject: [PATCH 1/5] Add IntegerRelation::rangeProduct

This is intended to match `isl::map`'s `flat_range_product`.
---
 .../Analysis/Presburger/IntegerRelation.h     | 15 ++++++++
 .../Analysis/Presburger/IntegerRelation.cpp   | 35 +++++++++++++++++++
 .../Presburger/IntegerRelationTest.cpp        | 14 ++++++++
 3 files changed, 64 insertions(+)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index b68262f09f485..ae213181f2c3b 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -707,6 +707,21 @@ class IntegerRelation {
   /// this for uniformity with `applyDomain`.
   void applyRange(const IntegerRelation &rel);
 
+  /// Let the relation `this` be R1, and the relation `rel` be R2. Requires
+  /// R1 and R2 to have the same domain.
+  ///
+  /// This operation computes the relation whose domain is the same as R1 and
+  /// whose range is the product of the ranges of R1 and R2, and whose
+  /// constraints are the conjunction of the constraints of R1 and R2 applied
+  /// to the relevant subspaces of the range.
+  ///
+  /// Example:
+  ///
+  /// R1: (i, j) -> k : f(i, j, k) = 0
+  /// R2: (i, j) -> l : g(i, j, l) = 0
+  /// R1.rangeProduct(R2): (i, j) -> (k, l) : f(i, j, k) = 0 and g(i, j, l) = 0
+  IntegerRelation rangeProduct(const IntegerRelation &rel);
+
   /// Given a relation `other: (A -> B)`, this operation merges the symbol and
   /// local variables and then takes the composition of `other` on `this: (B ->
   /// C)`. The resulting relation represents tuples of the form: `A -> C`.
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 631e085574fd0..87cec81e86b59 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -2481,6 +2481,41 @@ void IntegerRelation::applyDomain(const IntegerRelation &rel) {
 
 void IntegerRelation::applyRange(const IntegerRelation &rel) { compose(rel); }
 
+IntegerRelation IntegerRelation::rangeProduct(const IntegerRelation &rel) {
+  /// R1: (i, j) -> k : f(i, j, k) = 0
+  /// R2: (i, j) -> l : g(i, j, l) = 0
+  /// R1.rangeProduct(R2): (i, j) -> (k, l) : f(i, j, k) = 0 and g(i, j, l) = 0
+  assert(getNumDomainVars() == rel.getNumDomainVars() &&
+         "Range product is only defined for relations with equal domains");
+
+  // explicit copy of the context relation
+  IntegerRelation result = *this;
+  unsigned srcOffset = getVarKindOffset(VarKind::Range);
+  unsigned newNumRangeVars = rel.getNumRangeVars();
+
+  result.appendVar(VarKind::Range, newNumRangeVars);
+
+  for (unsigned i = 0; i < rel.getNumEqualities(); ++i) {
+    // Add a new equality that uses the new range variables.
+    // The old equality is a list of coefficients of the variables
+    // from `rel`, and so the range variables need to be shifted
+    // right by the number of range variables added to `result`.
+    SmallVector<DynamicAPInt> copy =
+        SmallVector<DynamicAPInt>(rel.getEquality(i));
+    copy.insert(copy.begin() + srcOffset, newNumRangeVars, DynamicAPInt(0));
+    result.addEquality(copy);
+  }
+
+  for (unsigned i = 0; i < rel.getNumInequalities(); ++i) {
+    SmallVector<DynamicAPInt> copy =
+        SmallVector<DynamicAPInt>(rel.getInequality(i));
+    copy.insert(copy.begin() + srcOffset, newNumRangeVars, DynamicAPInt(0));
+    result.addInequality(copy);
+  }
+
+  return result;
+}
+
 void IntegerRelation::printSpace(raw_ostream &os) const {
   space.print(os);
   os << getNumConstraints() << " constraints\n";
diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index 7df500bc9568a..dd8b9e3f03330 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -608,3 +608,17 @@ TEST(IntegerRelationTest, convertVarKindToLocal) {
   EXPECT_EQ(space.getId(VarKind::Symbol, 0), Identifier(&identifiers[3]));
   EXPECT_EQ(space.getId(VarKind::Symbol, 1), Identifier(&identifiers[4]));
 }
+
+TEST(IntegerRelationTest, rangeProduct) {
+  IntegerRelation r1 = parseRelationFromSet(
+      "(i, j, k) : (2*i + 3*k == 0, i >= 0, j >= 0, k >= 0)", 2);
+  IntegerRelation r2 = parseRelationFromSet(
+      "(i, j, l) : (4*i + 6*j + 9*l == 0, i >= 0, j >= 0, l >= 0)", 2);
+
+  IntegerRelation rangeProd = r1.rangeProduct(r2);
+  IntegerRelation expected = parseRelationFromSet(
+      "(i, j, k, l) : (2*i + 3*k == 0, 4*i + 6*j + 9*l == 0, i >= 0, j >= 0, k >= 0, l >= 0)", 2);
+
+  EXPECT_TRUE(expected.isEqual(rangeProd));
+}
+

>From 3600fb63abbc974e96b14c0fa36276b5aba86ec1 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun at users.noreply.github.com>
Date: Tue, 15 Jul 2025 13:43:15 -0700
Subject: [PATCH 2/5] formatting

---
 mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index dd8b9e3f03330..db24600111cee 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -616,9 +616,10 @@ TEST(IntegerRelationTest, rangeProduct) {
       "(i, j, l) : (4*i + 6*j + 9*l == 0, i >= 0, j >= 0, l >= 0)", 2);
 
   IntegerRelation rangeProd = r1.rangeProduct(r2);
-  IntegerRelation expected = parseRelationFromSet(
-      "(i, j, k, l) : (2*i + 3*k == 0, 4*i + 6*j + 9*l == 0, i >= 0, j >= 0, k >= 0, l >= 0)", 2);
+  IntegerRelation expected =
+      parseRelationFromSet("(i, j, k, l) : (2*i + 3*k == 0, 4*i + 6*j + 9*l == "
+                           "0, i >= 0, j >= 0, k >= 0, l >= 0)",
+                           2);
 
   EXPECT_TRUE(expected.isEqual(rangeProd));
 }
-

>From bfd023a50f87c790a41c2d8e08770e2d01b177e3 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun at users.noreply.github.com>
Date: Tue, 15 Jul 2025 13:51:07 -0700
Subject: [PATCH 3/5] address review comments

---
 .../include/mlir/Analysis/Presburger/IntegerRelation.h |  6 ++----
 mlir/lib/Analysis/Presburger/IntegerRelation.cpp       | 10 +++++-----
 2 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index ae213181f2c3b..ee401cca8f552 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -710,10 +710,8 @@ class IntegerRelation {
   /// Let the relation `this` be R1, and the relation `rel` be R2. Requires
   /// R1 and R2 to have the same domain.
   ///
-  /// This operation computes the relation whose domain is the same as R1 and
-  /// whose range is the product of the ranges of R1 and R2, and whose
-  /// constraints are the conjunction of the constraints of R1 and R2 applied
-  /// to the relevant subspaces of the range.
+  /// Let R3 be the rangeProduct of R1 and R2. Then x R3 (y, z) iff
+  /// (x R1 y and x R2 z).
   ///
   /// Example:
   ///
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 87cec81e86b59..ea207069852b4 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -2488,12 +2488,12 @@ IntegerRelation IntegerRelation::rangeProduct(const IntegerRelation &rel) {
   assert(getNumDomainVars() == rel.getNumDomainVars() &&
          "Range product is only defined for relations with equal domains");
 
-  // explicit copy of the context relation
+  // explicit copy of `this`
   IntegerRelation result = *this;
   unsigned srcOffset = getVarKindOffset(VarKind::Range);
-  unsigned newNumRangeVars = rel.getNumRangeVars();
+  unsigned numRelRangeVars = rel.getNumRangeVars();
 
-  result.appendVar(VarKind::Range, newNumRangeVars);
+  result.appendVar(VarKind::Range, numRelRangeVars);
 
   for (unsigned i = 0; i < rel.getNumEqualities(); ++i) {
     // Add a new equality that uses the new range variables.
@@ -2502,14 +2502,14 @@ IntegerRelation IntegerRelation::rangeProduct(const IntegerRelation &rel) {
     // right by the number of range variables added to `result`.
     SmallVector<DynamicAPInt> copy =
         SmallVector<DynamicAPInt>(rel.getEquality(i));
-    copy.insert(copy.begin() + srcOffset, newNumRangeVars, DynamicAPInt(0));
+    copy.insert(copy.begin() + srcOffset, numRelRangeVars, DynamicAPInt(0));
     result.addEquality(copy);
   }
 
   for (unsigned i = 0; i < rel.getNumInequalities(); ++i) {
     SmallVector<DynamicAPInt> copy =
         SmallVector<DynamicAPInt>(rel.getInequality(i));
-    copy.insert(copy.begin() + srcOffset, newNumRangeVars, DynamicAPInt(0));
+    copy.insert(copy.begin() + srcOffset, numRelRangeVars, DynamicAPInt(0));
     result.addInequality(copy);
   }
 

>From 1636a92950a02c848b824b054d63975d340fa8b2 Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun at users.noreply.github.com>
Date: Tue, 15 Jul 2025 14:26:01 -0700
Subject: [PATCH 4/5] fix an indexing bug and add more tests

---
 .../Analysis/Presburger/IntegerRelation.cpp   | 17 ++++++-----
 .../Presburger/IntegerRelationTest.cpp        | 30 +++++++++++++++++++
 2 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index ea207069852b4..f05067863d03c 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -2490,26 +2490,29 @@ IntegerRelation IntegerRelation::rangeProduct(const IntegerRelation &rel) {
 
   // explicit copy of `this`
   IntegerRelation result = *this;
-  unsigned srcOffset = getVarKindOffset(VarKind::Range);
+  unsigned relRangeVarStart = rel.getVarKindOffset(VarKind::Range);
   unsigned numRelRangeVars = rel.getNumRangeVars();
+  unsigned numThisRangeVars = getNumRangeVars();
 
   result.appendVar(VarKind::Range, numRelRangeVars);
 
+  // Copy each equality from `rel` and update the copy to account for range
+  // variables from `this`. The `rel` equality is a list of coefficients of the
+  // variables from `rel`, and so the range variables need to be shifted right
+  // by the number of `this` range variables.
   for (unsigned i = 0; i < rel.getNumEqualities(); ++i) {
-    // Add a new equality that uses the new range variables.
-    // The old equality is a list of coefficients of the variables
-    // from `rel`, and so the range variables need to be shifted
-    // right by the number of range variables added to `result`.
     SmallVector<DynamicAPInt> copy =
         SmallVector<DynamicAPInt>(rel.getEquality(i));
-    copy.insert(copy.begin() + srcOffset, numRelRangeVars, DynamicAPInt(0));
+    copy.insert(copy.begin() + relRangeVarStart, numThisRangeVars,
+                DynamicAPInt(0));
     result.addEquality(copy);
   }
 
   for (unsigned i = 0; i < rel.getNumInequalities(); ++i) {
     SmallVector<DynamicAPInt> copy =
         SmallVector<DynamicAPInt>(rel.getInequality(i));
-    copy.insert(copy.begin() + srcOffset, numRelRangeVars, DynamicAPInt(0));
+    copy.insert(copy.begin() + relRangeVarStart, numThisRangeVars,
+                DynamicAPInt(0));
     result.addInequality(copy);
   }
 
diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index db24600111cee..563e94d0c5e16 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -623,3 +623,33 @@ TEST(IntegerRelationTest, rangeProduct) {
 
   EXPECT_TRUE(expected.isEqual(rangeProd));
 }
+
+TEST(IntegerRelationTest, rangeProductMultdimRange) {
+  IntegerRelation r1 = parseRelationFromSet(
+      "(i, k) : (2*i + 3*k == 0, i >= 0, k >= 0)", 1);
+  IntegerRelation r2 = parseRelationFromSet(
+      "(i, l, m) : (4*i + 6*m + 9*l == 0, i >= 0, l >= 0, m >= 0)", 1);
+
+  IntegerRelation rangeProd = r1.rangeProduct(r2);
+  IntegerRelation expected =
+      parseRelationFromSet("(i, k, l, m) : (2*i + 3*k == 0, 4*i + 6*m + 9*l == "
+                           "0, i >= 0, k >= 0, l >= 0, m >= 0)",
+                           1);
+
+  EXPECT_TRUE(expected.isEqual(rangeProd));
+}
+
+TEST(IntegerRelationTest, rangeProductMultdimRangeSwapped) {
+  IntegerRelation r1 = parseRelationFromSet(
+      "(i, l, m) : (4*i + 6*m + 9*l == 0, i >= 0, l >= 0, m >= 0)", 1);
+  IntegerRelation r2 = parseRelationFromSet(
+      "(i, k) : (2*i + 3*k == 0, i >= 0, k >= 0)", 1);
+
+  IntegerRelation rangeProd = r1.rangeProduct(r2);
+  IntegerRelation expected =
+      parseRelationFromSet("(i, l, m, k) : (2*i + 3*k == 0, 4*i + 6*m + 9*l == "
+                           "0, i >= 0, k >= 0, l >= 0, m >= 0)",
+                           1);
+
+  EXPECT_TRUE(expected.isEqual(rangeProd));
+}

>From 7b96190813e3808f049ce32f2918b8ebda35746c Mon Sep 17 00:00:00 2001
From: Jeremy Kun <j2kun at users.noreply.github.com>
Date: Tue, 15 Jul 2025 20:55:40 -0700
Subject: [PATCH 5/5] formatting

---
 .../unittests/Analysis/Presburger/IntegerRelationTest.cpp | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index 563e94d0c5e16..39a74dbd7b18c 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -625,8 +625,8 @@ TEST(IntegerRelationTest, rangeProduct) {
 }
 
 TEST(IntegerRelationTest, rangeProductMultdimRange) {
-  IntegerRelation r1 = parseRelationFromSet(
-      "(i, k) : (2*i + 3*k == 0, i >= 0, k >= 0)", 1);
+  IntegerRelation r1 =
+      parseRelationFromSet("(i, k) : (2*i + 3*k == 0, i >= 0, k >= 0)", 1);
   IntegerRelation r2 = parseRelationFromSet(
       "(i, l, m) : (4*i + 6*m + 9*l == 0, i >= 0, l >= 0, m >= 0)", 1);
 
@@ -642,8 +642,8 @@ TEST(IntegerRelationTest, rangeProductMultdimRange) {
 TEST(IntegerRelationTest, rangeProductMultdimRangeSwapped) {
   IntegerRelation r1 = parseRelationFromSet(
       "(i, l, m) : (4*i + 6*m + 9*l == 0, i >= 0, l >= 0, m >= 0)", 1);
-  IntegerRelation r2 = parseRelationFromSet(
-      "(i, k) : (2*i + 3*k == 0, i >= 0, k >= 0)", 1);
+  IntegerRelation r2 =
+      parseRelationFromSet("(i, k) : (2*i + 3*k == 0, i >= 0, k >= 0)", 1);
 
   IntegerRelation rangeProd = r1.rangeProduct(r2);
   IntegerRelation expected =



More information about the Mlir-commits mailing list