[MLIR][Presburger] Introduce MaybeOptimum type to represent computed optima

This allows to differentiate between the cases where the optimum does not
exist due to being unbounded and due to the polytope being empty.

Reviewed By: Groverkss

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

Author: Superty

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

@@ -209,7 +209,8 @@ class IntegerPolyhedron : public PresburgerLocalSpace {
   /// constraints. Returns an empty optional if the polyhedron is empty or if
   /// the lexmin is unbounded. Symbols are not supported and will result in
   /// assert-failure.
-  Optional<SmallVector<Fraction, 8>> getRationalLexMin() const;
+  presburger_utils::MaybeOptimum<SmallVector<Fraction, 8>>
+  getRationalLexMin() const;
 
   /// Swap the posA^th identifier with the posB^th identifier.
   virtual void swapId(unsigned posA, unsigned posB);

mlir/include/mlir/Analysis/Presburger/Simplex.h

@@ -18,6 +18,7 @@
 #include "mlir/Analysis/Presburger/Fraction.h"
 #include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/Optional.h"
@@ -202,14 +203,6 @@ class SimplexBase {
   /// Add all the constraints from the given IntegerPolyhedron.
   void intersectIntegerPolyhedron(const IntegerPolyhedron &poly);
 
-  /// Returns the current sample point, which may contain non-integer (rational)
-  /// coordinates. Returns an empty optional when the tableau is empty.
-  ///
-  /// Also returns empty when the big M parameter is used and a variable
-  /// has a non-zero big M coefficient, meaning its value is infinite or
-  /// unbounded.
-  Optional<SmallVector<Fraction, 8>> getRationalSample() const;
-
   /// Print the tableau's internal state.
   void print(raw_ostream &os) const;
   void dump() const;
@@ -441,9 +434,18 @@ class LexSimplex : public SimplexBase {
   unsigned getSnapshot() { return SimplexBase::getSnapshotBasis(); }
 
   /// Return the lexicographically minimum rational solution to the constraints.
-  Optional<SmallVector<Fraction, 8>> getRationalLexMin();
+  presburger_utils::MaybeOptimum<SmallVector<Fraction, 8>> getRationalLexMin();
 
 protected:
+  /// Returns the current sample point, which may contain non-integer (rational)
+  /// coordinates. Returns an empty optimum when the tableau is empty.
+  ///
+  /// Returns an unbounded optimum when the big M parameter is used and a
+  /// variable has a non-zero big M coefficient, meaning its value is infinite
+  /// or unbounded.
+  presburger_utils::MaybeOptimum<SmallVector<Fraction, 8>>
+  getRationalSample() const;
+
   /// Undo the addition of the last constraint. This is only called while
   /// rolling back.
   void undoLastConstraint() final;
@@ -510,15 +512,16 @@ class Simplex : public SimplexBase {
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
-  Optional<Fraction> computeRowOptimum(Direction direction, unsigned row);
+  presburger_utils::MaybeOptimum<Fraction>
+  computeRowOptimum(Direction direction, unsigned row);
 
   /// Compute the maximum or minimum value of the given expression, depending on
   /// direction. Should not be called when the Simplex is empty.
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
-  Optional<Fraction> computeOptimum(Direction direction,
-                                    ArrayRef<int64_t> coeffs);
+  presburger_utils::MaybeOptimum<Fraction>
+  computeOptimum(Direction direction, ArrayRef<int64_t> coeffs);
 
   /// Returns whether the perpendicular of the specified constraint is a
   /// is a direction along which the polytope is bounded.
@@ -537,10 +540,10 @@ class Simplex : public SimplexBase {
   void detectRedundant();
 
   /// Returns a (min, max) pair denoting the minimum and maximum integer values
-  /// of the given expression. If either of the values is unbounded, an empty
-  /// optional is returned in its place. If the result has min > max then no
-  /// integer value exists.
-  std::pair<Optional<int64_t>, Optional<int64_t>>
+  /// of the given expression. If no integer value exists, both results will be
+  /// of kind Empty.
+  std::pair<presburger_utils::MaybeOptimum<int64_t>,
+            presburger_utils::MaybeOptimum<int64_t>>
   computeIntegerBounds(ArrayRef<int64_t> coeffs);
 
   /// Returns true if the polytope is unbounded, i.e., extends to infinity in
@@ -569,6 +572,10 @@ class Simplex : public SimplexBase {
   /// None.
   Optional<SmallVector<int64_t, 8>> getSamplePointIfIntegral() const;
 
+  /// Returns the current sample point, which may contain non-integer (rational)
+  /// coordinates. Returns an empty optional when the tableau is empty.
+  Optional<SmallVector<Fraction, 8>> getRationalSample() const;
+
 private:
   friend class GBRSimplex;
 
@@ -610,7 +617,8 @@ class Simplex : public SimplexBase {
   ///
   /// Returns a Fraction denoting the optimum, or a null value if no optimum
   /// exists, i.e., if the expression is unbounded in this direction.
-  Optional<Fraction> computeOptimum(Direction direction, Unknown &u);
+  presburger_utils::MaybeOptimum<Fraction> computeOptimum(Direction direction,
+                                                          Unknown &u);
 
   /// Mark the specified unknown redundant. This operation is added to the undo
   /// log and will be undone by rollbacks. The specified unknown must be in row

mlir/include/mlir/Analysis/Presburger/Utils.h

@@ -22,6 +22,67 @@ class IntegerPolyhedron;
 
 namespace presburger_utils {
 
+/// This class represents the result of operations optimizing something subject
+/// to some constraints. If the constraints were not satisfiable the, kind will
+/// be Empty. If the optimum is unbounded, the kind is Unbounded, and if the
+/// optimum is bounded, the kind will be Bounded and `optimum` holds the optimal
+/// value.
+enum class OptimumKind { Empty, Unbounded, Bounded };
+template <typename T>
+class MaybeOptimum {
+public:
+private:
+  OptimumKind kind = OptimumKind::Empty;
+  T optimum;
+
+public:
+  MaybeOptimum() = default;
+  MaybeOptimum(OptimumKind kind) : kind(kind) {
+    assert(kind != OptimumKind::Bounded &&
+           "Bounded optima should be constructed by specifying the optimum!");
+  }
+  MaybeOptimum(const T &optimum)
+      : kind(OptimumKind::Bounded), optimum(optimum) {}
+
+  OptimumKind getKind() const { return kind; }
+  bool isBounded() const { return kind == OptimumKind::Bounded; }
+  bool isUnbounded() const { return kind == OptimumKind::Unbounded; }
+  bool isEmpty() const { return kind == OptimumKind::Empty; }
+
+  Optional<T> getOptimumIfBounded() const { return optimum; }
+  const T &getBoundedOptimum() const {
+    assert(kind == OptimumKind::Bounded &&
+           "This should be called only for bounded optima");
+    return optimum;
+  }
+  T &getBoundedOptimum() {
+    assert(kind == OptimumKind::Bounded &&
+           "This should be called only for bounded optima");
+    return optimum;
+  }
+  const T &operator*() const { return getBoundedOptimum(); }
+  T &operator*() { return getBoundedOptimum(); }
+  const T *operator->() const { return &getBoundedOptimum(); }
+  T *operator->() { return &getBoundedOptimum(); }
+  bool operator==(const MaybeOptimum<T> &other) const {
+    if (kind != other.kind)
+      return false;
+    if (kind != OptimumKind::Bounded)
+      return true;
+    return optimum == other.optimum;
+  }
+
+  // Given f that takes a T and returns a U, convert this `MaybeOptimum<T>` to
+  // a `MaybeOptimum<U>` by applying `f` to the bounded optimum if it exists, or
+  // returning a MaybeOptimum of the same kind otherwise.
+  template <class Function>
+  auto map(const Function &f) const & -> MaybeOptimum<decltype(f(optimum))> {
+    if (kind == OptimumKind::Bounded)
+      return f(optimum);
+    return kind;
+  }
+};
+
 /// `ReprKind` enum is used to set the constraint type in `MaybeLocalRepr`.
 enum class ReprKind { Inequality, Equality, None };
 

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

@@ -72,14 +72,14 @@ bool IntegerPolyhedron::isSubsetOf(const IntegerPolyhedron &other) const {
   return PresburgerSet(*this).isSubsetOf(PresburgerSet(other));
 }
 
-Optional<SmallVector<Fraction, 8>>
+MaybeOptimum<SmallVector<Fraction, 8>>
 IntegerPolyhedron::getRationalLexMin() const {
   assert(getNumSymbolIds() == 0 && "Symbols are not supported!");
-  Optional<SmallVector<Fraction, 8>> maybeLexMin =
+  MaybeOptimum<SmallVector<Fraction, 8>> maybeLexMin =
       LexSimplex(*this).getRationalLexMin();
 
-  if (!maybeLexMin)
-    return {};
+  if (!maybeLexMin.isBounded())
+    return maybeLexMin;
 
   // The Simplex returns the lexmin over all the variables including locals. But
   // locals are not actually part of the space and should not be returned in the
@@ -1032,20 +1032,23 @@ Optional<uint64_t> IntegerPolyhedron::computeVolume() const {
   bool hasUnboundedId = false;
   for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i) {
     dim[i] = 1;
-    Optional<int64_t> min, max;
+    MaybeOptimum<int64_t> min, max;
     std::tie(min, max) = simplex.computeIntegerBounds(dim);
     dim[i] = 0;
 
+    assert((!min.isEmpty() && !max.isEmpty()) &&
+           "Polytope should be rationally non-empty!");
+
     // 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) {
+    if (min.isUnbounded() || max.isUnbounded()) {
       hasUnboundedId = true;
       continue;
     }
 
     // In this case there are no valid integer points and the volume is
     // definitely zero.
-    if (*min > *max)
+    if (min.getBoundedOptimum() > max.getBoundedOptimum())
       return 0;
 
     count *= (*max - *min + 1);