[Mlir-commits] [mlir] 7ce800b - [MLIR][Presburger] Fix bug in PresburgerSpace::convertVarKind (#67267)

[llvmlistbot at llvm.org]

Author: Kunwar Grover
Date: 2023-09-25T12:55:58+05:30
New Revision: 7ce800ba3197c6df866565b9be5709c8c2d64fc4

URL: https://github.com/llvm/llvm-project/commit/7ce800ba3197c6df866565b9be5709c8c2d64fc4
DIFF: https://github.com/llvm/llvm-project/commit/7ce800ba3197c6df866565b9be5709c8c2d64fc4.diff

LOG: [MLIR][Presburger] Fix bug in PresburgerSpace::convertVarKind (#67267)

This patch fixes a bug in PresburgerSpace::convertVarKind where the
identifiers were not moved properly due to offset being invalidated.

Added: 
    

Modified: 
    mlir/lib/Analysis/Presburger/PresburgerSpace.cpp
    mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp b/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp
index c4d01c551b43795..2ce5528365eb52f 100644
--- a/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp
+++ b/mlir/lib/Analysis/Presburger/PresburgerSpace.cpp
@@ -161,6 +161,26 @@ void PresburgerSpace::convertVarKind(VarKind srcKind, unsigned srcPos,
   assert(dstPos <= getNumVarKind(dstKind) &&
          "invalid position for destination variables");
 
+  // Move identifiers if `usingIds` and variables moved are not locals.
+  unsigned srcOffset = getVarKindOffset(srcKind) + srcPos;
+  unsigned dstOffset = getVarKindOffset(dstKind) + dstPos;
+  if (isUsingIds() && srcKind != VarKind::Local && dstKind != VarKind::Local) {
+    identifiers.insert(identifiers.begin() + dstOffset, num, Identifier());
+    // Update srcOffset if insertion of new elements invalidates it.
+    if (dstOffset < srcOffset)
+      srcOffset += num;
+    std::move(identifiers.begin() + srcOffset,
+              identifiers.begin() + srcOffset + num,
+              identifiers.begin() + dstOffset);
+    identifiers.erase(identifiers.begin() + srcOffset,
+                      identifiers.begin() + srcOffset + num);
+  } else if (isUsingIds() && srcKind != VarKind::Local) {
+    identifiers.erase(identifiers.begin() + srcOffset,
+                      identifiers.begin() + srcOffset + num);
+  } else if (isUsingIds() && dstKind != VarKind::Local) {
+    identifiers.insert(identifiers.begin() + dstOffset, num, Identifier());
+  }
+
   auto addVars = [&](VarKind kind, int num) {
     switch (kind) {
     case VarKind::Domain:
@@ -180,22 +200,6 @@ void PresburgerSpace::convertVarKind(VarKind srcKind, unsigned srcPos,
 
   addVars(srcKind, -(signed)num);
   addVars(dstKind, num);
-
-  // Move identifiers if `usingIds` and variables moved are not locals.
-  unsigned srcOffset = getVarKindOffset(srcKind) + srcPos;
-  unsigned dstOffset = getVarKindOffset(dstKind) + dstPos;
-  if (isUsingIds() && srcKind != VarKind::Local && dstKind != VarKind::Local) {
-    identifiers.insert(identifiers.begin() + dstOffset, num, Identifier());
-    for (unsigned i = 0; i < num; ++i)
-      identifiers[dstOffset + i] = identifiers[srcOffset + i];
-    identifiers.erase(identifiers.begin() + srcOffset,
-                      identifiers.begin() + srcOffset + num);
-  } else if (isUsingIds() && srcKind != VarKind::Local) {
-    identifiers.erase(identifiers.begin() + srcOffset,
-                      identifiers.begin() + srcOffset + num);
-  } else if (isUsingIds() && dstKind != VarKind::Local) {
-    identifiers.insert(identifiers.begin() + dstOffset, num, Identifier());
-  }
 }
 
 void PresburgerSpace::swapVar(VarKind kindA, VarKind kindB, unsigned posA,

diff  --git a/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp
index cb23174b939c38b..dd06d462f54bee7 100644
--- a/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/PresburgerSpaceTest.cpp
@@ -158,3 +158,24 @@ TEST(PresburgerSpaceTest, convertVarKindLocals) {
   EXPECT_FALSE(space.getId(VarKind::Range, 0).hasValue());
   EXPECT_FALSE(space.getId(VarKind::Range, 1).hasValue());
 }
+
+TEST(PresburgerSpaceTest, convertVarKind2) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(0, 2, 2, 0);
+  space.resetIds();
+
+  // Attach identifiers.
+  int identifiers[4] = {0, 1, 2, 3};
+  space.getId(VarKind::Range, 0) = Identifier(&identifiers[0]);
+  space.getId(VarKind::Range, 1) = Identifier(&identifiers[1]);
+  space.getId(VarKind::Symbol, 0) = Identifier(&identifiers[2]);
+  space.getId(VarKind::Symbol, 1) = Identifier(&identifiers[3]);
+
+  // Convert Range variables to symbols.
+  space.convertVarKind(VarKind::Range, 0, 2, VarKind::Symbol, 1);
+
+  // Check if the identifiers are moved to symbols.
+  EXPECT_EQ(space.getId(VarKind::Symbol, 0), Identifier(&identifiers[2]));
+  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[3]));
+}