[Mlir-commits] [mlir] [mlir][affine] Use value bound inference to determine minimum/maximum trip counts in loop analysis (PR #128113)

[Oleksandr Alex Zinenko]

================
@@ -223,20 +259,21 @@ std::optional<uint64_t> mlir::affine::getConstantTripCount(AffineForOp forOp) {
 
   if (!map)
     return std::nullopt;
+  return getKnownTripCountBound(map, operands, presburger::BoundType::LB);
+}
 
-  // Take the min if all trip counts are constant.
-  std::optional<uint64_t> tripCount;
-  for (auto resultExpr : map.getResults()) {
-    if (auto constExpr = dyn_cast<AffineConstantExpr>(resultExpr)) {
-      if (tripCount.has_value())
-        tripCount =
-            std::min(*tripCount, static_cast<uint64_t>(constExpr.getValue()));
-      else
-        tripCount = constExpr.getValue();
-    } else
-      return std::nullopt;
-  }
-  return tripCount;
+/// Returns the maximum trip count when the operand of forOp has a range. If the
+/// operand of forOp is a constant, the return value is the same as
+/// `getConstantTripCount`.
+std::optional<uint64_t>
+mlir::affine::getUpperBoundOnTripCount(AffineForOp forOp) {
+  SmallVector<Value, 4> operands;
----------------
ftynse wrote:

```suggestion
  SmallVector<Value> operands;
```

https://github.com/llvm/llvm-project/pull/128113