[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

34

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

50

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

[dcaballe]

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

34

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.
 FailureOr<OpFoldResult>
 reifyValueBound(OpBuilder &b, Location loc, presburger::BoundType type,
                 Value value, std::optional<int64_t> dim,
                 function_ref<bool(Value, std::optional<int64_t>)> stopCondition,
                 bool closedUB = false);

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/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,
                 function_ref<bool(Value, std::optional<int64_t>)> 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.
 ///
 /// 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/
-  /// 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);
 
-  /// Bound the given index-typed value by the given expression.
+  /// Bound the given integer-typed value by the given expression.
   void addBound(presburger::BoundType type, Value value, AffineExpr expr);
 
   /// Bound the given shaped value dimension by the given expression.
   void addBound(presburger::BoundType type, Value value, int64_t dim,
                 AffineExpr expr);
 
-  /// Bound the given index-typed value by a constant or SSA value.
+  /// Bound the given integer-typed value by a constant or SSA value.
   void addBound(presburger::BoundType type, Value value, OpFoldResult ofr);
 
   /// Bound the given shaped value dimension by a constant or SSA value.
   void addBound(presburger::BoundType type, Value value, int64_t dim,
                 OpFoldResult ofr);
 
-  /// 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.
-  static constexpr int64_t kIndexValue = -1;
+  /// 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 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);
 
   /// Return the column position of the given value/dimension. Asserts that the
   /// value/dimension exists in the constraint set.
-  int64_t getPos(Value value, int64_t dim = kIndexValue) const;
+  int64_t getPos(Value value, int64_t dim = kIntegerValue) const;
 
   /// Insert a value/dimension into the constraint set. If `isSymbol` is set to
   /// "false", a dimension is added.
   ///
   /// Note: There are certain affine restrictions wrt. dimensions. E.g., they
   /// cannot be multiplied. Furthermore, bounds can only be queried for
   /// dimensions but not for symbols.
   int64_t insert(ValueDim valueDim, bool isSymbol = true);

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 presburger::BoundType;
 
 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.addBound(BoundType::EQ, 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);
       AffineExpr openBound = bound + 1;
       cstr.addBound(BoundType::UB, value, openBound);
     }
   };
 };
 
 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(value.getType().cast<ShapedType>().isDynamicDim(*dim) &&
            "expected dynamic dim");
     if (value.getType().isa<RankedTensorType>()) {
       // A tensor dimension is used: generate a tensor.dim.

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

 using presburger::BoundType;
 
 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.addBound(BoundType::EQ, 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.addBound(BoundType::EQ, value, cstr.getExpr(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.addBound(BoundType::EQ, 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.addBound(BoundType::EQ, 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>() &&
       cstr.addBound(BoundType::EQ, 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.addBound(BoundType::EQ, value,
                   cstr.getExpr(dimOp.getSource(), *constIndex));
void populateBoundsForShapedValueDim(Operation *op, Value value, int64_t dim,
     assert(!type.isDynamicDim(dim) && "expected static dim");
     cstr.addBound(BoundType::EQ, value, dim,
                   cstr.getExpr(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.addBound(BoundType::EQ, value, cstr.getExpr(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.addBound(BoundType::LB, value, cstr.getExpr(forOp.getLowerBound()));
       cstr.addBound(BoundType::UB, value, cstr.getExpr(forOp.getUpperBound()));
       return;
     }

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

if (castOp.getResult().getType().isa<RankedTensorType>() &&
       cstr.addBound(BoundType::EQ, 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.addBound(BoundType::EQ, value,
                   cstr.getExpr(dimOp.getSource(), *constIndex));
AffineExpr expr = cstr.getExpr(padOp.getSource(), dim) +
                       cstr.getExpr(padOp.getMixedLowPad()[dim]) +
                       cstr.getExpr(padOp.getMixedHighPad()[dim]);
     cstr.addBound(BoundType::EQ, 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.addBound(BoundType::EQ, value, cstr.getExpr(tensorType.getRank()));
   }

mlir/lib/Interfaces/ValueBoundsOpInterface.cpp

 
 ValueBoundsConstraintSet::ValueBoundsConstraintSet(ValueDim valueDim)
     : builder(valueDim.first.getContext()) {
   insert(valueDim, /*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
 
 void ValueBoundsConstraintSet::addBound(BoundType type, int64_t pos,
                                         AffineExpr expr) {
   LogicalResult status = cstr.addBound(
       type, pos,
       AffineMap::get(cstr.getNumDimVars(), cstr.getNumSymbolVars(), expr));
   (void)status;
   assert(succeeded(status) && "failed to add bound to constraint system");
 }
 
 void ValueBoundsConstraintSet::addBound(BoundType type, Value value,
                                         AffineExpr expr) {
-  assert(value.getType().isIndex() && "expected index type");
+  assert(value.getType().isIntOrIndex() && "expected integer type");
   assert((value.isa<OpResult>() ||
           value.cast<BlockArgument>().getOwner()->isEntryBlock()) &&
          "unstructured control flow is not supported");
 
   addBound(type, getPos(value), expr);
 }
 
 void ValueBoundsConstraintSet::addBound(BoundType type, Value value,
                                         int64_t dim, AffineExpr expr) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
   assert((value.isa<OpResult>() ||
           value.cast<BlockArgument>().getOwner()->isEntryBlock()) &&
          "unstructured control flow is not supported");
 #endif // NDEBUG
 
   addBound(type, getPos(value, dim), expr);
 }
 
 void ValueBoundsConstraintSet::addBound(BoundType type, Value value,
                                         OpFoldResult ofr) {
-  assert(value.getType().isIndex() && "expected index type");
+  assert(value.getType().isIntOrIndex() && "expected integer type");
   addBound(type, getPos(value), getExpr(ofr));
 }
 
 void ValueBoundsConstraintSet::addBound(BoundType type, Value value,
                                         int64_t dim, OpFoldResult ofr) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
 #endif // NDEBUG
   addBound(type, getPos(value, dim), getExpr(ofr));
 }
 
 AffineExpr ValueBoundsConstraintSet::getExpr(Value value,
                                              std::optional<int64_t> dim) {
 #ifndef NDEBUG
   assertValidValueDim(value, dim);
 #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(valueDim);
-  int64_t pos = getPos(value, dim.value_or(kIndexValue));
+  int64_t pos = getPos(value, dim.value_or(kIntegerValue));
   return pos < cstr.getNumDimVars()
              ? builder.getAffineDimExpr(pos)
              : builder.getAffineSymbolExpr(pos - cstr.getNumDimVars());
 }
 
 AffineExpr ValueBoundsConstraintSet::getExpr(OpFoldResult ofr) {
   if (Value value = ofr.dyn_cast<Value>())
     return getExpr(value, /*dim=*/std::nullopt);
for (int64_t i = pos; i < positionToValueDim.size(); ++i)
     valueDimToPosition[positionToValueDim[i]] = i;
 
   worklist.push(pos);
   return pos;
 }
 
 int64_t ValueBoundsConstraintSet::getPos(Value value, int64_t dim) const {
 #ifndef NDEBUG
-  assertValidValueDim(value, dim == kIndexValue
+  assertValidValueDim(value, dim == kIntegerValue
                                  ? std::nullopt
                                  : std::make_optional<int64_t>(dim));
 #endif // NDEBUG
 
   auto it = valueDimToPosition.find(std::make_pair(value, dim));
   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 = 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 == 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(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 == 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.value_or(kIndexValue));
+  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);
 
   // Note: There are TODOs in the implementation of `getSliceBounds`. In such a
   // case, no lower/upper bound can be computed at the moment.
   // EQ, UB bounds: upper bound is needed.
   if ((type != BoundType::LB) &&
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(value.getType().cast<ShapedType>().isDynamicDim(dim) &&
            "expected dynamic dim");
     mapOperands.push_back(std::make_pair(value, dim));
   }
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(kIndexValue));
+  ValueDim valueDim = std::make_pair(value, dim.value_or(kIntegerValue));
   ValueBoundsConstraintSet cstr(valueDim);
-  int64_t pos = cstr.getPos(value, dim.value_or(kIndexValue));
+  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) {
           return cstr.cstr.getConstantBound64(type, pos).has_value();

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());
 }