[Mlir-commits] [mlir] [mlir] Extend affine.min/max ValueBoundsOpInterfaceImpls (PR #118840)

[Matthias Springer]

================
@@ -67,6 +67,27 @@ struct AffineMinOpInterface
           expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
       cstr.bound(value) <= bound;
     }
+    // Get all constant lower bounds, choose minimum, and set lower bound to it.
+    MLIRContext *ctx = op->getContext();
+    AffineMap map = minOp.getAffineMap();
+    SmallVector<Value> mapOperands = minOp.getOperands();
+    std::optional<int64_t> minBound;
+    for (AffineExpr expr : map.getResults()) {
+      auto exprMap =
+          AffineMap::get(map.getNumDims(), map.getNumSymbols(), expr, ctx);
+      ValueBoundsConstraintSet::Variable exprVar(exprMap, mapOperands);
+      FailureOr<int64_t> exprBound =
+          cstr.computeConstantBound(presburger::BoundType::LB, exprVar,
----------------
matthias-springer wrote:

Let me make sure I understand the problem.

Can you double check that my math is correct here?
```mlir
func.func @affine_min_const_lb(%a: index) -> (index) {
  %0 = affine.max affine_map<(d0) -> (d0, 0)>(%a)
  %1 = affine.min affine_map<(d0) -> (d0, 2)>(%0)
  %2 = "test.reify_bound"(%1) {type = "LB"}: (index) -> (index)
  return %2 : index
}

%0 >= %a
%0 >= 0
LB(%0) = {%a, 0}   // both are lower bounds

%1 <= 2
%1 <= %0
LB(%1) = min(LB(2), LB(%0))
       = min(2, {%a, 0})
       = {min(2, %a), min(2, 0)}
       = {min(2, %a), 0}   // both are lower bounds
```

When there are multiple LBs / UBs, the infrastructure may return either one of the two bounds. In practice, for `affine.min`/`affine.max`, it currently returns the constant bound if there is one (see `FlatLinearConstraints::getSliceBounds` for details).

The problem with `populateAndCompare` is that it compares two SSA values, whereas we have to compare the LBs of two SSA values here.

So we need something like `populateAndComputeBound`. This allows us to compute LBs for both `affine.min` operands in the same constraint set. These LBs can then be compared with the existing `populateAndCompare`.

`populateAndComputeBound` would be similar to `ValueBoundsConstraintSet::computeBound`, but reuses the existing constraint set. It's not a static function. Also, it does not project out any variables.

Can you explore this direction a bit? I think it's worth optimizing for performance here. I'm afraid, the current approach may have exponential runtime complexity in the worst case, as every time an `affine.max/min` op is visited, we start two new analyses for the remaining (not yet visited) IR.


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