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 @@ -110,7 +110,8 @@ void MultiAffineFunction::eliminateRedundantLocalId(unsigned posA, unsigned posB) { - output.addToColumn(posB, posA, /*scale=*/1); + unsigned localOffset = getIdKindOffset(IdKind::Local); + output.addToColumn(localOffset + posB, localOffset + posA, /*scale=*/1); IntegerPolyhedron::eliminateRedundantLocalId(posA, posB); } 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 @@ -157,19 +157,35 @@ } TEST(PWMAFunction, removeIdRangeRegressionTest) { - PWMAFunction pwafA = parsePWMAF( + PWMAFunction pwmafA = 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( + PWMAFunction pwmafB = 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)); + EXPECT_TRUE(pwmafA.isEqual(pwmafB)); +} + +TEST(PWMAFunction, eliminateRedundantLocalIdRegressionTest) { + PWMAFunction pwmafA = parsePWMAF( + /*numInputs=*/2, /*numOutputs=*/1, + { + {"(x, y) : (x - 2*(x floordiv 2) == 0, x - 2*y == 0)", + {{0, 1, 0, 0}}} // (0, 0) + }); + PWMAFunction pwmafB = parsePWMAF( + /*numInputs=*/2, /*numOutputs=*/1, + { + {"(x, y) : (x - 2*(x floordiv 2) == 0, x - 2*y == 0)", + {{1, -1, 0, 0}}} // (0, 0) + }); + EXPECT_TRUE(pwmafA.isEqual(pwmafB)); }