[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

1983

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

1988

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

[springerm]

mlir/lib/Analysis/AffineStructures.cpp

1988

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

1988

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

1954

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

1974

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

1954

Addressed in a subsequent commit.

1974

It's on my TODO list.

mlir/include/mlir/Analysis/AffineStructures.h

public:
   /// 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. 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 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);
 
   /// Returns the bound for the identifier at `pos` from the inequality at

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(););
   }
 }
 
+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> 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`.
+    // TODO: This should be an assertion. Fix `addDomainFromSliceMaps` and/or
+    // its callers to prevent invalid bounds from being added.
+    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
 FlatAffineConstraints::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");