[mlir][affine] addLowerOrUpperBound: Make map+operand composing optional

[springerm]

Summary

There are cases in which it is not desirable to fully compose the bound map with the operands when adding lower/upper bounds to a FlatAffineConstraints.

E.g., this is the case when bounds should be expressed in terms of the operands only (and not the operands' dependencies). This also makes addLowerOrUpperBound useable together with operands that are defined through semi-affine expressions.

Depends On D107220

[springerm]

update

mlir/include/mlir/Analysis/AffineStructures.h

public:
   /// 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.
+  /// first result expr. By default, the bound map is fully composed with the
+  /// operands before adding any bounds. This allows for bounds being expressed
+  /// in terms of values that are used by the operands.
   LogicalResult addLowerOrUpperBound(unsigned pos, AffineMap boundMap,
                                      ValueRange operands, bool eq,
-                                     bool lower = true);
+                                     bool lower = true,
+                                     bool composeMapAndOperands = true);
 
   /// 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,
public:
   void addId(IdKind kind, unsigned pos, Value id = nullptr);
 
   /// Add the specified values as a dim or symbol id depending on its nature, if
   /// it already doesn't exist in the system. `id' has to be either a terminal
   /// symbol or a loop IV, i.e., it cannot be the result affine.apply of any
   /// symbols or loop IVs. The identifier is added to the end of the existing
   /// dims or symbols. Additional information on the identifier is extracted
   /// from the IR and added to the constraint system.
-  void addInductionVarOrTerminalSymbol(Value id);
+  /// Note: If `allowNonTerminal`, any symbol (incl. potentially non-terminal
+  /// ones) is allowed.
+  void addInductionVarOrTerminalSymbol(Value id, bool allowNonTerminal = false);
 
   /// Composes the affine value map with this FlatAffineConstrains, adding the
   /// results of the map as dimensions at the front [0, vMap->getNumResults())
   /// and with the dimensions set to the equalities specified by the value map.
   /// Returns failure if the composition fails (when vMap is a semi-affine map).
   /// The vMap's operand Value's are used to look up the right positions in
   /// the FlatAffineConstraints with which to associate. The dimensional and
   /// symbolic operands of vMap should match 1:1 (in the same order) with those

mlir/lib/Analysis/AffineStructures.cpp

if (ids[i].hasValue() && getForInductionVarOwner(ids[i].getValue()))
       loopIVs.push_back(ids[i].getValue());
   }
   // Turn each symbol in 'loopIVs' into a dim identifier.
   for (auto iv : loopIVs) {
     turnSymbolIntoDim(this, iv);
   }
 }
 
-void FlatAffineConstraints::addInductionVarOrTerminalSymbol(Value id) {
+void FlatAffineConstraints::addInductionVarOrTerminalSymbol(
+    Value id, bool allowNonTerminal) {
   if (containsId(id))
     return;
 
   // Caller is expected to fully compose map/operands if necessary.
-  assert((isTopLevelValue(id) || isForInductionVar(id)) &&
+  assert((allowNonTerminal || isTopLevelValue(id) || isForInductionVar(id)) &&
          "non-terminal symbol / loop IV expected");
   // Outer loop IVs could be used in forOp's bounds.
   if (auto loop = getForInductionVarOwner(id)) {
     addDimId(getNumDimIds(), id);
     if (failed(this->addAffineForOpDomain(loop)))
       LLVM_DEBUG(
           loop.emitWarning("failed to add domain info to constraint system"));
     return;
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,
-                                            ValueRange boundOperands, bool eq,
-                                            bool lower) {
+LogicalResult FlatAffineConstraints::addLowerOrUpperBound(
+    unsigned pos, AffineMap boundMap, ValueRange boundOperands, bool eq,
+    bool lower, bool composeMapAndOperands) {
   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;
 
   // 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());
+  if (composeMapAndOperands)
     fullyComposeAffineMapAndOperands(&map, &operands);
   map = simplifyAffineMap(map);
   canonicalizeMapAndOperands(&map, &operands);
   for (auto operand : operands)
-    addInductionVarOrTerminalSymbol(operand);
+    addInductionVarOrTerminalSymbol(
+        operand, /*allowNonTerminal=*/!composeMapAndOperands);
 
   FlatAffineConstraints localVarCst;
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst))) {
     LLVM_DEBUG(llvm::dbgs() << "semi-affine expressions not yet supported\n");
     return failure();
   }