[mlir][NFC] Move offsets/sizes/strides helper to dialect utils and interface header

[springerm]

Summary

• Move foldDynamicIndexList to DialectUtils and simplify function.
• Move OpWithOffsetSizesAndStridesConstantArgumentFolder to ViewLikeInterface and add documentation.

mlir/include/mlir/Dialect/Arith/Utils/Utils.h

 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/Value.h"
 
 namespace mlir {
 
 /// Matches a ConstantIndexOp.
 detail::op_matcher<arith::ConstantIndexOp> matchConstantIndex();
 
-/// Returns `success` when any of the elements in `ofrs` was produced by
-/// arith::ConstantIndexOp. In that case the constant attribute replaces the
-/// Value. Returns `failure` when no folding happened.
-LogicalResult foldDynamicIndexList(Builder &b,
-                                   SmallVectorImpl<OpFoldResult> &ofrs);
-
 llvm::SmallBitVector getPositionsOfShapeOne(unsigned rank,
                                             ArrayRef<int64_t> shape);
 
-/// Pattern to rewrite a subview op with constant arguments.
-template <typename OpType, typename ResultTypeFunc, typename CastOpFunc>
-class OpWithOffsetSizesAndStridesConstantArgumentFolder final
-    : public OpRewritePattern<OpType> {
-public:
-  using OpRewritePattern<OpType>::OpRewritePattern;
-
-  LogicalResult matchAndRewrite(OpType op,
-                                PatternRewriter &rewriter) const override {
-    SmallVector<OpFoldResult> mixedOffsets(op.getMixedOffsets());
-    SmallVector<OpFoldResult> mixedSizes(op.getMixedSizes());
-    SmallVector<OpFoldResult> mixedStrides(op.getMixedStrides());
-
-    // No constant operands were folded, just return;
-    if (failed(foldDynamicIndexList(rewriter, mixedOffsets)) &&
-        failed(foldDynamicIndexList(rewriter, mixedSizes)) &&
-        failed(foldDynamicIndexList(rewriter, mixedStrides)))
-      return failure();
-
-    // Create the new op in canonical form.
-    ResultTypeFunc resultTypeFunc;
-    auto resultType =
-        resultTypeFunc(op, mixedOffsets, mixedSizes, mixedStrides);
-    if (!resultType)
-      return failure();
-    auto newOp =
-        rewriter.create<OpType>(op.getLoc(), resultType, op.getSource(),
-                                mixedOffsets, mixedSizes, mixedStrides);
-    CastOpFunc func;
-    func(rewriter, op, newOp);
-
-    return success();
-  }
-};
-
 /// Converts an OpFoldResult to a Value. Returns the fold result if it casts to
 /// a Value or creates a ConstantIndexOp if it casts to an IntegerAttribute.
 /// Other attribute types are not supported.
 Value getValueOrCreateConstantIndexOp(OpBuilder &b, Location loc,
                                       OpFoldResult ofr);
 
 /// Similar to the other overload, but converts multiple OpFoldResults into
 /// Values.

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

getValuesSortedByKey(ArrayRef<Attribute> keys, ArrayRef<Value> values,
                      llvm::function_ref<bool(Attribute, Attribute)> compare);
 SmallVector<OpFoldResult>
 getValuesSortedByKey(ArrayRef<Attribute> keys, ArrayRef<OpFoldResult> values,
                      llvm::function_ref<bool(Attribute, Attribute)> compare);
 SmallVector<int64_t>
 getValuesSortedByKey(ArrayRef<Attribute> keys, ArrayRef<int64_t> values,
                      llvm::function_ref<bool(Attribute, Attribute)> compare);
 
+/// Returns "success" when any of the elements in `ofrs` is a constant value. In
+/// that case the value is replaced by an attribute. Returns "failure" when no
+/// folding happened.
+LogicalResult foldDynamicIndexList(SmallVectorImpl<OpFoldResult> &ofrs);
+
 /// Return the number of iterations for a loop with a lower bound `lb`, upper
 /// bound `ub` and step `step`.
 std::optional<int64_t> constantTripCount(OpFoldResult lb, OpFoldResult ub,
                                          OpFoldResult step);
 
 } // namespace mlir
 
 #endif // MLIR_DIALECT_UTILS_STATICVALUEUTILS_H

mlir/include/mlir/Interfaces/ViewLikeInterface.h

 #ifndef MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
 #define MLIR_INTERFACES_VIEWLIKEINTERFACE_H_
 
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributes.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/OpImplementation.h"
+#include "mlir/IR/PatternMatch.h"
 
 namespace mlir {
 
 class OffsetSizeAndStrideOpInterface;
 
 namespace detail {
 
 LogicalResult verifyOffsetSizeAndStrideOp(OffsetSizeAndStrideOpInterface op);
 
 bool sameOffsetsSizesAndStrides(
     OffsetSizeAndStrideOpInterface a, OffsetSizeAndStrideOpInterface b,
     llvm::function_ref<bool(OpFoldResult, OpFoldResult)> cmp);
 
 } // namespace detail
 } // namespace mlir
 
 /// Include the generated interface declarations.
 #include "mlir/Interfaces/ViewLikeInterface.h.inc"
 
 namespace mlir {
 
+/// Pattern to rewrite dynamic offsets/sizes/strides of view/slice-like ops as
+/// constant arguments. This pattern assumes that the op has a suitable builder
+/// that takes a result type, a "source" operand and mixed offsets, sizes and
+/// strides.
+///
+/// `OpType` is the type of op to which this pattern is applied. `ResultTypeFn`
+/// returns the new result type of the op, based on the new offsets, sizes and
+/// strides. `CastOpFunc` is used to generate a cast op if the result type of
+/// the op has changed.
+template <typename OpType, typename ResultTypeFn, typename CastOpFunc>
+class OpWithOffsetSizesAndStridesConstantArgumentFolder final
+    : public OpRewritePattern<OpType> {
+public:
+  using OpRewritePattern<OpType>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpType op,
+                                PatternRewriter &rewriter) const override {
+    SmallVector<OpFoldResult> mixedOffsets(op.getMixedOffsets());
+    SmallVector<OpFoldResult> mixedSizes(op.getMixedSizes());
+    SmallVector<OpFoldResult> mixedStrides(op.getMixedStrides());
+
+    // No constant operands were folded, just return;
+    if (failed(foldDynamicIndexList(mixedOffsets)) &&
+        failed(foldDynamicIndexList(mixedSizes)) &&
+        failed(foldDynamicIndexList(mixedStrides)))
+      return failure();
+
+    // Create the new op in canonical form.
+    auto resultType =
+        ResultTypeFn()(op, mixedOffsets, mixedSizes, mixedStrides);
+    if (!resultType)
+      return failure();
+    auto newOp =
+        rewriter.create<OpType>(op.getLoc(), resultType, op.getSource(),
+                                mixedOffsets, mixedSizes, mixedStrides);
+    CastOpFunc()(rewriter, op, newOp);
+
+    return success();
+  }
+};
+
 /// Printer hook for custom directive in assemblyFormat.
 ///
 ///   custom<DynamicIndexList>($values, $integers)
 ///   custom<DynamicIndexList>($values, $integers, type($values))
 ///
 /// where `values` is of ODS type `Variadic<*>` and `integers` is of ODS
 /// type `I64ArrayAttr`. Prints a list with either (1) the static integer value
 /// in `integers` is `kDynamic` or (2) the next value otherwise. If `valueTypes`

mlir/lib/Dialect/Arith/Utils/Utils.cpp

 
 /// Matches a ConstantIndexOp.
 /// TODO: This should probably just be a general matcher that uses matchConstant
 /// and checks the operation for an index type.
 detail::op_matcher<arith::ConstantIndexOp> mlir::matchConstantIndex() {
   return detail::op_matcher<arith::ConstantIndexOp>();
 }
 
-// Returns `success` when any of the elements in `ofrs` was produced by
-// arith::ConstantIndexOp. In that case the constant attribute replaces the
-// Value. Returns `failure` when no folding happened.
-LogicalResult mlir::foldDynamicIndexList(Builder &b,
-                                         SmallVectorImpl<OpFoldResult> &ofrs) {
-  bool valuesChanged = false;
-  for (OpFoldResult &ofr : ofrs) {
-    if (ofr.is<Attribute>())
-      continue;
-    // Newly static, move from Value to constant.
-    if (auto cstOp = llvm::dyn_cast_if_present<Value>(ofr)
-                         .getDefiningOp<arith::ConstantIndexOp>()) {
-      ofr = b.getIndexAttr(cstOp.value());
-      valuesChanged = true;
-    }
-  }
-  return success(valuesChanged);
-}
-
 llvm::SmallBitVector mlir::getPositionsOfShapeOne(unsigned rank,
                                                   ArrayRef<int64_t> shape) {
   llvm::SmallBitVector dimsToProject(shape.size());
   for (unsigned pos = 0, e = shape.size(); pos < e && rank > 0; ++pos) {
     if (shape[pos] == 1) {
       dimsToProject.set(pos);
       --rank;
     }

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

 public:
   using OpRewritePattern<ForallOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(ForallOp op,
                                 PatternRewriter &rewriter) const override {
     SmallVector<OpFoldResult> mixedLowerBound(op.getMixedLowerBound());
     SmallVector<OpFoldResult> mixedUpperBound(op.getMixedUpperBound());
     SmallVector<OpFoldResult> mixedStep(op.getMixedStep());
-    if (failed(foldDynamicIndexList(rewriter, mixedLowerBound)) &&
-        failed(foldDynamicIndexList(rewriter, mixedUpperBound)) &&
-        failed(foldDynamicIndexList(rewriter, mixedStep)))
+    if (failed(foldDynamicIndexList(mixedLowerBound)) &&
+        failed(foldDynamicIndexList(mixedUpperBound)) &&
+        failed(foldDynamicIndexList(mixedStep)))
       return failure();
 
     rewriter.updateRootInPlace(op, [&]() {
       SmallVector<Value> dynamicLowerBound, dynamicUpperBound, dynamicStep;
       SmallVector<int64_t> staticLowerBound, staticUpperBound, staticStep;
       dispatchIndexOpFoldResults(mixedLowerBound, dynamicLowerBound,
                                  staticLowerBound);
       op.getDynamicLowerBoundMutable().assign(dynamicLowerBound);

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

public:
 
   LogicalResult matchAndRewrite(InsertOpTy insertSliceOp,
                                 PatternRewriter &rewriter) const override {
     SmallVector<OpFoldResult> mixedOffsets(insertSliceOp.getMixedOffsets());
     SmallVector<OpFoldResult> mixedSizes(insertSliceOp.getMixedSizes());
     SmallVector<OpFoldResult> mixedStrides(insertSliceOp.getMixedStrides());
 
     // No constant operands were folded, just return;
-    if (failed(foldDynamicIndexList(rewriter, mixedOffsets)) &&
-        failed(foldDynamicIndexList(rewriter, mixedSizes)) &&
-        failed(foldDynamicIndexList(rewriter, mixedStrides)))
+    if (failed(foldDynamicIndexList(mixedOffsets)) &&
+        failed(foldDynamicIndexList(mixedSizes)) &&
+        failed(foldDynamicIndexList(mixedStrides)))
       return failure();
 
     // Create the new op in canonical form.
     auto sourceType = ExtractSliceOp::inferCanonicalRankReducedResultType(
         insertSliceOp.getSourceType().getRank(), insertSliceOp.getDestType(),
         mixedOffsets, mixedSizes, mixedStrides);
     Value toInsert = insertSliceOp.getSource();
     if (sourceType != insertSliceOp.getSourceType()) {

mlir/lib/Dialect/Utils/StaticValueUtils.cpp

if (!ubConstant)
     return std::nullopt;
   std::optional<int64_t> stepConstant = getConstantIntValue(step);
   if (!stepConstant)
     return std::nullopt;
 
   return mlir::ceilDiv(*ubConstant - *lbConstant, *stepConstant);
 }
 
+LogicalResult foldDynamicIndexList(SmallVectorImpl<OpFoldResult> &ofrs) {
+  bool valuesChanged = false;
+  for (OpFoldResult &ofr : ofrs) {
+    if (ofr.is<Attribute>())
+      continue;
+    Attribute attr;
+    if (matchPattern(ofr.get<Value>(), m_Constant(&attr))) {
+      ofr = attr;
+      valuesChanged = true;
+    }
+  }
+  return success(valuesChanged);
+}
+
 } // namespace mlir