[MLIR][Presburger] Implement IntegerRelation::setId

[iambrj]

No description provided.

[llvmbot]

@llvm/pr-subscribers-mlir-presburger

@llvm/pr-subscribers-mlir

Author: Bharathi Ramana Joshi (iambrj)

Changes

---

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

3 Files Affected:

• (modified) mlir/include/mlir/Analysis/Presburger/IntegerRelation.h (+4)
• (modified) mlir/lib/Analysis/Presburger/IntegerRelation.cpp (+8)
• (modified) mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp (+37)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index 8e2c9fca0a17cb..00dd230e853991 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -122,6 +122,10 @@ class IntegerRelation {
   /// current space; this will result in an assert failure.
   void setSpaceExceptLocals(const PresburgerSpace &oSpace);
 
+  /// Set the ith identifier of the IntegerRelation's PresburgerSpace. The index
+  /// is relative to the kind of the variable.
+  void setId(VarKind kind, unsigned i, Identifier id);
+
   /// Returns a copy of the space without locals.
   PresburgerSpace getSpaceWithoutLocals() const {
     return PresburgerSpace::getRelationSpace(space.getNumDomainVars(),
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index f78e21ccd38eb9..3c7a0a8b8804eb 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -66,6 +66,14 @@ void IntegerRelation::setSpaceExceptLocals(const PresburgerSpace &oSpace) {
   space.insertVar(VarKind::Local, 0, newNumLocals);
 }
 
+void IntegerRelation::setId(VarKind kind, unsigned i, Identifier id) {
+  assert(space.isUsingIds() &&
+      "space must be using identifiers to set an identifier");
+  assert(kind != VarKind::Local && "local variables cannot have identifiers");
+  assert(i < space.getNumVarKind(kind) && "invalid variable index");
+  space.getId(kind, i) = id;
+}
+
 void IntegerRelation::append(const IntegerRelation &other) {
   assert(space.isEqual(other.getSpace()) && "Spaces must be equal.");
 
diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index dff092b3204bb3..0ef94be9d5d1be 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -451,3 +451,40 @@ TEST(IntegerRelationTest, mergeAndAlignCommonSuffixSymbols) {
   EXPECT_EQ(otherSpace.getId(VarKind::Range, 1),
             Identifier(&otherIdentifiers[3]));
 }
+
+TEST(IntegerRelationTest, setId) {
+  IntegerRelation rel = parseRelationFromSet(
+      "(x, y, z)[A, B, C, D] : (x + A - C - y + D - z >= 0)", 2);
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 4, 0);
+  space.resetIds();
+
+  // Attach identifiers.
+  int identifiers[7] = {'x', 'y', 'z', 'A', 'B', 'C', 'D'};
+  space.getId(VarKind::Domain, 0) = Identifier(&identifiers[0]);
+  space.getId(VarKind::Domain, 1) = Identifier(&identifiers[1]);
+  space.getId(VarKind::Range, 0) = Identifier(&identifiers[2]);
+  space.getId(VarKind::Symbol, 0) = Identifier(&identifiers[3]);
+  space.getId(VarKind::Symbol, 1) = Identifier(&identifiers[4]);
+  space.getId(VarKind::Symbol, 2) = Identifier(&identifiers[5]);
+  space.getId(VarKind::Symbol, 3) = Identifier(&identifiers[6]);
+  rel.setSpace(space);
+
+  rel.getSpace().dump();
+  int newIdentifiers[3] = {1, 2, 3};
+  rel.setId(VarKind::Domain, 1, Identifier(&newIdentifiers[0]));
+  rel.setId(VarKind::Range, 0, Identifier(&newIdentifiers[1]));
+  rel.setId(VarKind::Symbol, 2, Identifier(&newIdentifiers[2]));
+  rel.getSpace().dump();
+
+  space = rel.getSpace();
+  // Check that new identifiers are set correctly.
+  EXPECT_EQ(space.getId(VarKind::Domain, 1), Identifier(&newIdentifiers[0]));
+  EXPECT_EQ(space.getId(VarKind::Range, 0), Identifier(&newIdentifiers[1]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 2), Identifier(&newIdentifiers[2]));
+  // Check that old identifier are not changed.
+  EXPECT_EQ(space.getId(VarKind::Domain, 0), Identifier(&identifiers[0]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 0), Identifier(&identifiers[3]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 1), Identifier(&identifiers[4]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 3), Identifier(&identifiers[6]));
+}
+

[Groverkss]

LGTM

[Groverkss]

"Set the identifier for the ith variable of the specified kind... "