diff --git a/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h b/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h --- a/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h +++ b/mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h @@ -277,6 +277,11 @@ /// otherwise. Optional> findIntegerSample() const; + /// Compute an overapproximation of the number of integer points in the + /// polyhedron. Symbol ids are currently not supported. If the computed + /// overapproximation is infinite, an empty optional is returned. + Optional computeVolume() const; + /// Returns true if the given point satisfies the constraints, or false /// otherwise. /// diff --git a/mlir/include/mlir/Analysis/Presburger/PresburgerSet.h b/mlir/include/mlir/Analysis/Presburger/PresburgerSet.h --- a/mlir/include/mlir/Analysis/Presburger/PresburgerSet.h +++ b/mlir/include/mlir/Analysis/Presburger/PresburgerSet.h @@ -99,6 +99,15 @@ /// any of the Polys in the union are unbounded. bool findIntegerSample(SmallVectorImpl &sample); + /// Compute an overapproximation of the number of integer points in the + /// polyhedron. Symbol ids are currently not supported. If the computed + /// overapproximation is infinite, an empty optional is returned. + /// + /// This currently just sums up the overapproximations of the volumes of the + /// disjuncts, so the approximation might be far from the true volume in the + /// case when there is a lot of overlap between disjuncts. + Optional computeVolume() const; + /// Simplifies the representation of a PresburgerSet. /// /// In particular, removes all Polys which are subsets of other Polys in the diff --git a/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp b/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp --- a/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp +++ b/mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp @@ -1065,6 +1065,60 @@ equalities.resizeVertically(pos); } +Optional IntegerPolyhedron::computeVolume() const { + assert(getNumSymbolIds() == 0 && "Symbols are not yet supported!"); + + Simplex simplex(*this); + // If the polytope is rationally empty, there are certainly no integer + // points. + if (simplex.isEmpty()) + return 0; + + // Just find the maximum and minimum integer value of each non-local id + // separately, thus finding the number of integer values each such id can + // take. Multiplying these together gives a valid overapproximation of the + // number of integer points in the polyhedron. The result this gives is + // equivalent to projecting (rationally) the polyhedron onto its non-local ids + // and returning the number of integer points in a minimal axis-parallel + // hyperrectangular overapproximation of that. + // + // We also handle the special case where one dimension is unbounded and + // another dimension can take no integer values. In this case, the volume is + // zero. + // + // If there is no such empty dimension, if any dimension is unbounded we + // just return the result as unbounded. + uint64_t count = 1; + SmallVector dim(getNumIds() + 1); + bool hasUnboundedId = false; + for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i) { + dim[i] = 1; + Optional min, max; + std::tie(min, max) = simplex.computeIntegerBounds(dim); + dim[i] = 0; + + // One of the dimensions is unbounded. Note this fact. We will return + // unbounded if none of the other dimensions makes the volume zero. + if (!min || !max) { + hasUnboundedId = true; + continue; + } + + // In this case there are no valid integer points and the volume is + // definitely zero. + if (*min > *max) + return 0; + + count *= (*max - *min + 1); + } + + if (count == 0) + return 0; + if (hasUnboundedId) + return {}; + return count; +} + void IntegerPolyhedron::eliminateRedundantLocalId(unsigned posA, unsigned posB) { assert(posA < getNumLocalIds() && "Invalid local id position"); diff --git a/mlir/lib/Analysis/Presburger/PresburgerSet.cpp b/mlir/lib/Analysis/Presburger/PresburgerSet.cpp --- a/mlir/lib/Analysis/Presburger/PresburgerSet.cpp +++ b/mlir/lib/Analysis/Presburger/PresburgerSet.cpp @@ -385,6 +385,20 @@ return false; } +Optional PresburgerSet::computeVolume() const { + assert(getNumSymbolIds() == 0 && "Symbols are not yet supported!"); + // The sum of the volumes of the disjuncts is a valid overapproximation of the + // volume of their union, even if they overlap. + uint64_t result = 0; + for (const IntegerPolyhedron &poly : integerPolyhedrons) { + Optional volume = poly.computeVolume(); + if (!volume) + return {}; + result += *volume; + } + return result; +} + PresburgerSet PresburgerSet::coalesce() const { PresburgerSet newSet = PresburgerSet::getEmptySet(getNumDimIds(), getNumSymbolIds()); diff --git a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp --- a/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp +++ b/mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp @@ -1118,4 +1118,66 @@ expectNoRationalLexMin(parsePoly("(x) : (2*x >= 0, -x - 1 >= 0)", &context)); } +static void +expectComputedVolumeIsValidOverapprox(const IntegerPolyhedron &poly, + Optional trueVolume, + Optional resultBound) { + expectComputedVolumeIsValidOverapprox(poly.computeVolume(), trueVolume, + resultBound); +} + +TEST(IntegerPolyhedronTest, computeVolume) { + MLIRContext context; + + // 0 <= x <= 3 + 1/3, -5.5 <= y <= 2 + 3/5, 3 <= z <= 1.75. + // i.e. 0 <= x <= 3, -5 <= y <= 2, 3 <= z <= 3 + 1/4. + // So volume is 4 * 8 * 1 = 32. + expectComputedVolumeIsValidOverapprox( + parsePoly("(x, y, z) : (x >= 0, -3*x + 10 >= 0, 2*y + 11 >= 0," + "-5*y + 13 >= 0, z - 3 >= 0, -4*z + 13 >= 0)", + &context), + /*trueVolume=*/32ull, /*resultBound=*/32ull); + + // Same as above but y has bounds 2 + 1/5 <= y <= 2 + 3/5. So the volume is + // zero. + expectComputedVolumeIsValidOverapprox( + parsePoly("(x, y, z) : (x >= 0, -3*x + 10 >= 0, 5*y - 11 >= 0," + "-5*y + 13 >= 0, z - 3 >= 0, -4*z + 13 >= 0)", + &context), + /*trueVolume=*/0ull, /*resultBound=*/0ull); + + // Now x is unbounded below but y still has no integer values. + expectComputedVolumeIsValidOverapprox( + parsePoly("(x, y, z) : (-3*x + 10 >= 0, 5*y - 11 >= 0," + "-5*y + 13 >= 0, z - 3 >= 0, -4*z + 13 >= 0)", + &context), + /*trueVolume=*/0ull, /*resultBound=*/0ull); + + // A diamond shape, 0 <= x + y <= 10, 0 <= x - y <= 10, + // with vertices at (0, 0), (5, 5), (5, 5), (10, 0). + // x and y can take 11 possible values so result computed is 11*11 = 121. + expectComputedVolumeIsValidOverapprox( + parsePoly("(x, y) : (x + y >= 0, -x - y + 10 >= 0, x - y >= 0," + "-x + y + 10 >= 0)", + &context), + /*trueVolume=*/61ull, /*resultBound=*/121ull); + + // Effectively the same diamond as above; constrain the variables to be even + // and double the constant terms of the constraints. The algorithm can't + // eliminate locals exactly, so the result is an overapproximation by + // computing that x and y can take 21 possible values so result is 21*21 = + // 441. + expectComputedVolumeIsValidOverapprox( + parsePoly("(x, y) : (x + y >= 0, -x - y + 20 >= 0, x - y >= 0," + " -x + y + 20 >= 0, x - 2*(x floordiv 2) == 0," + "y - 2*(y floordiv 2) == 0)", + &context), + /*trueVolume=*/61ull, /*resultBound=*/441ull); + + // Unbounded polytope. + expectComputedVolumeIsValidOverapprox( + parsePoly("(x, y) : (2*x - y >= 0, y - 3*x >= 0)", &context), + /*trueVolume=*/{}, /*resultBound=*/{}); +} + } // namespace mlir diff --git a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp --- a/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp +++ b/mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp @@ -696,4 +696,64 @@ expectCoalesce(1, set); } +static void +expectComputedVolumeIsValidOverapprox(const PresburgerSet &set, + Optional trueVolume, + Optional resultBound) { + expectComputedVolumeIsValidOverapprox(set.computeVolume(), trueVolume, + resultBound); +} + +TEST(SetTest, computeVolume) { + MLIRContext context; + // Diamond with vertices at (0, 0), (5, 5), (5, 5), (10, 0). + PresburgerSet diamond( + parsePoly("(x, y) : (x + y >= 0, -x - y + 10 >= 0, x - y >= 0, -x + y + " + "10 >= 0)", + &context)); + expectComputedVolumeIsValidOverapprox(diamond, + /*trueVolume=*/61ull, + /*resultBound=*/121ull); + + // Diamond with vertices at (-5, 0), (0, -5), (0, 5), (5, 0). + PresburgerSet shiftedDiamond(parsePoly( + "(x, y) : (x + y + 5 >= 0, -x - y + 5 >= 0, x - y + 5 >= 0, -x + y + " + "5 >= 0)", + &context)); + expectComputedVolumeIsValidOverapprox(shiftedDiamond, + /*trueVolume=*/61ull, + /*resultBound=*/121ull); + + // Diamond with vertices at (-5, 0), (5, -10), (5, 10), (15, 0). + PresburgerSet biggerDiamond(parsePoly( + "(x, y) : (x + y + 5 >= 0, -x - y + 15 >= 0, x - y + 5 >= 0, -x + y + " + "15 >= 0)", + &context)); + expectComputedVolumeIsValidOverapprox(biggerDiamond, + /*trueVolume=*/221ull, + /*resultBound=*/441ull); + + // There is some overlap between diamond and shiftedDiamond. + expectComputedVolumeIsValidOverapprox(diamond.unionSet(shiftedDiamond), + /*trueVolume=*/104ull, + /*resultBound=*/242ull); + + // biggerDiamond subsumes both the small ones. + expectComputedVolumeIsValidOverapprox( + diamond.unionSet(shiftedDiamond).unionSet(biggerDiamond), + /*trueVolume=*/221ull, + /*resultBound=*/683ull); + + // Unbounded polytope. + PresburgerSet unbounded( + parsePoly("(x, y) : (2*x - y >= 0, y - 3*x >= 0)", &context)); + expectComputedVolumeIsValidOverapprox(unbounded, /*trueVolume=*/{}, + /*resultBound=*/{}); + + // Union of unbounded with bounded is unbounded. + expectComputedVolumeIsValidOverapprox(unbounded.unionSet(diamond), + /*trueVolume=*/{}, + /*resultBound=*/{}); +} + } // namespace mlir diff --git a/mlir/unittests/Analysis/Presburger/Utils.h b/mlir/unittests/Analysis/Presburger/Utils.h --- a/mlir/unittests/Analysis/Presburger/Utils.h +++ b/mlir/unittests/Analysis/Presburger/Utils.h @@ -14,6 +14,8 @@ #define MLIR_UNITTESTS_ANALYSIS_PRESBURGER_UTILS_H #include "../../Dialect/Affine/Analysis/AffineStructuresParser.h" +#include "mlir/Analysis/Presburger/IntegerPolyhedron.h" +#include "mlir/Analysis/Presburger/PresburgerSet.h" #include "mlir/IR/MLIRContext.h" #include @@ -27,6 +29,31 @@ EXPECT_TRUE(succeeded(poly)); return *poly; } + +/// lhs and rhs represent non-negative integers or positive infinity. The +/// infinity case corresponds to when the Optional is empty. +static bool infinityOrUInt64LE(Optional lhs, Optional rhs) { + // No constraint. + if (!rhs) + return true; + // Finite rhs provided so lhs has to be finite too. + if (!lhs) + return false; + return *lhs <= *rhs; +} + +/// Expect that the computed volume is a valid overapproximation of +/// the true volume `trueVolume`, while also being at least as good an +/// approximation as `resultBound`. +static void +expectComputedVolumeIsValidOverapprox(Optional computedVolume, + Optional trueVolume, + Optional resultBound) { + assert(infinityOrUInt64LE(trueVolume, resultBound) && + "can't expect result to be less than the true volume"); + EXPECT_TRUE(infinityOrUInt64LE(trueVolume, computedVolume)); + EXPECT_TRUE(infinityOrUInt64LE(computedVolume, resultBound)); +} } // namespace mlir #endif // MLIR_UNITTESTS_ANALYSIS_PRESBURGER_UTILS_H