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

Reviewed By: Groverkss, grosser

Differential Revision: https://reviews.llvm.org/D109892
This commit is contained in:
Arjun P 2021-09-17 00:57:11 +05:30
parent 6607bd9fd8
commit 44db07f11f
3 changed files with 58 additions and 1 deletions

View File

@ -294,6 +294,10 @@ public:
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);

View File

@ -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

View File

@ -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, {}, {});