[MLIR][Presburger] Support computing volumes via hyperrectangular overapproximation

Add support for computing an overapproximation of the number of integer points
in a polyhedron. The returned result is actually the number of integer points
one gets by computing the "rational shadow" obtained by projecting out the
local IDs, finding the minimal axis-parallel hyperrectangular approximation
of the shadow, and returning the number of integer points in that. This does
not currently support symbols.

Reviewed By: Groverkss

Differential Revision: https://reviews.llvm.org/D119228

mlir/include/mlir/Analysis/Presburger/IntegerPolyhedron.h

@@ -277,6 +277,11 @@ public:
   /// otherwise.
   Optional<SmallVector<int64_t, 8>> 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<uint64_t> computeVolume() const;
+
   /// Returns true if the given point satisfies the constraints, or false
   /// otherwise.
   ///

mlir/include/mlir/Analysis/Presburger/PresburgerSet.h

@@ -99,6 +99,15 @@ public:
   /// any of the Polys in the union are unbounded.
   bool findIntegerSample(SmallVectorImpl<int64_t> &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<uint64_t> computeVolume() const;
+
   /// Simplifies the representation of a PresburgerSet.
   ///
   /// In particular, removes all Polys which are subsets of other Polys in the

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

@@ -1065,6 +1065,60 @@ void IntegerPolyhedron::removeRedundantConstraints() {
   equalities.resizeVertically(pos);
 }
 
+Optional<uint64_t> 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<int64_t, 8> dim(getNumIds() + 1);
+  bool hasUnboundedId = false;
+  for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i) {
+    dim[i] = 1;
+    Optional<int64_t> 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");

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

@@ -385,6 +385,20 @@ bool PresburgerSet::findIntegerSample(SmallVectorImpl<int64_t> &sample) {
   return false;
 }
 
+Optional<uint64_t> 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<uint64_t> volume = poly.computeVolume();
+    if (!volume)
+      return {};
+    result += *volume;
+  }
+  return result;
+}
+
 PresburgerSet PresburgerSet::coalesce() const {
   PresburgerSet newSet =
       PresburgerSet::getEmptySet(getNumDimIds(), getNumSymbolIds());

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

@@ -1118,4 +1118,66 @@ TEST(IntegerPolyhedronTest, getRationalLexMin) {
   expectNoRationalLexMin(parsePoly("(x) : (2*x >= 0, -x - 1 >= 0)", &context));
 }
 
+static void
+expectComputedVolumeIsValidOverapprox(const IntegerPolyhedron &poly,
+                                      Optional<uint64_t> trueVolume,
+                                      Optional<uint64_t> 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

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

@@ -696,4 +696,64 @@ TEST(SetTest, coalesceContainedEqComplex) {
   expectCoalesce(1, set);
 }
 
+static void
+expectComputedVolumeIsValidOverapprox(const PresburgerSet &set,
+                                      Optional<uint64_t> trueVolume,
+                                      Optional<uint64_t> 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

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 <gtest/gtest.h>
@@ -27,6 +29,31 @@ static IntegerPolyhedron parsePoly(StringRef str, MLIRContext *context) {
   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<uint64_t> lhs, Optional<uint64_t> 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<uint64_t> computedVolume,
+                                      Optional<uint64_t> trueVolume,
+                                      Optional<uint64_t> 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