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

[arjunp]

Add a comment.

[Groverkss]

Added some comments.

mlir/include/mlir/Analysis/AffineStructures.h

297

Please use triple / doc comments here.

mlir/lib/Analysis/AffineStructures.cpp

2338–2339

An assert would be useful here as end - begin could overflow if not checked and could cause problems.

2342–2343

Assert here, same reason as previous comment.

[arjunp]

Address Kunwar's comments.

[arjunp]

If begin >= end just return, in line with the behaviour of removeIdRange.

[Groverkss]

LGTM

[grosser]

This looks also good from my side. Given that this is a very simple and non-algorithmic change you can probably go ahead with committing this.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   void removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit);
 
   /// Removes the specified identifier from the system.
   void removeId(unsigned pos);
 
   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);
 
   /// Changes the partition between dimensions and symbols. Depending on the new
   /// symbol count, either a chunk of trailing dimensional identifiers becomes
   /// symbols, or some of the leading symbols become dimensions.
   void setDimSymbolSeparation(unsigned newSymbolCount);

mlir/lib/Analysis/AffineStructures.cpp

 void FlatAffineConstraints::removeEquality(unsigned pos) {
   equalities.removeRow(pos);
 }
 
 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
 /// symbols is also treated like a constant. Returns -1 if such an equality
 /// could not be found.
 static int findEqualityToConstant(const FlatAffineConstraints &cst,
                                   unsigned pos, bool symbolic = false) {
   assert(pos < cst.getNumIds() && "invalid position");

mlir/unittests/Analysis/AffineStructuresTest.cpp

 
 #include <gmock/gmock.h>
 #include <gtest/gtest.h>
 
 #include <numeric>
 
 namespace mlir {
 
+using testing::ElementsAre;
+
 enum class TestFunction { Sample, Empty };
 
 /// If fn is TestFunction::Sample (default):
 /// If hasSample is true, check that findIntegerSample returns a valid sample
 /// for the FlatAffineConstraints fac.
 /// If hasSample is false, check that findIntegerSample returns None.
 ///
 /// If fn is TestFunction::Empty, check that isIntegerEmpty returns the
FlatAffineConstraints fac2 =
                                  {0, 1, -2}  // y >= 2 (redundant).
                              },
                              {{1, -1, 0}}); // x == y.
   fac2.removeRedundantConstraints();
 
   // 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 =
       makeFACFromConstraints(3, {},
                              {{1, -1, 0, 0},   // x == y.
                               {1, 0, -1, 0},   // x == z.
                               {0, 1, -1, 0}}); // y == z.
   fac3.removeRedundantConstraints();
TEST(FlatAffineConstraintsTest, addConstantLowerBound) {
   EXPECT_EQ(fac.atIneq(0, 2), -1);
 
   fac.addBound(FlatAffineConstraints::LB, {1, 2, 3}, 1);
   EXPECT_EQ(fac.atIneq(1, 0), 1);
   EXPECT_EQ(fac.atIneq(1, 1), 2);
   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, {}, {});
 
   fac.addInequality({1, 0});
   EXPECT_EQ(fac.atIneq(0, 0), 1);
   EXPECT_EQ(fac.atIneq(0, 1), 0);
 
   fac.clearConstraints();