[Mlir-commits] [mlir] [mlir][affine] Update getSliceBounds to allow multi-result upper bound maps (PR #219369)

[llvmlistbot at llvm.org]

================
@@ -1791,21 +1792,35 @@ mlir::affine::computeSliceUnion(ArrayRef<Operation *> opsA,
   return SliceComputationResult::Success;
 }
 
-// TODO: extend this to handle multiple result maps.
+/// Returns the number of iterations the slice bounded below by `lbMap` and
+/// above by `ubMap` runs for, where that is a constant.
+///
+/// An upper bound of several results is the min of them, so each result taken
+/// against the lower bound bounds the count from above and the smallest of
+/// those that comes out constant is the tightest constant bound there is. A
+/// tiled loop clamped at the end of the data has exactly this shape --
+/// `min(%i * 64 + 64, 1000)` over `%i * 64` -- where the tile-relative result
+/// gives the 64 and the extent gives nothing constant at all.
 static std::optional<uint64_t> getConstDifference(AffineMap lbMap,
                                                   AffineMap ubMap) {
-  assert(lbMap.getNumResults() == 1 && "expected single result bound map");
-  assert(ubMap.getNumResults() == 1 && "expected single result bound map");
+  assert(lbMap.getNumResults() == 1 && "expected single result lower bound");
+  assert(ubMap.getNumResults() >= 1 && "expected at least one upper bound");
   assert(lbMap.getNumDims() == ubMap.getNumDims());
   assert(lbMap.getNumSymbols() == ubMap.getNumSymbols());
   AffineExpr lbExpr(lbMap.getResult(0));
-  AffineExpr ubExpr(ubMap.getResult(0));
-  auto loopSpanExpr = simplifyAffineExpr(ubExpr - lbExpr, lbMap.getNumDims(),
-                                         lbMap.getNumSymbols());
-  auto cExpr = dyn_cast<AffineConstantExpr>(loopSpanExpr);
-  if (!cExpr)
-    return std::nullopt;
-  return cExpr.getValue();
+  std::optional<uint64_t> tripCount;
+  for (AffineExpr ubExpr : ubMap.getResults()) {
+    AffineExpr loopSpanExpr = simplifyAffineExpr(
+        ubExpr - lbExpr, lbMap.getNumDims(), lbMap.getNumSymbols());
+    auto cExpr = dyn_cast<AffineConstantExpr>(loopSpanExpr);
+    if (!cExpr)
+      continue;
+    if (cExpr.getValue() < 0)
+      return 0;
+    tripCount =
+        std::min(tripCount.value_or(UINT64_MAX), (uint64_t)cExpr.getValue());
----------------
lipracer wrote:

Two more nits on this line, feel free to ignore:
- `(uint64_t)` → `static_cast<uint64_t>`, which is what the rest of mlir/lib uses.
- `UINT64_MAX` → `std::numeric_limits<uint64_t>::max()`; `UINT64_MAX` doesn't
  appear anywhere else under mlir/.

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