[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,

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,
-                                            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);
 
   FlatAffineConstraints localVarCst;
   std::vector<SmallVector<int64_t, 8>> flatExprs;
   if (failed(getFlattenedAffineExprs(map, &flatExprs, &localVarCst))) {