[mlir][Interfaces] ValueBoundsOpInterface: Support IntegerTypes

[springerm]

Summary

Both IndexType and IntegerTypes are now supported. This is needed to compute bounds for condition values of arith.select and scf.if.

Depends On: D143910

[springerm]

rebase

[dcaballe]

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

35

To void this complexity, could the user just get the default type and convert it at their wish?

51

Ok, some of my questions in previous commits are answered now :)

[dcaballe]

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

35

Sorry if I'm missing something trivial but I'm still not convinced that this conversion should be part of the reify logic. Couldn't we just reify the value using the default type and let the conversion to the user?

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.
+/// Reify a bound for the given integer/index-typed value or shape
+/// size in terms of the owning op's operands. `dim` must be `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,
                                         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 integer/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 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.
 ///
 /// Example:
 /// %0 = arith.addi %a, %b : index
 /// %1 = arith.addi %0, %c : index
 ///

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

 //===- Transforms.h - Arith Transforms --------------------------*- C++ -*-===//
 //
 // 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_TRANSFORMS_TRANSFORMS_H
 #define MLIR_DIALECT_ARITH_TRANSFORMS_TRANSFORMS_H
 
+#include "mlir/IR/Types.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 #include "mlir/Support/LogicalResult.h"
 
 namespace mlir {
 class Location;
 class OpBuilder;
 class OpFoldResult;
+class Type;
 class Value;
 
 namespace presburger {
 enum class BoundType;
 } // namespace presburger
 
 namespace arith {
 
-/// 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.
+/// Reify a bound for the given integer/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 integer/index-typed.
+///
+/// The type of the reified result (IntegerType or IndexType) can be specified.
+/// 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,
-                                        bool closedUB = false);
+                                        bool closedUB = false,
+                                        Type resultType = {});
 
-/// 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 integer/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 integer/index-typed.
+///
+/// The type of the reified result (IntegerType or IndexType) can be specified.
+/// 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,
                 ValueBoundsConstraintSet::StopConditionFn stopCondition,
-                bool closedUB = false);
+                bool closedUB = false, Type resultType = {});
 
 } // namespace arith
 } // namespace mlir
 
 #endif // MLIR_DIALECT_ARITH_TRANSFORMS_TRANSFORMS_H

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.h

 
 #include <queue>
 
 namespace mlir {
 
 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 index-typed
+/// 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.
private:
 
     ValueBoundsConstraintSet &cstr;
     Value value;
     std::optional<int64_t> dim = std::nullopt;
   };
 
 public:
   /// The stop condition when traversing the backward slice of a shaped value/
-  /// index-type value. The traversal continues until the stop condition
+  /// 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>)>;
 
-  /// Compute a bound for the given index-typed value or shape dimension size.
+  /// 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 index-type SSA value
+  /// 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 index-typed. The bound
+  /// `dim` must be `nullopt` if and only if `value` is integer-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.
public:
   /// Compute a bound in that is independent of the of the values in
   /// `independencies`.
   static LogicalResult
   computeIndependentBound(AffineMap &resultMap, ValueDimList &mapOperands,
                           presburger::BoundType type, Value value,
                           std::optional<int64_t> dim, ValueRange independencies,
                           bool closedUB = false);
 
-  /// Compute a constant bound for the given index-typed value or shape
+  /// Compute a constant bound for the given integer-typed value or shape
   /// dimension size.
   ///
-  /// `dim` must be `nullopt` if and only if `value` is index-typed. This
+  /// `dim` must be `nullopt` if and only if `value` is integer-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,
                        bool closedUB = false);
 
-  /// Add a bound for the given index-typed value or shaped value. This function
-  /// returns a builder that adds the bound.
+  /// Add a bound for the given integer-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.
+  /// Return an expression that represents the given integer-typed value or
+  /// shaped value dimension. If this value/dimension was not used so far, it is
+  /// added to the worklist.
   ///
-  /// `dim` must be `nullopt` if and only if the given value is of index type.
+  /// `dim` must be `nullopt` if and only if the given value is of integer type.
   AffineExpr getExpr(Value value, std::optional<int64_t> dim = std::nullopt);
 
-  /// Return an expression that represents a constant or index-typed SSA value.
-  /// In case of a value, if this value was not used so far, it is added to the
-  /// worklist.
+  /// 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);
 
 protected:
-  /// Dimension identifier to indicate a value is index-typed. This is used for
-  /// internal data structures/API only.
-  static constexpr int64_t kIndexValue = -1;
+  /// Dimension identifier to indicate a value is integer-typed. This is used
+  /// for internal data structures/API only.
+  static constexpr int64_t kIntegerValue = -1;
 
-  /// An index-typed value or the dimension of a shaped-type value.
+  /// An integer-typed value or the dimension of a shaped-type value.
   using ValueDim = std::pair<Value, int64_t>;
 
   ValueBoundsConstraintSet(Value value, std::optional<int64_t> dim);
 
-  /// Iteratively process all elements on the worklist until an index-typed
+  /// Iteratively process all elements on the worklist until an integer-typed
   /// value or shaped value meets `stopCondition`. Such values are not processed
   /// any further.
   void processWorklist(StopConditionFn stopCondition);
 
   /// Bound the given column in the underlying constraint set by the given
   /// expression.
   void addBound(presburger::BoundType type, int64_t pos, AffineExpr expr);
 

mlir/include/mlir/Interfaces/ValueBoundsOpInterface.td

 //===-- ValueBoundsOpInterface.td - Value Bounds -----------*- tablegen -*-===//
 //
 // 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 VALUEBOUNDSOPINTERFACE
 #define VALUEBOUNDSOPINTERFACE
 
 include "mlir/IR/OpBase.td"
 
 def ValueBoundsOpInterface : OpInterface<"ValueBoundsOpInterface"> {
   let description = [{
-    This interface allows operations with index-typed and/or shaped value-typed
-    results/block arguments to specify range bounds. These bounds are stored in
-    a constraint set. The constraint set can then be queried to compute bounds
-    in terms of other values that are stored in the constraint set.
+    This interface allows operations with integer/index-typed and/or shaped
+    value-typed results/block arguments to specify range bounds. These bounds
+    are stored in a constraint set. The constraint set can then be queried to
+    compute bounds in terms of other values that are stored in the constraint
+    set.
   }];
   let cppNamespace = "::mlir";
   let methods = [
     InterfaceMethod<
       /*desc=*/[{
-        Populate the constraint set with bounds for the given index-typed
-        value.
+        Populate the constraint set with bounds for the given
+        integer/index-typed value.
 
         Note: If `value` is a block argument, it must belong to an entry block
         of a region. Unstructured control flow graphs are not supported at the
         moment.
       }],
       /*retType=*/"void",
-      /*methodName=*/"populateBoundsForIndexValue",
+      /*methodName=*/"populateBoundsForIntegerValue",
       /*args=*/(ins "::mlir::Value":$value,
                     "::mlir::ValueBoundsConstraintSet &":$cstr),
       /*methodBody=*/"",
       /*defaultImplementation=*/[{
         llvm_unreachable("not implemented");
       }]
     >,
     InterfaceMethod<

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

 using namespace mlir;
 
 namespace mlir {
 namespace {
 
 struct AffineApplyOpInterface
     : public ValueBoundsOpInterface::ExternalModel<AffineApplyOpInterface,
                                                    AffineApplyOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto applyOp = cast<AffineApplyOp>(op);
     assert(value == applyOp.getResult() && "invalid value");
     assert(applyOp.getAffineMap().getNumResults() == 1 &&
            "expected single result");
 
     // Align affine map result with dims/symbols in the constraint set.
     AffineExpr expr = applyOp.getAffineMap().getResult(0);
     SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
         applyOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
     SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
         applyOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
     AffineExpr bound =
         expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
     cstr.bound(value) == bound;
   }
 };
 
 struct AffineMinOpInterface
     : public ValueBoundsOpInterface::ExternalModel<AffineMinOpInterface,
                                                    AffineMinOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto minOp = cast<AffineMinOp>(op);
     assert(value == minOp.getResult() && "invalid value");
 
     // Align affine map results with dims/symbols in the constraint set.
     for (AffineExpr expr : minOp.getAffineMap().getResults()) {
       SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
           minOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
       SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(
           minOp.getSymbolOperands(), [&](Value v) { return cstr.getExpr(v); }));
       AffineExpr bound =
           expr.replaceDimsAndSymbols(dimReplacements, symReplacements);
       cstr.bound(value) <= bound;
     }
   };
 };
 
 struct AffineMaxOpInterface
     : public ValueBoundsOpInterface::ExternalModel<AffineMaxOpInterface,
                                                    AffineMaxOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto maxOp = cast<AffineMaxOp>(op);
     assert(value == maxOp.getResult() && "invalid value");
 
     // Align affine map results with dims/symbols in the constraint set.
     for (AffineExpr expr : maxOp.getAffineMap().getResults()) {
       SmallVector<AffineExpr> dimReplacements = llvm::to_vector(llvm::map_range(
           maxOp.getDimOperands(), [&](Value v) { return cstr.getExpr(v); }));
       SmallVector<AffineExpr> symReplacements = llvm::to_vector(llvm::map_range(

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

OpFoldResult mlir::reifyValueBound(
     ArrayRef<std::pair<Value, std::optional<int64_t>>> mapOperands) {
   // Materialize tensor.dim/memref.dim ops.
   SmallVector<Value> operands;
   for (auto valueDim : mapOperands) {
     Value value = valueDim.first;
     std::optional<int64_t> dim = valueDim.second;
 
     if (!dim.has_value()) {
-      // This is an index-typed value.
-      assert(value.getType().isIndex() && "expected index type");
+      // This is an integer-typed value.
+      assert(value.getType().isIntOrIndex() && "expected integer type");
       operands.push_back(value);
       continue;
     }
 
     assert(cast<ShapedType>(value.getType()).isDynamicDim(*dim) &&
            "expected dynamic dim");
     if (isa<RankedTensorType>(value.getType())) {
       // A tensor dimension is used: generate a tensor.dim.

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

 using namespace mlir;
 
 namespace mlir {
 namespace arith {
 namespace {
 
 struct AddIOpInterface
     : public ValueBoundsOpInterface::ExternalModel<AddIOpInterface, AddIOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto addIOp = cast<AddIOp>(op);
     assert(value == addIOp.getResult() && "invalid value");
 
     cstr.bound(value) ==
         cstr.getExpr(addIOp.getLhs()) + cstr.getExpr(addIOp.getRhs());
   }
 };
 
 struct ConstantOpInterface
     : public ValueBoundsOpInterface::ExternalModel<ConstantOpInterface,
                                                    ConstantOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto constantOp = cast<ConstantOp>(op);
     assert(value == constantOp.getResult() && "invalid value");
 
     if (auto attr = constantOp.getValue().dyn_cast<IntegerAttr>())
       cstr.bound(value) == attr.getInt();
   }
 };
 
 struct SubIOpInterface
     : public ValueBoundsOpInterface::ExternalModel<SubIOpInterface, SubIOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto subIOp = cast<SubIOp>(op);
     assert(value == subIOp.getResult() && "invalid value");
 
     cstr.bound(value) ==
         cstr.getExpr(subIOp.getLhs()) - cstr.getExpr(subIOp.getRhs());
   }
 };
 
 struct MulIOpInterface
     : public ValueBoundsOpInterface::ExternalModel<MulIOpInterface, MulIOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto mulIOp = cast<MulIOp>(op);
     assert(value == mulIOp.getResult() && "invalid value");
 
     cstr.bound(value) ==
         cstr.getExpr(mulIOp.getLhs()) * cstr.getExpr(mulIOp.getRhs());
   }
 };
 

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

 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Interfaces/ValueBoundsOpInterface.h"
 
 using namespace mlir;
 using namespace mlir::arith;
 
 /// Build Arith IR for the given affine map and its operands.
-static Value buildArithValue(OpBuilder &b, Location loc, AffineMap map,
-                             ValueRange operands) {
+static Value buildArithValue(OpBuilder &b, Location loc, Type type,
+                             AffineMap map, ValueRange operands) {
   assert(map.getNumResults() == 1 && "multiple results not supported yet");
   std::function<Value(AffineExpr)> buildExpr = [&](AffineExpr e) -> Value {
     switch (e.getKind()) {
     case AffineExprKind::Constant:
-      return b.create<ConstantIndexOp>(loc,
-                                       e.cast<AffineConstantExpr>().getValue());
+      return b.create<ConstantOp>(
+          loc, type,
+          b.getIntegerAttr(type, e.cast<AffineConstantExpr>().getValue()));
     case AffineExprKind::DimId:
       return operands[e.cast<AffineDimExpr>().getPosition()];
     case AffineExprKind::SymbolId:
       return operands[e.cast<AffineSymbolExpr>().getPosition() +
                       map.getNumDims()];
     case AffineExprKind::Add: {
       auto binaryExpr = e.cast<AffineBinaryOpExpr>();
       return b.create<AddIOp>(loc, buildExpr(binaryExpr.getLHS()),
std::function<Value(AffineExpr)> buildExpr = [&](AffineExpr e) -> Value {
     }
     default:
       llvm_unreachable("unknown AffineExpr");
     }
   };
   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) {
+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) {
     // Reify in terms of SSA values that are different from `value`.
     return v != value;
   };
-  return reifyValueBound(b, loc, type, value, dim, stopCondition, closedUB);
+  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) {
+    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;
   ValueDimList mapOperands;
   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;
 
     if (!dim.has_value()) {
-      // This is an index-typed value.
-      assert(value.getType().isIndex() && "expected index type");
+      // This is an integer-typed value.
+      assert(value.getType().isIntOrIndex() && "expected integer type");
+      if (value.getType() != resultType)
+        value = b.create<arith::IndexCastOp>(loc, resultType, value);
       operands.push_back(value);
       continue;
     }
 
     assert(value.getType().cast<ShapedType>().isDynamicDim(*dim) &&
            "expected dynamic dim");
     if (value.getType().isa<RankedTensorType>()) {
       // A tensor dimension is used: generate a tensor.dim.
-      operands.push_back(b.create<tensor::DimOp>(loc, value, *dim));
+      Value dimSize = b.create<tensor::DimOp>(loc, value, *dim);
+      if (!resultType.isIndex())
+        dimSize = b.create<arith::IndexCastOp>(loc, resultType, dimSize);
+      operands.push_back(dimSize);
     } else if (value.getType().isa<MemRefType>()) {
       // A memref dimension is used: generate a memref.dim.
-      operands.push_back(b.create<memref::DimOp>(loc, value, *dim));
+      Value dimSize = b.create<memref::DimOp>(loc, value, *dim);
+      if (!resultType.isIndex())
+        dimSize = b.create<arith::IndexCastOp>(loc, resultType, dimSize);
+      operands.push_back(dimSize);
     } else {
       llvm_unreachable("cannot generate DimOp for unsupported shaped type");
     }
   }
 
   // Check for special cases where no arith ops are needed.
   if (boundMap.isSingleConstant()) {
     // Bound is a constant: return an IntegerAttr.
     return static_cast<OpFoldResult>(
         b.getIndexAttr(boundMap.getSingleConstantResult()));
   }
   // No arith ops are needed if the bound is a single SSA value.
   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 Arith ops.
-  return static_cast<OpFoldResult>(buildArithValue(b, loc, boundMap, operands));
+  return static_cast<OpFoldResult>(
+      buildArithValue(b, loc, resultType, boundMap, operands));
 }

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

if (castOp.getResult().getType().isa<MemRefType>() &&
         castOp.getSource().getType().isa<MemRefType>()) {
       cstr.bound(value)[dim] == cstr.getExpr(castOp.getSource(), dim);
     }
   }
 };
 
 struct DimOpInterface
     : public ValueBoundsOpInterface::ExternalModel<DimOpInterface, DimOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto dimOp = cast<DimOp>(op);
     assert(value == dimOp.getResult() && "invalid value");
 
     auto constIndex = dimOp.getConstantIndex();
     if (!constIndex.has_value())
       return;
     cstr.bound(value) == cstr.getExpr(dimOp.getSource(), *constIndex);
   }
void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
     auto type = getGlobalOp.getType();
     assert(!type.isDynamicDim(dim) && "expected static dim");
     cstr.bound(value)[dim] == type.getDimSize(dim);
   }
 };
 
 struct RankOpInterface
     : public ValueBoundsOpInterface::ExternalModel<RankOpInterface, RankOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto rankOp = cast<RankOp>(op);
     assert(value == rankOp.getResult() && "invalid value");
 
     auto memrefType = rankOp.getMemref().getType().dyn_cast<MemRefType>();
     if (!memrefType)
       return;
     cstr.bound(value) == memrefType.getRank();
   }

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

if (auto dimExpr = bound.getResult(0).dyn_cast<AffineDimExpr>()) {
       int64_t idx = symExpr.getPosition() + bound.getNumDims();
       singleValue = boundOperands[idx].first;
       singleDim = boundOperands[idx].second;
     }
     if (singleValue == iterArg && singleDim == dim)
       addEqBound();
   }
 
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto forOp = cast<ForOp>(op);
 
     if (value == forOp.getInductionVar()) {
       // TODO: Take into account step size.
       cstr.bound(value) >= forOp.getLowerBound();
       cstr.bound(value) < forOp.getUpperBound();
       return;
     }

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

if (castOp.getResult().getType().isa<RankedTensorType>() &&
         castOp.getSource().getType().isa<RankedTensorType>()) {
       cstr.bound(value)[dim] == cstr.getExpr(castOp.getSource(), dim);
     }
   }
 };
 
 struct DimOpInterface
     : public ValueBoundsOpInterface::ExternalModel<DimOpInterface, DimOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto dimOp = cast<DimOp>(op);
     assert(value == dimOp.getResult() && "invalid value");
 
     auto constIndex = dimOp.getConstantIndex();
     if (!constIndex.has_value())
       return;
     cstr.bound(value) == cstr.getExpr(dimOp.getSource(), *constIndex);
   }
AffineExpr expr = cstr.getExpr(padOp.getSource(), dim) +
                       cstr.getExpr(padOp.getMixedLowPad()[dim]) +
                       cstr.getExpr(padOp.getMixedHighPad()[dim]);
     cstr.bound(value)[dim] == expr;
   }
 };
 
 struct RankOpInterface
     : public ValueBoundsOpInterface::ExternalModel<RankOpInterface, RankOp> {
-  void populateBoundsForIndexValue(Operation *op, Value value,
+  void populateBoundsForIntegerValue(Operation *op, Value value,
                                      ValueBoundsConstraintSet &cstr) const {
     auto rankOp = cast<RankOp>(op);
     assert(value == rankOp.getResult() && "invalid value");
 
     auto tensorType = rankOp.getTensor().getType().dyn_cast<RankedTensorType>();
     if (!tensorType)
       return;
     cstr.bound(value) == tensorType.getRank();
   }

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

 ValueBoundsConstraintSet::ValueBoundsConstraintSet(Value value,
                                                    std::optional<int64_t> dim)
     : builder(value.getContext()) {
   insert(value, dim, /*isSymbol=*/false);
 }
 
 #ifndef NDEBUG
 static void assertValidValueDim(Value value, std::optional<int64_t> dim) {
-  if (value.getType().isIndex()) {
+  if (value.getType().isIntOrIndex()) {
     assert(!dim.has_value() && "invalid dim value");
   } else if (auto shapedType = value.getType().dyn_cast<ShapedType>()) {
     assert(*dim >= 0 && "invalid dim value");
     if (shapedType.hasRank())
       assert(*dim < shapedType.getRank() && "invalid dim value");
   } else {
     llvm_unreachable("unsupported type");
   }

 #endif // NDEBUG
 
   auto shapedType = value.getType().dyn_cast<ShapedType>();
   if (shapedType) {
     // Static dimension: return constant directly.
     if (shapedType.hasRank() && !shapedType.isDynamicDim(*dim))
       return builder.getAffineConstantExpr(shapedType.getDimSize(*dim));
   } else {
-    // Constant index value: return directly.
+    // Constant integer value: return directly.
     if (auto constInt = getConstantIntValue(value))
       return builder.getAffineConstantExpr(*constInt);
   }
 
   // Dynamic value: add to constraint set.
-  ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
+  ValueDim valueDim = std::make_pair(value, dim.value_or(kIntegerValue));
   if (valueDimToPosition.find(valueDim) == valueDimToPosition.end())
     (void)insert(value, dim);
   int64_t pos = getPos(value, dim);
   return pos < cstr.getNumDimVars()
              ? builder.getAffineDimExpr(pos)
              : builder.getAffineSymbolExpr(pos - cstr.getNumDimVars());
 }
 

 
 int64_t ValueBoundsConstraintSet::insert(Value value,
                                          std::optional<int64_t> dim,
                                          bool isSymbol) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
 #endif // NDEBUG
 
-  ValueDim valueDim = std::make_pair(value, dim.value_or(kIndexValue));
+  ValueDim valueDim = std::make_pair(value, dim.value_or(kIntegerValue));
   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, e = positionToValueDim.size(); i < e; ++i)
     valueDimToPosition[positionToValueDim[i]] = i;
 
   worklist.push(pos);
   return pos;
 }
 
 int64_t ValueBoundsConstraintSet::getPos(Value value,
                                          std::optional<int64_t> dim) const {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
   assert((value.isa<OpResult>() ||
           value.cast<BlockArgument>().getOwner()->isEntryBlock()) &&
          "unstructured control flow is not supported");
 #endif // NDEBUG
 
-  auto it =
-      valueDimToPosition.find(std::make_pair(value, dim.value_or(kIndexValue)));
+  auto it = valueDimToPosition.find(
+      std::make_pair(value, dim.value_or(kIntegerValue)));
   assert(it != valueDimToPosition.end() && "expected mapped entry");
   return it->second;
 }
 
 static Operation *getOwnerOfValue(Value value) {
   if (auto bbArg = value.dyn_cast<BlockArgument>())
     return bbArg.getOwner()->getParentOp();
   return value.getDefiningOp();
 }
 
 void ValueBoundsConstraintSet::processWorklist(StopConditionFn 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 != kIndexValue) {
+    if (dim != kIntegerValue) {
       auto shapedType = cast<ShapedType>(value.getType());
       if (shapedType.hasRank() && !shapedType.isDynamicDim(dim)) {
         bound(value)[dim] == getExpr(shapedType.getDimSize(dim));
         continue;
       }
     }
 
     // Do not process any further if the stop condition is met.
-    auto maybeDim = dim == kIndexValue ? std::nullopt : std::make_optional(dim);
+    auto maybeDim =
+        dim == kIntegerValue ? std::nullopt : std::make_optional(dim);
     if (stopCondition(value, maybeDim))
       continue;
 
     // Query `ValueBoundsOpInterface` for constraints. New items may be added to
     // the worklist.
     auto valueBoundsOp =
         dyn_cast<ValueBoundsOpInterface>(getOwnerOfValue(value));
     if (!valueBoundsOp)
       continue;
-    if (dim == kIndexValue) {
-      valueBoundsOp.populateBoundsForIndexValue(value, *this);
+    if (dim == kIntegerValue) {
+      valueBoundsOp.populateBoundsForIntegerValue(value, *this);
     } else {
       valueBoundsOp.populateBoundsForShapedValueDim(value, dim, *this);
     }
   }
 }
 
 void ValueBoundsConstraintSet::projectOut(int64_t pos) {
   assert(pos >= 0 && pos < positionToValueDim.size() && "invalid position");
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));
+  ValueDim valueDim = std::make_pair(value, dim.value_or(kIntegerValue));
   ValueBoundsConstraintSet cstr(value, dim);
   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 == kIndexValue ? std::nullopt : std::make_optional(p.second);
+        p.second == kIntegerValue ? std::nullopt : std::make_optional(p.second);
     return !stopCondition(p.first, maybeDim);
   });
 
   // Compute lower and upper bounds for `valueDim`.
   int64_t pos = cstr.getPos(value, dim);
   SmallVector<AffineMap> lb(1), ub(1);
   cstr.cstr.getSliceBounds(pos, 1, value.getContext(), &lb, &ub,
                            /*getClosedUB=*/true);
if (isDim) {
       replacementDims.push_back(b.getAffineDimExpr(numDims++));
     } else {
       replacementSymbols.push_back(b.getAffineSymbolExpr(numSymbols++));
     }
 
     ValueBoundsConstraintSet::ValueDim valueDim = cstr.positionToValueDim[i];
     Value value = valueDim.first;
     int64_t dim = valueDim.second;
-    if (dim == ValueBoundsConstraintSet::kIndexValue) {
-      // An index-type value is used: can be used directly in the affine.apply
+    if (dim == ValueBoundsConstraintSet::kIntegerValue) {
+      // An integer-type value is used: can be used directly in the affine.apply
       // op.
-      assert(value.getType().isIndex() && "expected index type");
+      assert(value.getType().isIntOrIndex() && "expected integer type");
       mapOperands.push_back(std::make_pair(value, std::nullopt));
       continue;
     }
 
     assert(cast<ShapedType>(value.getType()).isDynamicDim(dim) &&
            "expected dynamic dim");
     mapOperands.push_back(std::make_pair(value, dim));
   }

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

 }
 
 // -----
 
 // CHECK: #[[$map:.*]] = affine_map<()[s0] -> (s0 * 5)>
 // CHECK-LABEL: func @arith_muli(
 //  CHECK-SAME:     %[[a:.*]]: index
 //       CHECK:   %[[apply:.*]] = affine.apply #[[$map]]()[%[[a]]]
-//       CHECK:   return %[[apply]]
-func.func @arith_muli(%a: index) -> index {
+//       CHECK:   %[[apply2:.*]] = affine.apply #[[$map]]()[%[[a]]]
+//       CHECK:   %[[casted:.*]] = arith.index_cast %[[apply2]] : index to i32
+//       CHECK:   return %[[apply]], %[[casted]]
+
+// CHECK-ARITH-LABEL: func @arith_muli(
+//  CHECK-ARITH-SAME:     %[[a:.*]]: index
+//       CHECK-ARITH:   %[[c5:.*]] = arith.constant 5 : index
+//       CHECK-ARITH:   %[[mul:.*]] = arith.muli %[[a]], %[[c5]]
+//       CHECK-ARITH:   %[[a_casted:.*]] = arith.index_cast %[[a]] : index to i32
+//       CHECK-ARITH:   %[[c5_i32:.*]] = arith.constant 5 : i32
+//       CHECK-ARITH:   %[[mul_i32:.*]] = arith.muli %[[a_casted]], %[[c5_i32]]
+//       CHECK-ARITH:   return %[[mul]], %[[mul_i32]]
+func.func @arith_muli(%a: index) -> (index, i32) {
   %0 = arith.constant 5 : index
   %1 = arith.muli %0, %a : index
-  %2 = "test.reify_bound"(%1) : (index) -> (index)
-  return %2 : index
+  %bound = "test.reify_bound"(%1) : (index) -> (index)
+  %bound_i32 = "test.reify_bound"(%1) : (index) -> (i32)
+  return %bound, %bound_i32 : index, i32
 }
 
 // -----
 
 // CHECK-LABEL: func @arith_const()
 //       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
 }

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

 static LogicalResult testReifyValueBounds(func::FuncOp funcOp,
                                           bool reifyToFuncArgs,
                                           bool useArithOps) {
   IRRewriter rewriter(funcOp.getContext());
   WalkResult result = funcOp.walk([&](Operation *op) {
     // Look for test.reify_bound ops.
     if (op->getName().getStringRef() == "test.reify_bound") {
       if (op->getNumOperands() != 1 || op->getNumResults() != 1 ||
-          !op->getResultTypes()[0].isIndex()) {
+          -!op->getResultTypes()[0].isIntOrIndex()) {
         op->emitOpError("invalid op");
         return WalkResult::skip();
       }
       Value value = op->getOperand(0);
-      if (value.getType().isa<IndexType>() !=
+      if (value.getType().isIntOrIndex() !=
           !op->hasAttrOfType<IntegerAttr>("dim")) {
         // Op should have "dim" attribute if and only if the operand is an
-        // index-typed value.
+        // integer/index-typed value.
         op->emitOpError("invalid op");
         return WalkResult::skip();
       }
 
       // Get bound type.
       std::string boundTypeStr = "EQ";
       if (auto boundTypeAttr = op->getAttrOfType<StringAttr>("type"))
         boundTypeStr = boundTypeAttr.str();
       auto boundType = parseBoundType(boundTypeStr);
       if (failed(boundType)) {
         op->emitOpError("invalid op");
         return WalkResult::interrupt();
       }
 
       // Get shape dimension (if any).
-      auto dim = value.getType().isIndex()
+      auto dim = value.getType().isIntOrIndex()
                      ? std::nullopt
                      : 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
if (op->getName().getStringRef() == "test.reify_bound") {
 
       // Reify value bound
       rewriter.setInsertionPointAfter(op);
       FailureOr<OpFoldResult> reified = failure();
       if (constant) {
         auto reifiedConst = ValueBoundsConstraintSet::computeConstantBound(
             *boundType, value, dim, /*stopCondition=*/nullptr);
         if (succeeded(reifiedConst))
-          reified =
-              FailureOr<OpFoldResult>(rewriter.getIndexAttr(*reifiedConst));
+          reified = FailureOr<OpFoldResult>(
+              rewriter.getIntegerAttr(op->getResultTypes()[0], *reifiedConst));
       } else {
         if (useArithOps) {
-          reified = arith::reifyValueBound(rewriter, op->getLoc(), *boundType,
-                                           value, dim, stopCondition);
+          reified = arith::reifyValueBound(
+              rewriter, op->getLoc(), *boundType, value, dim, stopCondition,
+              /*closedUB=*/false, op->getResultTypes()[0]);
         } else {
           reified = reifyValueBound(rewriter, op->getLoc(), *boundType, value,
                                     dim, 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>()) {
+      if (Value val = reified->dyn_cast<Value>()) {
+        if (val.getType() != op->getResultTypes()[0])
+          val = rewriter.create<arith::IndexCastOp>(
+              op->getLoc(), op->getResultTypes()[0], val);
         rewriter.replaceOp(op, val);
         return WalkResult::skip();
       }
-      Value constOp = rewriter.create<arith::ConstantIndexOp>(
-          op->getLoc(), reified->get<Attribute>().cast<IntegerAttr>().getInt());
+      Value constOp = rewriter.create<arith::ConstantOp>(
+          op->getLoc(), op->getResultTypes()[0],
+          reified->get<Attribute>().cast<IntegerAttr>());
       rewriter.replaceOp(op, constOp);
       return WalkResult::skip();
     }
     return WalkResult::advance();
   });
   return failure(result.wasInterrupted());
 }