[mlir][NFC] Use `getConstantIntValue` instead of casting to `ConstantIndexOp`

[springerm]

Summary

getConstantIntValue extracts constant values from all constant-like ops, not just arith::ConstantIndexOp.

Depends On: D154343

mlir/lib/Conversion/SCFToGPU/SCFToGPU.cpp

struct AffineLoopToGpuConverter {
   SmallVector<Value, 6> lbs;
   // Induction variables of the loops mapped to blocks or threads.
   SmallVector<Value, 6> ivs;
   // Steps of the loops mapped to blocks or threads.
   SmallVector<Value, 6> steps;
 };
 } // namespace
 
-// Return true if the value is obviously a constant "one".
-static bool isConstantOne(Value value) {
-  if (auto def = value.getDefiningOp<arith::ConstantIndexOp>())
-    return def.value() == 1;
-  return false;
-}
-
 // Collect ranges, bounds, steps and induction variables in preparation for
 // mapping a loop nest of depth "numLoops" rooted at "forOp" to a GPU kernel.
 // This may fail if the IR for computing loop bounds cannot be constructed, for
 // example if an affine loop uses semi-affine maps. Return the last loop to be
 // mapped on success, std::nullopt on failure.
 std::optional<AffineForOp>
 AffineLoopToGpuConverter::collectBounds(AffineForOp forOp, unsigned numLoops) {
   OpBuilder builder(forOp.getOperation());
   dims.reserve(numLoops);
   lbs.reserve(numLoops);
   ivs.reserve(numLoops);
   steps.reserve(numLoops);
   AffineForOp currentLoop = forOp;
   for (unsigned i = 0; i < numLoops; ++i) {
     Value lowerBound = getOrEmitLowerBound(currentLoop, builder);
     Value upperBound = getOrEmitUpperBound(currentLoop, builder);
     if (!lowerBound || !upperBound) {
       return std::nullopt;
     }
 
     Value range = builder.create<arith::SubIOp>(currentLoop.getLoc(),
                                                 upperBound, lowerBound);
     Value step = getOrCreateStep(currentLoop, builder);
-    if (!isConstantOne(step))
+    if (getConstantIntValue(step) != static_cast<int64_t>(1))
       range = builder.create<arith::DivSIOp>(currentLoop.getLoc(), range, step);
     dims.push_back(range);
 
     lbs.push_back(lowerBound);
     ivs.push_back(currentLoop.getInductionVar());
     steps.push_back(step);
 
     if (i != numLoops - 1)
void AffineLoopToGpuConverter::createLaunch(AffineForOp rootForOp,
   auto *lbArgumentIt = lbs.begin();
   auto *stepArgumentIt = steps.begin();
   for (const auto &en : llvm::enumerate(ivs)) {
     Value id =
         en.index() < numBlockDims
             ? getDim3Value(launchOp.getBlockIds(), en.index())
             : getDim3Value(launchOp.getThreadIds(), en.index() - numBlockDims);
     Value step = steps[en.index()];
-    if (!isConstantOne(step))
+    if (getConstantIntValue(step) != static_cast<int64_t>(1))
       id = builder.create<arith::MulIOp>(rootForOp.getLoc(), step, id);
 
     Value ivReplacement =
         builder.create<arith::AddIOp>(rootForOp.getLoc(), *lbArgumentIt, id);
     en.value().replaceAllUsesWith(ivReplacement);
     std::advance(lbArgumentIt, 1);
     std::advance(stepArgumentIt, 1);
   }

mlir/lib/Dialect/Affine/Analysis/AffineStructures.cpp

 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/Analysis/AffineStructures.h"
 #include "mlir/Analysis/Presburger/LinearTransform.h"
 #include "mlir/Analysis/Presburger/Simplex.h"
 #include "mlir/Analysis/Presburger/Utils.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Affine/IR/AffineValueMap.h"
-#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/MathExtras.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallPtrSet.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/Support/Debug.h"
if (failed(this->addAffineParallelOpDomain(parallel)))
       LLVM_DEBUG(parallel.emitWarning(
           "failed to add domain info to constraint system"));
     return;
   }
 
   // Add top level symbol.
   appendSymbolVar(val);
   // Check if the symbol is a constant.
-  if (auto constOp = val.getDefiningOp<arith::ConstantIndexOp>())
+  if (std::optional<int64_t> constOp = getConstantIntValue(val))
     addBound(BoundType::EQ, val, constOp.value());
 }
 
 LogicalResult
 FlatAffineValueConstraints::addAffineForOpDomain(AffineForOp forOp) {
   unsigned pos;
   // Pre-condition for this method.
   if (!findVar(forOp.getInductionVar(), &pos)) {

mlir/lib/Dialect/Affine/Analysis/Utils.cpp

ComputationSliceState::getAsConstraints(FlatAffineValueConstraints *cst) const {
 
   // Add loop bound constraints for values which are loop IVs of the destination
   // of fusion and equality constraints for symbols which are constants.
   for (unsigned i = numDims, end = values.size(); i < end; ++i) {
     Value value = values[i];
     assert(cst->containsVar(value) && "value expected to be present");
     if (isValidSymbol(value)) {
       // Check if the symbol is a constant.
-      if (auto cOp = value.getDefiningOp<arith::ConstantIndexOp>())
+      if (std::optional<int64_t> cOp = getConstantIntValue(value))
         cst->addBound(BoundType::EQ, value, cOp.value());
     } else if (auto loop = getForInductionVarOwner(value)) {
       if (failed(cst->addAffineForOpDomain(loop)))
         return failure();
     }
   }
 
   // Add slices bounds on 'ivs' using maps 'lbs'/'ubs' with 'lbOperands[0]'

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

return builder.create<AffineForOp>(loc, lb, ub, step,
                                      /*iterArgs=*/std::nullopt, bodyBuilderFn);
 }
 
 /// Creates an affine loop from the bounds that may or may not be constants.
 static AffineForOp
 buildAffineLoopFromValues(OpBuilder &builder, Location loc, Value lb, Value ub,
                           int64_t step,
                           AffineForOp::BodyBuilderFn bodyBuilderFn) {
-  auto lbConst = lb.getDefiningOp<arith::ConstantIndexOp>();
-  auto ubConst = ub.getDefiningOp<arith::ConstantIndexOp>();
+  std::optional<int64_t> lbConst = getConstantIntValue(lb);
+  std::optional<int64_t> ubConst = getConstantIntValue(ub);
   if (lbConst && ubConst)
     return buildAffineLoopFromConstants(builder, loc, lbConst.value(),
                                         ubConst.value(), step, bodyBuilderFn);
   return builder.create<AffineForOp>(loc, lb, builder.getDimIdentityMap(), ub,
                                      builder.getDimIdentityMap(), step,
                                      /*iterArgs=*/std::nullopt, bodyBuilderFn);
 }
 

mlir/lib/Dialect/GPU/IR/GPUDialect.cpp

 
 /// Folding of memref.dim(gpu.alloc(%size), %idx) -> %size similar to
 /// `memref::AllocOp`.
 struct SimplifyDimOfAllocOp : public OpRewritePattern<memref::DimOp> {
   using OpRewritePattern<memref::DimOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(memref::DimOp dimOp,
                                 PatternRewriter &rewriter) const override {
-    auto index = dimOp.getIndex().getDefiningOp<arith::ConstantIndexOp>();
+    std::optional<int64_t> index = dimOp.getConstantIndex();
     if (!index)
       return failure();
 
     auto memrefType = llvm::dyn_cast<MemRefType>(dimOp.getSource().getType());
     if (!memrefType || !memrefType.isDynamicDim(index.value()))
       return failure();
 
     auto alloc = dimOp.getSource().getDefiningOp<AllocOp>();

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

static Value allocBuffer(ImplicitLocOpBuilder &b,
                          std::optional<unsigned> alignment = std::nullopt) {
   auto width = layout.getTypeSize(elementType);
 
   IntegerAttr alignmentAttr;
   if (alignment.has_value())
     alignmentAttr = b.getI64IntegerAttr(alignment.value());
 
   // Static buffer.
-  if (auto cst = allocSize.getDefiningOp<arith::ConstantIndexOp>()) {
+  if (std::optional<int64_t> cst = getConstantIntValue(allocSize)) {
     auto staticBufferType =
         MemRefType::get(width * cst.value(), b.getIntegerType(8));
     if (options.useAlloca) {
       return b.create<memref::AllocaOp>(staticBufferType, ValueRange{},
                                         alignmentAttr);
     }
     return b.create<memref::AllocOp>(staticBufferType, ValueRange{},
                                      alignmentAttr);

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

 
 using namespace mlir;
 using namespace mlir::affine;
 using namespace mlir::linalg;
 using namespace mlir::scf;
 
 #define DEBUG_TYPE "linalg-tiling"
 
-static bool isZero(OpFoldResult v) {
-  if (!v)
-    return false;
-  if (auto attr = llvm::dyn_cast_if_present<Attribute>(v)) {
-    IntegerAttr intAttr = dyn_cast<IntegerAttr>(attr);
-    return intAttr && intAttr.getValue().isZero();
-  }
-  if (auto cst = v.get<Value>().getDefiningOp<arith::ConstantIndexOp>())
-    return cst.value() == 0;
-  return false;
-}
-
 std::tuple<SmallVector<Range, 4>, LoopIndexToRangeIndexMap>
 mlir::linalg::makeTiledLoopRanges(RewriterBase &b, Location loc, AffineMap map,
                                   ArrayRef<OpFoldResult> allShapeSizes,
                                   ArrayRef<OpFoldResult> allTileSizes) {
   assert(allTileSizes.size() == map.getNumResults());
   // Apply `map` to get shape sizes in loop order.
   SmallVector<OpFoldResult> shapeSizes =
       makeComposedFoldedMultiResultAffineApply(b, loc, map, allShapeSizes);
   SmallVector<OpFoldResult> tileSizes(allTileSizes.begin(), allTileSizes.end());
 
   // Traverse the tile sizes, which are in loop order, erase zeros everywhere.
   LoopIndexToRangeIndexMap loopIndexToRangeIndex;
   for (int idx = 0, e = tileSizes.size(), zerosCount = 0; idx < e; ++idx) {
-    if (isZero(tileSizes[idx - zerosCount])) {
+    if (getConstantIntValue(tileSizes[idx - zerosCount]) ==
+        static_cast<int64_t>(0)) {
       shapeSizes.erase(shapeSizes.begin() + idx - zerosCount);
       tileSizes.erase(tileSizes.begin() + idx - zerosCount);
       ++zerosCount;
       continue;
     }
     loopIndexToRangeIndex[idx] = idx - zerosCount;
   }
 

 tileLinalgOpImpl(RewriterBase &b, LinalgOp op, ArrayRef<OpFoldResult> tileSizes,
                  const LinalgTilingOptions &options) {
   OpBuilder::InsertionGuard g(b);
 
   auto nLoops = op.getNumLoops();
   // Initial tile sizes may be too big, only take the first nLoops.
   tileSizes = tileSizes.take_front(nLoops);
 
-  if (llvm::all_of(tileSizes, isZero)) {
+  if (llvm::all_of(tileSizes, [](OpFoldResult ofr) {
+        return getConstantIntValue(ofr) == static_cast<int64_t>(0);
+      })) {
     TiledLinalgOp tiledOp;
     tiledOp.op = cast<LinalgOp>(b.clone(*op.getOperation()));
     tiledOp.tensorResults.assign(tiledOp.op->result_begin(),
                                  tiledOp.op->result_end());
     return tiledOp;
   }
 
   // 1. Build the tiled loop ranges.

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

void GenerateLoopNest<AffineForOp>::doit(
   assert(iterArgInitValues.empty() && "unexpected AffineForOp init values");
   SmallVector<Value, 4> lbs, ubs, steps;
   unpackRanges(b, loc, loopRanges, lbs, ubs, steps);
 
   // Affine loops require constant steps.
   SmallVector<int64_t, 4> constantSteps;
   constantSteps.reserve(steps.size());
   for (Value v : steps) {
-    auto op = v.getDefiningOp<arith::ConstantIndexOp>();
-    assert(op && "Affine loops require constant steps");
-    constantSteps.push_back(op.value());
+    auto constVal = getConstantIntValue(v);
+    assert(constVal.has_value() && "Affine loops require constant steps");
+    constantSteps.push_back(constVal.value());
   }
 
   affine::buildAffineLoopNest(b, loc, lbs, ubs, constantSteps,
                               [&](OpBuilder &b, Location loc, ValueRange ivs) {
                                 bodyBuilderFn(b, loc, ivs,
                                               linalgOp->getOperands());
                               });
 }

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

LogicalResult ParallelOp::verify() {
   // number of elements in lowerBound, upperBound and step are the same.
   Operation::operand_range stepValues = getStep();
   if (stepValues.empty())
     return emitOpError(
         "needs at least one tuple element for lowerBound, upperBound and step");
 
   // Check whether all constant step values are positive.
   for (Value stepValue : stepValues)
-    if (auto cst = stepValue.getDefiningOp<arith::ConstantIndexOp>())
-      if (cst.value() <= 0)
+    if (auto cst = getConstantIntValue(stepValue))
+      if (*cst <= 0)
         return emitOpError("constant step operand must be positive");
 
   // Check that the body defines the same number of block arguments as the
   // number of tuple elements in step.
   Block *body = getBody();
   if (body->getNumArguments() != stepValues.size())
     return emitOpError() << "expects the same number of induction variables: "
                          << body->getNumArguments()

mlir/lib/Dialect/SCF/Utils/Utils.cpp

LogicalResult mlir::loopUnrollByFactor(
   OpBuilder boundsBuilder(forOp);
   IRRewriter rewriter(forOp.getContext());
   auto loc = forOp.getLoc();
   Value step = forOp.getStep();
   Value upperBoundUnrolled;
   Value stepUnrolled;
   bool generateEpilogueLoop = true;
 
-  auto lbCstOp = forOp.getLowerBound().getDefiningOp<arith::ConstantIndexOp>();
-  auto ubCstOp = forOp.getUpperBound().getDefiningOp<arith::ConstantIndexOp>();
-  auto stepCstOp = forOp.getStep().getDefiningOp<arith::ConstantIndexOp>();
+  std::optional<int64_t> lbCstOp = getConstantIntValue(forOp.getLowerBound());
+  std::optional<int64_t> ubCstOp = getConstantIntValue(forOp.getUpperBound());
+  std::optional<int64_t> stepCstOp = getConstantIntValue(forOp.getStep());
   if (lbCstOp && ubCstOp && stepCstOp) {
     // Constant loop bounds computation.
     int64_t lbCst = lbCstOp.value();
     int64_t ubCst = ubCstOp.value();
     int64_t stepCst = stepCstOp.value();
     assert(lbCst >= 0 && ubCst >= 0 && stepCst >= 0 &&
            "expected positive loop bounds and step");
     int64_t tripCount = mlir::ceilDiv(ubCst - lbCst, stepCst);
if (lbCstOp && ubCstOp && stepCstOp) {
     int64_t stepUnrolledCst = stepCst * unrollFactor;
 
     // Create constant for 'upperBoundUnrolled' and set epilogue loop flag.
     generateEpilogueLoop = upperBoundUnrolledCst < ubCst;
     if (generateEpilogueLoop)
       upperBoundUnrolled = boundsBuilder.create<arith::ConstantIndexOp>(
           loc, upperBoundUnrolledCst);
     else
-      upperBoundUnrolled = ubCstOp;
+      upperBoundUnrolled = forOp.getUpperBound();
 
     // Create constant for 'stepUnrolled'.
     stepUnrolled = stepCst == stepUnrolledCst
                        ? step
                        : boundsBuilder.create<arith::ConstantIndexOp>(
                              loc, stepUnrolledCst);
   } else {
     // Dynamic loop bounds computation.

 /// conversion back to the original loop induction value inside the given Block.
 static LoopParams normalizeLoop(OpBuilder &boundsBuilder,
                                 OpBuilder &insideLoopBuilder, Location loc,
                                 Value lowerBound, Value upperBound, Value step,
                                 Value inductionVar) {
   // Check if the loop is already known to have a constant zero lower bound or
   // a constant one step.
   bool isZeroBased = false;
-  if (auto ubCst = lowerBound.getDefiningOp<arith::ConstantIndexOp>())
+  if (auto ubCst = getConstantIntValue(lowerBound))
     isZeroBased = ubCst.value() == 0;
 
   bool isStepOne = false;
-  if (auto stepCst = step.getDefiningOp<arith::ConstantIndexOp>())
+  if (auto stepCst = getConstantIntValue(step))
     isStepOne = stepCst.value() == 1;
 
   // Compute the number of iterations the loop executes: ceildiv(ub - lb, step)
   // assuming the step is strictly positive.  Update the bounds and the step
   // of the loop to go from 0 to the number of iterations, if necessary.
   if (isZeroBased && isStepOne)
     return {/*lowerBound=*/lowerBound, /*upperBound=*/upperBound,
             /*step=*/step};

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

add_mlir_dialect_library(MLIRSparseTensorDialect
   DEPENDS
   MLIRSparseTensorAttrDefsIncGen
   MLIRSparseTensorOpsIncGen
   MLIRSparseTensorTypesIncGen
 
   LINK_LIBS PUBLIC
   MLIRArithDialect
   MLIRDialect
+  MLIRDialectUtils
   MLIRIR
   MLIRInferTypeOpInterface
   MLIRSupport
   MLIRSparseTensorEnums
   )

mlir/lib/Dialect/SparseTensor/IR/SparseTensorDialect.cpp

 
 #include "Detail/DimLvlMapParser.h"
 
 #include "mlir/Dialect/SparseTensor/IR/SparseTensor.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h"
 #include "mlir/Dialect/SparseTensor/IR/SparseTensorType.h"
 
 #include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "llvm/ADT/TypeSwitch.h"
 #include "llvm/Support/FormatVariadic.h"
 

 
 void PushBackOp::build(OpBuilder &builder, OperationState &result,
                        Value curSize, Value inBuffer, Value value) {
   build(builder, result, curSize, inBuffer, value, Value());
 }
 
 LogicalResult PushBackOp::verify() {
   if (Value n = getN()) {
-    auto nValue = dyn_cast_or_null<arith::ConstantIndexOp>(n.getDefiningOp());
+    std::optional<int64_t> nValue = getConstantIntValue(n);
     if (nValue && nValue.value() < 1)
       return emitOpError("n must be not less than 1");
   }
   return success();
 }
 
 LogicalResult CompressOp::verify() {
   const auto stt = getSparseTensorType(getTensor());
RETURN_FAILURE_IF_FAILED(verifyNumBlockArgs(this, formula, "select",
                                               TypeRange{inputType}, boolType))
   return success();
 }
 
 LogicalResult SortOp::verify() {
   if (getXs().empty())
     return emitError("need at least one xs buffer.");
 
-  auto n = getN().getDefiningOp<arith::ConstantIndexOp>();
+  std::optional<int64_t> n = getConstantIntValue(getN());
 
   Type xtp = getMemRefType(getXs().front()).getElementType();
   auto checkTypes = [&](ValueRange operands,
                         bool checkEleType = true) -> LogicalResult {
     for (Value opnd : operands) {
       auto mtp = getMemRefType(opnd);
       const DynSize sh = mtp.getShape()[0];
       // We can't check the size of dynamic dimension at compile-time, but all
       // xs and ys should have a dimension not less than n at runtime.
       if (n && !ShapedType::isDynamic(sh) && sh < n.value())
         return emitError(llvm::formatv("xs and ys need to have a dimension >= n"
                                        ": {0} < {1}",
                                        sh, n.value()));
 
       if (checkEleType && xtp != mtp.getElementType())
         return emitError("mismatch xs element types");
     }
     return success();
   };
   RETURN_FAILURE_IF_FAILED(checkTypes(getXs()))
   return n ? checkTypes(getYs(), false) : success();
 }
 
 LogicalResult SortCooOp::verify() {
-  auto cn = getN().getDefiningOp<arith::ConstantIndexOp>();
+  std::optional<int64_t> cn = getConstantIntValue(getN());
   // We can't check the size of the buffers when n or buffer dimensions aren't
   // compile-time constants.
   if (!cn)
     return success();
 
   uint64_t n = cn.value();
   uint64_t nx = 1;
   if (auto nxAttr = getNxAttr()) {

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
   build(builder, result, source, getVectorSubscriptAttr(builder, position));
 }
 
 // Convenience builder which assumes the values are constant indices.
 void vector::ExtractOp::build(OpBuilder &builder, OperationState &result,
                               Value source, ValueRange position) {
   SmallVector<int64_t, 4> positionConstants =
       llvm::to_vector<4>(llvm::map_range(position, [](Value pos) {
-        return pos.getDefiningOp<arith::ConstantIndexOp>().value();
+        return getConstantIntValue(pos).value();
       }));
   build(builder, result, source, positionConstants);
 }
 
 LogicalResult
 ExtractOp::inferReturnTypes(MLIRContext *, std::optional<Location>,
                             ValueRange operands, DictionaryAttr attributes,
                             OpaqueProperties properties, RegionRange,
void InsertOp::build(OpBuilder &builder, OperationState &result, Value source,
   result.addAttribute(getPositionAttrStrName(), positionAttr);
 }
 
 // Convenience builder which assumes the values are constant indices.
 void InsertOp::build(OpBuilder &builder, OperationState &result, Value source,
                      Value dest, ValueRange position) {
   SmallVector<int64_t, 4> positionConstants =
       llvm::to_vector<4>(llvm::map_range(position, [](Value pos) {
-        return pos.getDefiningOp<arith::ConstantIndexOp>().value();
+        return getConstantIntValue(pos).value();
       }));
   build(builder, result, source, dest, positionConstants);
 }
 
 LogicalResult InsertOp::verify() {
   auto positionAttr = getPosition().getValue();
   auto destVectorType = getDestVectorType();
   if (positionAttr.size() > static_cast<unsigned>(destVectorType.getRank()))
for (unsigned i = 0; i < numFixedVectorSizes; ++i) {
     if (fixedVectorSizes[i] != outputVectorType.getShape()[index])
       return emitError("fixed vector size must match output vector for dim ")
              << i;
   }
 
   // If all shape operands are produced by constant ops, verify that product
   // of dimensions for input/output shape match.
   auto isDefByConstant = [](Value operand) {
-    return isa_and_nonnull<arith::ConstantIndexOp>(operand.getDefiningOp());
+    return getConstantIntValue(operand).has_value();
   };
   if (llvm::all_of(getInputShape(), isDefByConstant) &&
       llvm::all_of(getOutputShape(), isDefByConstant)) {
     int64_t numInputElements = 1;
     for (auto operand : getInputShape())
-      numInputElements *=
-          cast<arith::ConstantIndexOp>(operand.getDefiningOp()).value();
+      numInputElements *= getConstantIntValue(operand).value();
     int64_t numOutputElements = 1;
     for (auto operand : getOutputShape())
-      numOutputElements *=
-          cast<arith::ConstantIndexOp>(operand.getDefiningOp()).value();
+      numOutputElements *= getConstantIntValue(operand).value();
     if (numInputElements != numOutputElements)
       return emitError("product of input and output shape sizes must match");
   }
   return success();
 }
 
 void ReshapeOp::getFixedVectorSizes(SmallVectorImpl<int64_t> &results) {
   populateFromInt64AttrArray(getFixedVectorSizes(), results);

 
 template <typename TransferOp>
 static bool isInBounds(TransferOp op, int64_t resultIdx, int64_t indicesIdx) {
   // TODO: support more aggressive createOrFold on:
   // `op.indices()[indicesIdx] + vectorType < dim(op.source(), indicesIdx)`
   if (op.getShapedType().isDynamicDim(indicesIdx))
     return false;
   Value index = op.getIndices()[indicesIdx];
-  auto cstOp = index.getDefiningOp<arith::ConstantIndexOp>();
-  if (!cstOp)
+  std::optional<int64_t> cstOp = getConstantIntValue(index);
+  if (!cstOp.has_value())
     return false;
 
   int64_t sourceSize = op.getShapedType().getDimSize(indicesIdx);
   int64_t vectorSize = op.getVectorType().getDimSize(resultIdx);
 
   return cstOp.value() + vectorSize <= sourceSize;
 }
 
static LogicalResult foldReadInitWrite(TransferWriteOp write,
   // Vector types must be the same.
   if (read.getVectorType() != write.getVectorType())
     return failure();
   // Vector and Tensor shapes must match.
   if (read.getVectorType().getShape() != rankedTensorType.getShape())
     return failure();
   // If any index is nonzero.
   auto isNotConstantZero = [](Value v) {
-    auto cstOp = v.getDefiningOp<arith::ConstantIndexOp>();
-    return !cstOp || cstOp.value() != 0;
+    auto cstOp = getConstantIntValue(v);
+    return !cstOp.has_value() || cstOp.value() != 0;
   };
   if (llvm::any_of(read.getIndices(), isNotConstantZero) ||
       llvm::any_of(write.getIndices(), isNotConstantZero))
     return failure();
   // Success.
   results.push_back(read.getSource());
   return success();
 }

 class CreateMaskFolder final : public OpRewritePattern<CreateMaskOp> {
 public:
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(CreateMaskOp createMaskOp,
                                 PatternRewriter &rewriter) const override {
     // Return if any of 'createMaskOp' operands are not defined by a constant.
     auto isNotDefByConstant = [](Value operand) {
-      return !isa_and_nonnull<arith::ConstantIndexOp>(operand.getDefiningOp());
+      return !getConstantIntValue(operand).has_value();
     };
     if (llvm::any_of(createMaskOp.getOperands(), isNotDefByConstant))
       return failure();
 
     // CreateMaskOp for scalable vectors can be folded only if all dimensions
     // are negative or zero.
     if (auto vType = llvm::dyn_cast<VectorType>(createMaskOp.getType())) {
       if (vType.isScalable())
         for (auto opDim : createMaskOp.getOperands()) {
           APInt intVal;
           if (matchPattern(opDim, m_ConstantInt(&intVal)) &&
               intVal.isStrictlyPositive())
             return failure();
         }
     }
 
     // Gather constant mask dimension sizes.
     SmallVector<int64_t, 4> maskDimSizes;
     maskDimSizes.reserve(createMaskOp->getNumOperands());
     for (auto [operand, maxDimSize] : llvm::zip_equal(
              createMaskOp.getOperands(), createMaskOp.getType().getShape())) {
-      Operation *defOp = operand.getDefiningOp();
-      int64_t dimSize = cast<arith::ConstantIndexOp>(defOp).value();
+      int64_t dimSize = getConstantIntValue(operand).value();
       dimSize = std::min(dimSize, maxDimSize);
       // If one of dim sizes is zero, set all dims to zero.
       if (dimSize <= 0) {
         maskDimSizes.assign(createMaskOp.getType().getRank(), 0);
         break;
       }
       maskDimSizes.push_back(dimSize);
     }

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

 /// Returns a copy of `shape` without unit dims.
 static SmallVector<int64_t> getReducedShape(ArrayRef<int64_t> shape) {
   SmallVector<int64_t> reducedShape;
   llvm::copy_if(shape, std::back_inserter(reducedShape),
                 [](int64_t dimSize) { return dimSize != 1; });
   return reducedShape;
 }
 
-/// Returns true if all values are `arith.constant 0 : index`
-static bool isZero(Value v) {
-  auto cst = v.getDefiningOp<arith::ConstantIndexOp>();
-  return cst && cst.value() == 0;
-}
-
 namespace {
 
 /// Rewrites `vector.transfer_read` ops where the source has unit dims, by
 /// inserting a memref.subview dropping those unit dims. The vector shapes are
 /// also reduced accordingly.
 class TransferReadDropUnitDimsPattern
     : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern::OpRewritePattern;
LogicalResult matchAndRewrite(vector::TransferReadOp transferReadOp,
     int reducedRank = getReducedRank(sourceType.getShape());
     if (reducedRank == sourceType.getRank())
       return failure();
     // Check if the reduced vector shape matches the reduced source shape.
     // Otherwise, this case is not supported yet.
     int vectorReducedRank = getReducedRank(vectorType.getShape());
     if (reducedRank != vectorReducedRank)
       return failure();
-    if (llvm::any_of(transferReadOp.getIndices(),
-                     [](Value v) { return !isZero(v); }))
+    if (llvm::any_of(transferReadOp.getIndices(), [](Value v) {
+          return getConstantIntValue(v) != static_cast<int64_t>(0);
+        }))
       return failure();
     Value reducedShapeSource =
         rankReducingSubviewDroppingUnitDims(rewriter, loc, source);
     Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     SmallVector<Value> zeros(reducedRank, c0);
     auto identityMap = rewriter.getMultiDimIdentityMap(reducedRank);
     auto reducedVectorType = VectorType::get(
         getReducedShape(vectorType.getShape()), vectorType.getElementType());
LogicalResult matchAndRewrite(vector::TransferWriteOp transferWriteOp,
     int reducedRank = getReducedRank(sourceType.getShape());
     if (reducedRank == sourceType.getRank())
       return failure();
     // Check if the reduced vector shape matches the reduced destination shape.
     // Otherwise, this case is not supported yet.
     int vectorReducedRank = getReducedRank(vectorType.getShape());
     if (reducedRank != vectorReducedRank)
       return failure();
-    if (llvm::any_of(transferWriteOp.getIndices(),
-                     [](Value v) { return !isZero(v); }))
+    if (llvm::any_of(transferWriteOp.getIndices(), [](Value v) {
+          return getConstantIntValue(v) != static_cast<int64_t>(0);
+        }))
       return failure();
     Value reducedShapeSource =
         rankReducingSubviewDroppingUnitDims(rewriter, loc, source);
     Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     SmallVector<Value> zeros(reducedRank, c0);
     auto identityMap = rewriter.getMultiDimIdentityMap(reducedRank);
     VectorType reducedVectorType = VectorType::get(
         getReducedShape(vectorType.getShape()), vectorType.getElementType());

 /// the truncated indices where `firstDimToCollapse` is now the innermost dim.
 static LogicalResult
 checkAndCollapseInnerZeroIndices(ValueRange indices, int64_t firstDimToCollapse,
                                  SmallVector<Value> &outIndices) {
   int64_t rank = indices.size();
   if (firstDimToCollapse >= rank)
     return failure();
   for (int64_t i = firstDimToCollapse; i < rank; ++i) {
-    arith::ConstantIndexOp cst =
-        indices[i].getDefiningOp<arith::ConstantIndexOp>();
+    std::optional<int64_t> cst = getConstantIntValue(indices[i]);
     if (!cst || cst.value() != 0)
       return failure();
   }
   outIndices = indices;
   outIndices.resize(firstDimToCollapse + 1);
   return success();
 }
 

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

cc_library(
     hdrs = [
         "include/mlir/Dialect/SparseTensor/IR/SparseTensor.h",
         "include/mlir/Dialect/SparseTensor/IR/SparseTensorStorageLayout.h",
         "include/mlir/Dialect/SparseTensor/IR/SparseTensorType.h",
     ],
     includes = ["include"],
     deps = [
         ":ArithDialect",
+        ":DialectUtils",
         ":IR",
         ":InferTypeOpInterface",
         ":SparseTensorAttrDefsIncGen",
         ":SparseTensorEnums",
         ":SparseTensorOpsIncGen",
         ":SparseTensorTypesIncGen",
         "//llvm:Support",
     ],