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

[Bharathi Ramana Joshi]

https://github.com/iambrj created https://github.com/llvm/llvm-project/pull/74407

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

>From 5a22d974d849fe236ae583a178c0dc33bf88d722 Mon Sep 17 00:00:00 2001
From: iambrj <joshibharathiramana at gmail.com>
Date: Tue, 5 Dec 2023 08:35:50 +0530
Subject: [PATCH] [MLIR][Presburger] Fix IntegerRelation::swapVar not swapping
 identifiers

This commit fixes a bug where identifiers were not swapped when doing a
IntegerRelation::swapVar.
---
 .../Analysis/Presburger/IntegerRelation.cpp   |  6 +++
 .../Presburger/IntegerRelationTest.cpp        | 40 +++++++++++++++++++
 2 files changed, 46 insertions(+)

diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 3724df5abccca..0109384f1689d 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 287f7c7c56549..f390296da648d 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)));
+}