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

[arjunp]

Summary

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.

[Groverkss]

mlir/lib/Analysis/Presburger/Simplex.cpp

1478–1479 ↗

(On Diff #406758)

Maybe add an assert here with this comment.

mlir/unittests/Analysis/Presburger/PresburgerSetTest.cpp

699–723

You can put these into Utils.h after https://reviews.llvm.org/D119194 is landed

[arjunp]

Thanks for the quick review!

mlir/lib/Analysis/Presburger/Simplex.cpp

1478–1479 ↗

(On Diff #406758)

There is already an assert, inside Optional::operator*

[bondhugula]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

390–392

You may want to state in the doc comment that the approximation could be far away from a tight one depending on the overlap.

mlir/lib/Analysis/Presburger/Simplex.cpp

1476–1481 ↗

(On Diff #406758)

Are these changes related?

[arjunp]

Move common functions into utils.h

[arjunp]

Rebase

[arjunp]

mlir/lib/Analysis/Presburger/Simplex.cpp

1476–1481 ↗

(On Diff #406758)

Yeah, they're needed because we had to change the return type of computeIntegerBounds. That can be separated out into another commit though, which I've done now

[jpienaar]

mlir/unittests/Analysis/Presburger/Utils.h

49

Getting error about this function being unused. Can it be removed?

[arjunp]

mlir/unittests/Analysis/Presburger/Utils.h

49

It can't be removed. It's called from function called expectComputedVolumeIsValidOverapprox in IntegerPolyhedronTest.cpp and another with same name in PresburgerSetTest.cpp. Where does the error occur?

[arjunp]

mlir/unittests/Analysis/Presburger/Utils.h

49

This has apparently been fixed

https://reviews.llvm.org/rG78eeda7529e7

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

public:
   /// Find an integer sample point satisfying the constraints using a
   /// branch and bound algorithm with generalized basis reduction, with some
   /// additional processing using Simplex for unbounded sets.
   ///
   /// Returns an integer sample point if one exists, or an empty Optional
   /// 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.
   ///
   /// Note: currently, if the polyhedron contains local ids, the values of
   /// the local ids must also be provided.
   bool containsPoint(ArrayRef<int64_t> point) const;
 
   /// Find equality and pairs of inequality contraints identified by their

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

public:
   /// Return true if all the sets in the union are known to be integer empty
   /// false otherwise.
   bool isIntegerEmpty() const;
 
   /// Find an integer sample from the given set. This should not be called if
   /// 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
   /// union.
   PresburgerSet coalesce() const;
 
 private:
   /// Construct an empty PresburgerSet.

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

void IntegerPolyhedron::removeRedundantConstraints() {
   for (unsigned r = 0, e = getNumEqualities(); r < e; r++) {
     if (!(simplex.isMarkedRedundant(numIneqs + 2 * r) &&
           simplex.isMarkedRedundant(numIneqs + 2 * r + 1)))
       copyEquality(r, pos++);
   }
   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");
   assert(posB < getNumLocalIds() && "Invalid local id position");
 
   unsigned localOffset = getIdKindOffset(IdKind::Local);
   posA += localOffset;
   posB += localOffset;

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

for (const IntegerPolyhedron &poly : integerPolyhedrons) {
     if (Optional<SmallVector<int64_t, 8>> opt = poly.findIntegerSample()) {
       sample = std::move(*opt);
       return true;
     }
   }
   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());
   llvm::SmallBitVector isRedundant(getNumPolys());
 
   for (unsigned i = 0, e = integerPolyhedrons.size(); i < e; ++i) {
     if (isRedundant[i])
       continue;

mlir/unittests/Analysis/Presburger/IntegerPolyhedronTest.cpp

expectNoRationalLexMin(
       parsePoly("(x, y, z) : (2*x + 5*y + 8*z - 10 >= 0,"
                 "2*x + 10*y + 8*z - 10 >= 0, 2*x + 5*y + 10*z - 10 >= 0)",
                 &context));
 
   // The set is empty.
   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

PresburgerSet set = parsePresburgerSetFromPolyStrings(
       {
           "(x,y) : (x - 2 == 0, x - y == 0)",
           "(x,y) : (x - 1 >= 0, -x + 2 >= 0, x - y == 0)",
       },
       &context);
   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

 //===- Utils.h - Utils for Presburger Tests ---------------------*- C++ -*-===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file defines helper functions for Presburger unittests.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_UNITTESTS_ANALYSIS_PRESBURGER_UTILS_H
 #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>
 
 namespace mlir {
 /// Parses a IntegerPolyhedron from a StringRef. It is expected that the
 /// string represents a valid IntegerSet, otherwise it will violate a gtest
 /// assertion.
 static IntegerPolyhedron parsePoly(StringRef str, MLIRContext *context) {
   FailureOr<IntegerPolyhedron> poly = parseIntegerSetToFAC(str, 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