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

[llvmlistbot at llvm.org]

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

>From 34e3e0ceb50163826ab7610827d1b7a8845da880 Mon Sep 17 00:00:00 2001
From: Max Dawkins <max.dawkins at gmail.com>
Date: Wed, 27 Nov 2024 14:30:37 -0600
Subject: [PATCH] [mlir] Extend affine.min/max ValueBoundsOpInterfaceImpls

Signed-off-by: Max Dawkins <max.dawkins at gmail.com>
---
 .../Affine/IR/ValueBoundsOpInterfaceImpl.cpp  | 42 +++++++++++++++++++
 .../value-bounds-op-interface-impl.mlir       | 26 ++++++++++++
 2 files changed, 68 insertions(+)

diff --git a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
index e93b99b4f49866d..f49691b20108e34 100644
--- a/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Affine/IR/ValueBoundsOpInterfaceImpl.cpp
@@ -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,
+                                    /*stopCondition=*/nullptr);
+      // If any LB cannot be computed, then the total LB cannot be known.
+      if (failed(exprBound))
+        return;
+      if (!minBound.has_value() || exprBound.value() < minBound.value())
+        minBound = exprBound.value();
+    }
+    if (!minBound.has_value())
+      return;
+    cstr.bound(value) >= minBound.value();
   };
 };
 
@@ -88,6 +109,27 @@ struct AffineMaxOpInterface
           expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
       cstr.bound(value) >= bound;
     }
+    // Get all constant upper bounds, choose maximum, and set upper bound to it.
+    MLIRContext *ctx = op->getContext();
+    AffineMap map = maxOp.getAffineMap();
+    SmallVector<Value> mapOperands = maxOp.getOperands();
+    std::optional<int64_t> maxBound;
+    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::UB, exprVar,
+          /*stopCondition=*/nullptr, /*closedUB=*/true);
+      // If any UB cannot be computed, then the total UB cannot be known.
+      if (failed(exprBound))
+        return;
+      if (!maxBound.has_value() || exprBound.value() > maxBound.value())
+        maxBound = exprBound.value();
+    }
+    if (!maxBound.has_value())
+      return;
+    cstr.bound(value) <= maxBound.value();
   };
 };
 
diff --git a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
index a4310b91a37b3dd..72af500e623a5f8 100644
--- a/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
+++ b/mlir/test/Dialect/Affine/value-bounds-op-interface-impl.mlir
@@ -38,6 +38,19 @@ func.func @affine_max_ub(%a: index) -> (index) {
 
 // -----
 
+// CHECK-LABEL: func @affine_max_const_ub(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK:   return %[[c5]]
+func.func @affine_max_const_ub(%a: index) -> (index) {
+  %0 = affine.min affine_map<(d0) -> (d0, 4)>(%a)
+  %1 = affine.max affine_map<(d0) -> (d0, 2)>(%0)
+  %2 = "test.reify_bound"(%1) {type = "UB"}: (index) -> (index)
+  return %2 : index
+}
+
+// -----
+
 // CHECK-LABEL: func @affine_min_ub(
 //  CHECK-SAME:     %[[a:.*]]: index
 //       CHECK:   %[[c3:.*]] = arith.constant 3 : index
@@ -61,6 +74,19 @@ func.func @affine_min_lb(%a: index) -> (index) {
 
 // -----
 
+// CHECK-LABEL: func @affine_min_const_lb(
+//  CHECK-SAME:     %[[a:.*]]: index
+//       CHECK:   %[[c0:.*]] = arith.constant 0 : index
+//       CHECK:   return %[[c0]]
+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
+}
+
+// -----
+
 // CHECK-LABEL: func @composed_affine_apply(
 //       CHECK:   %[[cst:.*]] = arith.constant -8 : index
 //       CHECK:   return %[[cst]]