[Mlir-commits] [mlir] 8d7c979 - [MLIR][Presburger] Fix IntegerRelation::swapVar not swapping identifiers (#74407)

[llvmlistbot at llvm.org]

Author: Bharathi Ramana Joshi
Date: 2023-12-13T22:47:19+05:30
New Revision: 8d7c9798153fbaa285f73585a1d87c8d00de0030

URL: https://github.com/llvm/llvm-project/commit/8d7c9798153fbaa285f73585a1d87c8d00de0030
DIFF: https://github.com/llvm/llvm-project/commit/8d7c9798153fbaa285f73585a1d87c8d00de0030.diff

LOG: [MLIR][Presburger] Fix IntegerRelation::swapVar not swapping identifiers (#74407)

This commit fixes a bug where identifiers were not swapped when doing a
IntegerRelation::swapVar.

Added: 
    

Modified: 
    mlir/lib/Analysis/Presburger/IntegerRelation.cpp
    mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 3724df5abcccaa..0109384f1689dd 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -449,6 +449,12 @@ void IntegerRelation::swapVar(unsigned posA, unsigned posB) {
   if (posA == posB)
     return;
 
+  VarKind kindA = space.getVarKindAt(posA);
+  VarKind kindB = space.getVarKindAt(posB);
+  unsigned relativePosA = posA - getVarKindOffset(kindA);
+  unsigned relativePosB = posB - getVarKindOffset(kindB);
+  space.swapVar(kindA, kindB, relativePosA, relativePosB);
+
   inequalities.swapColumns(posA, posB);
   equalities.swapColumns(posA, posB);
 }

diff  --git a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
index 287f7c7c56549f..f390296da648d2 100644
--- a/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
+++ b/mlir/unittests/Analysis/Presburger/IntegerRelationTest.cpp
@@ -8,6 +8,7 @@
 
 #include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "Parser.h"
+#include "mlir/Analysis/Presburger/PresburgerSpace.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 
 #include <gmock/gmock.h>
@@ -167,3 +168,42 @@ TEST(IntegerRelationTest, symbolicLexmax) {
   EXPECT_TRUE(lexmax3.unboundedDomain.isIntegerEmpty());
   EXPECT_TRUE(lexmax3.lexopt.isEqual(expectedLexmax3));
 }
+
+TEST(IntegerRelationTest, swapVar) {
+  PresburgerSpace space = PresburgerSpace::getRelationSpace(2, 1, 2, 0);
+  space.resetIds();
+
+  int identifiers[6] = {0, 1, 2, 3, 4};
+
+  // Attach identifiers to domain identifiers.
+  space.getId(VarKind::Domain, 0) = Identifier(&identifiers[0]);
+  space.getId(VarKind::Domain, 1) = Identifier(&identifiers[1]);
+
+  // Attach identifiers to range identifiers.
+  space.getId(VarKind::Range, 0) = Identifier(&identifiers[2]);
+
+  // Attach identifiers to symbol identifiers.
+  space.getId(VarKind::Symbol, 0) = Identifier(&identifiers[3]);
+  space.getId(VarKind::Symbol, 1) = Identifier(&identifiers[4]);
+
+  IntegerRelation rel =
+      parseRelationFromSet("(x, y, z)[N, M] : (z - x - y == 0, x >= 0, N - x "
+                           ">= 0, y >= 0, M - y >= 0)",
+                           2);
+  rel.setSpace(space);
+  // Swap (Domain 0, Range 0)
+  rel.swapVar(0, 2);
+  // Swap (Domain 1, Symbol 1)
+  rel.swapVar(1, 4);
+
+  PresburgerSpace swappedSpace = rel.getSpace();
+
+  EXPECT_TRUE(swappedSpace.getId(VarKind::Domain, 0)
+                  .isEqual(space.getId(VarKind::Range, 0)));
+  EXPECT_TRUE(swappedSpace.getId(VarKind::Domain, 1)
+                  .isEqual(space.getId(VarKind::Symbol, 1)));
+  EXPECT_TRUE(swappedSpace.getId(VarKind::Range, 0)
+                  .isEqual(space.getId(VarKind::Domain, 0)));
+  EXPECT_TRUE(swappedSpace.getId(VarKind::Symbol, 1)
+                  .isEqual(space.getId(VarKind::Domain, 1)));
+}