[mlir][Arith][WIP] ValueBoundsOpInterface: Add support for `arith.select`/`scf.if`

[springerm]

Summary

Depends On: D146870

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

 class AffineApplyOp;
 class AffineMap;
 class Location;
 class OpBuilder;
 class OpFoldResult;
 class RewritePatternSet;
 class RewriterBase;
 class Value;
+class ValueBoundsConstraintSet;
 
 namespace presburger {
 enum class BoundType;
 } // namespace presburger
 
 /// Populate patterns that expand affine index operations into more fundamental
 /// operations (not necessarily restricted to Affine dialect).
 void populateAffineExpandIndexOpsPatterns(RewritePatternSet &patterns);

 /// `nullopt` if and only if `value` is integer/index-typed. The reified result
 /// has index type.
 ///
 /// 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,
-                function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
+                function_ref<bool(Value, std::optional<int64_t>,
+                                  const ValueBoundsConstraintSet &)>
+                    stopCondition,
                 bool closedUB = false);
 
 /// Reify an already computed bound with Affine dialect ops.
 OpFoldResult
 reifyValueBound(OpBuilder &b, Location loc, AffineMap boundMap,
                 ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_AFFINE_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h

 //===- ValueBoundsOpInterfaceImpl.h - Impl. of ValueBoundsOpInterface -----===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_ARITH_IR_VALUEBOUNDSOPINTERFACEIMPL_H
 #define MLIR_DIALECT_ARITH_IR_VALUEBOUNDSOPINTERFACEIMPL_H
 
+#include <optional>
+
 namespace mlir {
 class DialectRegistry;
+class Value;
+class ValueBoundsConstraintSet;
 
 namespace arith {
+/// Populate bounds for an arith.select-like op.
+void populateSelectLikeBounds(Value value, std::optional<int64_t> dim,
+                              Value condition, Value trueValue,
+                              Value falseValue, ValueBoundsConstraintSet &cstr);
+
 void registerValueBoundsOpInterfaceExternalModels(DialectRegistry &registry);
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_IR_VALUEBOUNDSOPINTERFACEIMPL_H

mlir/include/mlir/Dialect/Arith/Transforms/Transforms.h

 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 class Location;
 class OpBuilder;
 class OpFoldResult;
 class Type;
 class Value;
+class ValueBoundsConstraintSet;
 
 namespace presburger {
 enum class BoundType;
 } // namespace presburger
 
 namespace arith {
 
 /// Reify a bound for the given integer/index-typed value or shape dimension

 /// By default, it is the same type as the input. (IndexType in case of a
 /// dimension).
 ///
 /// 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,
-                function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
+                function_ref<bool(Value, std::optional<int64_t>,
+                                  const ValueBoundsConstraintSet &cstr)>
+                    stopCondition,
                 bool closedUB = false, Type resultType = {});
 
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

 #include "mlir/IR/Value.h"
 #include "mlir/Interfaces/DestinationStyleOpInterface.h"
 #include "llvm/ADT/SetVector.h"
 
 #include <queue>
 
 namespace mlir {
 
+class ValueBoundsConstraintSet;
+
 using ValueDimList = SmallVector<std::pair<Value, std::optional<int64_t>>>;
 
 /// A helper class to be used with `ValueBoundsOpInterface`. This class stores a
 /// constraint system and mapping of constrained variables to integer-typed
 /// values or dimension sizes of shaped values.
 ///
 /// Note: "Integer type" refers to both integer and index types.
 ///
 /// Interface implementations of `ValueBoundsOpInterface` use `addBounds` to
 /// insert constraints about their results and/or region block arguments into
 /// the constraint set in the form of an AffineExpr. When a bound should be
 /// expressed in terms of another value/dimension, `getExpr` can be used to
 /// retrieve an AffineExpr that represents the specified value/dimension.
 ///
 /// When a value/dimension is retrieved for the first time through `getExpr`,
 /// it is added to an internal worklist. See `computeBound` for more details.
 ///
 /// Note: Any modification of the IR invalides the data stored in this class.
 class ValueBoundsConstraintSet {
 public:
   /// The stop condition when traversing the backward slice of a shaped value/
   /// integer-type value. The traversal continues until the stop condition
   /// evaluates to "true" for a value.
-  using StopConditionFn = function_ref<bool(Value, std::optional<int64_t>)>;
+  using StopConditionFn = function_ref<bool(
+      Value, std::optional<int64_t>, const ValueBoundsConstraintSet &cstr)>;
 
   /// Compute a bound for the given integer-typed value or shape dimension size.
   /// The computed bound is stored in `resultMap`. The operands of the bound are
   /// stored in `mapOperands`. An operand is either an integer-type SSA value
   /// or a shaped value and a dimension.
   ///
   /// `dim` must be `nullopt` if and only if `value` is integer-typed. The bound
   /// is computed in terms of values/dimensions for which `stopCondition`
public:
   /// Return an expression that represents a constant or integer-typed SSA
   /// value. In case of a value, if this value was not used so far, it is added
   /// to the worklist.
   AffineExpr getExpr(OpFoldResult ofr);
 
   /// Return an expression that represents a constant.
   AffineExpr getExpr(int64_t constant);
 
+  AffineExpr getAlignedExpr(AffineExpr expr, int64_t numDims,
+                            int64_t numSymbols, ValueDimList valueDims);
+
+  StopConditionFn getStopCondition() const;
+
 protected:
   /// Dimension identifier to indicate a value is integer-typed.
   static constexpr int64_t kIntegerValue = -1;
 
   using ValueDim = std::pair<Value, int64_t>;
   ValueBoundsConstraintSet(ValueDim valueDim);
 
   /// Iteratively process all elements on the worklist until an integer-typed
protected:
   /// Worklist of values/shape dimensions that have not been processed yet.
   std::queue<int64_t> worklist;
 
   /// Constraint system of equalities and inequalities.
   FlatLinearConstraints cstr;
 
   /// Builder for constructing affine expressions.
   Builder builder;
+
+  StopConditionFn stopCondition;
 };
 
 } // namespace mlir
 
 #include "mlir/Interfaces/ValueBoundsOpInterface.h.inc"
 
 namespace mlir {
 

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

 #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) {
-  auto stopCondition = [&](Value v, std::optional<int64_t> d) {
+  auto stopCondition = [&](Value v, std::optional<int64_t> d,
+                           const ValueBoundsConstraintSet &cstr) {
     // 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,
-    function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
+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>,
+                                        const ValueBoundsConstraintSet &)>
+                          stopCondition,
                       bool closedUB) {
   // Compute bound.
   AffineMap boundMap;
   ValueDimList mapOperands;
   if (failed(ValueBoundsConstraintSet::computeBound(
           boundMap, mapOperands, type, value, dim, stopCondition, closedUB)))
     return failure();
 
   // Reify bound.

mlir/lib/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.cpp

 //===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 using presburger::BoundType;
 
+/// Try to align the first map with the second one. I.e., replace dims/symbols
+/// in such a way that both have the same ValueDims at the same dim/symbol
+/// positions.
+static FailureOr<AffineMap> computeAlignedExpr(Builder &b, AffineMap m1,
+                                               ValueDimList v1, AffineMap m2,
+                                               ValueDimList v2) {
+  if (v1.size() != v2.size())
+    return failure();
+  SmallVector<AffineExpr> dimReplacements, symbolReplacements;
+  for (int i = 0; i < m1.getNumDims(); ++i) {
+    auto it = llvm::find(v2, v1[i]);
+    if (!it)
+      return failure();
+    int64_t pos = std::distance(v2.begin(), it);
+    if (pos < m2.getNumDims()) {
+      dimReplacements.push_back(b.getAffineDimExpr(pos));
+    } else {
+      dimReplacements.push_back(b.getAffineSymbolExpr(pos - m2.getNumDims()));
+    }
+  }
+  for (int i = m1.getNumDims(); i < m1.getNumDims() + m1.getNumSymbols(); ++i) {
+    auto it = llvm::find(v2, v1[i]);
+    if (!it)
+      return failure();
+    int64_t pos = std::distance(v2.begin(), it);
+    if (pos < m2.getNumDims()) {
+      symbolReplacements.push_back(b.getAffineDimExpr(pos));
+    } else {
+      symbolReplacements.push_back(
+          b.getAffineSymbolExpr(pos - m2.getNumDims()));
+    }
+  }
+
+  SmallVector<AffineExpr> results =
+      llvm::to_vector(llvm::map_range(m1.getResults(), [&](AffineExpr e) {
+        return e.replaceDimsAndSymbols(dimReplacements, symbolReplacements);
+      }));
+  return AffineMap::get(m2.getNumDims(), m2.getNumSymbols(), results,
+                        b.getContext());
+}
+
+/// Return "true" if the two given bounds are equivalent.
+static bool isEquivalentBound(AffineMap m1, ValueDimList v1, AffineMap m2,
+                              ValueDimList v2) {
+  Builder b(m1.getContext());
+
+  // Only maps with one result are supported at the moment.
+  if (m1.getNumResults() != 1 || m2.getNumResults() != 1)
+    return false;
+
+  // Align the first bound with the second one.
+  auto alignedMap = computeAlignedExpr(b, m1, v1, m2, v2);
+  if (failed(alignedMap))
+    return false;
+
+  // Compute "bound1 - bounds" and simplify. Both bounds are equivalent if this
+  // expression simplifies to "0".
+  AffineExpr sub = alignedMap->getResult(0) - m2.getResult(0);
+  AffineExpr simplified =
+      simplifyAffineExpr(sub, m2.getNumDims(), m2.getNumSymbols());
+  return simplified == b.getAffineConstantExpr(0);
+}
+
+void mlir::arith::populateSelectLikeBounds(Value value,
+                                           std::optional<int64_t> dim,
+                                           Value condition, Value trueValue,
+                                           Value falseValue,
+                                           ValueBoundsConstraintSet &cstr) {
+  auto addBound = [&](BoundType type, AffineExpr expr) {
+    if (dim.has_value()) {
+      cstr.addBound(type, value, *dim, expr);
+    } else {
+      cstr.addBound(type, value, expr);
+    }
+  };
+
+  // Case 1: Constant condition.
+  FailureOr<int64_t> cond = ValueBoundsConstraintSet::computeConstantBound(
+      presburger::BoundType::EQ, condition);
+  if (succeeded(cond)) {
+    if (*cond == 0) {
+      addBound(BoundType::EQ, cstr.getExpr(falseValue, dim));
+      return;
+    }
+    assert(*cond == 1 && "expected 'true'");
+    addBound(BoundType::EQ, cstr.getExpr(trueValue, dim));
+    return;
+  }
+
+  // Case 2: Both values have a constant bound.
+  auto findConstantBound = [&](BoundType type) {
+    auto trueBound =
+        ValueBoundsConstraintSet::computeConstantBound(type, trueValue, dim);
+    if (succeeded(trueBound)) {
+      auto falseBound =
+          ValueBoundsConstraintSet::computeConstantBound(type, falseValue, dim);
+      if (succeeded(falseBound)) {
+        // LB: Take the smaller bound.
+        // UB: Take the larger bound.
+        addBound(type, cstr.getExpr(type == BoundType::LB
+                                        ? std::min(*trueBound, *falseBound)
+                                        : std::max(*trueBound, *falseBound)));
+        return true;
+      }
+    }
+    return false;
+  };
+  if (findConstantBound(BoundType::LB) && findConstantBound(BoundType::UB))
+    return;
+
+  // Case 3: Both values have the same bound.
+  AffineMap mapFalse, mapTrue;
+  ValueDimList mapFalseOperands, mapTrueOperands;
+  LogicalResult status = ValueBoundsConstraintSet::computeBound(
+      mapFalse, mapFalseOperands, BoundType::EQ, falseValue, dim,
+      cstr.getStopCondition());
+  if (failed(status) || mapFalse.getNumResults() != 1)
+    return;
+  status = ValueBoundsConstraintSet::computeBound(mapTrue, mapTrueOperands,
+                                                  BoundType::EQ, trueValue, dim,
+                                                  mapFalseOperands);
+  if (failed(status) || mapTrue.getNumResults() != 1)
+    return;
+  if (isEquivalentBound(mapFalse, mapFalseOperands, mapTrue, mapTrueOperands))
+    addBound(BoundType::EQ,
+             cstr.getAlignedExpr(mapFalse.getResult(0), mapFalse.getNumDims(),
+                                 mapFalse.getNumSymbols(), mapFalseOperands));
+}
+
 namespace mlir {
 namespace arith {
 namespace {
 
 struct AddIOpInterface
     : public ValueBoundsOpInterface::ExternalModel<AddIOpInterface, AddIOp> {
   void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
void populateBoundsForIntegerValue(Operation *op, Value value,
     assert(value == mulIOp.getResult() && "invalid value");
 
     cstr.addBound(BoundType::EQ, value,
                   cstr.getExpr(mulIOp.getLhs()) *
                       cstr.getExpr(mulIOp.getRhs()));
   }
 };
 
+struct SelectOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<SelectOpInterface,
+                                                   SelectOp> {
+  void populateBoundsForIntegerValue(Operation *op, Value value,
+                                     ValueBoundsConstraintSet &cstr) const {
+    auto selectOp = cast<SelectOp>(op);
+    assert(value == selectOp.getResult() && "invalid value");
+    populateSelectLikeBounds(value, /*dim=*/std::nullopt,
+                             selectOp.getCondition(), selectOp.getTrueValue(),
+                             selectOp.getFalseValue(), cstr);
+  }
+
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto selectOp = cast<SelectOp>(op);
+    assert(value == selectOp.getResult() && "invalid value");
+    populateSelectLikeBounds(value, dim, selectOp.getCondition(),
+                             selectOp.getTrueValue(), selectOp.getFalseValue(),
+                             cstr);
+  }
+};
+
 } // namespace
 } // namespace arith
 } // namespace mlir
 
 void mlir::arith::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, arith::ArithDialect *dialect) {
     arith::AddIOp::attachInterface<arith::AddIOpInterface>(*ctx);
     arith::ConstantOp::attachInterface<arith::ConstantOpInterface>(*ctx);
+    arith::SelectOp::attachInterface<arith::SelectOpInterface>(*ctx);
     arith::SubIOp::attachInterface<arith::SubIOpInterface>(*ctx);
     arith::MulIOp::attachInterface<arith::MulIOpInterface>(*ctx);
   });
 }

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

std::function<Value(AffineExpr)> buildExpr = [&](AffineExpr e) -> Value {
     }
   };
   return buildExpr(map.getResult(0));
 }
 
 FailureOr<OpFoldResult> mlir::arith::reifyValueBound(
     OpBuilder &b, Location loc, presburger::BoundType type, Value value,
     std::optional<int64_t> dim, bool closedUB, Type resultType) {
-  auto stopCondition = [&](Value v, std::optional<int64_t> d) {
+  auto stopCondition = [&](Value v, std::optional<int64_t> d,
+                           const ValueBoundsConstraintSet &cstr) {
     // Reify in terms of SSA values that are different from `value`.
     return v != value;
   };
   return reifyValueBound(b, loc, type, value, dim, stopCondition, closedUB,
                          resultType);
 }
 
 FailureOr<OpFoldResult> mlir::arith::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, Type resultType) {
+    ValueBoundsConstraintSet::StopConditionFn stopCondition, bool closedUB,
+    Type resultType) {
   // Default result type if not specified: same as value.
   if (!resultType) {
     resultType = dim.has_value() ? b.getIndexType() : value.getType();
   }
   assert(resultType.isIntOrIndex() && "invalid result type");
 
   // Compute bound.
   AffineMap boundMap;

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

HoistingAnalysis::getPackedTensorSizes(RewriterBase &rewriter, Location loc) {
   // 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(
         rewriter, loc, presburger::BoundType::UB, forOp.getUpperBound(),
         /*dim=*/std::nullopt, /*stopCondition=*/
-        [&](Value v, std::optional<int64_t> d) {
+        [&](Value v, std::optional<int64_t> d,
+            const ValueBoundsConstraintSet &cstr) {
           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/lib/Dialect/SCF/IR/CMakeLists.txt

 add_mlir_dialect_library(MLIRSCFDialect
   SCF.cpp
   DeviceMappingInterface.cpp
   ValueBoundsOpInterfaceImpl.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${MLIR_MAIN_INCLUDE_DIR}/mlir/SCF
 
   DEPENDS
   MLIRSCFOpsIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
+  MLIRArithValueBoundsOpInterfaceImpl
   MLIRBufferizationDialect
   MLIRControlFlowDialect
   MLIRIR
   MLIRLoopLikeInterface
   MLIRSideEffectInterfaces
   MLIRValueBoundsOpInterface
   )
 

mlir/lib/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.cpp

 //===- ValueBoundsOpInterfaceImpl.cpp - Impl. of ValueBoundsOpInterface ---===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/SCF/IR/ValueBoundsOpInterfaceImpl.h"
 
+#include "mlir/Dialect/Arith/IR/ValueBoundsOpInterfaceImpl.h"
 #include "mlir/Dialect/SCF/IR/SCF.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 using presburger::BoundType;
 
 namespace mlir {
 namespace scf {
if (yieldedValue == iterArg) {
       return;
     }
 
     // Compute EQ bound for yielded value.
     AffineMap bound;
     ValueDimList boundOperands;
     LogicalResult status = ValueBoundsConstraintSet::computeBound(
         bound, boundOperands, BoundType::EQ, yieldedValue, dim,
-        [&](Value v, std::optional<int64_t> d) {
+        [&](Value v, std::optional<int64_t> d,
+            const ValueBoundsConstraintSet &cstr) {
           // Stop when reaching a block argument of the loop body.
           if (auto bbArg = v.dyn_cast<BlockArgument>())
             return bbArg.getOwner()->getParentOp() == forOp;
           // Stop when reaching a value that is defined outside of the loop. It
           // is impossible to reach an iter_arg from there.
           Operation *op = v.getDefiningOp();
           return forOp.getLoopBody().findAncestorOpInRegion(*op) == nullptr;
         });
struct ForOpInterface
   void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
                                        ValueBoundsConstraintSet &cstr) const {
     auto forOp = cast<ForOp>(op);
     // Handle iter_args and OpResults.
     populateIterArgBounds(forOp, value, dim, cstr);
   }
 };
 
+struct IfOpInterface
+    : public ValueBoundsOpInterface::ExternalModel<IfOpInterface, IfOp> {
+  void populateBoundsForIntegerValue(Operation *op, Value value,
+                                     ValueBoundsConstraintSet &cstr) const {
+    auto ifOp = cast<IfOp>(op);
+    Value condition = ifOp.getCondition();
+    Value thenValue =
+        ifOp.thenYield().getOperand(value.cast<OpResult>().getResultNumber());
+    Value elseValue =
+        ifOp.elseYield().getOperand(value.cast<OpResult>().getResultNumber());
+    arith::populateSelectLikeBounds(value, /*dim=*/std::nullopt, condition,
+                                    thenValue, elseValue, cstr);
+  }
+
+  void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
+                                       ValueBoundsConstraintSet &cstr) const {
+    auto ifOp = cast<IfOp>(op);
+    Value condition = ifOp.getCondition();
+    Value thenValue =
+        ifOp.thenYield().getOperand(value.cast<OpResult>().getResultNumber());
+    Value elseValue =
+        ifOp.elseYield().getOperand(value.cast<OpResult>().getResultNumber());
+    arith::populateSelectLikeBounds(value, dim, condition, thenValue, elseValue,
+                                    cstr);
+  }
+};
+
 } // namespace
 } // namespace scf
 } // namespace mlir
 
 void mlir::scf::registerValueBoundsOpInterfaceExternalModels(
     DialectRegistry &registry) {
   registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
     scf::ForOp::attachInterface<scf::ForOpInterface>(*ctx);
+    scf::IfOp::attachInterface<scf::IfOpInterface>(*ctx);
   });
 }

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

 //===- ValueBoundsOpInterface.cpp - Value Bounds  -------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "llvm/ADT/APSInt.h"
+#include "llvm/ADT/ScopeExit.h"
 
 using namespace mlir;
 using presburger::BoundType;
 using presburger::VarKind;
 
 namespace mlir {
 #include "mlir/Interfaces/ValueBoundsOpInterface.cpp.inc"
 } // namespace mlir
ValueBoundsConstraintSet::ValueBoundsConstraintSet(ValueDim valueDim)
   insert(valueDim, /*isSymbol=*/false);
 }
 
 #ifndef NDEBUG
 static void assertValidValueDim(Value value, std::optional<int64_t> dim) {
   if (value.getType().isIntOrIndex()) {
     assert(!dim.has_value() && "invalid dim value");
   } else if (auto shapedType = value.getType().dyn_cast<ShapedType>()) {
+    assert(dim.has_value() && "invalid dim value");
     assert(*dim >= 0 && "invalid dim value");
     if (shapedType.hasRank())
       assert(*dim < shapedType.getRank() && "invalid dim value");
   } else {
     llvm_unreachable("unsupported type");
   }
 }
 #endif // NDEBUG
AffineExpr ValueBoundsConstraintSet::getExpr(OpFoldResult ofr) {
   assert(constInt.has_value() && "expected Integer constant");
   return builder.getAffineConstantExpr(*constInt);
 }
 
 AffineExpr ValueBoundsConstraintSet::getExpr(int64_t constant) {
   return builder.getAffineConstantExpr(constant);
 }
 
+AffineExpr ValueBoundsConstraintSet::getAlignedExpr(AffineExpr expr,
+                                                    int64_t numDims,
+                                                    int64_t numSymbols,
+                                                    ValueDimList valueDims) {
+  SmallVector<AffineExpr> dimReplacements, symbolReplacements;
+  for (int i = 0; i < numDims; ++i)
+    dimReplacements.push_back(getExpr(valueDims[i].first, valueDims[i].second));
+  for (int i = numDims; i < numDims + numSymbols; ++i)
+    symbolReplacements.push_back(
+        getExpr(valueDims[i].first, valueDims[i].second));
+  return expr.replaceDimsAndSymbols(dimReplacements, symbolReplacements);
+}
+
+ValueBoundsConstraintSet::StopConditionFn
+ValueBoundsConstraintSet::getStopCondition() const {
+  assert(stopCondition && "stop condition not set");
+  return stopCondition;
+}
+
 int64_t ValueBoundsConstraintSet::insert(ValueDim valueDim, bool isSymbol) {
   assert((valueDimToPosition.find(valueDim) == valueDimToPosition.end()) &&
          "already mapped");
   int64_t pos = isSymbol ? cstr.appendVar(VarKind::Symbol)
                          : cstr.appendVar(VarKind::SetDim);
   positionToValueDim.insert(positionToValueDim.begin() + pos, valueDim);
   // Update reverse mapping.
   for (int64_t i = pos; i < positionToValueDim.size(); ++i)

 
 static Operation *getOwnerOfValue(Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>())
     return bbArg.getOwner()->getParentOp();
   return value.getDefiningOp();
 }
 
 void ValueBoundsConstraintSet::processWorklist(StopConditionFn stopCondition) {
+  auto resetStopCondition =
+      llvm::make_scope_exit([&]() { this->stopCondition = nullptr; });
+  this->stopCondition = stopCondition;
+
   while (!worklist.empty()) {
     int64_t pos = worklist.front();
     worklist.pop();
     ValueDim valueDim = positionToValueDim[pos];
     Value value = valueDim.first;
     int64_t dim = valueDim.second;
 
     // Check for static dim size.
     if (dim != kIntegerValue) {
       auto shapedType = value.getType().cast<ShapedType>();
       if (shapedType.hasRank() && !shapedType.isDynamicDim(dim)) {
         addBound(BoundType::EQ, value, dim,
                  builder.getAffineConstantExpr(shapedType.getDimSize(dim)));
         continue;
       }
     }
 
     // Do not process any further if the stop condition is met.
     auto maybeDim =
         dim == kIntegerValue ? std::nullopt : std::make_optional(dim);
-    if (stopCondition(value, maybeDim))
+    if (stopCondition(value, maybeDim, *this))
       continue;
 
     // Query `ValueBoundsOpInterface` for constraints. New items may be added to
     // the worklist.
     auto valueBoundsOp =
         dyn_cast<ValueBoundsOpInterface>(getOwnerOfValue(value));
     if (!valueBoundsOp)
       continue;

 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
     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)) {
+  // 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(kIntegerValue));
+  ValueBoundsConstraintSet cstr(valueDim);
+  if (stopCondition(value, dim, cstr)) {
     // 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(kIntegerValue));
-  ValueBoundsConstraintSet cstr(valueDim);
   cstr.processWorklist(stopCondition);
 
   // Project out all variables (apart from `valueDim`) that do not match the
   // stop condition.
   cstr.projectOut([&](ValueDim p) {
     // Do not project out `valueDim`.
     if (valueDim == p)
       return false;
     auto maybeDim =
         p.second == kIntegerValue ? std::nullopt : std::make_optional(p.second);
-    return !stopCondition(p.first, maybeDim);
+    return !stopCondition(p.first, maybeDim, cstr);
   });
 
   // Compute lower and upper bounds for `valueDim`.
   int64_t pos = cstr.getPos(value, dim.value_or(kIntegerValue));
   SmallVector<AffineMap> lb(1), ub(1);
   cstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lb, &ub,
                            /*getClosedUB=*/true);
 

 }
 
 LogicalResult ValueBoundsConstraintSet::computeBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
     Value value, std::optional<int64_t> dim, ValueDimList dependencies,
     bool closedUB) {
   return computeBound(
       resultMap, mapOperands, type, value, dim,
-      [&](Value v, std::optional<int64_t> d) {
+      [&](Value v, std::optional<int64_t> d,
+          const ValueBoundsConstraintSet &cstr) {
         return llvm::is_contained(dependencies, std::make_pair(v, d));
       },
       closedUB);
 }
 
 LogicalResult ValueBoundsConstraintSet::computeIndependentBound(
     AffineMap &resultMap, ValueDimList &mapOperands, presburger::BoundType type,
     Value value, std::optional<int64_t> dim, ValueRange independencies,
while (!worklist.empty()) {
       worklist.append(op->getOperands().begin(), op->getOperands().end());
     }
     return true;
   };
 
   // Reify bounds in terms of any independent values.
   return computeBound(
       resultMap, mapOperands, type, value, dim,
-      [&](Value v, std::optional<int64_t> d) { return isIndependent(v); },
+      [&](Value v, std::optional<int64_t> d,
+          const ValueBoundsConstraintSet &cstr) { return isIndependent(v); },
       closedUB);
 }
 
 FailureOr<int64_t> ValueBoundsConstraintSet::computeConstantBound(
     presburger::BoundType type, Value value, std::optional<int64_t> dim,
     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.
   ValueDim valueDim = std::make_pair(value, dim.value_or(kIntegerValue));
   ValueBoundsConstraintSet cstr(valueDim);
   int64_t pos = cstr.getPos(value, dim.value_or(kIntegerValue));
   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) {
+        /*stopCondition=*/[&](Value v, std::optional<int64_t> dim,
+                              const ValueBoundsConstraintSet &cstr) {
           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 + ubAdjustment : *bound;
   return failure();
 }

mlir/test/Dialect/Arith/value-bounds-op-interface-impl.mlir

 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
 //       CHECK:   %[[c5:.*]] = arith.constant 5 : index
 //       CHECK:   return %[[c5]]
 func.func @arith_const() -> index {
   %c5 = arith.constant 5 : index
   %0 = "test.reify_bound"(%c5) : (index) -> (index)
   return %0 : index
 }
+
+// -----
+
+// CHECK-LABEL: func @arith_select_const(
+//  CHECK-SAME:     %[[a:.*]]: index, %[[b:.*]]: index
+//       CHECK:   return %[[b]]
+func.func @arith_select_const(%a: index, %b: index) -> index {
+  %c0 = arith.constant 0 : i1
+  %0 = arith.select %c0, %a, %b : index
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
+// CHECK-LABEL: func @arith_select_same_bound(
+//  CHECK-SAME:     %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: i1
+//       CHECK:   return %[[a]]
+func.func @arith_select_same_bound(%a: index, %b: index, %c: i1) -> index {
+  %0 = arith.select %c, %a, %a : index
+  %1 = "test.reify_bound"(%0) : (index) -> (index)
+  return %1 : index
+}
+
+// -----
+
+//       CHECK: #[[$map:.*]] = affine_map<()[s0, s1] -> (s0 + s1 * 2)>
+// CHECK-LABEL: func @arith_select_same_bound_2(
+//  CHECK-SAME:     %[[a:.*]]: index, %[[b:.*]]: index, %[[c:.*]]: i1
+//       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[b]], %[[a]]]
+//       CHECK:   return %[[apply]]
+func.func @arith_select_same_bound_2(%a: index, %b: index, %c: i1) -> index {
+  %c1 = arith.constant 1 : index
+  %c2 = arith.constant 2 : index
+  %0 = arith.addi %a, %b : index
+  %1 = arith.addi %0, %a : index
+  %2 = arith.muli %c2, %a : index
+  %3 = arith.addi %2, %b : index
+
+  %selected = arith.select %c, %1, %3 : index
+  %bound = "test.reify_bound"(%selected) {reify_to_func_args}
+      : (index) -> (index)
+  return %bound : index
+}
+
+// -----
+
+// CHECK-LABEL: func @arith_select_merge_const_bounds(
+//   CHECK-DAG:   %[[c5:.*]] = arith.constant 5 : index
+//   CHECK-DAG:   %[[c9:.*]] = arith.constant 9 : index
+//       CHECK:   return %[[c5]], %[[c9]]
+func.func @arith_select_merge_const_bounds(%a: tensor<5xf32>, %b: tensor<8xf32>,
+                                           %c: i1) -> (index, index) {
+  %0 = tensor.cast %a : tensor<5xf32> to tensor<?xf32>
+  %1 = tensor.cast %b : tensor<8xf32> to tensor<?xf32>
+
+  %selected = arith.select %c, %0, %1 : tensor<?xf32>
+  %lb = "test.reify_bound"(%selected) {reify_to_func_args, type = "LB", dim = 0}
+      : (tensor<?xf32>) -> (index)
+  %ub = "test.reify_bound"(%selected) {reify_to_func_args, type = "UB", dim = 0}
+      : (tensor<?xf32>) -> (index)
+  return %lb, %ub : index, index
+}

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

if (op->getName().getStringRef() == "test.reify_bound") {
                      : std::make_optional<int64_t>(
                            op->getAttrOfType<IntegerAttr>("dim").getInt());
 
       // Check if a constant was requested.
       bool constant = op->hasAttr("constant");
 
       // Prepare stop condition. By default, reify in terms of the op's
       // operands. No stop condition is used when a constant was requested.
-      std::function<bool(Value, std::optional<int64_t>)> stopCondition =
-          [&](Value v, std::optional<int64_t> d) {
+      std::function<bool(Value, std::optional<int64_t>,
+                         const ValueBoundsConstraintSet &cstr)>
+          stopCondition = [&](Value v, std::optional<int64_t> d,
+                              const ValueBoundsConstraintSet &cstr) {
             // Reify in terms of SSA values that are different from `value`.
             return v != value;
           };
-      if (reifyToFuncArgs) {
+      if (op->hasAttr("reify_to_func_args") || reifyToFuncArgs) {
         // Reify in terms of function block arguments.
-        stopCondition = stopCondition = [](Value v, std::optional<int64_t> d) {
+        stopCondition = [](Value v, std::optional<int64_t> d,
+                           const ValueBoundsConstraintSet &cstr) {
           auto bbArg = v.dyn_cast<BlockArgument>();
           if (!bbArg)
             return false;
           return isa<FunctionOpInterface>(
               bbArg.getParentBlock()->getParentOp());
         };
       }
 

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

hdrs = glob(
         [
             "include/mlir/Dialect/SCF/IR/*.h",
         ],
     ),
     includes = ["include"],
     deps = [
         ":ArithDialect",
         ":ArithUtils",
+        ":ArithValueBoundsOpInterfaceImpl",
         ":BufferizationDialect",
         ":ControlFlowDialect",
         ":ControlFlowInterfaces",
         ":FuncDialect",
         ":IR",
         ":InferTypeOpInterface",
         ":LoopLikeInterface",
         ":MemRefDialect",