[mlir][Analysis][NFC] FlatAffineConstraints: Use BoundType enum in functions

[springerm]

fix test

[nicolasvasilache]

Nice, thanks!

mlir/include/mlir/Analysis/AffineStructures.h

public:
   inline ArrayRef<int64_t> getEquality(unsigned idx) const {
     return equalities.getRow(idx);
   }
 
   inline ArrayRef<int64_t> getInequality(unsigned idx) const {
     return inequalities.getRow(idx);
   }
 
-  /// Adds a lower or an upper bound for the identifier at the specified
-  /// position with constraints being drawn from the specified bound map. If
-  /// `eq` is true, add a single equality equal to the bound map's first result
-  /// expr.
+  /// The type of bound: equal, lower bound or upper bound.
+  enum BoundType { EQ, LB, UB };
+
+  /// Adds a bound for the identifier at the specified position with constraints
+  /// being drawn from the specified bound map. In case of an EQ bound, the
+  /// bound map is expected to have exactly one result. In case of a LB/UB, the
+  /// bound map may have more than one result, for each of which an inequality
+  /// is added.
   /// Note: The dimensions/symbols of this FlatAffineConstraints must match the
   /// dimensions/symbols of the affine map.
-  LogicalResult addLowerOrUpperBound(unsigned pos, AffineMap boundMap, bool eq,
-                                     bool lower = true);
+  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap);
+
+  /// Adds a constant bound for the specified identifier.
+  void addBound(BoundType type, unsigned pos, int64_t value);
+
+  /// Adds a constant bound for the specified expression.
+  void addBound(BoundType type, ArrayRef<int64_t> expr, int64_t value);
 
   /// Returns the constraint system as an integer set. Returns a null integer
   /// set if the system has no constraints, or if an integer set couldn't be
   /// constructed as a result of a local variable's explicit representation not
   /// being known and such a local variable appearing in any of the constraints.
   IntegerSet getAsIntegerSet(MLIRContext *context) const;
 
   /// Computes the lower and upper bounds of the first `num` dimensional
   /// identifiers (starting at `offset`) as an affine map of the remaining
   /// identifiers (dimensional and symbolic). This method is able to detect
   /// identifiers as floordiv's and mod's of affine expressions of other
   /// identifiers with respect to (positive) constants. Sets bound map to a
   /// null AffineMap if such a bound can't be found (or yet unimplemented).
   void getSliceBounds(unsigned offset, unsigned num, MLIRContext *context,
                       SmallVectorImpl<AffineMap> *lbMaps,
                       SmallVectorImpl<AffineMap> *ubMaps);
 
   /// Adds an inequality (>= 0) from the coefficients specified in `inEq`.
   void addInequality(ArrayRef<int64_t> inEq);
   /// Adds an equality from the coefficients specified in `eq`.
   void addEquality(ArrayRef<int64_t> eq);
 
-  /// Adds a constant lower bound constraint for the specified identifier.
-  void addConstantLowerBound(unsigned pos, int64_t lb);
-  /// Adds a constant upper bound constraint for the specified identifier.
-  void addConstantUpperBound(unsigned pos, int64_t ub);
-
   /// Adds a new local identifier as the floordiv of an affine function of other
   /// identifiers, the coefficients of which are provided in `dividend` and with
   /// respect to a positive constant `divisor`. Two constraints are added to the
   /// system to capture equivalence with the floordiv:
   /// q = dividend floordiv c    <=>   c*q <= dividend <= c*q + c - 1.
   void addLocalFloorDiv(ArrayRef<int64_t> dividend, int64_t divisor);
 
-  /// Adds a constant lower bound constraint for the specified expression.
-  void addConstantLowerBound(ArrayRef<int64_t> expr, int64_t lb);
-  /// Adds a constant upper bound constraint for the specified expression.
-  void addConstantUpperBound(ArrayRef<int64_t> expr, int64_t ub);
-
-  /// Sets the identifier at the specified position to a constant.
-  void setIdToConstant(unsigned pos, int64_t val);
-
   /// Swap the posA^th identifier with the posB^th identifier.
   virtual void swapId(unsigned posA, unsigned posB);
 
   /// Add identifiers of the specified kind - specified positions are relative
   /// to the kind of identifier. The coefficient column corresponding to the
   /// added identifier is initialized to zero.
   void addDimId(unsigned pos);
   void addSymbolId(unsigned pos);
public:
   /// symbolic identifiers, *lb = [1, 0, 1], lbDivisor = 32. See comments at
   /// function definition for examples.
   Optional<int64_t> getConstantBoundOnDimSize(
       unsigned pos, SmallVectorImpl<int64_t> *lb = nullptr,
       int64_t *boundFloorDivisor = nullptr,
       SmallVectorImpl<int64_t> *ub = nullptr, unsigned *minLbPos = nullptr,
       unsigned *minUbPos = nullptr) const;
 
-  /// Returns the constant lower bound for the pos^th identifier if there is
-  /// one; None otherwise.
-  Optional<int64_t> getConstantLowerBound(unsigned pos) const;
-
-  /// Returns the constant upper bound for the pos^th identifier if there is
-  /// one; None otherwise.
-  Optional<int64_t> getConstantUpperBound(unsigned pos) const;
+  /// Returns the constant bound for the pos^th identifier if there is one;
+  /// None otherwise.
+  // TODO: Support EQ bounds.
+  Optional<int64_t> getConstantBound(BoundType type, unsigned pos) const;
 
   /// Gets the lower and upper bound of the `offset` + `pos`th identifier
   /// treating [0, offset) U [offset + num, symStartPos) as dimensions and
   /// [symStartPos, getNumDimAndSymbolIds) as symbols, and `pos` lies in
   /// [0, num). The multi-dimensional maps in the returned pair represent the
   /// max and min of potentially multiple affine expressions. The upper bound is
   /// exclusive. `localExprs` holds pre-computed AffineExpr's for all local
   /// identifiers in the system.
public:
   /// are collected from the IntegerSet attached to the given `affine.if`
   /// instance argument (`ifOp`). It is asserted that:
   /// 1) The IntegerSet of the given `affine.if` instance should not contain
   /// semi-affine expressions,
   /// 2) The columns of the constraint system created from `ifOp` should match
   /// the columns in the current one regarding numbers and values.
   void addAffineIfOpDomain(AffineIfOp ifOp);
 
-  /// Adds a lower or an upper bound for the identifier at the specified
-  /// position with constraints being drawn from the specified bound map and
-  /// operands. If `eq` is true, add a single equality equal to the bound map's
-  /// first result expr.
-  LogicalResult addLowerOrUpperBound(unsigned pos, AffineMap boundMap,
-                                     ValueRange operands, bool eq,
-                                     bool lower = true);
-  using FlatAffineConstraints::addLowerOrUpperBound;
+  /// Adds a bound for the identifier at the specified position with constraints
+  /// being drawn from the specified bound map and operands. In case of an
+  /// EQ bound, the  bound map is expected to have exactly one result. In case
+  /// of a LB/UB, the bound map may have more than one result, for each of which
+  /// an inequality is added.
+  LogicalResult addBound(BoundType type, unsigned pos, AffineMap boundMap,
+                         ValueRange operands);
+
+  /// Adds a constant bound for the identifier associated with the given Value.
+  void addBound(BoundType type, Value val, int64_t value);
+
+  using FlatAffineConstraints::addBound;
 
   /// Returns the bound for the identifier at `pos` from the inequality at
   /// `ineqPos` as a 1-d affine value map (affine map + operands). The returned
   /// affine value map can either be a lower bound or an upper bound depending
   /// on the sign of atIneq(ineqPos, pos). Asserts if the row at `ineqPos` does
   /// not involve the `pos`th identifier.
   void getIneqAsAffineValueMap(unsigned pos, unsigned ineqPos,
                                AffineValueMap &vmap,
                                MLIRContext *context) const;
 
   /// Adds slice lower bounds represented by lower bounds in `lbMaps` and upper
   /// bounds in `ubMaps` to each identifier in the constraint system which has
   /// a value in `values`. Note that both lower/upper bounds share the same
   /// operand list `operands`.
   /// This function assumes `values.size` == `lbMaps.size` == `ubMaps.size`.
   /// Note that both lower/upper bounds use operands from `operands`.
   LogicalResult addSliceBounds(ArrayRef<Value> values,
                                ArrayRef<AffineMap> lbMaps,
                                ArrayRef<AffineMap> ubMaps,
                                ArrayRef<Value> operands);
 
-  /// Sets the identifier corresponding to the specified Value `value` to a
-  /// constant. Asserts if the `value` is not found.
-  void setIdToConstant(Value value, int64_t val);
-  using FlatAffineConstraints::setIdToConstant;
-
   /// Looks up the position of the identifier with the specified Value. Returns
   /// true if found (false otherwise). `pos` is set to the (column) position of
   /// the identifier.
   bool findId(Value val, unsigned *pos) const;
 
   /// Returns true if an identifier with the specified Value exists, false
   /// otherwise.
   bool containsId(Value val) const;

mlir/lib/Analysis/AffineAnalysis.cpp

addMemRefAccessConstraints(const AffineValueMap &srcAccessMap,
   auto addEqForConstOperands = [&](ArrayRef<Value> operands) {
     for (unsigned i = 0, e = operands.size(); i < e; ++i) {
       if (isForInductionVar(operands[i]))
         continue;
       auto symbol = operands[i];
       assert(isValidSymbol(symbol));
       // Check if the symbol is a constant.
       if (auto cOp = symbol.getDefiningOp<ConstantIndexOp>())
-        dependenceDomain->setIdToConstant(valuePosMap.getSymPos(symbol),
-                                          cOp.getValue());
+        dependenceDomain->addBound(FlatAffineConstraints::EQ,
+                                   valuePosMap.getSymPos(symbol),
     }
   };
 
   // Add equality constraints for any src symbols defined by constant ops.
   addEqForConstOperands(srcOperands);
   // Add equality constraints for any dst symbols defined by constant ops.
   addEqForConstOperands(dstOperands);
 
static void computeDirectionVector(
   // Eliminate all variables other than the direction variables just added.
   dependenceDomain->projectOut(numCommonLoops, numIdsToEliminate);
 
   // Scan each common loop variable column and set direction vectors based
   // on eliminated constraint system.
   dependenceComponents->resize(numCommonLoops);
   for (unsigned j = 0; j < numCommonLoops; ++j) {
     (*dependenceComponents)[j].op = commonLoops[j].getOperation();
-    auto lbConst = dependenceDomain->getConstantLowerBound(j);
+    auto lbConst =
+        dependenceDomain->getConstantBound(FlatAffineConstraints::LB, j);
     (*dependenceComponents)[j].lb =
         lbConst.getValueOr(std::numeric_limits<int64_t>::min());
-    auto ubConst = dependenceDomain->getConstantUpperBound(j);
+    auto ubConst =
+        dependenceDomain->getConstantBound(FlatAffineConstraints::UB, j);
     (*dependenceComponents)[j].ub =
         ubConst.getValueOr(std::numeric_limits<int64_t>::max());
   }
 }
 
 // Populates 'accessMap' with composition of AffineApplyOps reachable from
 // indices of MemRefAccess.
 void MemRefAccess::getAccessMap(AffineValueMap *accessMap) const {

mlir/lib/Analysis/AffineStructures.cpp

if (failed(this->addAffineForOpDomain(loop)))
       LLVM_DEBUG(
           loop.emitWarning("failed to add domain info to constraint system"));
     return;
   }
   // Add top level symbol.
   addSymbolId(getNumSymbolIds(), val);
   // Check if the symbol is a constant.
   if (auto constOp = val.getDefiningOp<ConstantIndexOp>())
-    setIdToConstant(val, constOp.getValue());
+    addBound(BoundType::EQ, val, constOp.getValue());
 }
 
 LogicalResult
 FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) {
   unsigned pos;
   // Pre-condition for this method.
   if (!findId(forOp.getInductionVar(), &pos)) {
     assert(false && "Value not found");
else {
       eq.back() -= lb;
       // For the local var just added above.
       eq[getNumCols() - 2] = -step;
       addEquality(eq);
     }
   }
 
   if (forOp.hasConstantLowerBound()) {
-    addConstantLowerBound(pos, forOp.getConstantLowerBound());
+    addBound(BoundType::LB, pos, forOp.getConstantLowerBound());
   } else {
     // Non-constant lower bound case.
-    if (failed(addLowerOrUpperBound(pos, forOp.getLowerBoundMap(),
-                                    forOp.getLowerBoundOperands(),
-                                    /*eq=*/false, /*lower=*/true)))
+    if (failed(addBound(BoundType::LB, pos, forOp.getLowerBoundMap(),
+                        forOp.getLowerBoundOperands())))
       return failure();
   }
 
   if (forOp.hasConstantUpperBound()) {
-    addConstantUpperBound(pos, forOp.getConstantUpperBound() - 1);
+    addBound(BoundType::UB, pos, forOp.getConstantUpperBound() - 1);
     return success();
   }
   // Non-constant upper bound case.
-  return addLowerOrUpperBound(pos, forOp.getUpperBoundMap(),
-                              forOp.getUpperBoundOperands(),
-                              /*eq=*/false, /*lower=*/false);
+  return addBound(BoundType::UB, pos, forOp.getUpperBoundMap(),
+                  forOp.getUpperBoundOperands());
 }
 
 LogicalResult
 FlatAffineValueConstraints::addDomainFromSliceMaps(ArrayRef<AffineMap> lbMaps,
                                                    ArrayRef<AffineMap> ubMaps,
                                                    ArrayRef<Value> operands) {
   assert(lbMaps.size() == ubMaps.size());
   assert(lbMaps.size() <= getNumDimIds());
if (lbMap && ubMap && lbMap.getNumResults() == 1 &&
       if (failed(addAffineForOpDomain(loop)))
         return failure();
       continue;
     }
 
     // This slice refers to a loop that doesn't exist in the IR yet. Add its
     // bounds to the system assuming its dimension identifier position is the
     // same as the position of the loop in the loop nest.
-    if (lbMap && failed(addLowerOrUpperBound(i, lbMap, operands, /*eq=*/false,
-                                             /*lower=*/true)))
+    if (lbMap && failed(addBound(BoundType::LB, i, lbMap, operands)))
       return failure();
-
-    if (ubMap && failed(addLowerOrUpperBound(i, ubMap, operands, /*eq=*/false,
-                                             /*lower=*/false)))
+    if (ubMap && failed(addBound(BoundType::UB, i, ubMap, operands)))
       return failure();
   }
   return success();
 }
 
 void FlatAffineValueConstraints::addAffineIfOpDomain(AffineIfOp ifOp) {
   // Create the base constraints from the integer set attached to ifOp.
   FlatAffineValueConstraints cst(ifOp.getIntegerSet());
for (unsigned curEquality = 0, numEqualities = cst.getNumEqualities();
     // TODO: Handle AffineSymbolExpr as well. There is no reason to restrict it
     // to dims themselves.
     auto dimExpr = dividendExpr.dyn_cast<AffineDimExpr>();
     if (!dimExpr)
       continue;
 
     // Express `id_r` as `id_n % divisor` and store the expression in `memo`.
     if (quotientCount >= 1) {
-      auto ub = cst.getConstantUpperBound(dimExpr.getPosition());
+      auto ub = cst.getConstantBound(FlatAffineConstraints::BoundType::UB,
+                                     dimExpr.getPosition());
       // If `id_n` has an upperbound that is less than the divisor, mod can be
       // eliminated altogether.
       if (ub.hasValue() && ub.getValue() < divisor)
         memo[pos] = dimExpr;
       else
         memo[pos] = dimExpr % divisor;
       // If a unique quotient `id_q` was seen, it can be expressed as
       // `id_n floordiv divisor`.
void FlatAffineConstraints::getSliceBounds(unsigned offset, unsigned num,
   do {
     changed = false;
     // Identify yet unknown identifiers as constants or mod's / floordiv's of
     // other identifiers if possible.
     for (unsigned pos = 0; pos < getNumIds(); pos++) {
       if (memo[pos])
         continue;
 
-      auto lbConst = getConstantLowerBound(pos);
-      auto ubConst = getConstantUpperBound(pos);
+      auto lbConst = getConstantBound(BoundType::LB, pos);
+      auto ubConst = getConstantBound(BoundType::UB, pos);
       if (lbConst.hasValue() && ubConst.hasValue()) {
         // Detect equality to a constant.
         if (lbConst.getValue() == ubConst.getValue()) {
           memo[pos] = getAffineConstantExpr(lbConst.getValue(), context);
           changed = true;
           continue;
         }
 
if (expr) {
       // If the above fails, we'll just use the constant lower bound and the
       // constant upper bound (if they exist) as the slice bounds.
       // TODO: being conservative for the moment in cases that
       // lead to multiple bounds - until getConstDifference in LoopFusion.cpp is
       // fixed (b/126426796).
       if (!lbMap || lbMap.getNumResults() > 1) {
         LLVM_DEBUG(llvm::dbgs()
                    << "WARNING: Potentially over-approximating slice lb\n");
-        auto lbConst = getConstantLowerBound(pos + offset);
+        auto lbConst = getConstantBound(BoundType::LB, pos + offset);
         if (lbConst.hasValue()) {
           lbMap = AffineMap::get(
               numMapDims, numMapSymbols,
               getAffineConstantExpr(lbConst.getValue(), context));
         }
       }
       if (!ubMap || ubMap.getNumResults() > 1) {
         LLVM_DEBUG(llvm::dbgs()
                    << "WARNING: Potentially over-approximating slice ub\n");
-        auto ubConst = getConstantUpperBound(pos + offset);
+        auto ubConst = getConstantBound(BoundType::UB, pos + offset);
         if (ubConst.hasValue()) {
           (ubMap) = AffineMap::get(
               numMapDims, numMapSymbols,
               getAffineConstantExpr(ubConst.getValue() + 1, context));
         }
       }
     }
     LLVM_DEBUG(llvm::dbgs()
if (localCst.getNumLocalIds() > 0) {
     // Dimensions of localCst and this constraint set match. Append localCst to
     // this constraint set.
     append(localCst);
   }
 
   return success();
 }
 
-LogicalResult FlatAffineConstraints::addLowerOrUpperBound(unsigned pos,
-                                                          AffineMap boundMap,
-                                                          bool eq, bool lower) {
+LogicalResult FlatAffineConstraints::addBound(BoundType type, unsigned pos,
+                                              AffineMap boundMap) {
   assert(boundMap.getNumDims() == getNumDimIds() && "dim mismatch");
   assert(boundMap.getNumSymbols() == getNumSymbolIds() && "symbol mismatch");
   assert(pos < getNumDimAndSymbolIds() && "invalid position");
 
   // Equality follows the logic of lower bound except that we add an equality
   // instead of an inequality.
-  assert((!eq || boundMap.getNumResults() == 1) && "single result expected");
-  if (eq)
-    lower = true;
+  assert((type != BoundType::EQ || boundMap.getNumResults() == 1) &&
+         "single result expected");
+  bool lower = type == BoundType::LB || type == BoundType::EQ;
 
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   if (failed(flattenAlignedMapAndMergeLocals(boundMap, &flatExprs)))
     return failure();
   assert(flatExprs.size() == boundMap.getNumResults());
 
   // Add one (in)equality for each result.
   for (const auto &flatExpr : flatExprs) {
for (const auto &flatExpr : flatExprs) {
     for (unsigned i = boundMap.getNumInputs(); i < end; i++, j++) {
       ineq[j] = lower ? -flatExpr[i] : flatExpr[i];
     }
     // Constant term.
     ineq[getNumCols() - 1] =
         lower ? -flatExpr[flatExpr.size() - 1]
               // Upper bound in flattenedExpr is an exclusive one.
               : flatExpr[flatExpr.size() - 1] - 1;
-    eq ? addEquality(ineq) : addInequality(ineq);
+    type == BoundType::EQ ? addEquality(ineq) : addInequality(ineq);
   }
 
   return success();
 }
 
 AffineMap
 FlatAffineValueConstraints::computeAlignedMap(AffineMap map,
                                               ValueRange operands) const {
AffineMap alignedMap =
       alignAffineMapWithValues(map, operands, dims, syms, newSymsPtr);
   // All symbols are already part of this FlatAffineConstraints.
   assert(syms.size() == newSymsPtr->size() && "unexpected new/missing symbols");
   assert(std::equal(syms.begin(), syms.end(), newSymsPtr->begin()) &&
          "unexpected new/missing symbols");
   return alignedMap;
 }
 
-LogicalResult FlatAffineValueConstraints::addLowerOrUpperBound(
-    unsigned pos, AffineMap boundMap, ValueRange boundOperands, bool eq,
-    bool lower) {
+LogicalResult FlatAffineValueConstraints::addBound(BoundType type, unsigned pos,
+                                                   AffineMap boundMap,
+                                                   ValueRange boundOperands) {
   // Fully compose map and operands; canonicalize and simplify so that we
   // transitively get to terminal symbols or loop IVs.
   auto map = boundMap;
   SmallVector<Value, 4> operands(boundOperands.begin(), boundOperands.end());
   fullyComposeAffineMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
   canonicalizeMapAndOperands(&map, &operands);
   for (auto operand : operands)
     addInductionVarOrTerminalSymbol(operand);
-  return addLowerOrUpperBound(pos, computeAlignedMap(map, operands), eq, lower);
+  return addBound(type, pos, computeAlignedMap(map, operands));
 }
 
 // Adds slice lower bounds represented by lower bounds in 'lbMaps' and upper
 // bounds in 'ubMaps' to each value in `values' that appears in the constraint
 // system. Note that both lower/upper bounds share the same operand list
 // 'operands'.
 // This function assumes 'values.size' == 'lbMaps.size' == 'ubMaps.size', and
 // skips any null AffineMaps in 'lbMaps' or 'ubMaps'.
for (unsigned i = 0, e = lbMaps.size(); i < e; ++i) {
     AffineMap ubMap = ubMaps[i];
     assert(!lbMap || lbMap.getNumInputs() == operands.size());
     assert(!ubMap || ubMap.getNumInputs() == operands.size());
 
     // Check if this slice is just an equality along this dimension.
     if (lbMap && ubMap && lbMap.getNumResults() == 1 &&
         ubMap.getNumResults() == 1 &&
         lbMap.getResult(0) + 1 == ubMap.getResult(0)) {
-      if (failed(addLowerOrUpperBound(pos, lbMap, operands, /*eq=*/true,
-                                      /*lower=*/true)))
+      if (failed(addBound(BoundType::EQ, pos, lbMap, operands)))
         return failure();
       continue;
     }
 
     // If lower or upper bound maps are null or provide no results, it implies
     // that the source loop was not at all sliced, and the entire loop will be a
     // part of the slice.
     if (lbMap && lbMap.getNumResults() != 0 && ubMap &&
         ubMap.getNumResults() != 0) {
-      if (failed(addLowerOrUpperBound(pos, lbMap, operands, /*eq=*/false,
-                                      /*lower=*/true)))
+      if (failed(addBound(BoundType::LB, pos, lbMap, operands)))
         return failure();
-      if (failed(addLowerOrUpperBound(pos, ubMap, operands, /*eq=*/false,
-                                      /*lower=*/false)))
+      if (failed(addBound(BoundType::UB, pos, ubMap, operands)))
         return failure();
     } else {
       auto loop = getForInductionVarOwner(values[i]);
       if (failed(this->addAffineForOpDomain(loop)))
         return failure();
     }
   }
   return success();
 }
 
 void FlatAffineConstraints::addEquality(ArrayRef<int64_t> eq) {
   assert(eq.size() == getNumCols());
   unsigned row = equalities.appendExtraRow();
   for (unsigned i = 0, e = eq.size(); i < e; ++i)
     equalities(row, i) = eq[i];
 }
 
 void FlatAffineConstraints::addInequality(ArrayRef<int64_t> inEq) {
   assert(inEq.size() == getNumCols());
   unsigned row = inequalities.appendExtraRow();
   for (unsigned i = 0, e = inEq.size(); i < e; ++i)
     inequalities(row, i) = inEq[i];
 }
 
-void FlatAffineConstraints::addConstantLowerBound(unsigned pos, int64_t lb) {
+void FlatAffineConstraints::addBound(BoundType type, unsigned pos,
+                                     int64_t value) {
   assert(pos < getNumCols());
+  if (type == BoundType::EQ) {
+    unsigned row = equalities.appendExtraRow();
+    equalities(row, pos) = 1;
+    equalities(row, getNumCols() - 1) = -value;
+  } else {
     unsigned row = inequalities.appendExtraRow();
-  inequalities(row, pos) = 1;
-  inequalities(row, getNumCols() - 1) = -lb;
+    inequalities(row, pos) = type == BoundType::LB ? 1 : -1;
+    inequalities(row, getNumCols() - 1) =
+        type == BoundType::LB ? -value : value;
   }
-
-void FlatAffineConstraints::addConstantUpperBound(unsigned pos, int64_t ub) {
-  assert(pos < getNumCols());
-  unsigned row = inequalities.appendExtraRow();
-  inequalities(row, pos) = -1;
-  inequalities(row, getNumCols() - 1) = ub;
 }
 
-void FlatAffineConstraints::addConstantLowerBound(ArrayRef<int64_t> expr,
-                                                  int64_t lb) {
-  addInequality(expr);
-  inequalities(inequalities.getNumRows() - 1, getNumCols() - 1) += -lb;
-}
-
-void FlatAffineConstraints::addConstantUpperBound(ArrayRef<int64_t> expr,
-                                                  int64_t ub) {
+void FlatAffineConstraints::addBound(BoundType type, ArrayRef<int64_t> expr,
+                                     int64_t value) {
+  assert(type != BoundType::EQ && "EQ not implemented");
   assert(expr.size() == getNumCols());
   unsigned row = inequalities.appendExtraRow();
   for (unsigned i = 0, e = expr.size(); i < e; ++i)
-    inequalities(row, i) = -expr[i];
-  inequalities(inequalities.getNumRows() - 1, getNumCols() - 1) += ub;
+    inequalities(row, i) = type == BoundType::LB ? expr[i] : -expr[i];
+  inequalities(inequalities.getNumRows() - 1, getNumCols() - 1) +=
+      type == BoundType::LB ? -value : value;
 }
 
 /// Adds a new local identifier as the floordiv of an affine function of other
 /// identifiers, the coefficients of which are provided in 'dividend' and with
 /// respect to a positive constant 'divisor'. Two constraints are added to the
 /// system to capture equivalence with the floordiv.
 ///      q = expr floordiv c    <=>   c*q <= expr <= c*q + c - 1.
 void FlatAffineConstraints::addLocalFloorDiv(ArrayRef<int64_t> dividend,

 
 void FlatAffineConstraints::setDimSymbolSeparation(unsigned newSymbolCount) {
   assert(newSymbolCount <= numDims + numSymbols &&
          "invalid separation position");
   numDims = numDims + numSymbols - newSymbolCount;
   numSymbols = newSymbolCount;
 }
 
-/// Sets the specified identifier to a constant value.
-void FlatAffineConstraints::setIdToConstant(unsigned pos, int64_t val) {
-  equalities.resizeVertically(equalities.getNumRows() + 1);
-  unsigned row = equalities.getNumRows() - 1;
-  equalities(row, pos) = 1;
-  equalities(row, getNumCols() - 1) = -val;
-}
-
-/// Sets the specified identifier to a constant value; asserts if the id is not
-/// found.
-void FlatAffineValueConstraints::setIdToConstant(Value value, int64_t val) {
+void FlatAffineValueConstraints::addBound(BoundType type, Value val,
+                                          int64_t value) {
   unsigned pos;
-  if (!findId(value, &pos))
+  if (!findId(val, &pos))
     // This is a pre-condition for this method.
     assert(0 && "id not found");
-  setIdToConstant(pos, val);
+  addBound(type, pos, value);
 }
 
 void FlatAffineConstraints::removeEquality(unsigned pos) {
   equalities.removeRow(pos);
 }
 
 void FlatAffineConstraints::removeInequality(unsigned pos) {
   inequalities.removeRow(pos);
for (unsigned r = 0, e = getNumInequalities(); r < e; r++) {
     } else {
       if (minOrMaxConst == None || boundConst < minOrMaxConst)
         minOrMaxConst = boundConst;
     }
   }
   return minOrMaxConst;
 }
 
-Optional<int64_t>
-FlatAffineConstraints::getConstantLowerBound(unsigned pos) const {
+Optional<int64_t> FlatAffineConstraints::getConstantBound(BoundType type,
+                                                          unsigned pos) const {
+  assert(type != BoundType::EQ && "EQ not implemented");
   FlatAffineConstraints tmpCst(*this);
+  if (type == BoundType::LB)
     return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/true>(pos);
-}
-
-Optional<int64_t>
-FlatAffineConstraints::getConstantUpperBound(unsigned pos) const {
-  FlatAffineConstraints tmpCst(*this);
   return tmpCst.computeConstantLowerOrUpperBound</*isLower=*/false>(pos);
 }
 
 // A simple (naive and conservative) check for hyper-rectangularity.
 bool FlatAffineConstraints::isHyperRectangular(unsigned pos,
                                                unsigned num) const {
   assert(pos < getNumCols() - 1);
   // Check for two non-zero coefficients in the range [pos, pos + sum).
if (res == BoundCmpResult::Less || res == BoundCmpResult::Equal) {
       // i.e., i >= expr floordiv div <=> i >= (expr - div + 1) ceildiv div <=>
       // div * i >= expr - div + 1.
       minLb.back() -= lbFloorDivisor - 1;
     } else if (res == BoundCmpResult::Greater) {
       minLb = otherLb;
       minLb.back() -= otherLbFloorDivisor - 1;
     } else {
       // Uncomparable - check for constant lower/upper bounds.
-      auto constLb = getConstantLowerBound(d);
-      auto constOtherLb = otherCst.getConstantLowerBound(d);
+      auto constLb = getConstantBound(BoundType::LB, d);
+      auto constOtherLb = otherCst.getConstantBound(BoundType::LB, d);
       if (!constLb.hasValue() || !constOtherLb.hasValue())
         return failure();
       std::fill(minLb.begin(), minLb.end(), 0);
       minLb.back() = std::min(constLb.getValue(), constOtherLb.getValue());
     }
 
     // Do the same for ub's but max of upper bounds. Identify max.
     auto uRes = compareBounds(ub, otherUb);
     if (uRes == BoundCmpResult::Greater || uRes == BoundCmpResult::Equal) {
       maxUb = ub;
     } else if (uRes == BoundCmpResult::Less) {
       maxUb = otherUb;
     } else {
       // Uncomparable - check for constant lower/upper bounds.
-      auto constUb = getConstantUpperBound(d);
-      auto constOtherUb = otherCst.getConstantUpperBound(d);
+      auto constUb = getConstantBound(BoundType::UB, d);
+      auto constOtherUb = otherCst.getConstantBound(BoundType::UB, d);
       if (!constUb.hasValue() || !constOtherUb.hasValue())
         return failure();
       std::fill(maxUb.begin(), maxUb.end(), 0);
       maxUb.back() = std::max(constUb.getValue(), constOtherUb.getValue());
     }
 
     std::fill(newLb.begin(), newLb.end(), 0);
     std::fill(newUb.begin(), newUb.end(), 0);

mlir/lib/Analysis/Utils.cpp

ComputationSliceState::getAsConstraints(FlatAffineValueConstraints *cst) {
   // Add loop bound constraints for values which are loop IVs of the destination
   // of fusion and equality constraints for symbols which are constants.
   for (unsigned i = numDims, end = values.size(); i < end; ++i) {
     Value value = values[i];
     assert(cst->containsId(value) && "value expected to be present");
     if (isValidSymbol(value)) {
       // Check if the symbol is a constant.
       if (auto cOp = value.getDefiningOp<ConstantIndexOp>())
-        cst->setIdToConstant(value, cOp.getValue());
+        cst->addBound(FlatAffineConstraints::EQ, value, cOp.getValue());
     } else if (auto loop = getForInductionVarOwner(value)) {
       if (failed(cst->addAffineForOpDomain(loop)))
         return failure();
     }
   }
 
   // Add slices bounds on 'ivs' using maps 'lbs'/'ubs' with 'lbOperands[0]'
   LogicalResult ret = cst->addSliceBounds(ivs, lbs, ubs, lbOperands[0]);
Optional<int64_t> MemRefRegion::getConstantBoundingSizeAndShape(
 
   // Use a copy of the region constraints that has upper/lower bounds for each
   // memref dimension with static size added to guard against potential
   // over-approximation from projection or union bounding box. We may not add
   // this on the region itself since they might just be redundant constraints
   // that will need non-trivials means to eliminate.
   FlatAffineConstraints cstWithShapeBounds(cst);
   for (unsigned r = 0; r < rank; r++) {
-    cstWithShapeBounds.addConstantLowerBound(r, 0);
+    cstWithShapeBounds.addBound(FlatAffineConstraints::LB, r, 0);
     int64_t dimSize = memRefType.getDimSize(r);
     if (ShapedType::isDynamic(dimSize))
       continue;
-    cstWithShapeBounds.addConstantUpperBound(r, dimSize - 1);
+    cstWithShapeBounds.addBound(FlatAffineConstraints::UB, r, dimSize - 1);
   }
 
   // Find a constant upper bound on the extent of this memref region along each
   // dimension.
   int64_t numElements = 1;
   int64_t diffConstant;
   int64_t lbDivisor;
   for (unsigned d = 0; d < rank; d++) {
if (auto loop = getForInductionVarOwner(operand)) {
         return failure();
     } else {
       // Has to be a valid symbol.
       auto symbol = operand;
       assert(isValidSymbol(symbol));
       // Check if the symbol is a constant.
       if (auto *op = symbol.getDefiningOp()) {
         if (auto constOp = dyn_cast<ConstantIndexOp>(op)) {
-          cst.setIdToConstant(symbol, constOp.getValue());
+          cst.addBound(FlatAffineConstraints::EQ, symbol, constOp.getValue());
         }
       }
     }
   }
 
   // Add lower/upper bounds on loop IVs using bounds from 'sliceState'.
   if (sliceState != nullptr) {
     // Add dim and symbol slice operands.
LogicalResult MemRefRegion::compute(Operation *op, unsigned loopDepth,
   assert(cst.getNumDimIds() == rank && "unexpected MemRefRegion format");
 
   // Add upper/lower bounds for each memref dimension with static size
   // to guard against potential over-approximation from projection.
   // TODO: Support dynamic memref dimensions.
   if (addMemRefDimBounds) {
     auto memRefType = memref.getType().cast<MemRefType>();
     for (unsigned r = 0; r < rank; r++) {
-      cst.addConstantLowerBound(/*pos=*/r, /*lb=*/0);
+      cst.addBound(FlatAffineConstraints::LB, /*pos=*/r, /*value=*/0);
       if (memRefType.isDynamicDim(r))
         continue;
-      cst.addConstantUpperBound(/*pos=*/r, memRefType.getDimSize(r) - 1);
+      cst.addBound(FlatAffineConstraints::UB, /*pos=*/r,
+                   memRefType.getDimSize(r) - 1);
     }
   }
   cst.removeTrivialRedundancy();
 
   LLVM_DEBUG(llvm::dbgs() << "Memory region:\n");
   LLVM_DEBUG(cst.dump());
   return success();
 }
for (unsigned r = 0; r < rank; r++) {
     // feasible. If it is, there is at least one point out of bounds.
     SmallVector<int64_t, 4> ineq(rank + 1, 0);
     int64_t dimSize = loadOrStoreOp.getMemRefType().getDimSize(r);
     // TODO: handle dynamic dim sizes.
     if (dimSize == -1)
       continue;
 
     // Check for overflow: d_i >= memref dim size.
-    ucst.addConstantLowerBound(r, dimSize);
+    ucst.addBound(FlatAffineConstraints::LB, r, dimSize);
     outOfBounds = !ucst.isEmpty();
     if (outOfBounds && emitError) {
       loadOrStoreOp.emitOpError()
           << "memref out of upper bound access along dimension #" << (r + 1);
     }
 
     // Check for a negative index.
     FlatAffineConstraints lcst(*region.getConstraints());
     std::fill(ineq.begin(), ineq.end(), 0);
     // d_i <= -1;
-    lcst.addConstantUpperBound(r, -1);
+    lcst.addBound(FlatAffineConstraints::UB, r, -1);
     outOfBounds = !lcst.isEmpty();
     if (outOfBounds && emitError) {
       loadOrStoreOp.emitOpError()
           << "memref out of lower bound access along dimension #" << (r + 1);
     }
   }
   return failure(outOfBounds);
 }

mlir/lib/Transforms/Utils/LoopUtils.cpp

static bool getFullMemRefAsRegion(Operation *op, unsigned numParamLoopIVs,
   extractForInductionVars(ivs, &symbols);
   regionCst->reset(rank, numParamLoopIVs, 0);
   regionCst->setValues(rank, rank + numParamLoopIVs, symbols);
 
   // Memref dim sizes provide the bounds.
   for (unsigned d = 0; d < rank; d++) {
     auto dimSize = memRefType.getDimSize(d);
     assert(dimSize > 0 && "filtered dynamic shapes above");
-    regionCst->addConstantLowerBound(d, 0);
-    regionCst->addConstantUpperBound(d, dimSize - 1);
+    regionCst->addBound(FlatAffineConstraints::LB, d, 0);
+    regionCst->addBound(FlatAffineConstraints::UB, d, dimSize - 1);
   }
   return true;
 }
 
 /// Performs explicit copying for the contiguous sequence of operations in the
 /// block iterator range [`begin', `end'), where `end' can't be past the
 /// terminator of the block (since additional operations are potentially
 /// inserted right before `end`. Returns the total size of fast memory space

mlir/lib/Transforms/Utils/Utils.cpp

MemRefType mlir::normalizeMemRefType(MemRefType memrefType, OpBuilder b,
   // with the right shape and an identity layout map.
   ArrayRef<int64_t> shape = memrefType.getShape();
   // FlatAffineConstraint may later on use symbolicOperands.
   FlatAffineConstraints fac(rank, numSymbolicOperands);
   SmallVector<unsigned, 4> memrefTypeDynDims;
   for (unsigned d = 0; d < rank; ++d) {
     // Use constraint system only in static dimensions.
     if (shape[d] > 0) {
-      fac.addConstantLowerBound(d, 0);
-      fac.addConstantUpperBound(d, shape[d] - 1);
+      fac.addBound(FlatAffineConstraints::LB, d, 0);
+      fac.addBound(FlatAffineConstraints::UB, d, shape[d] - 1);
     } else {
       memrefTypeDynDims.emplace_back(d);
     }
   }
   // We compose this map with the original index (logical) space to derive
   // the upper bounds for the new index space.
   AffineMap layoutMap = layoutMaps.front();
   unsigned newRank = layoutMap.getNumResults();
   if (failed(fac.composeMatchingMap(layoutMap)))
     return memrefType;
   // TODO: Handle semi-affine maps.
   // Project out the old data dimensions.
   fac.projectOut(newRank, fac.getNumIds() - newRank - fac.getNumLocalIds());
   SmallVector<int64_t, 4> newShape(newRank);
   for (unsigned d = 0; d < newRank; ++d) {
     // Check if each dimension of normalized memrefType is dynamic.
     bool isDynDim = isNormalizedMemRefDynamicDim(
         d, layoutMap, memrefTypeDynDims, b.getContext());
     if (isDynDim) {
       newShape[d] = -1;
     } else {
       // The lower bound for the shape is always zero.
-      auto ubConst = fac.getConstantUpperBound(d);
+      auto ubConst = fac.getConstantBound(FlatAffineConstraints::UB, d);
       // For a static memref and an affine map with no symbols, this is
       // always bounded.
       assert(ubConst.hasValue() && "should always have an upper bound");
       if (ubConst.getValue() < 0)
         // This is due to an invalid map that maps to a negative space.
         return memrefType;
       // If dimension of new memrefType is dynamic, the value is -1.
       newShape[d] = ubConst.getValue() + 1;

mlir/unittests/Analysis/AffineStructuresTest.cpp

TEST(FlatAffineConstraintsTest, removeRedundantConstraintsTest) {
   for (unsigned i = 0; i < 3; ++i) {
     // Ensure that the removed constraint was the redundant constraint [3].
     EXPECT_NE(fac5.getInequality(i), ArrayRef<int64_t>(redundantConstraint));
   }
 }
 
 TEST(FlatAffineConstraintsTest, addConstantUpperBound) {
   FlatAffineConstraints fac = makeFACFromConstraints(2, {}, {});
-  fac.addConstantUpperBound(0, 1);
+  fac.addBound(FlatAffineConstraints::UB, 0, 1);
   EXPECT_EQ(fac.atIneq(0, 0), -1);
   EXPECT_EQ(fac.atIneq(0, 1), 0);
   EXPECT_EQ(fac.atIneq(0, 2), 1);
 
-  fac.addConstantUpperBound({1, 2, 3}, 1);
+  fac.addBound(FlatAffineConstraints::UB, {1, 2, 3}, 1);
   EXPECT_EQ(fac.atIneq(1, 0), -1);
   EXPECT_EQ(fac.atIneq(1, 1), -2);
   EXPECT_EQ(fac.atIneq(1, 2), -2);
 }
 
 TEST(FlatAffineConstraintsTest, addConstantLowerBound) {
   FlatAffineConstraints fac = makeFACFromConstraints(2, {}, {});
-  fac.addConstantLowerBound(0, 1);
+  fac.addBound(FlatAffineConstraints::LB, 0, 1);
   EXPECT_EQ(fac.atIneq(0, 0), 1);
   EXPECT_EQ(fac.atIneq(0, 1), 0);
   EXPECT_EQ(fac.atIneq(0, 2), -1);
 
-  fac.addConstantLowerBound({1, 2, 3}, 1);
+  fac.addBound(FlatAffineConstraints::LB, {1, 2, 3}, 1);
   EXPECT_EQ(fac.atIneq(1, 0), 1);
   EXPECT_EQ(fac.atIneq(1, 1), 2);
   EXPECT_EQ(fac.atIneq(1, 2), 2);
 }
 
 TEST(FlatAffineConstraintsTest, clearConstraints) {
   FlatAffineConstraints fac = makeFACFromConstraints(1, {}, {});