[mlir][linalg] Remove `linalg::getUpperBoundForIndex` util

[springerm]

Summary

Use reifyValueBound instead, which is more general not hard-coded to a specific list supported ops.

Also add a closedUB parameter to the ValueBoundsOpInterface API.

Depends On: D147769

[springerm]

rebase

[springerm]

rebase

[dcaballe]

Mask this as NFC?

[mravishankar]

In D146356#4250059, @dcaballe wrote:

Mask this as NFC?

I don't think you can mark it as NFC if it is dropping an API method

[springerm]

This change cannot be landed yet. There are certain inputs where it triggers an assertion in FlatLinearValueConstraints.cpp:

LLVM_DEBUG(llvm::dbgs()
           << "composition unimplemented for semi-affine maps\n");

[springerm]

In D146356#4250592, @springerm wrote:

This change cannot be landed yet. There are certain inputs where it triggers an assertion in FlatLinearValueConstraints.cpp:

LLVM_DEBUG(llvm::dbgs()
           << "composition unimplemented for semi-affine maps\n");

Actually, it can be landed. We have the same limitation in linalg::getUpperBoundForIndex. We just need better error handling: D147769

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

 /// 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. LB and EQ bounds are closed, UB bounds are open.
+/// `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);
+                                        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. LB and EQ bounds are
-/// closed, UB bounds are open.
+/// `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.
 ///
 /// 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);
+                ValueBoundsConstraintSet::StopConditionFn stopCondition,
+                bool closedUB = false);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Linalg/Utils/Utils.h

 Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 OpFoldResult createFoldedDimOp(OpBuilder &b, Location loc, Value source,
                                int64_t dim);
 
 /// Given an operation, retrieves the value of each dynamic dimension through
 /// constructing the necessary DimOp operators.
 SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b);
 
-/// Computes an upper bound for the result `value` of an index computation.
-/// Translates AffineMinOps and AffineApplyOps along the use-def chains of the
-/// index computation to affine constraints and projects out intermediate
-/// values. The method sets `boundMap` to an affine map that given
-/// `boundOperands` evaluates to an upper bound for the index computation.
-///
-/// If constantRequired is true, only returns the constant bounds (potentially
-/// over-approximating) and fails when not possible.
-///
-/// Example:
-/// ```
-/// %dim0 = dim %tensor, %c0
-/// %dim1 = dim %tensor, %c1
-/// %0 = affine.min affine.map<(d0) -> (40, d0)> (%dim0)
-/// %1 = affine.apply affine.map<(d0, d1) -> (d0 + d1)> (%0, %dim1)
-/// ```
-/// getUpperBoundForIndex(%1, boundMap, boundOperands)
-/// set the output parameters to:
-/// - boundMap = affine.map<(d0) -> (d0 + 40)>
-/// - boundOperands = [%dim1]
-void getUpperBoundForIndex(Value value, AffineMap &boundMap,
-                           SmallVectorImpl<Value> &boundOperands,
-                           bool constantRequired = false);
-
-/// Returns a constant upper bound for the result `value` of an index
-/// computation. Calls `getUpperBoundForIndex` and returns a constant upper
-/// bound if the result of `boundMap` is a constant expression and failure
-/// otherwise.
-///
-/// Example:
-/// ```
-/// %0 = affine.min affine.map<(d0) -> (40, d0)> (%d0)
-/// %1 = affine.apply affine.map<(d0) -> (d0 + 2)> (%0)
-/// ```
-/// getConstantUpperBoundForIndex(%1) returns 42
-/// (boundsMap = affine.map<() -> (42)>)
-FailureOr<int64_t> getConstantUpperBoundForIndex(Value value);
-
 /// Create a tensor::PadOp that pads `source` to the size of the statically
 /// sized `type` whose static sizes are assumed to be greater than the dynamic
 /// `source` size. The padding introduces trailing `pad` values until the
 /// target size is met. If `source` is defined by one or more LinalgOps that
 /// have been padded with the same value and sizes, return their padded result
 /// instead of creating a tensor::PadOp.
 ///
 /// Example:

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

public:
   /// or a shaped value and a dimension.
   ///
   /// `dim` must be `nullopt` if and only if `value` is index-typed. The bound
   /// is computed in terms of values/dimensions for which `stopCondition`
   /// evaluates to "true". To that end, the backward slice (reverse use-def
   /// chain) of the given value is visited in a worklist-driven manner and the
   /// constraint set is populated according to `ValueBoundsOpInterface` for each
   /// visited value.
+  ///
+  /// By default, lower/equal bounds are closed and upper bounds are open. If
+  /// `closedUB` is set to "true", upper bounds are also closed.
   static LogicalResult computeBound(AffineMap &resultMap,
                                     ValueDimList &mapOperands,
                                     presburger::BoundType type, Value value,
                                     std::optional<int64_t> dim,
-                                    StopConditionFn stopCondition);
+                                    StopConditionFn stopCondition,
+                                    bool closedUB = false);
 
   /// Compute a bound in terms of the values/dimensions in `dependencies`. The
   /// computed bound consists of only constant terms and dependent values (or
   /// dimension sizes thereof).
   static LogicalResult computeBound(AffineMap &resultMap,
                                     ValueDimList &mapOperands,
                                     presburger::BoundType type, Value value,
                                     std::optional<int64_t> dim,
-                                    ValueDimList dependencies);
+                                    ValueDimList dependencies,
+                                    bool closedUB = false);
 
   /// Compute a constant bound for the given index-typed value or shape
   /// dimension size.
   ///
   /// `dim` must be `nullopt` if and only if `value` is index-typed. This
   /// function traverses the backward slice of the given value in a
   /// worklist-driven manner until `stopCondition` evaluates to "true". The
   /// constraint set is populated according to `ValueBoundsOpInterface` for each
   /// visited value. (No constraints are added for values for which the stop
   /// condition evaluates to "true".)
   ///
   /// The stop condition is optional: If none is specified, the backward slice
   /// is traversed in a breadth-first manner until a constant bound could be
   /// computed.
+  ///
+  /// By default, lower/equal bounds are closed and upper bounds are open. If
+  /// `closedUB` is set to "true", upper bounds are also closed.
   static FailureOr<int64_t>
   computeConstantBound(presburger::BoundType type, Value value,
                        std::optional<int64_t> dim = std::nullopt,
-                       StopConditionFn stopCondition = nullptr);
+                       StopConditionFn stopCondition = nullptr,
+                       bool closedUB = false);
 
   /// Add a bound for the given index-typed value or shaped value. This function
   /// returns a builder that adds the bound.
   BoundBuilder bound(Value value) { return BoundBuilder(*this, value); }
 
   /// Return an expression that represents the given index-typed value or shaped
   /// value dimension. If this value/dimension was not used so far, it is added
   /// to the worklist.

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) {
+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 except 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);
+  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,
-    function_ref<bool(Value, std::optional<int64_t>)> stopCondition) {
+    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)))
+  if (failed(ValueBoundsConstraintSet::computeBound(
+          boundMap, mapOperands, type, value, dim, stopCondition, closedUB)))
     return failure();
 
   // Materialize tensor.dim/memref.dim ops.
   SmallVector<Value> operands;
   for (auto valueDim : mapOperands) {
     Value value = valueDim.first;
     std::optional<int64_t> dim = valueDim.second;
 

mlir/lib/Dialect/Linalg/Transforms/CMakeLists.txt

add_mlir_dialect_library(MLIRLinalgTransforms
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/Dialect/Linalg
 
   DEPENDS
   MLIRLinalgPassIncGen
 
   LINK_LIBS PUBLIC
   MLIRAffineDialect
+  MLIRAffineTransforms
   MLIRAffineUtils
   MLIRAnalysis
   MLIRArithDialect
   MLIRArithTransforms
   MLIRArithUtils
   MLIRBufferizationDialect
   MLIRBufferizationTransforms
   MLIRComplexDialect
add_mlir_dialect_library(MLIRLinalgTransforms
   MLIRSCFUtils
   MLIRPass
   MLIRSparseTensorDialect
   MLIRTensorDialect
   MLIRTensorTilingInterfaceImpl
   MLIRTensorTransforms
   MLIRTransforms
   MLIRTransformUtils
+  MLIRValueBoundsOpInterface
   MLIRVectorDialect
   MLIRVectorTransforms
   MLIRVectorUtils
   MLIRX86VectorTransforms
   MLIRVectorToSCF
 )

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

 //===- HoistPadding.cpp - Hoisting for tensor::PadOp ----------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements functions concerned with hoisting padding operations.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Analysis/Presburger/IntegerRelation.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Affine/Transforms/Transforms.h"
 #include "mlir/Dialect/Func/IR/FuncOps.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/IR/AsmState.h"
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`.
-    AffineMap boundMap;
-    SmallVector<Value> boundOperands;
-    getUpperBoundForIndex(forOp.getUpperBound(), boundMap, boundOperands);
-    Value ubVal =
-        rewriter.createOrFold<AffineMinOp>(loc, boundMap, boundOperands);
+    FailureOr<OpFoldResult> loopUb = reifyValueBound(
+        rewriter, loc, presburger::BoundType::UB, forOp.getUpperBound(),
+        /*dim=*/std::nullopt, /*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);
+        },
+        /*closedUB=*/true);
+    assert(succeeded(loopUb) && "could not get upper bound");
+    Value ubVal = getValueOrCreateConstantIndexOp(rewriter, loc, *loopUb);
+
     // Compute the maximal packing loop length as (ub - lb).ceilDiv(step) and
     // store the result to `dynamicTensorSizes`.
     // TODO: instead of using the lower bound of `forOp` directly, implement a
     // lower bound computation similar to the upper bound computation.
     AffineExpr lb, ub, step;
     bindDims(rewriter.getContext(), lb, ub);
     bindSymbols(rewriter.getContext(), step);
     Value res = rewriter.createOrFold<AffineApplyOp>(

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

 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/ImplicitLocOpBuilder.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/CommandLine.h"
 #include "llvm/Support/Debug.h"
 
for (const auto &en : llvm::enumerate(subView.getOrCreateRanges(b, loc))) {
     LLVM_DEBUG(llvm::dbgs() << "Extract tightest: " << rangeValue.size << "\n");
     Value size;
     if (auto attr = rangeValue.size.dyn_cast<Attribute>()) {
       size = getValueOrCreateConstantIndexOp(b, loc, rangeValue.size);
     } else {
       Value materializedSize =
           getValueOrCreateConstantIndexOp(b, loc, rangeValue.size);
       FailureOr<int64_t> upperBound =
-          getConstantUpperBoundForIndex(materializedSize);
+          ValueBoundsConstraintSet::computeConstantBound(
+              presburger::BoundType::UB, materializedSize, /*dim=*/std::nullopt,
+              /*stopCondition=*/nullptr, /*closedUB=*/true);
       size = failed(upperBound)
                  ? materializedSize
                  : b.create<arith::ConstantIndexOp>(loc, *upperBound);
     }
     LLVM_DEBUG(llvm::dbgs() << "Extracted tightest: " << size << "\n");
     fullSizes.push_back(size);
     partialSizes.push_back(
         b.createOrFold<memref::DimOp>(loc, subView, resultDimIdx++));

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

 #include "mlir/Dialect/Tensor/IR/TensorTilingInterfaceImpl.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Utils/StructuredOpsUtils.h"
 #include "mlir/Dialect/Vector/IR/VectorOps.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 #include "llvm/ADT/ScopeExit.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/raw_ostream.h"
 #include <type_traits>
for (const auto &en : enumerate(mixedSizes)) {
     // If the size is an attribute add it directly to `paddedShape`.
     if (en.value().is<Attribute>()) {
       paddedShape[shapeIdx++] =
           en.value().get<Attribute>().dyn_cast<IntegerAttr>().getInt();
       continue;
     }
     // Otherwise, try to compute a constant upper bound for the size value.
     FailureOr<int64_t> upperBound =
-        getConstantUpperBoundForIndex(en.value().get<Value>());
+        ValueBoundsConstraintSet::computeConstantBound(
+            presburger::BoundType::UB, en.value().get<Value>(),
+            /*dim=*/std::nullopt, /*stopCondition=*/nullptr, /*closedUB=*/true);
     if (failed(upperBound)) {
       LLVM_DEBUG(DBGS() << "count not compute a bonding box for padding");
       return rewriter.notifyMatchFailure(
           opToPad, "count not compute a bonding box for padding");
     }
     paddedShape[shapeIdx++] = *upperBound;
   }
   assert(shapeIdx == static_cast<int64_t>(shape.size()) &&

mlir/lib/Dialect/Linalg/Utils/Utils.cpp

SmallVector<Value, 4> getDynOperands(Location loc, Value val, OpBuilder &b) {
   auto shapedType = val.getType().cast<ShapedType>();
   for (const auto &dim : llvm::enumerate(shapedType.getShape())) {
     if (dim.value() == ShapedType::kDynamic)
       dynOperands.push_back(createOrFoldDimOp(b, loc, val, dim.index()));
   }
   return dynOperands;
 }
 
-void getUpperBoundForIndex(Value value, AffineMap &boundMap,
-                           SmallVectorImpl<Value> &boundOperands,
-                           bool constantRequired) {
-  // Initialize `boundMap` and `boundOperands` to the identity returning
-  // `value`. This combination is the default result of the method if no
-  // simplification is possible.
-  assert(value.getType().isIndex() && "expect value to have index type");
-  boundMap = AffineMap::getMultiDimIdentityMap(1, value.getContext());
-  boundOperands.assign({value});
-  canonicalizeMapAndOperands(&boundMap, &boundOperands);
-
-  // Continue only if there is an affine index computation to simplify.
-  Operation *definingOp = value.getDefiningOp();
-  if (!definingOp || !isa<AffineApplyOp, AffineMinOp>(definingOp))
-    return;
-
-  // Get the backward slice containing the affine index computation.
-  SetVector<Operation *> backwardSlice;
-  getBackwardSlice(definingOp, &backwardSlice, [](Operation *op) {
-    return isa<AffineApplyOp, AffineMinOp>(op);
-  });
-  backwardSlice.insert(definingOp);
-
-  // Setup a system of affine constraints that describe the index computation.
-  FlatAffineValueConstraints constraints;
-
-  // Helper to find or create an identifier for the given value.
-  auto findOrCreateId = [&](Value value) {
-    if (!constraints.containsVar(value)) {
-      constraints.appendDimVar(value);
-      return true;
-    }
-    unsigned pos;
-    constraints.findVar(value, &pos);
-    return pos < constraints.getNumDimVars();
-  };
-  // Helper to get the position for the given value.
-  auto getPosition = [&](Value value) {
-    unsigned pos;
-    bool exists = constraints.findVar(value, &pos);
-    (void)exists;
-    assert(exists && "expect to find the identifier");
-    return pos;
-  };
-
-  // Add the affine operations in `backwardSlice` to the constraints.
-  for (Operation *op : llvm::reverse(backwardSlice)) {
-    // Add an identifier for all op results and operands.
-    if (!(llvm::all_of(op->getResults(), findOrCreateId) &&
-          llvm::all_of(op->getOperands(), findOrCreateId)))
-      return;
-
-    // Add AffineApplyOps to the constraints.
-    if (auto applyOp = dyn_cast<AffineApplyOp>(op)) {
-      AffineMap map = constraints.computeAlignedMap(applyOp.getAffineMap(),
-                                                    applyOp.getOperands());
-      if (failed(constraints.addBound(BoundType::EQ,
-                                      getPosition(applyOp.getResult()), map)))
-        return;
-      continue;
-    }
-    // Add AffineMinOps to the constraints.
-    auto minOp = cast<AffineMinOp>(op);
-    AffineMap map = constraints.computeAlignedMap(minOp.getAffineMap(),
-                                                  minOp.getOperands());
-    if (failed(constraints.addBound(BoundType::UB,
-                                    getPosition(minOp.getResult()), map,
-                                    /*isClosedBound=*/true)))
-      return;
-  }
-
-  // Obtain an upper bound for the affine index computation by projecting out
-  // all temporary results and expressing the upper bound for `value` in terms
-  // of the terminals of the index computation.
-  unsigned pos = getPosition(value);
-  if (constantRequired) {
-    auto ubConst = constraints.getConstantBound64(BoundType::UB, pos);
-    if (!ubConst)
-      return;
-
-    boundMap = AffineMap::getConstantMap(*ubConst, value.getContext());
-    return;
-  }
-
-  SmallVector<AffineMap> lowerBounds(1), upperBounds(1);
-  constraints.getSliceBounds(pos, 1, value.getContext(), &lowerBounds,
-                             &upperBounds,
-                             /*getClosedUB=*/true);
-  // Verify `upperBounds[0]` is valid and has at least one result.
-  if (!upperBounds[0] || upperBounds[0].getNumResults() == 0)
-    return;
-
-  // Set `boundMap` and `boundOperands` to the computed upper bound.
-  boundMap = upperBounds[0];
-  constraints.getAllValues(&boundOperands);
-  erase_value(boundOperands, value);
-  canonicalizeMapAndOperands(&boundMap, &boundOperands);
-}
-
-FailureOr<int64_t> getConstantUpperBoundForIndex(Value value) {
-  // Compute an upper bound for `value`.
-  AffineMap boundMap;
-  SmallVector<Value> boundOperands;
-  getUpperBoundForIndex(value, boundMap, boundOperands,
-                        /*constantRequired=*/true);
-
-  // Search the results of `boundMap` for constant upper bounds.
-  SmallVector<int64_t> constantBounds;
-  for (AffineExpr result : boundMap.getResults())
-    if (auto constExpr = result.dyn_cast<AffineConstantExpr>())
-      constantBounds.push_back(constExpr.getValue());
-
-  // Return the minimal upper bound or failure if none is found.
-  if (constantBounds.empty())
-    return failure();
-  return *std::min_element(constantBounds.begin(), constantBounds.end());
-}
-
 Value makeComposedPadHighOp(OpBuilder &b, Location loc, RankedTensorType type,
                             Value source, Value pad, bool nofold) {
   // Exit if `source` is not defined by an ExtractSliceOp.
   auto sliceOp = source.getDefiningOp<tensor::ExtractSliceOp>();
   if (!sliceOp)
     return tensor::createPadHighOp(type, source, pad, nofold, loc, b);
 
   // Search the `source` use-def chain for padded LinalgOps.

mlir/lib/Dialect/Tensor/IR/CMakeLists.txt

add_mlir_dialect_library(MLIRTensorTilingInterfaceImpl
   MLIRIR
   MLIRLinalgDialect
   MLIRLinalgUtils
   MLIRSCFDialect
   MLIRSupport
   MLIRTensorDialect
   MLIRTensorUtils
   MLIRTilingInterface
+  MLIRValueBoundsOpInterface
   )

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

 #include "mlir/Dialect/Arith/Utils/Utils.h"
 #include "mlir/Dialect/Linalg/IR/Linalg.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Utils/Utils.h"
 #include "mlir/Dialect/Utils/IndexingUtils.h"
 #include "mlir/Interfaces/TilingInterface.h"
+#include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 namespace {
 
 struct PadOpTiling : public TilingInterface::ExternalModel<PadOpTiling, PadOp> {
 
static UnpackTileDimInfo getUnpackTileDimInfo(OpBuilder &b, UnPackOp unpackOp,
   AffineBuilder ab(b, loc);
   AffineExpr dim0, dim1, sym0;
   bindDims(b.getContext(), dim0, dim1);
   bindSymbols(b.getContext(), sym0);
 
   OpFoldResult innerTileSize = dimAndTileMapping[tileDim];
 
   info.isAlignedToInnerTileSize = false;
-  FailureOr<int64_t> cstSize = linalg::getConstantUpperBoundForIndex(
-      getValueOrCreateConstantIndexOp(b, loc, tileSize));
+  FailureOr<int64_t> cstSize = ValueBoundsConstraintSet::computeConstantBound(
+      presburger::BoundType::UB,
+      getValueOrCreateConstantIndexOp(b, loc, tileSize), /*dim=*/std::nullopt,
+      /*stopCondition=*/nullptr, /*closedUB=*/true);
   std::optional<int64_t> cstInnerSize = getConstantIntValue(innerTileSize);
   if (!failed(cstSize) && cstInnerSize) {
     if (*cstSize % *cstInnerSize == 0)
       info.isAlignedToInnerTileSize = true;
 
     // If the tiling size equals to the inner tiling size, the outer dims are
     // always 1.
     if (*cstInnerSize == *cstSize) {

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

while (nextPos < static_cast<int64_t>(positionToValueDim.size())) {
     } else {
       ++nextPos;
     }
   }
 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
-    Value value, std::optional<int64_t> dim, StopConditionFn stopCondition) {
+    Value value, std::optional<int64_t> dim, StopConditionFn stopCondition,
+    bool closedUB) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
   assert(!stopCondition(value, dim) &&
          "stop condition should not be satisfied for starting point");
 #endif // NDEBUG
 
+  int64_t ubAdjustment = closedUB ? 0 : 1;
   Builder b(value.getContext());
   mapOperands.clear();
 
   if (stopCondition(value, dim)) {
     // Special case: If the stop condition is satisfied for the input
     // value/dimension, directly return it.
     mapOperands.push_back(std::make_pair(value, dim));
+    AffineExpr bound = b.getAffineDimExpr(0);
+    if (type == BoundType::UB)
+      bound = bound + ubAdjustment;
     resultMap = AffineMap::get(/*dimCount=*/1, /*symbolCount=*/0,
                                b.getAffineDimExpr(0));
     return success();
   }
 
   // Process the backward slice of `value` (i.e., reverse use-def chain) until
   // `stopCondition` is met.
   ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
#endif // NDEBUG
   // EQ bound: lower and upper bound must match.
   if (type == BoundType::EQ && ub[0] != lb[0])
     return failure();
 
   AffineMap bound;
   if (type == BoundType::EQ || type == BoundType::LB) {
     bound = lb[0];
   } else {
-    // Computed UB is a closed bound. Turn into an open bound.
+    // Computed UB is a closed bound.
     bound = AffineMap::get(ub[0].getNumDims(), ub[0].getNumSymbols(),
-                           ub[0].getResult(0) + 1);
+                           ub[0].getResult(0) + ubAdjustment);
   }
 
   // Gather all SSA values that are used in the computed bound.
   assert(cstr.cstr.getNumDimAndSymbolVars() == cstr.positionToValueDim.size() &&
          "inconsistent mapping state");
   SmallVector<AffineExpr> replacementDims, replacementSymbols;
   int64_t numDims = 0, numSymbols = 0;
   for (int64_t i = 0; i < cstr.cstr.getNumDimAndSymbolVars(); ++i) {
#endif // NDEBUG
 
   resultMap = bound.replaceDimsAndSymbols(replacementDims, replacementSymbols,
                                           numDims, numSymbols);
   return success();
 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
-    Value value, std::optional<int64_t> dim, ValueDimList dependencies) {
-  return computeBound(resultMap, mapOperands, type, value, dim,
-                      [&](Value v, std::optional<int64_t> d) {
-                        return llvm::is_contained(dependencies,
-                                                  std::make_pair(v, d));
-                      });
+    Value value, std::optional<int64_t> dim, ValueDimList dependencies,
+    bool closedUB) {
+  return computeBound(
+      resultMap, mapOperands, type, value, dim,
+        return llvm::is_contained(dependencies, std::make_pair(v, d));
+      },
+      closedUB);
 }
 
 FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
     presburger::BoundType type, Value value, std::optional<int64_t> dim,
-    StopConditionFn stopCondition) {
+    StopConditionFn stopCondition, bool closedUB) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
 #endif // NDEBUG
 
   // Process the backward slice of `value` (i.e., reverse use-def chain) until
   // `stopCondition` is met.
   ValueBoundsConstraintSet cstr(value, dim);
   int64_t pos = cstr.getPos(value, dim);
   if (stopCondition) {
     cstr.processWorklist(stopCondition);
   } else {
     // No stop condition specified: Keep adding constraints until a bound could
     // be computed.
     cstr.processWorklist(
         /*stopCondition=*/[&](Value v, std::optional<int64_t> dim) {
           return cstr.cstr.getConstantBound64(type, pos).has_value();
         });
   }
 
   // Compute constant bound for `valueDim`.
+  int64_t ubAdjustment = closedUB ? 0 : 1;
   if (auto bound = cstr.cstr.getConstantBound64(type, pos))
-    return type == BoundType::UB ? *bound + 1 : *bound;
+    return type == BoundType::UB ? *bound + ubAdjustment : *bound;
   return failure();
 }
 
 ValueBoundsConstraintSet::BoundBuilder &
 ValueBoundsConstraintSet::BoundBuilder::operator[](int64_t dim) {
   assert(!this->dim.has_value() && "dim was already set");
   this->dim = dim;
 #ifndef NDEBUG

utils/bazel/llvm-project-overlay/mlir/BUILD.bazel

deps = [
         ":DialectUtils",
         ":IR",
         ":LinalgDialect",
         ":LinalgUtils",
         ":SCFDialect",
         ":TensorDialect",
         ":TensorUtils",
         ":TilingInterface",
+        ":ValueBoundsOpInterface",
         "//llvm:Support",
     ],
 )
 
 cc_library(
     name = "TensorUtils",
     srcs = ["lib/Dialect/Tensor/Utils/Utils.cpp"],
     hdrs = ["include/mlir/Dialect/Tensor/Utils/Utils.h"],
hdrs = [
         "include/mlir/Dialect/Linalg/Passes.h",
     ] + glob([
         "include/mlir/Dialect/Linalg/Transforms/*.h",
     ]),
     includes = ["include"],
     deps = [
         ":AffineAnalysis",
         ":AffineDialect",
+        ":AffineTransforms",
         ":AffineUtils",
         ":Analysis",
         ":ArithDialect",
         ":ArithTransforms",
         ":ArithUtils",
         ":BufferizationDialect",
         ":BufferizationTransforms",
         ":ComplexDialect",
deps = [
         ":Support",
         ":TensorDialect",
         ":TensorTilingInterfaceImpl",
         ":TensorTransforms",
         ":TensorUtils",
         ":TilingInterface",
         ":TransformUtils",
         ":Transforms",
+        ":ValueBoundsOpInterface",
         ":VectorDialect",
         ":VectorToSCF",
         ":VectorTransforms",
         ":VectorUtils",
         ":X86VectorTransforms",
         "//llvm:Support",
     ],
 )