[mlir][Analysis] Add FlatAffineConstraints::addLowerOrUpperBound

[springerm]

Summary

This function overload is similar to the existing FlatAffineConstraints::addLowerOrUpperBound. It constrains a dimension based on an affine map. However, in contrast to the other overloading, it does not attempt to align dimensions/symbols of the affine map with the dimensions/symbols of the constraint set. Instead, dimensions/symbols are expected to already be aligned.

Depends On D107726

[ftynse]

mlir/lib/Analysis/AffineStructures.cpp

2055

Let's omit the number of stack elements in SmallVector unless there is a good reason to choose a specific number.

2060

If it's invalid, should this be an assertion instead of continue?

[springerm]

mlir/lib/Analysis/AffineStructures.cpp

2060

I think so. But some existing code/test cases would fail. E.g., FlatAffineConstraints::addDomainFromSliceMaps adds invalid bounds here and relies on the fact that the other (non-invalid) inequalities are being added:

// 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)))
  return failure();

Looks like a bug in addDomainFromSliceMaps or one of its callers. Maybe better to fix it now. Taking a look...

[springerm]

mlir/lib/Analysis/AffineStructures.cpp

2060

I don't know enough about the code base (in particular ComputationSliceState) to fix this. Putting a TODO in here for now, and continuing with the original refactoring.

[springerm]

rebase

[nicolasvasilache]

Great to see this refactoring chain, some fly-by comments.

mlir/lib/Analysis/AffineStructures.cpp

2025

I find this somewhat unfortunate.
How about addBound(unsigned pos, AffineMap boundMap, BoundType type) where enum class BoundType { EQ, LB, UB }; ?

2045

While we are refactoring and cleaning up, I would drop all these one off functions and have a single, batched:
addIds(IdType idType, int pos, int numIds = 1);
where enum class IdType { Dim, Sym, Local };

The all this could be simply rewritten as:

if (...) {
  addIds(IdType::Local, 0,  localCst.getNumLocalIds());
  localCst.addIds(IdType::Local, localCst.getNumLocalIds(), getNumLocalIds());
  append(localCst);
}

[springerm]

mlir/lib/Analysis/AffineStructures.cpp

2025

Addressed in a subsequent commit.

2045

It's on my TODO list.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   // 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 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.
+  /// 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);
+
   /// 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.
public:
 
   /// 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;
 
   /// 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,

mlir/lib/Analysis/AffineStructures.cpp

LLVM_DEBUG(llvm::dbgs()
                << "lb map for pos = " << Twine(pos + offset) << ", expr: ");
     LLVM_DEBUG(lbMap.dump(););
     LLVM_DEBUG(llvm::dbgs()
                << "ub map for pos = " << Twine(pos + offset) << ", expr: ");
     LLVM_DEBUG(ubMap.dump(););
   }
 }
 
-// TODO: Add FlatAffineConstraints::addLowerOrUpperBound
+LogicalResult FlatAffineConstraints::addLowerOrUpperBound(unsigned pos,
+                                                          AffineMap boundMap,
+                                                          bool eq, bool lower) {
+  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;
+
+  std::vector<SmallVector<int64_t, 8>> flatExprs;
+  FlatAffineConstraints localCst;
+  if (failed(getFlattenedAffineExprs(boundMap, &flatExprs, &localCst))) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "composition unimplemented for semi-affine maps\n");
+    return failure();
+  }
+  assert(flatExprs.size() == boundMap.getNumResults());
+
+  // Add localCst information.
+  if (localCst.getNumLocalIds() > 0) {
+    unsigned numLocalIds = getNumLocalIds();
+    // Insert local dims of localCst at the beginning.
+    for (unsigned l = 0, e = localCst.getNumLocalIds(); l < e; ++l)
+      addLocalId(0);
+    // Insert local dims of `this` at the end of localCst.
+    for (unsigned l = 0; l < numLocalIds; ++l)
+      localCst.addLocalId(localCst.getNumLocalIds());
+    // Dimensions of localCst and this constraint set match. Append localCst to
+    // this constraint set.
+    append(localCst);
+  }
+
+  // Add one (in)equality for each result.
+  for (const auto &flatExpr : flatExprs) {
+    SmallVector<int64_t, 4> ineq(getNumCols(), 0);
+    // Dims and symbols.
+    for (unsigned j = 0, e = boundMap.getNumInputs(); j < e; j++) {
+      ineq[j] = lower ? -flatExpr[j] : flatExpr[j];
+    }
+    // Invalid bound: pos appears in `boundMap`.
+    if (ineq[pos] != 0)
+      continue;
+    ineq[pos] = lower ? 1 : -1;
+    // Local vars common to eq and localCst are at the beginning.
+    unsigned j = getNumDimIds() + getNumSymbolIds();
+    unsigned end = flatExpr.size() - 1;
+    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);
+  }
+
+  return success();
+}
 
 LogicalResult FlatAffineValueConstraints::addLowerOrUpperBound(
     unsigned pos, AffineMap boundMap, ValueRange boundOperands, bool eq,
     bool lower) {
   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");

 void FlatAffineConstraints::print(raw_ostream &os) const {
   assert(hasConsistentState());
   os << "\nConstraints (" << getNumDimIds() << " dims, " << getNumSymbolIds()
      << " symbols, " << getNumLocalIds() << " locals), (" << getNumConstraints()
      << " constraints)\n";
   os << "(";
   for (unsigned i = 0, e = getNumIds(); i < e; i++) {
     if (auto *valueConstrSet =
             dynamic_cast<const FlatAffineValueConstraints *>(this)) {
       if (valueConstrSet->hasIdValue(i))
         os << "Value ";
       else
         os << "None ";
     } else {
       os << "None ";
     }
   }
void FlatAffineConstraints::removeTrivialRedundancy() {
 
   // TODO: consider doing this for equalities as well, but probably not worth
   // the savings.
 }
 
 void FlatAffineConstraints::clearAndCopyFrom(
     const FlatAffineConstraints &other) {
   if (auto *otherValueSet =
           dynamic_cast<const FlatAffineValueConstraints *>(&other))
     assert(!otherValueSet->hasIdValues() &&
            "cannot copy associated Values into FlatAffineConstraints");
   *this = other;
 }
 
 void FlatAffineValueConstraints::clearAndCopyFrom(
     const FlatAffineConstraints &other) {
   if (auto *otherValueSet =
           dynamic_cast<const FlatAffineValueConstraints *>(&other)) {
     *this = *otherValueSet;
   } else {
     *static_cast<FlatAffineConstraints *>(this) = other;
     ids.clear();
     ids.resize(numIds, None);
   }
 }