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

llvmlistbot at llvm.org llvmlistbot at llvm.org
Thu Dec 5 09:35:27 PST 2024


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

The ValueBoundsOpInterface implementation for affine.min and affine.max only supported bounding on one side (LB for affine.max and UB for affine.min). However, it is possible to be smarter about the bounding in some cases, and bound values on both sides for these ops. An affine.min op result can be lower bounded if a lower bound can be computed for every result expression in its affine map. In this case, the new lower bound can be set as the minimum of the lower bounds of all result expressions. The converse can be done for affine.max, computing the maximum of all result expression bounds. This PR adds this logic to the ValueBoundsOpInterface implementations of affine.min and affine.max.

>From a1ccd9e3d03684e5e7500a0069394d78d2fbc366 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 82a9fb0d490882..a575c0603aee70 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 935c08aceff548..8014a9c3d087ae 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]]



More information about the Mlir-commits mailing list