[MLIR] The return type in the `computeSingleVarRepr` function is modified to include equality expressions.

[pashu123]

Summary

Earlier computeSingleVarRepr was returning a pair of upper bound and
lower bound indices of the inequality contraints that can be expressed
as a floordiv of an affine function. The equality expression can also be
expressed as a floordiv but contains only one index and hence the LocalRepr
class is introduced to facilitate this.

[pashu123]

Clang format.

[Groverkss]

mlir/lib/Analysis/AffineStructures.cpp

859–860 ↗

(On Diff #400383)

Until support for equalities is added I would recommend returning false here in case of equality also.

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

839

You can add using namespace presburger_utils in this file and use ReprKind directly.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

224–225

The name "maybePair" does not make sense now. Please change it to repr or something similar.

mlir/lib/Analysis/Presburger/Utils.cpp

148

I would recommend add using namespace presburge_utils in this file and using LocalRepr directly.

[arjunp]

Since LocalRepr is nullable, a better name for it might be MaybeLocalRepr.

[pashu123]

Addressing comments.

[pashu123]

Changing LocalRepr to MaybeLocalRepr.

[pashu123]

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

224–225

Since we are asserting for inequality any way below, I think maybeInequality makes sense.

[pashu123]

Minor change.

[Groverkss]

LGTM

[mehdi_amini]

mlir/lib/Analysis/Presburger/Utils.cpp

157

It's not clear that this is guaranteed to be initialized before the return statement.
I think you can just add braces here: MaybeLocalRepr repr{};

[bondhugula]

Please rewrite the commit title to use active: eg:
[MLIR] Update return type for computeSingleVarRepr` to include equalities`
(much more compact)
Edit: Looks like this was already merged.

[pashu123]

In D117430#3268031, @bondhugula wrote:

Please rewrite the commit title to use active: eg:
[MLIR] Update return type for computeSingleVarRepr` to include equalities`
(much more compact)
Edit: Looks like this was already merged.

Next time I will take care of this. Thanks.

mlir/lib/Analysis/Presburger/Utils.cpp

157

@mehdi_amini I have the required changes in this patch: https://reviews.llvm.org/D117959

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

 //===- IntegerPolyhedron.h - MLIR IntegerPolyhedron Class -----*- 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
 //
 //===----------------------------------------------------------------------===//
 //
 // A class to represent an integer polyhedron.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H
 #define MLIR_ANALYSIS_PRESBURGER_INTEGERPOLYHEDRON_H
 
 #include "mlir/Analysis/Presburger/Matrix.h"
+#include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 
 /// An integer polyhedron is the set of solutions to a list of affine
 /// constraints over n integer-valued variables/identifiers. Affine constraints
 /// can be inequalities or equalities in the form:
 ///
public:
   /// which an explicit representation for each local variable can be computed.
   /// The pairs are stored as indices of upperbound, lowerbound inequalities. If
   /// no such pair can be found, it is stored as llvm::None.
   ///
   /// The dividends of the explicit representations are stored in `dividends`
   /// and the denominators in `denominators`. If no explicit representation
   /// could be found for the `i^th` local identifier, `denominators[i]` is set
   /// to 0.
-  void getLocalReprs(
-      std::vector<SmallVector<int64_t, 8>> &dividends,
+  void getLocalReprs(std::vector<SmallVector<int64_t, 8>> &dividends,
                      SmallVector<unsigned, 4> &denominators,
-      std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
-  void getLocalReprs(
-      std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const;
+                     std::vector<presburger_utils::MaybeLocalRepr> &repr) const;
+  void getLocalReprs(std::vector<presburger_utils::MaybeLocalRepr> &repr) const;
   void getLocalReprs(std::vector<SmallVector<int64_t, 8>> &dividends,
                      SmallVector<unsigned, 4> &denominators) const;
 
   /// The type of bound: equal, lower bound or upper bound.
   enum BoundType { EQ, LB, UB };
 
   /// Adds a constant bound for the specified identifier.
   void addBound(BoundType type, unsigned pos, int64_t value);

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

 #include "mlir/Support/LLVM.h"
 
 namespace mlir {
 
 class IntegerPolyhedron;
 
 namespace presburger_utils {
 
+/// `ReprKind` enum is used to set the constraint type in `MaybeLocalRepr`.
+enum class ReprKind { Inequality, Equality, None };
+
+/// `MaybeLocalRepr` contains the indices of the contraints that can be
+/// expressed as a floordiv of an affine function. If it's an `equality`
+/// contraint `equalityIdx` is set, in case of `inequality` the `lowerBoundIdx`
+/// and `upperBoundIdx` is set. By default the kind attribute is set to None.
+struct MaybeLocalRepr {
+  ReprKind kind = ReprKind::None;
+  union {
+    unsigned equalityIdx;
+    struct {
+      unsigned lowerBoundIdx, upperBoundIdx;
+    } inEqualityPair;
+  } repr;
+};
+
 /// Check if the pos^th identifier can be expressed as a floordiv of an affine
 /// function of other identifiers (where the divisor is a positive constant).
 /// `foundRepr` contains a boolean for each identifier indicating if the
 /// explicit representation for that identifier has already been computed.
 /// Returns the upper and lower bound inequalities using which the floordiv
 /// can be computed. If the representation could be computed, `dividend` and
 /// `denominator` are set. If the representation could not be computed,
 /// `llvm::None` is returned.
-Optional<std::pair<unsigned, unsigned>>
-computeSingleVarRepr(const IntegerPolyhedron &cst, ArrayRef<bool> foundRepr,
-                     unsigned pos, SmallVector<int64_t, 8> &dividend,
+MaybeLocalRepr computeSingleVarRepr(const IntegerPolyhedron &cst,
+                                    ArrayRef<bool> foundRepr, unsigned pos,
+                                    SmallVector<int64_t, 8> &dividend,
                                     unsigned &divisor);
 
 } // namespace presburger_utils
 } // namespace mlir
 
 #endif // MLIR_ANALYSIS_PRESBURGER_UTILS_H

mlir/lib/Analysis/Presburger/IntegerPolyhedron.cpp

 #include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/Support/Debug.h"
 
 #define DEBUG_TYPE "presburger"
 
 using namespace mlir;
+using namespace presburger_utils;
 using llvm::SmallDenseMap;
 using llvm::SmallDenseSet;
 
 std::unique_ptr<IntegerPolyhedron> IntegerPolyhedron::clone() const {
   return std::make_unique<IntegerPolyhedron>(*this);
 }
 
 void IntegerPolyhedron::reset(unsigned numReservedInequalities,
bool IntegerPolyhedron::containsPoint(ArrayRef<int64_t> point) const {
   }
   for (unsigned i = 0, e = getNumInequalities(); i < e; ++i) {
     if (valueAt(getInequality(i), point) < 0)
       return false;
   }
   return true;
 }
 
-void IntegerPolyhedron::getLocalReprs(
-    std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
+void IntegerPolyhedron::getLocalReprs(std::vector<MaybeLocalRepr> &repr) const {
   std::vector<SmallVector<int64_t, 8>> dividends(getNumLocalIds());
   SmallVector<unsigned, 4> denominators(getNumLocalIds());
   getLocalReprs(dividends, denominators, repr);
 }
 
 void IntegerPolyhedron::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,
     SmallVector<unsigned, 4> &denominators) const {
-  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
-      getNumLocalIds());
+  std::vector<MaybeLocalRepr> repr(getNumLocalIds());
   getLocalReprs(dividends, denominators, repr);
 }
 
 void IntegerPolyhedron::getLocalReprs(
     std::vector<SmallVector<int64_t, 8>> &dividends,
     SmallVector<unsigned, 4> &denominators,
-    std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> &repr) const {
+    std::vector<MaybeLocalRepr> &repr) const {
 
   repr.resize(getNumLocalIds());
   dividends.resize(getNumLocalIds());
   denominators.resize(getNumLocalIds());
 
   SmallVector<bool, 8> foundRepr(getNumIds(), false);
   for (unsigned i = 0, e = getNumDimAndSymbolIds(); i < e; ++i)
     foundRepr[i] = true;
 
   unsigned divOffset = getNumDimAndSymbolIds();
   bool changed;
   do {
     // Each time changed is true, at end of this iteration, one or more local
     // vars have been detected as floor divs.
     changed = false;
     for (unsigned i = 0, e = getNumLocalIds(); i < e; ++i) {
       if (!foundRepr[i + divOffset]) {
-        if (auto res = presburger_utils::computeSingleVarRepr(
-                *this, foundRepr, divOffset + i, dividends[i],
-                denominators[i])) {
+        auto res = computeSingleVarRepr(*this, foundRepr, divOffset + i,
+                                        dividends[i], denominators[i]);
+        if (res.kind == ReprKind::Inequality) {
           foundRepr[i + divOffset] = true;
-          repr[i] = res;
+          repr[i].kind = ReprKind::Inequality;
+          repr[i].repr.inEqualityPair = {res.repr.inEqualityPair.lowerBoundIdx,
+                                         res.repr.inEqualityPair.upperBoundIdx};
           changed = true;
         }
       }
     }
   } while (changed);
 
   // Set 0 denominator for identifiers for which no division representation
   // could be found.
   for (unsigned i = 0, e = repr.size(); i < e; ++i)
-    if (!repr[i].hasValue())
+    if (repr[i].kind == ReprKind::None)
       denominators[i] = 0;
 }
 
 /// Tightens inequalities given that we are dealing with integer spaces. This is
 /// analogous to the GCD test but applied to inequalities. The constant term can
 /// be reduced to the preceding multiple of the GCD of the coefficients, i.e.,
 ///  64*i - 100 >= 0  =>  64*i - 128 >= 0 (since 'i' is an integer). This is a
 /// fast method - linear in the number of coefficients.

mlir/lib/Analysis/Presburger/PresburgerSet.cpp

 //===- Set.cpp - MLIR PresburgerSet Class ---------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/PresburgerSet.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
+#include "mlir/Analysis/Presburger/Utils.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
+using namespace presburger_utils;
 
 PresburgerSet::PresburgerSet(const IntegerPolyhedron &poly)
     : nDim(poly.getNumDimIds()), nSym(poly.getNumSymbolIds()) {
   unionPolyInPlace(poly);
 }
 
 unsigned PresburgerSet::getNumPolys() const {
   return integerPolyhedrons.size();
b.removeIdRange(IntegerPolyhedron::IdKind::Local, bInitNumLocals,
                     b.getNumLocalIds());
     b.removeInequalityRange(bInitNumIneqs, b.getNumInequalities());
     b.removeEqualityRange(bInitNumEqs, b.getNumEqualities());
     simplex.rollback(initialSnapshot);
   };
 
   // Find out which inequalities of sI correspond to division inequalities for
   // the local variables of sI.
-  std::vector<llvm::Optional<std::pair<unsigned, unsigned>>> repr(
-      sI.getNumLocalIds());
+  std::vector<MaybeLocalRepr> repr(sI.getNumLocalIds());
   sI.getLocalReprs(repr);
 
   // Add sI's locals to b, after b's locals. Also add b's locals to sI, before
   // sI's locals.
   b.mergeLocalIds(sI);
 
   // Mark which inequalities of sI are division inequalities and add all such
   // inequalities to b.
   llvm::SmallBitVector isDivInequality(sI.getNumInequalities());
-  for (Optional<std::pair<unsigned, unsigned>> &maybePair : repr) {
-    assert(maybePair &&
+  for (MaybeLocalRepr &maybeInequality : repr) {
+    assert(maybeInequality.kind == ReprKind::Inequality &&
            "Subtraction is not supported when a representation of the local "
            "variables of the subtrahend cannot be found!");
+    auto lb = maybeInequality.repr.inEqualityPair.lowerBoundIdx;
+    auto ub = maybeInequality.repr.inEqualityPair.upperBoundIdx;
 
-    b.addInequality(sI.getInequality(maybePair->first));
-    b.addInequality(sI.getInequality(maybePair->second));
+    b.addInequality(sI.getInequality(lb));
+    b.addInequality(sI.getInequality(ub));
 
-    assert(maybePair->first != maybePair->second &&
+    assert(lb != ub &&
            "Upper and lower bounds must be different inequalities!");
-    isDivInequality[maybePair->first] = true;
-    isDivInequality[maybePair->second] = true;
+    isDivInequality[lb] = true;
+    isDivInequality[ub] = true;
   }
 
   unsigned offset = simplex.getNumConstraints();
   unsigned numLocalsAdded = b.getNumLocalIds() - bInitNumLocals;
   simplex.appendVariable(numLocalsAdded);
 
   unsigned snapshotBeforeIntersect = simplex.getSnapshot();
   simplex.intersectIntegerPolyhedron(sI);

mlir/lib/Analysis/Presburger/Utils.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Analysis/Presburger/IntegerPolyhedron.h"
 #include "mlir/Support/LogicalResult.h"
 #include "mlir/Support/MathExtras.h"
 
 using namespace mlir;
+using namespace presburger_utils;
 
 /// Normalize a division's `dividend` and the `divisor` by their GCD. For
 /// example: if the dividend and divisor are [2,0,4] and 4 respectively,
 /// they get normalized to [1,0,2] and 2.
 static void normalizeDivisionByGCD(SmallVectorImpl<int64_t> &dividend,
                                    unsigned &divisor) {
   if (divisor == 0 || dividend.empty())
     return;

 /// Check if the pos^th identifier can be expressed as a floordiv of an affine
 /// function of other identifiers (where the divisor is a positive constant).
 /// `foundRepr` contains a boolean for each identifier indicating if the
 /// explicit representation for that identifier has already been computed.
 /// Returns the upper and lower bound inequalities using which the floordiv can
 /// be computed. If the representation could be computed, `dividend` and
 /// `denominator` are set. If the representation could not be computed,
 /// `llvm::None` is returned.
-Optional<std::pair<unsigned, unsigned>> presburger_utils::computeSingleVarRepr(
+MaybeLocalRepr presburger_utils::computeSingleVarRepr(
     const IntegerPolyhedron &cst, ArrayRef<bool> foundRepr, unsigned pos,
     SmallVector<int64_t, 8> &dividend, unsigned &divisor) {
   assert(pos < cst.getNumIds() && "invalid position");
   assert(foundRepr.size() == cst.getNumIds() &&
          "Size of foundRepr does not match total number of variables");
 
   SmallVector<unsigned, 4> lbIndices, ubIndices;
   cst.getLowerAndUpperBoundIndices(pos, &lbIndices, &ubIndices);
+  MaybeLocalRepr repr;
 
   for (unsigned ubPos : ubIndices) {
     for (unsigned lbPos : lbIndices) {
       // Attempt to get divison representation from ubPos, lbPos.
       if (failed(getDivRepr(cst, pos, ubPos, lbPos, dividend, divisor)))
         continue;
 
       // Check if the inequalities depend on a variable for which
for (unsigned lbPos : lbIndices) {
 
       // Expression can't be constructed as it depends on a yet unknown
       // identifier.
       // TODO: Visit/compute the identifiers in an order so that this doesn't
       // happen. More complex but much more efficient.
       if (c < f)
         continue;
 
-      return std::make_pair(ubPos, lbPos);
+      repr.kind = ReprKind::Inequality;
+      repr.repr.inEqualityPair = {ubPos, lbPos};
+      return repr;
     }
   }
-
-  return llvm::None;
+  return repr;
 }

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 
 #define DEBUG_TYPE "affine-structures"
 
 using namespace mlir;
+using namespace presburger_utils;
 
 namespace {
 
 // See comments for SimpleAffineExprFlattener.
 // An AffineExprFlattener extends a SimpleAffineExprFlattener by recording
 // constraint information associated with mod's, floordiv's, and ceildiv's
 // in FlatAffineConstraints 'localVarCst'.
 struct AffineExprFlattener : public SimpleAffineExprFlattener {
static bool detectAsFloorDiv(const FlatAffineConstraints &cst, unsigned pos,
   // Get upper-lower bound pair for this variable.
   SmallVector<bool, 8> foundRepr(cst.getNumIds(), false);
   for (unsigned i = 0, e = cst.getNumIds(); i < e; ++i)
     if (exprs[i])
       foundRepr[i] = true;
 
   SmallVector<int64_t, 8> dividend;
   unsigned divisor;
-  auto ulPair = presburger_utils::computeSingleVarRepr(cst, foundRepr, pos,
-                                                       dividend, divisor);
+  auto ulPair = computeSingleVarRepr(cst, foundRepr, pos, dividend, divisor);
 
   // No upper-lower bound pair found for this var.
-  if (!ulPair)
+  if (ulPair.kind == ReprKind::None || ulPair.kind == ReprKind::Equality)
     return false;
 
   // Construct the dividend expression.
   auto dividendExpr = getAffineConstantExpr(dividend.back(), context);
   for (unsigned c = 0, f = cst.getNumIds(); c < f; c++)
     if (dividend[c] != 0)
       dividendExpr = dividendExpr + dividend[c] * exprs[c];