[mlir][affine] Split `reifyValueBound` in two functions

[springerm]

Summary

There are now two entry points. One for shaped values and one for index-typed values. This addresses a comment in D146524.

[mravishankar]

I don't think an API deprecation can be marked NFC

[dcaballe]

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

82

expect -> except? (2x)

mlir/include/mlir/Dialect/Affine/Transforms/Transforms.h

 /// AffineApplyOps. This can be used right before lowering affine ops to arith
 /// to exhibit more opportunities for CSE and LICM.
 /// Return the sink AffineApplyOp on success or failure if `op` does not
 /// decompose into smaller AffineApplyOps.
 /// Note that this can be undone by canonicalization which tries to
 /// maximally compose chains of AffineApplyOps.
 FailureOr<AffineApplyOp> decompose(RewriterBase &rewriter, AffineApplyOp op);
 
-/// Reify a bound for the given index-typed value or shape dimension size in
-/// terms of the owning op's operands. `dim` must be `nullopt` if and only if
-/// `value` is index-typed.
-///
-/// By default, lower/equal bounds are closed and upper bounds are open. If
-/// `closedUB` is set to "true", upper bounds are also closed.
-FailureOr<OpFoldResult> reifyValueBound(OpBuilder &b, Location loc,
-                                        presburger::BoundType type, Value value,
-                                        std::optional<int64_t> dim,
-                                        bool closedUB = false);
-
-/// Reify a bound for the given index-typed value or shape dimension size in
-/// terms of SSA values for which `stopCondition` is met. `dim` must be
-/// `nullopt` if and only if `value` is index-typed.
+/// Reify a bound for the given index-typed value in terms of SSA values for
+/// which `stopCondition` is met. If no stop condition is specified, reify in
+/// terms of the operands of the owner op.
 ///
 /// By default, lower/equal bounds are closed and upper bounds are open. If
 /// `closedUB` is set to "true", upper bounds are also closed.
 ///
 /// Example:
 /// %0 = arith.addi %a, %b : index
 /// %1 = arith.addi %0, %c : index
 ///
 /// * If `stopCondition` evaluates to "true" for %0 and %c, "%0 + %c" is an EQ
 ///   bound for %1.
 /// * If `stopCondition` evaluates to "true" for %a, %b and %c, "%a + %b + %c"
 ///   is an EQ bound for %1.
 /// * Otherwise, if the owners of %a, %b or %c do not implement the
 ///   ValueBoundsOpInterface, no bound can be computed.
-FailureOr<OpFoldResult>
-reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
-                Value value, std::optional<int64_t> dim,
-                ValueBoundsConstraintSet::StopConditionFn stopCondition,
-                bool closedUB = false);
+FailureOr<OpFoldResult> reifyIndexValueBound(
+    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
+    ValueBoundsConstraintSet::StopConditionFn stopCondition = nullptr,
+    bool closedUB = false);
+
+/// Reify a bound for the specified dimension of the given shaped value in terms
+/// of SSA values for which `stopCondition` is met. If no stop condition is
+/// specified, reify in terms of the operands of the owner op.
+///
+/// By default, lower/equal bounds are closed and upper bounds are open. If
+/// `closedUB` is set to "true", upper bounds are also closed.
+FailureOr<OpFoldResult> reifyShapedValueDimBound(
+    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
+    int64_t dim,
+    ValueBoundsConstraintSet::StopConditionFn stopCondition = nullptr,
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/lib/Dialect/Affine/Transforms/ReifyValueBounds.cpp

 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 
-FailureOr<OpFoldResult>
-mlir::reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
-                      Value value, std::optional<int64_t> dim, bool closedUB) {
-  // We are trying to reify a bound for `value`. Construct a stop condition that
-  // evaluates to "true" for any SSA value expect for `value`. I.e., the bound
-  // will be computed in terms of any SSA values expect for `value`. The first
-  // such values are operands of the owner of `value`.
-  auto stopCondition = [&](Value v, std::optional<int64_t> d) {
-    // Reify in terms of SSA values that are different from `value`.
-    return v != value;
-  };
-  return reifyValueBound(b, loc, type, value, dim, stopCondition, closedUB);
-}
-
-FailureOr<OpFoldResult> mlir::reifyValueBound(
-    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
-    std::optional<int64_t> dim,
+static FailureOr<OpFoldResult>
+reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
+                Value value, std::optional<int64_t> dim,
                 function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
                 bool closedUB) {
   // Compute bound.
   AffineMap boundMap;
   ValueDimList mapOperands;
   if (failed(ValueBoundsConstraintSet::computeBound(
           boundMap, mapOperands, type, value, dim, stopCondition, closedUB)))
     return failure();
 
   // Materialize tensor.dim/memref.dim ops.
if (auto expr = boundMap.getResult(0).dyn_cast<AffineDimExpr>())
     return static_cast<OpFoldResult>(operands[expr.getPosition()]);
   if (auto expr = boundMap.getResult(0).dyn_cast<AffineSymbolExpr>())
     return static_cast<OpFoldResult>(
         operands[expr.getPosition() + boundMap.getNumDims()]);
   // General case: build affine.apply op.
   return static_cast<OpFoldResult>(
       b.create<AffineApplyOp>(loc, boundMap, operands).getResult());
 }
+
+FailureOr<OpFoldResult> mlir::reifyShapedValueDimBound(
+    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
+    int64_t dim, ValueBoundsConstraintSet::StopConditionFn stopCondition,
+    bool closedUB) {
+  auto reifyToOperands = [&](Value v, std::optional<int64_t> d) {
+    // We are trying to reify a bound for `value` in terms of the owning op's
+    // operands. Construct a stop condition that evaluates to "true" for any SSA
+    // value expect for `value`. I.e., the bound will be computed in terms of
+    // any SSA values expect for `value`. The first such values are operands of
+    // the owner of `value`.
+    return v != value;
+  };
+  return reifyValueBound(b, loc, type, value, dim,
+                         stopCondition ? stopCondition : reifyToOperands,
+                         closedUB);
+}
+
+FailureOr<OpFoldResult> mlir::reifyIndexValueBound(
+    OpBuilder &b, Location loc, presburger::BoundType type, Value value,
+    ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB) {
+  auto reifyToOperands = [&](Value v, std::optional<int64_t> d) {
+    return v != value;
+  };
+  return reifyValueBound(b, loc, type, value, /*dim=*/std::nullopt,
+                         stopCondition ? stopCondition : reifyToOperands,
+                         closedUB);
+}

mlir/lib/Dialect/Linalg/Transforms/HoistPadding.cpp

HoistPaddingAnalysis::getPackedTensorSizes(RewriterBase &rewriter,
 
   // Upper bound the packing loop lengths to size the packed tensor. Taking
   // upper bounds can make the sizes of the packed tensor independent of the
   // enclosing loops. This independence is a prerequisite for reusing the same
   // buffer for all enclosing loop iterations and hoisting its allocation out
   // of the enclosing loops.
   for (auto forOp : packingLoops) {
     // Compute an upper bound `ubVal` for the upper bound of `forOp`.
-    FailureOr<OpFoldResult> loopUb = reifyValueBound(
+    FailureOr<OpFoldResult> loopUb = reifyIndexValueBound(
         rewriter, loc, presburger::BoundType::UB, forOp.getUpperBound(),
-        /*dim=*/std::nullopt, /*stopCondition=*/
+        /*stopCondition=*/
         [&](Value v, std::optional<int64_t> d) {
           if (v == forOp.getUpperBound())
             return false;
           // Compute a bound that is independent of any affine op results.
           Operation *op = v.getDefiningOp();
           if (!op)
             return true;
           return !isa<AffineMinOp, AffineMaxOp, AffineApplyOp>(op);

mlir/test/lib/Dialect/Affine/TestReifyValueBounds.cpp

if (op->getName().getStringRef() == "test.reify_bound" ||
       FailureOr<OpFoldResult> reified = failure();
       if (constant) {
         auto reifiedConst = ValueBoundsConstraintSet::computeConstantBound(
             *boundType, value, dim, /*stopCondition=*/nullptr);
         if (succeeded(reifiedConst))
           reified =
               FailureOr<OpFoldResult>(rewriter.getIndexAttr(*reifiedConst));
       } else {
-        reified = reifyValueBound(rewriter, op->getLoc(), *boundType, value,
-                                  dim, stopCondition);
+        if (dim) {
+          reified = reifyShapedValueDimBound(rewriter, op->getLoc(), *boundType,
+                                             value, *dim, stopCondition);
+        } else {
+          reified = reifyIndexValueBound(rewriter, op->getLoc(), *boundType,
+                                         value, stopCondition);
+        }
       }
       if (failed(reified)) {
         op->emitOpError("could not reify bound");
         return WalkResult::interrupt();
       }
 
       // Replace the op with the reified bound.
       if (auto val = reified->dyn_cast<Value>()) {