[mlir][NFC] Cleanup: Move helper functions to StaticValueUtils

[springerm]

Summary

Reduce code duplication: Move various helper functions, that are duplicated in TensorDialect, MemRefDialect, LinalgDialect, StandardDialect, into a new StaticValueUtils.cpp.

Depends On D104676

[rriddle]

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

11

Why does this depend on Standard ops?

40

nit: Drop the llvm:: here.

49

Drop the llvm:: here and several other places.

52–53

Use m_Constant here, and then you can drop the dependency on StandardOps.

59–62

Is this method useful? It can be represented with just ==.

63–64

76

[springerm]

update CMakeFiles

[springerm]

no change

[springerm]

no change

[springerm]

address comments

[springerm]

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

59–62

Nice, I didn't think of that.

[rriddle]

LGTM, thanks!

mlir/include/mlir/Dialect/Linalg/IR/LinalgOps.td

272 ↗

(On Diff #353827)

Is the mlir:: necessary on these?

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

20

nit: Remove the newline here.

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

12

nit: Drop the newline here.

15

nit: We traditionally keep llvm::None spelled out. Optional and SmallVector should already be exported in the mlir namespace (in mlir/Support/LLVM.h)

41

nit: Spell out auto here.

mlir/include/mlir/Dialect/Linalg/Transforms/Transforms.h

 //===- Transforms.h - Linalg transformations as patterns --------*- 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 DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_
 #define DIALECT_LINALG_TRANSFORMS_TRANSFORMS_H_
 
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/SCF/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Transforms/Bufferize.h"
 #include "llvm/ADT/SmallBitVector.h"
 #include "llvm/ADT/SmallSet.h"
 
 namespace mlir {

mlir/include/mlir/Dialect/StandardOps/IR/Ops.h

 bool applyCmpPredicate(CmpIPredicate predicate, const APInt &lhs,
                        const APInt &rhs);
 
 /// Compute `lhs` `pred` `rhs`, where `pred` is one of the known floating point
 /// comparison predicates.
 bool applyCmpPredicate(CmpFPredicate predicate, const APFloat &lhs,
                        const APFloat &rhs);
 
-/// If ofr is a constant integer, i.e., an IntegerAttr or a ConstantOp with an
-/// IntegerAttr, return the integer.
-llvm::Optional<int64_t> getConstantIntValue(OpFoldResult ofr);
-
-/// Return true if ofr and value are the same integer.
-/// Ignore integer bitwidth and type mismatch that come from the fact there is
-/// no IndexAttr and that IndexType has no bitwidth.
-bool isEqualConstantInt(OpFoldResult ofr, int64_t value);
-
-/// Return true if ofr1 and ofr2 are the same integer constant attribute values
-/// or the same SSA value.
-/// Ignore integer bitwitdh and type mismatch that come from the fact there is
-/// no IndexAttr and that IndexType have no bitwidth.
-bool isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2);
-
 /// Returns the identity value attribute associated with an AtomicRMWKind op.
 Attribute getIdentityValueAttr(AtomicRMWKind kind, Type resultType,
                                OpBuilder &builder, Location loc);
 
 /// Returns the identity value associated with an AtomicRMWKind op.
 Value getIdentityValue(AtomicRMWKind op, Type resultType, OpBuilder &builder,
                        Location loc);
 
 /// Returns the value obtained by applying the reduction operation kind
 /// associated with a binary AtomicRMWKind op to `lhs` and `rhs`.
 Value getReductionOp(AtomicRMWKind op, OpBuilder &builder, Location loc,
                      Value lhs, Value rhs);
 
 } // end namespace mlir
 
 #endif // MLIR_DIALECT_IR_STANDARDOPS_IR_OPS_H

mlir/include/mlir/Dialect/Utils/StaticValueUtils.h

+//===- StaticValueUtils.h - Utilities for static values ---------*- 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
+//
+//===----------------------------------------------------------------------===//
+//
+// This header file defines utilities for dealing with static values, e.g.,
+// converting back and forth between Value and OpFoldResult. Such functionality
+// is used in multiple dialects.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef THIRD_PARTY_LLVM_LLVM_PROJECT_MLIR_INCLUDE_MLIR_DIALECT_UTILS_STATICVALUEUTILS_H_
+#define THIRD_PARTY_LLVM_LLVM_PROJECT_MLIR_INCLUDE_MLIR_DIALECT_UTILS_STATICVALUEUTILS_H_
+
+#include "mlir/IR/OpDefinition.h"
+
+#include "llvm/ADT/SmallVector.h"
+
+namespace mlir {
+
+/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if
+/// it is a Value or into `staticVec` if it is an IntegerAttr.
+/// In the case of a Value, a copy of the `sentinel` value is also pushed to
+/// `staticVec`. This is useful to extract mixed static and dynamic entries that
+/// come from an AttrSizedOperandSegments trait.
+void dispatchIndexOpFoldResult(OpFoldResult ofr,
+                               SmallVectorImpl<Value> &dynamicVec,
+                               SmallVectorImpl<int64_t> &staticVec,
+                               int64_t sentinel);
+
+/// Helper function to dispatch multiple OpFoldResults into either the
+/// `dynamicVec` (for Values) or into `staticVec` (for IntegerAttrs).
+/// In the case of a Value, a copy of the `sentinel` value is also pushed to
+/// `staticVec`. This is useful to extract mixed static and dynamic entries that
+/// come from an AttrSizedOperandSegments trait.
+void dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
+                                SmallVectorImpl<Value> &dynamicVec,
+                                SmallVectorImpl<int64_t> &staticVec,
+                                int64_t sentinel);
+
+/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
+SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr);
+
+/// If ofr is a constant integer, i.e., an IntegerAttr or a ConstantOp with an
+/// IntegerAttr, return the integer.
+llvm::Optional<int64_t> getConstantIntValue(OpFoldResult ofr);
+
+/// Return true if ofr and value are the same integer.
+/// Ignore integer bitwidth and type mismatch that come from the fact there is
+/// no IndexAttr and that IndexType has no bitwidth.
+bool isEqualConstantInt(OpFoldResult ofr, int64_t value);
+
+/// Return true if ofr1 and ofr2 are the same integer constant attribute values
+/// or the same SSA value.
+/// Ignore integer bitwitdh and type mismatch that come from the fact there is
+/// no IndexAttr and that IndexType have no bitwidth.
+bool isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2);
+
+} // namespace mlir
+
+#endif // THIRD_PARTY_LLVM_LLVM_PROJECT_MLIR_INCLUDE_MLIR_DIALECT_UTILS_STATICVALUEUTILS_H_

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 #include "mlir/Analysis/DataLayoutAnalysis.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Dialect/LLVMIR/FunctionCallUtils.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Math/IR/Math.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/DataLayoutInterfaces.h"
nDVectorIterate(vectorTypeInfo, rewriter, [&](ArrayAttr position) {
       desc = rewriter.create<LLVM::InsertValueOp>(loc, llvmNDVectorTy, desc, v,
                                                   position);
     });
     rewriter.replaceOp(splatOp, desc);
     return success();
   }
 };
 
-/// Helper function extracts int64_t from the assumedArrayAttr of IntegerAttr.
-static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
-  return llvm::to_vector<4>(
-      llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
-        return a.cast<IntegerAttr>().getInt();
-      }));
-}
-
 /// Conversion pattern that transforms a subview op into:
 ///   1. An `llvm.mlir.undef` operation to create a memref descriptor
 ///   2. Updates to the descriptor to introduce the data ptr, offset, size
 ///      and stride.
 /// The subview op is replaced by the descriptor.
 struct SubViewOpLowering : public ConvertOpToLLVMPattern<memref::SubViewOp> {
   using ConvertOpToLLVMPattern<memref::SubViewOp>::ConvertOpToLLVMPattern;
 

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

 
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineExprVisitor.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Parser.h"
 
 #include "llvm/ADT/DenseMap.h"
return llvm::to_vector<4>(
       llvm::map_range(valueOrAttrVec, [&](OpFoldResult value) -> Value {
         if (auto attr = value.dyn_cast<Attribute>())
           return b.create<ConstantIndexOp>(loc,
                                            attr.cast<IntegerAttr>().getInt());
         return value.get<Value>();
       }));
 }
 
-/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if
-/// it is a Value or into `staticVec` if it is an IntegerAttr.
-/// In the case of a Value, a copy of the `sentinel` value is also pushed to
-/// `staticVec`. This is useful to extract mixed static and dynamic entries that
-/// come from an AttrSizedOperandSegments trait.
-static void dispatchIndexOpFoldResult(OpFoldResult ofr,
-                                      SmallVectorImpl<Value> &dynamicVec,
-                                      SmallVectorImpl<int64_t> &staticVec,
-                                      int64_t sentinel) {
-  if (auto v = ofr.dyn_cast<Value>()) {
-    dynamicVec.push_back(v);
-    staticVec.push_back(sentinel);
-    return;
-  }
-  APInt apInt = ofr.dyn_cast<Attribute>().cast<IntegerAttr>().getValue();
-  staticVec.push_back(apInt.getSExtValue());
-}
-
 /// This is a common class used for patterns of the form
 /// ```
 ///    someop(memrefcast(%src)) -> someop(%src)
 /// ```
 /// It folds the source of the memref.cast into the root operation directly.
 static LogicalResult foldMemRefCast(Operation *op) {
   bool folded = false;
   for (OpOperand &operand : op->getOpOperands()) {
LogicalResult InitTensorOp::reifyReturnTypeShapesPerResultDim(
   reifiedReturnShapes.emplace_back(std::move(shapes));
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // PadTensorOp
 //===----------------------------------------------------------------------===//
 
-/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
-static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
-  return llvm::to_vector<4>(
-      llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
-        return a.cast<IntegerAttr>().getInt();
-      }));
-}
-
 static LogicalResult verify(PadTensorOp op) {
   auto sourceType = op.source().getType().cast<RankedTensorType>();
   auto resultType = op.result().getType().cast<RankedTensorType>();
   auto expectedType = PadTensorOp::inferResultType(
       sourceType, extractFromI64ArrayAttr(op.static_low()),
       extractFromI64ArrayAttr(op.static_high()));
   for (int i = 0, e = sourceType.getRank(); i < e; ++i) {
     if (resultType.getDimSize(i) == expectedType.getDimSize(i))

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

 //  canonical strided layout of the proper rank. This ensures compatibility with
 //  expected layouts after transformations. Combinations of memref.cast +
 //  canonicalization are responsible for clean ups.
 
 #include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Passes.h"
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/BufferUtils.h"
 
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/EquivalenceClasses.h"
 #include "llvm/ADT/ScopeExit.h"

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

 add_mlir_dialect_library(MLIRMemRef
   MemRefDialect.cpp
   MemRefOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/inlude/mlir/Dialect/MemRefDialect
 
   DEPENDS
   MLIRStandardOpsIncGen
   MLIRMemRefOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRDialect
+  MLIRDialectUtils
   MLIRInferTypeOpInterface
   MLIRIR
   MLIRMemRefUtils
   MLIRStandard
   MLIRTensor
   MLIRViewLikeInterface
 )

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

 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/MemRef/IR/MemRef.h"
 #include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Interfaces/InferTypeOpInterface.h"
 #include "mlir/Interfaces/ViewLikeInterface.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::memref;
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *MemRefDialect::materializeConstant(OpBuilder &builder,
                                               Attribute value, Type type,
                                               Location loc) {
   return builder.create<mlir::ConstantOp>(loc, type, value);
 }
 
-/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
-static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
-  return llvm::to_vector<4>(
-      llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
-        return a.cast<IntegerAttr>().getInt();
-      }));
-}
-
-/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if
-/// it is a Value or into `staticVec` if it is an IntegerAttr.
-/// In the case of a Value, a copy of the `sentinel` value is also pushed to
-/// `staticVec`. This is useful to extract mixed static and dynamic entries that
-/// come from an AttrSizedOperandSegments trait.
-static void dispatchIndexOpFoldResult(OpFoldResult ofr,
-                                      SmallVectorImpl<Value> &dynamicVec,
-                                      SmallVectorImpl<int64_t> &staticVec,
-                                      int64_t sentinel) {
-  if (auto v = ofr.dyn_cast<Value>()) {
-    dynamicVec.push_back(v);
-    staticVec.push_back(sentinel);
-    return;
-  }
-  APInt apInt = ofr.dyn_cast<Attribute>().cast<IntegerAttr>().getValue();
-  staticVec.push_back(apInt.getSExtValue());
-}
-
-static void dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
-                                       SmallVectorImpl<Value> &dynamicVec,
-                                       SmallVectorImpl<int64_t> &staticVec,
-                                       int64_t sentinel) {
-  for (auto ofr : ofrs)
-    dispatchIndexOpFoldResult(ofr, dynamicVec, staticVec, sentinel);
-}
-
 //===----------------------------------------------------------------------===//
 // Common canonicalization pattern support logic
 //===----------------------------------------------------------------------===//
 
 /// This is a common class used for patterns of the form
 /// "someop(memrefcast) -> someop".  It folds the source of any memref.cast
 /// into the root operation directly.
 static LogicalResult foldMemRefCast(Operation *op, Value inner = nullptr) {

mlir/lib/Dialect/StandardOps/IR/Ops.cpp

 #include "llvm/Support/FormatVariadic.h"
 #include "llvm/Support/raw_ostream.h"
 
 // Pull in all enum type definitions and utility function declarations.
 #include "mlir/Dialect/StandardOps/IR/OpsEnums.cpp.inc"
 
 using namespace mlir;
 
-/// If ofr is a constant integer, i.e., an IntegerAttr or a ConstantOp with an
-/// IntegerAttr, return the integer.
-llvm::Optional<int64_t> mlir::getConstantIntValue(OpFoldResult ofr) {
-  Attribute attr = ofr.dyn_cast<Attribute>();
-  // Note: isa+cast-like pattern allows writing the condition below as 1 line.
-  if (!attr && ofr.get<Value>().getDefiningOp<ConstantOp>())
-    attr = ofr.get<Value>().getDefiningOp<ConstantOp>().getValue();
-  if (auto intAttr = attr.dyn_cast_or_null<IntegerAttr>())
-    return intAttr.getValue().getSExtValue();
-  return llvm::None;
-}
-
-/// Return true if ofr and value are the same integer.
-/// Ignore integer bitwidth and type mismatch that come from the fact there is
-/// no IndexAttr and that IndexType has no bitwidth.
-bool mlir::isEqualConstantInt(OpFoldResult ofr, int64_t value) {
-  auto ofrValue = getConstantIntValue(ofr);
-  return ofrValue && *ofrValue == value;
-}
-
-/// Return true if ofr1 and ofr2 are the same integer constant attribute values
-/// or the same SSA value.
-/// Ignore integer bitwidth and type mismatch that come from the fact there is
-/// no IndexAttr and that IndexType has no bitwidth.
-bool mlir::isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2) {
-  auto cst1 = getConstantIntValue(ofr1), cst2 = getConstantIntValue(ofr2);
-  if (cst1 && cst2 && *cst1 == *cst2)
-    return true;
-  auto v1 = ofr1.dyn_cast<Value>(), v2 = ofr2.dyn_cast<Value>();
-  return v1 && v2 && v1 == v2;
-}
-
 //===----------------------------------------------------------------------===//
 // StandardOpsDialect Interfaces
 //===----------------------------------------------------------------------===//
 namespace {
 /// This class defines the interface for handling inlining with standard
 /// operations.
 struct StdInlinerInterface : public DialectInlinerInterface {
   using DialectInlinerInterface::DialectInlinerInterface;

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

 add_mlir_dialect_library(MLIRTensor
   TensorDialect.cpp
   TensorOps.cpp
 
   ADDITIONAL_HEADER_DIRS
   ${PROJECT_SOURCE_DIR}/include/mlir/Dialect/Tensor
 
   DEPENDS
   MLIRTensorOpsIncGen
 
   LINK_COMPONENTS
   Core
 
   LINK_LIBS PUBLIC
   MLIRCastInterfaces
+  MLIRDialectUtils
   MLIRIR
   MLIRSideEffectInterfaces
   MLIRSupport
   MLIRStandard
   MLIRViewLikeInterface
   )

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

 //===----------------------------------------------------------------------===//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/StandardOps/Utils/Utils.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
static LogicalResult verify(ReshapeOp op) {
   }
   return success();
 }
 
 //===----------------------------------------------------------------------===//
 // ExtractSliceOp
 //===----------------------------------------------------------------------===//
 
-/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if
-/// it is a Value or into `staticVec` if it is an IntegerAttr.
-/// In the case of a Value, a copy of the `sentinel` value is also pushed to
-/// `staticVec`. This is useful to extract mixed static and dynamic entries that
-/// come from an AttrSizedOperandSegments trait.
-static void dispatchIndexOpFoldResult(OpFoldResult ofr,
-                                      SmallVectorImpl<Value> &dynamicVec,
-                                      SmallVectorImpl<int64_t> &staticVec,
-                                      int64_t sentinel) {
-  if (auto v = ofr.dyn_cast<Value>()) {
-    dynamicVec.push_back(v);
-    staticVec.push_back(sentinel);
-    return;
-  }
-  APInt apInt = ofr.dyn_cast<Attribute>().cast<IntegerAttr>().getValue();
-  staticVec.push_back(apInt.getSExtValue());
-}
-
-static void dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
-                                       SmallVectorImpl<Value> &dynamicVec,
-                                       SmallVectorImpl<int64_t> &staticVec,
-                                       int64_t sentinel) {
-  for (auto ofr : ofrs)
-    dispatchIndexOpFoldResult(ofr, dynamicVec, staticVec, sentinel);
-}
-
 /// An extract_slice op result type can be fully inferred from the source type
 /// and the static representation of offsets, sizes and strides. Special
 /// sentinels encode the dynamic case.
 Type ExtractSliceOp::inferResultType(RankedTensorType sourceRankedTensorType,
                                      ArrayRef<int64_t> leadingStaticOffsets,
                                      ArrayRef<int64_t> leadingStaticSizes,
                                      ArrayRef<int64_t> leadingStaticStrides) {
   // An extract_slice op may specify only a leading subset of offset/sizes/
   // strides in which case we complete with offset=0, sizes from memref type and
   // strides=1.
   unsigned rank = sourceRankedTensorType.getRank();
   assert(leadingStaticSizes.size() <= rank &&
          "unexpected leadingStaticSizes overflow");
   auto staticSizes = llvm::to_vector<4>(leadingStaticSizes);
   unsigned numTrailingSizes = rank - staticSizes.size();
   llvm::append_range(staticSizes, sourceRankedTensorType.getShape().take_back(
                                       numTrailingSizes));
   return RankedTensorType::get(staticSizes,
                                sourceRankedTensorType.getElementType());
 }
 
-/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
-static SmallVector<int64_t, 4> extractFromI64ArrayAttr(Attribute attr) {
-  return llvm::to_vector<4>(
-      llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
-        return a.cast<IntegerAttr>().getInt();
-      }));
-}
-
 Type ExtractSliceOp::inferResultType(
     RankedTensorType sourceRankedTensorType,
     ArrayRef<OpFoldResult> leadingStaticOffsets,
     ArrayRef<OpFoldResult> leadingStaticSizes,
     ArrayRef<OpFoldResult> leadingStaticStrides) {
   SmallVector<int64_t> staticOffsets, staticSizes, staticStrides;
   SmallVector<Value> dynamicOffsets, dynamicSizes, dynamicStrides;
   dispatchIndexOpFoldResults(leadingStaticOffsets, dynamicOffsets,

mlir/lib/Dialect/Utils/CMakeLists.txt

 add_mlir_library(MLIRDialectUtils
   StructuredOpsUtils.cpp
+  StaticValueUtils.cpp
 
   LINK_LIBS PUBLIC
   MLIRIR
 )

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

+//===- StaticValueUtils.cpp - Utilities for dealing with static values ----===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/Dialect/Utils/StaticValueUtils.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
+
+/// Helper function to dispatch an OpFoldResult into either the `dynamicVec` if
+/// it is a Value or into `staticVec` if it is an IntegerAttr.
+/// In the case of a Value, a copy of the `sentinel` value is also pushed to
+/// `staticVec`. This is useful to extract mixed static and dynamic entries that
+/// come from an AttrSizedOperandSegments trait.
+void mlir::dispatchIndexOpFoldResult(OpFoldResult ofr,
+                                     SmallVectorImpl<Value> &dynamicVec,
+                                     SmallVectorImpl<int64_t> &staticVec,
+                                     int64_t sentinel) {
+  if (auto v = ofr.dyn_cast<Value>()) {
+    dynamicVec.push_back(v);
+    staticVec.push_back(sentinel);
+    return;
+  }
+  APInt apInt = ofr.dyn_cast<Attribute>().cast<IntegerAttr>().getValue();
+  staticVec.push_back(apInt.getSExtValue());
+}
+
+void mlir::dispatchIndexOpFoldResults(ArrayRef<OpFoldResult> ofrs,
+                                      SmallVectorImpl<Value> &dynamicVec,
+                                      SmallVectorImpl<int64_t> &staticVec,
+                                      int64_t sentinel) {
+  for (auto ofr : ofrs)
+    dispatchIndexOpFoldResult(ofr, dynamicVec, staticVec, sentinel);
+}
+
+/// Extract int64_t values from the assumed ArrayAttr of IntegerAttr.
+llvm::SmallVector<int64_t, 4> mlir::extractFromI64ArrayAttr(Attribute attr) {
+  return llvm::to_vector<4>(
+      llvm::map_range(attr.cast<ArrayAttr>(), [](Attribute a) -> int64_t {
+        return a.cast<IntegerAttr>().getInt();
+      }));
+}
+
+/// If ofr is a constant integer, i.e., an IntegerAttr or a ConstantOp with an
+/// IntegerAttr, return the integer.
+llvm::Optional<int64_t> mlir::getConstantIntValue(OpFoldResult ofr) {
+  Attribute attr = ofr.dyn_cast<Attribute>();
+  // Note: isa+cast-like pattern allows writing the condition below as 1 line.
+  if (!attr && ofr.get<Value>().getDefiningOp<ConstantOp>())
+    attr = ofr.get<Value>().getDefiningOp<ConstantOp>().getValue();
+  if (auto intAttr = attr.dyn_cast_or_null<IntegerAttr>())
+    return intAttr.getValue().getSExtValue();
+  return llvm::None;
+}
+
+/// Return true if ofr and value are the same integer.
+/// Ignore integer bitwidth and type mismatch that come from the fact there is
+/// no IndexAttr and that IndexType has no bitwidth.
+bool mlir::isEqualConstantInt(OpFoldResult ofr, int64_t value) {
+  auto ofrValue = getConstantIntValue(ofr);
+  return ofrValue && *ofrValue == value;
+}
+
+/// Return true if ofr1 and ofr2 are the same integer constant attribute values
+/// or the same SSA value.
+/// Ignore integer bitwidth and type mismatch that come from the fact there is
+/// no IndexAttr and that IndexType has no bitwidth.
+bool mlir::isEqualConstantIntOrValue(OpFoldResult ofr1, OpFoldResult ofr2) {
+  auto cst1 = getConstantIntValue(ofr1), cst2 = getConstantIntValue(ofr2);
+  if (cst1 && cst2 && *cst1 == *cst2)
+    return true;
+  auto v1 = ofr1.dyn_cast<Value>(), v2 = ofr2.dyn_cast<Value>();
+  return v1 && v2 && v1 == v2;
+}