diff --git a/mlir/lib/Analysis/Presburger/PWMAFunction.cpp b/mlir/lib/Analysis/Presburger/PWMAFunction.cpp --- a/mlir/lib/Analysis/Presburger/PWMAFunction.cpp +++ b/mlir/lib/Analysis/Presburger/PWMAFunction.cpp @@ -99,7 +99,7 @@ void MultiAffineFunction::removeIdRange(IdKind kind, unsigned idStart, unsigned idLimit) { output.removeColumns(idStart + getIdKindOffset(kind), idLimit - idStart); - IntegerPolyhedron::removeIdRange(idStart, idLimit); + IntegerPolyhedron::removeIdRange(kind, idStart, idLimit); } void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA, diff --git a/mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp b/mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp --- a/mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp +++ b/mlir/unittests/Analysis/Presburger/PWMAFunctionTest.cpp @@ -140,3 +140,21 @@ EXPECT_THAT(*nonNegPWAF.valueAt({2, -3}), ElementsAre(-1, -1)); EXPECT_FALSE(nonNegPWAF.valueAt({-2, -3}).hasValue()); } + +TEST(PWMAFunction, removeIdRangeRegressionTest) { + PWMAFunction pwafA = parsePWMAF( + /*numInputs=*/2, /*numOutputs=*/1, + { + {"(x, y) : (x == 0, y == 0, x - 2*(x floordiv 2) == 0, y - 2*(y " + "floordiv 2) == 0)", + {{0, 0, 0, 0, 0}}} // (0, 0) + }); + PWMAFunction pwafB = parsePWMAF( + /*numInputs=*/2, /*numOutputs=*/1, + { + {"(x, y) : (x - 11*y == 0, 11*x - y == 0, x - 2*(x floordiv 2) == 0, " + "y - 2*(y floordiv 2) == 0)", + {{0, 0, 0, 0, 0}}} // (0, 0) + }); + EXPECT_TRUE(pwafA.isEqual(pwafB)); +}