From b68e78cea65faef45342a2af420b32877413e9c5 Mon Sep 17 00:00:00 2001
From: Arjun P <arjunpitchanathan@gmail.com>
Date: Thu, 17 Mar 2022 22:25:40 +0000
Subject: [PATCH] [MLIR][Prebsurger] Add IntegerRelation::intersect supporting
 locals properly

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D122149
---
 mlir/include/mlir/Analysis/Presburger/IntegerRelation.h | 4 ++++
 mlir/lib/Analysis/Presburger/IntegerRelation.cpp        | 7 +++++++
 mlir/lib/Analysis/Presburger/PresburgerRelation.cpp     | 8 +++-----
 3 files changed, 14 insertions(+), 5 deletions(-)

diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
index f8f06b8..fa14949 100644
--- a/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
+++ b/mlir/include/mlir/Analysis/Presburger/IntegerRelation.h
@@ -95,6 +95,10 @@ public:
   /// intersection with no simplification of any sort attempted.
   void append(const IntegerRelation &other);
 
+  /// Return the intersection of the two sets.
+  /// If there are locals, they will be merged.
+  IntegerRelation intersect(IntegerRelation other) const;
+
   /// Return whether `this` and `other` are equal. This is integer-exact
   /// and somewhat expensive, since it uses the integer emptiness check
   /// (see IntegerRelation::findIntegerSample()).
diff --git a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
index 1917d9a..3f647a0 100644
--- a/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/IntegerRelation.cpp
@@ -52,6 +52,13 @@ void IntegerRelation::append(const IntegerRelation &other) {
   }
 }
 
+IntegerRelation IntegerRelation::intersect(IntegerRelation other) const {
+  IntegerRelation result = *this;
+  result.mergeLocalIds(other);
+  result.append(other);
+  return result;
+}
+
 bool IntegerRelation::isEqual(const IntegerRelation &other) const {
   assert(PresburgerLocalSpace::isEqual(other) && "Spaces must be equal.");
   return PresburgerRelation(*this).isEqual(PresburgerRelation(other));
diff --git a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
index 6f035aa..37a3b78 100644
--- a/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
+++ b/mlir/lib/Analysis/Presburger/PresburgerRelation.cpp
@@ -97,11 +97,9 @@ PresburgerRelation::intersect(const PresburgerRelation &set) const {
                             getNumSymbolIds());
   for (const IntegerRelation &csA : integerRelations) {
     for (const IntegerRelation &csB : set.integerRelations) {
-      IntegerRelation csACopy = csA, csBCopy = csB;
-      csACopy.mergeLocalIds(csBCopy);
-      csACopy.append(csBCopy);
-      if (!csACopy.isEmpty())
-        result.unionInPlace(csACopy);
+      IntegerRelation intersection = csA.intersect(csB);
+      if (!intersection.isEmpty())
+        result.unionInPlace(intersection);
     }
   }
   return result;
-- 
2.7.4