[Mlir-commits] [mlir] [MLIR][Presburger] Preserve identifiers in IntegerRelation::convertVa… (PR #67909)

[Bharathi Ramana Joshi]

https://github.com/iambrj updated https://github.com/llvm/llvm-project/pull/67909

>From da5741232944e27c5595bd4fb264b3dcc9b5c391 Mon Sep 17 00:00:00 2001
From: iambrj <joshibharathiramana at gmail.com>
Date: Sun, 1 Oct 2023 11:25:59 +0300
Subject: [PATCH] [MLIR][Presburger] Preserve identifiers in
 IntegerRelation::convertVarKind

---
 .../Analysis/Presburger/IntegerRelation.cpp   |  8 ++++
 .../Presburger/IntegerRelationTest.cpp        | 46 +++++++++++++++++++
 2 files changed, 54 insertions(+)

diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index be764bd7c9176b9..1ed33d7034d9733 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -1325,6 +1325,10 @@ void IntegerRelation::convertVarKind(VarKind srcKind, unsigned varStart,
   if (varStart >= varLimit)
     return;
 
+  // Save the space as the insert/delete vars operations affect the identifier
+  // information in the space.
+  PresburgerSpace oldSpace = space;
+
   // Append new local variables corresponding to the dimensions to be converted.
   unsigned convertCount = varLimit - varStart;
   unsigned newVarsBegin = insertVar(dstKind, pos, convertCount);
@@ -1342,6 +1346,10 @@ void IntegerRelation::convertVarKind(VarKind srcKind, unsigned varStart,
 
   // Complete the move by deleting the initially occupied columns.
   removeVarRange(srcKind, varStart, varLimit);
+
+  // Update the space while preserving identifiers.
+  space = oldSpace;
+  space.convertVarKind(srcKind, varStart, varLimit - varStart, dstKind, pos);
 }
 
 void IntegerRelation::addBound(BoundType type, unsigned pos,
diff --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index 287f7c7c56549ff..3e94fd02936ded3 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -167,3 +167,49 @@ TEST(IntegerRelationTest, symbolicLexmax) {
   EXPECT_TRUE(lexmax3.unboundedDomain.isIntegerEmpty());
   EXPECT_TRUE(lexmax3.lexopt.isEqual(expectedLexmax3));
 }
+
+TEST(IntegerRelationTest, convertVarKind) {
+  PresburgerSpace space = PresburgerSpace::getSetSpace(3, 3, 0);
+  space.resetIds();
+
+  // Attach identifiers.
+  int identifiers[6] = {0, 1, 2, 3, 4, 5};
+  space.getId(VarKind::SetDim, 0) = Identifier(&identifiers[0]);
+  space.getId(VarKind::SetDim, 1) = Identifier(&identifiers[1]);
+  space.getId(VarKind::SetDim, 2) = 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]);
+
+  // Cannot call parseIntegerRelation to test convertVarKind as
+  // parseIntegerRelation uses convertVarKind.
+  IntegerRelation rel = parseIntegerPolyhedron(
+      // 0  1  2  3  4  5
+      "(x, y, a)[U, V, W] : (x - U == 0, y + a - W == 0, U - V >= 0,"
+      "y - a >= 0)");
+  rel.setSpace(space);
+
+  // Make a few kind conversions.
+  rel.convertVarKind(VarKind::Symbol, 1, 2, VarKind::Domain, 0);
+  rel.convertVarKind(VarKind::Range, 2, 3, VarKind::Domain, 0);
+  rel.convertVarKind(VarKind::Range, 0, 2, VarKind::Symbol, 1);
+  rel.convertVarKind(VarKind::Domain, 1, 2, VarKind::Range, 0);
+  rel.convertVarKind(VarKind::Domain, 0, 1, VarKind::Range, 1);
+
+  space = rel.getSpace();
+
+  // Expected rel.
+  IntegerRelation expectedRel = parseIntegerPolyhedron(
+      "(V, a)[U, x, y, W] : (x - U == 0, y + a - W == 0, U - V >= 0,"
+      "y - a >= 0)");
+  expectedRel.setSpace(space);
+
+  EXPECT_TRUE(rel.isEqual(expectedRel));
+
+  EXPECT_EQ(space.getId(VarKind::SetDim, 0), Identifier(&identifiers[4]));
+  EXPECT_EQ(space.getId(VarKind::SetDim, 1), Identifier(&identifiers[2]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 0), Identifier(&identifiers[3]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 1), Identifier(&identifiers[0]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 2), Identifier(&identifiers[1]));
+  EXPECT_EQ(space.getId(VarKind::Symbol, 3), Identifier(&identifiers[5]));
+}