[Mlir-commits] [mlir] 44db07f - [MLIR] AffineStructures: support removing a range of constraints at once

[Arjun P]

Author: Arjun P
Date: 2021-09-17T16:27:48+05:30
New Revision: 44db07f11f00b2115e277dd89d830e2abde684ab

URL: https://github.com/llvm/llvm-project/commit/44db07f11f00b2115e277dd89d830e2abde684ab
DIFF: https://github.com/llvm/llvm-project/commit/44db07f11f00b2115e277dd89d830e2abde684ab.diff

LOG: [MLIR] AffineStructures: support removing a range of constraints at once

Reviewed By: Groverkss, grosser

Differential Revision: https://reviews.llvm.org/D109892

Added: 
    

Modified: 
    mlir/include/mlir/Analysis/AffineStructures.h
    mlir/lib/Analysis/AffineStructures.cpp
    mlir/unittests/Analysis/AffineStructuresTest.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/include/mlir/Analysis/AffineStructures.h b/mlir/include/mlir/Analysis/AffineStructures.h
index b5676186d87d..589858815b5b 100644
--- a/mlir/include/mlir/Analysis/AffineStructures.h
+++ b/mlir/include/mlir/Analysis/AffineStructures.h
@@ -294,6 +294,10 @@ class FlatAffineConstraints {
   void removeEquality(unsigned pos);
   void removeInequality(unsigned pos);
 
+  /// Remove the (in)equalities at positions [start, end).
+  void removeEqualityRange(unsigned start, unsigned end);
+  void removeInequalityRange(unsigned start, unsigned end);
+
   /// Sets the `values.size()` identifiers starting at `po`s to the specified
   /// values and removes them.
   void setAndEliminate(unsigned pos, ArrayRef<int64_t> values);

diff  --git a/mlir/lib/Analysis/AffineStructures.cpp b/mlir/lib/Analysis/AffineStructures.cpp
index 4453e8531884..36aedfb26ac6 100644
--- a/mlir/lib/Analysis/AffineStructures.cpp
+++ b/mlir/lib/Analysis/AffineStructures.cpp
@@ -2335,6 +2335,19 @@ void FlatAffineConstraints::removeInequality(unsigned pos) {
   inequalities.removeRow(pos);
 }
 
+void FlatAffineConstraints::removeEqualityRange(unsigned begin, unsigned end) {
+  if (begin >= end)
+    return;
+  equalities.removeRows(begin, end - begin);
+}
+
+void FlatAffineConstraints::removeInequalityRange(unsigned begin,
+                                                  unsigned end) {
+  if (begin >= end)
+    return;
+  inequalities.removeRows(begin, end - begin);
+}
+
 /// Finds an equality that equates the specified identifier to a constant.
 /// Returns the position of the equality row. If 'symbolic' is set to true,
 /// symbols are also treated like a constant, i.e., an affine function of the

diff  --git a/mlir/unittests/Analysis/AffineStructuresTest.cpp b/mlir/unittests/Analysis/AffineStructuresTest.cpp
index d5a88b684b9e..f74a78ddf5fa 100644
--- a/mlir/unittests/Analysis/AffineStructuresTest.cpp
+++ b/mlir/unittests/Analysis/AffineStructuresTest.cpp
@@ -17,6 +17,8 @@
 
 namespace mlir {
 
+using testing::ElementsAre;
+
 enum class TestFunction { Sample, Empty };
 
 /// If fn is TestFunction::Sample (default):
@@ -461,7 +463,7 @@ TEST(FlatAffineConstraintsTest, removeRedundantConstraintsTest) {
   // The second inequality is redundant and should have been removed. The
   // remaining inequality should be the first one.
   EXPECT_EQ(fac2.getNumInequalities(), 1u);
-  EXPECT_THAT(fac2.getInequality(0), testing::ElementsAre(1, 0, -3));
+  EXPECT_THAT(fac2.getInequality(0), ElementsAre(1, 0, -3));
   EXPECT_EQ(fac2.getNumEqualities(), 1u);
 
   FlatAffineConstraints fac3 =
@@ -575,6 +577,44 @@ TEST(FlatAffineConstraintsTest, addConstantLowerBound) {
   EXPECT_EQ(fac.atIneq(1, 2), 2);
 }
 
+TEST(FlatAffineConstraintsTest, removeInequality) {
+  FlatAffineConstraints fac =
+      makeFACFromConstraints(1, {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 4}}, {});
+
+  fac.removeInequalityRange(0, 0);
+  EXPECT_EQ(fac.getNumInequalities(), 5u);
+
+  fac.removeInequalityRange(1, 3);
+  EXPECT_EQ(fac.getNumInequalities(), 3u);
+  EXPECT_THAT(fac.getInequality(0), ElementsAre(0, 0));
+  EXPECT_THAT(fac.getInequality(1), ElementsAre(3, 3));
+  EXPECT_THAT(fac.getInequality(2), ElementsAre(4, 4));
+
+  fac.removeInequality(1);
+  EXPECT_EQ(fac.getNumInequalities(), 2u);
+  EXPECT_THAT(fac.getInequality(0), ElementsAre(0, 0));
+  EXPECT_THAT(fac.getInequality(1), ElementsAre(4, 4));
+}
+
+TEST(FlatAffineConstraintsTest, removeEquality) {
+  FlatAffineConstraints fac =
+      makeFACFromConstraints(1, {}, {{0, 0}, {1, 1}, {2, 2}, {3, 3}, {4, 4}});
+
+  fac.removeEqualityRange(0, 0);
+  EXPECT_EQ(fac.getNumEqualities(), 5u);
+
+  fac.removeEqualityRange(1, 3);
+  EXPECT_EQ(fac.getNumEqualities(), 3u);
+  EXPECT_THAT(fac.getEquality(0), ElementsAre(0, 0));
+  EXPECT_THAT(fac.getEquality(1), ElementsAre(3, 3));
+  EXPECT_THAT(fac.getEquality(2), ElementsAre(4, 4));
+
+  fac.removeEquality(1);
+  EXPECT_EQ(fac.getNumEqualities(), 2u);
+  EXPECT_THAT(fac.getEquality(0), ElementsAre(0, 0));
+  EXPECT_THAT(fac.getEquality(1), ElementsAre(4, 4));
+}
+
 TEST(FlatAffineConstraintsTest, clearConstraints) {
   FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});