[mlir] Fold Arithmetic::ConstantOp and Tensor::ExtractSliceOp.

[okkwon]

Summary

Fold ExtractSliceOp when the source is a constant.

[rriddle]

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

1366

Drop trivial braces here.

1368

Drop trivial braces here.

1379

Can you use longer names here? I and E aren't clear to me here.

1411

Drop the llvm:: here. Also do we need the hard coded 6?

1420

Please use m_Constant instead of hardcoding constant operations.

1424–1437

Can you format here?

[okkwon]

Addressed the comments.

[okkwon]

[mlir] Fold Arithmetic::ConstantOp and Tensor::ExtractSliceOp. │

│

Fold ExtractSliceOp when the source is a constant.

[okkwon]

apply clang-format

[mravishankar]

Actually, I was under the impression that you were handling the splat case first before moving on to the non-splat case. In reality, we probably do not want to do such a transformation for the non-splat case. Typically these constants are very large (several MBs). This would end up reading these values element by element in the compiler. That leads to a huge increase in compilation time. Still there is a value of doing that as a preprocessing step within the compiler to (a) reduce runtime costs, and (b) also reduce the binary size of the compiled model.
So one way this has been done in IREE is to extract out such computations as a pre-processing step, compile it, and run it using IREE, and use the result of the run as an input to the rest of the computation.
Long story short, we need to do such transformations only for either splat constants, or "small constants".

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

1396

Probably better to move this out of here and check much earlier. Then you can make the function just take the offsets, strides, and sizes directly as arguments.

1428

I have a broader comment on the constants that need to be folded this way on the patch itself.
In addition to that, I'd consider doing this only if the constant op has a single use. If the whole constant is used anyway, then taking slice of it will eventually just be doing a pointer offset. The value of this transformation is that if the constant has a single use, then you could just discard the unused values.

1433

Nit: Instead of templating on the iterator type, can we template on the Attribute type directly , i.e. DenseIntElementsAttr/DenseFPElementsAttr?

[okkwon]

Addressed Mahesh's comment by folding the single use case.

[okkwon]

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

1396

This will make the code in foldConstant() lengthy so I would keep this as is.

1428

This is done. Thank you for your thoughtful comments!

1433

I tried to use the Attribute only, but there is a difficulty to get the underlying concrete element type solely from the attribute since it is an implementation detail. Let's keep it as is.

[mravishankar]

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

1352–1353

You also need to have a size control here to avoid compilation time explosion.

1352–1353

Related to the comment below, if this method took the offsets, size and strides directly, say

static Attribute sliceDenseElementsAttr(Value constant, ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes, ArrayRef<int64_t> strides) {
  ...
}

then the logic here is agnostic to the fact that these came from a tensor.extract_slice. That separation of concerns seems cleaner to me.

1396

See comment about. Doing this processing in the matchAndRewrite method will separate the actual slice handling of the constant and getting the values of offsets, sizes and strides from the extract_slice op. You could just check for the fold conditions and also retrieve the values directly in the ExtractSliceOp::fold(...) method makes it clear when the folding is done.

[okkwon]

For some unknown reason, arg diff --update only pushed one commit instead of two.
Now the change for the single use case is correctly updated.

[okkwon]

Control the size of constant folding. When the size > 1024 (hard-coded now),
the constant folding won't be applied.

[rriddle]

Can you re-upload with the full diff?

[okkwon]

hoist various checks to fold() to clean up the code

[rriddle]

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

1353

When is the result not a ShapedType?

1357–1358

1358–1359

[okkwon]

Updated to use cast to get ShapedType.

[okkwon]

reuse resultType

[okkwon]

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

1352–1353

Hoisted the offsets, sizes, and strides calculation to fold(). Now the fold function directly calls sliceElements().

1353

I will use cast. The type should always be ShapedType. Thanks!

[okkwon]

Resolve the patch error

[okkwon]

rebase onto main

[okkwon]

Fix a missing change for IterTy and ElemTy

[mravishankar]

Looks good overall. Just a nit, and maybe some more lit tests?

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

1403

Nit : s/Sice/Since

1405

A better TODO is "If the size of the constant folded is to be controlled, move this out of folding and make it a separate pattern which can accept an option to control the size"

mlir/test/Dialect/Tensor/canonicalize.mlir

394 ↗

(On Diff #408654)

Sorry should have mentioned earlier. Might be worth adding more tests, like having a larger constant and taking an interior out of it, etc.

[rriddle]

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

1365

Spell out auto here.

1386

attr is guaranteed to be non-null here. Also, you can use SplatElementsAttr in place of DenseElementsAttr above and drop this call to isSplat.

1391

Why go to the defining op instead of op.source().hasOneUse()?

1410

nit: Spell out auto here.

1427

1429

Spell out auto here.

[mehdi_amini]

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

1391

I'm not thrilled by this heuristic, I don't think it the right thing to handle this kind of thing in such an ad-hoc way here.

In particular because if there are two uses for this constant and both of them can be folded but both of them are guarded the same way, then we don't fold anything...

1405

Right, but a corollary of the TODO is also that we shouldn't have an ad-hoc threshold like this here, because why this operation and no other? How is the threshold suitable in general?

1428

This is unusual for the folder to look

[okkwon]

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

1391

Thanks for the comment! Even though the pattern caught by the current code looks very limited, it is still incremental enhancement. I am going to add more patterns. I am also very interested in seeing what is the most common pattern regarding the constant folding for the real applications.

[mehdi_amini]

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

1391

I commenting on the guard itself, I'm fine with adding the pattern, I object to the way it is guarded though: both by this "hasOneUse" and the arbitrary "kConstantFoldingMaxNumElements".

[okkwon]

add two more unit tests

[okkwon]

I have added two more unit tests.

I will revisit the change to address Mehdi's comments. Thanks River, Mehdi, and Mahesh for the comments!

[mravishankar]

I see that you changed the condition to check for splat constants, but the same handling does not apply for splat constants. There is no need to take slices, etc (actually that might even be wrong, not sure)

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

1391

There is a difference between splat constants and non-splat constants. The reasoning for non-splat constants is, say you have a constant of type tensor<4x2xf32>. If there are multiple uses of this constant (potentially accessing different slices of the constant, or the whole constant), you will need to store the constant values in the binary somewhere. Having a new constant of smaller shape is just adding to binary size without any value. You might as well leave the tensor.extract_slice as is. Eventually it will just become a memref.subview followed by load. So when there are multiple uses it makes sense to leave the constant as is. if there is a single use, it means that the values outside of the slice are dead and you can just remove those values. It saves binary size. I agree with you about the kConstantFoldingMaxNumElements is fairly arbitrary.

The same reasoning does not hold for splat constants though. So now that this change is targeting splat constants, i'd also suggest dropping both these heuristics.

1395

IIUC, constants cannot have a dynamic shape.

1405

Replied above with more context on this.

1428

I am not sure I understand the comment here.

1437

All of this is not necessary for a splat constant. There is only a single value. You just need to change the type of the constant as dictated by the offsets, sizes and strides.

[mehdi_amini]

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

1391

So now that this change is targeting splat constants

Right now the code is targeting only non-splat unless I am mis-reading?

[okkwon]

Right. It does not handle the splat case yet.

[mravishankar]

My bad. Saw the patch too late in the night for me and misread the conditions. Will sync with Okwan to figure out next steps.

[okkwon]

Add control function support

[okkwon]

The code is updated with the control function approach. Please take a look.

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

1391

I used the control function approach to make the decision user-controllable.

1395

I believe this comment is outdated and the code shows a

[mravishankar]

Looks good to me.

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

1208

Nit: Combine this and next checks into one

if (!sourceType.hasStaticShape() && !resultType.hasStaticShape()) return failure();

1220

I am not sure we really need this case?

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

68

This will assert if result type is not static. In this case the constant verification enforces that the result of a constant is statically shaped. So maybe its fine, but FYI.

[okkwon]

Address Mahesh's comments

[okkwon]

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

1208

nit: the logic should be ||. Updated.

1220

Are you referring to if (count == 0)? shape itself is used below. I will move it down.

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

68

Yeah, the control function is just another control over whether the folding should happen or not.
The sanity check for the folding itself is done in the folding function.

[bondhugula]

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

12 ↗

(On Diff #411476)

Please use forward declaration instead of including headers: https://llvm.org/docs/CodingStandards.html#include-as-little-as-possible

32 ↗

(On Diff #411476)

Is it disabled -> Disable ...

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

16–19

Trim headers -- you don't need all of these.

1160

Missing doc comment.

1162–1186

Code comments are completely missing.

1243–1255

Likewise.

1266

Doc comment.

mlir/test/Dialect/Tensor/fold-constant-extractslice.mlir

1 ↗

(On Diff #411476)

Nit: extractslice -> extract-slice

[okkwon]

Thanks Uday for the comments. Please take another look.

mlir/include/mlir/Dialect/Tensor/Transforms/Transforms.h

12 ↗

(On Diff #411476)

Unfortunately, ExtractSliceOp is used in the default value and it needs its concrete type, so it is not avoidable.

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

1243–1255

Please elaborate. I put a simple comment over newAttr.

[okkwon]

Addressed Uday's comments.

[bondhugula]

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

1247

All comments are to be terminated with a full stop.

[okkwon]

Merged to origin/main.

[mehdi_amini]

Try to keep the last line of the commit message with Differential Revision: https://reviews.llvm.org/D119605 ; not only will this close the revision, but also it'll allow folks in the future to find the discussion in the review from a git log.

(also the title of the revision / commit could have been updated since this evolved into not hooking up into the fold of the operations)

[okkwon]

Reopening: The push to origin/main caused a build break due to a cyclic header inclusion between Tensor.h and Transforms.h. The issue didn't occur in my local machine, but it broke the testing system.

[okkwon]

Fix the bazel build error and rewrite the commit message.

mlir/include/mlir/Dialect/Tensor/IR/Tensor.h

 
 /// Create a rank-reducing InsertSliceOp @[0 .. 0] with strides [1 .. 1] and
 /// appropriate sizes (i.e. `dest.getSizes()`). The result is a new tensor with
 /// rank increased to that of `dest`, obtained by inserting `tensor` into `dest`
 /// at the canonical [0 .. 0] position.
 Value createCanonicalRankReducingInsertSliceOp(OpBuilder &b, Location loc,
                                                Value tensor, Value dest);
 
+/// Function to control the folding of constant and extract slice
+using ControlConstantExtractSliceFusionFn = std::function<bool(ExtractSliceOp)>;
+
+/// Patterns to fold the extract slice op with its constant operand
+void populateFoldConstantExtractSlicePatterns(
+    RewritePatternSet &patterns,
+    const ControlConstantExtractSliceFusionFn &controlFn =
+        [](ExtractSliceOp op) {
+          // Disable by default because the folding can generate a large
+          // constant tensor, which would affect the compile time and storage.
+          return false;
+        });
+
 } // namespace tensor
 } // namespace mlir
 
 #endif // MLIR_DIALECT_TENSOR_IR_TENSOR_H_

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/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/Arithmetic/Utils/Utils.h"
-#include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Utils/ReshapeOpsUtils.h"
 #include "mlir/Dialect/Utils/StaticValueUtils.h"
 #include "mlir/IR/BlockAndValueMapping.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/BuiltinAttributeInterfaces.h"
 #include "mlir/IR/Matchers.h"
-#include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ADT/SmallBitVector.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 /// Materialize a single constant operation from a given attribute value with
 /// the desired resultant type.
 Operation *TensorDialect::materializeConstant(OpBuilder &builder,
                                               Attribute value, Type type,
Value newSlice = rewriter.create<ExtractSliceOp>(
         sliceOp.getLoc(), resultType, castOp.source(), sliceOp.offsets(),
         sliceOp.sizes(), sliceOp.strides(), sliceOp.static_offsets(),
         sliceOp.static_sizes(), sliceOp.static_strides());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(sliceOp, sliceOp.getType(),
                                                 newSlice);
     return success();
   }
 };
+
+/// Slice elements from `values` into `outValues`. `counts` represents the
+/// numbers of elements to stride in the original values for each dimension.
+/// The output values can be used to construct a DenseElementsAttr.
+template <typename IterTy, typename ElemTy>
+static void sliceElements(IterTy values, ArrayRef<int64_t> counts,
+                          ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes,
+                          ArrayRef<int64_t> strides,
+                          llvm::SmallVectorImpl<ElemTy> *outValues) {
+  assert(offsets.size() == sizes.size());
+  assert(offsets.size() == strides.size());
+  if (offsets.empty())
+    return;
+
+  int64_t offset = offsets.front();
+  int64_t size = sizes.front();
+  int64_t stride = strides.front();
+  if (offsets.size() == 1) {
+    for (int64_t i = 0; i < size; ++i, offset += stride)
+      outValues->push_back(*(values + offset));
+
+    return;
+  }
+
+  for (int64_t i = 0; i < size; ++i, offset += stride) {
+    auto begin = values + offset * counts.front();
+    sliceElements<IterTy, ElemTy>(begin, counts.drop_front(),
+                                  offsets.drop_front(), sizes.drop_front(),
+                                  strides.drop_front(), outValues);
+  }
+}
+
+/// Fold arith.constant and tensor.extract_slice into arith.constant. The folded
+/// operation might introduce more constant data; Users can control their
+/// heuristics by the control function.
+class ConstantOpExtractSliceFolder final
+    : public OpRewritePattern<ExtractSliceOp> {
+public:
+  using OpRewritePattern<ExtractSliceOp>::OpRewritePattern;
+
+  ConstantOpExtractSliceFolder(MLIRContext *context,
+                               ControlConstantExtractSliceFusionFn controlFn)
+      : OpRewritePattern<ExtractSliceOp>(context),
+        controlFn(std::move(controlFn)) {}
+
+  LogicalResult matchAndRewrite(ExtractSliceOp op,
+                                PatternRewriter &rewriter) const override {
+    DenseElementsAttr attr;
+    if (!matchPattern(op.source(), m_Constant(&attr)))
+      return failure();
+
+    // A constant splat is handled by fold().
+    if (attr.isSplat())
+      return failure();
+
+    // Dynamic result shape is not supported.
+    auto sourceType = op.source().getType().cast<ShapedType>();
+    auto resultType = op.result().getType().cast<ShapedType>();
+    if (!sourceType.hasStaticShape() || !resultType.hasStaticShape())
+      return failure();
+
+    // Customized control over the folding.
+    if (!controlFn(op))
+      return failure();
+
+    int64_t count = sourceType.getNumElements();
+    if (count == 0)
+      return failure();
+
+    // Check if there are any dynamic parts, which are not supported.
+    auto offsets = extractFromI64ArrayAttr(op.static_offsets());
+    if (llvm::is_contained(offsets, ShapedType::kDynamicStrideOrOffset))
+      return failure();
+    auto sizes = extractFromI64ArrayAttr(op.static_sizes());
+    if (llvm::is_contained(sizes, ShapedType::kDynamicSize))
+      return failure();
+    auto strides = extractFromI64ArrayAttr(op.static_strides());
+    if (llvm::is_contained(strides, ShapedType::kDynamicStrideOrOffset))
+      return failure();
+
+    // Compute the stride for each dimension.
+    SmallVector<int64_t> counts;
+    ArrayRef<int64_t> shape = sourceType.getShape();
+    counts.reserve(shape.size());
+    for (int64_t v : shape) {
+      count = count / v;
+      counts.push_back(count);
+    }
+
+    // New attribute constructed by the sliced values.
+    DenseElementsAttr newAttr;
+
+    if (auto elems = attr.dyn_cast<DenseIntElementsAttr>()) {
+      SmallVector<APInt> outValues;
+      outValues.reserve(sourceType.getNumElements());
+      sliceElements<DenseElementsAttr::IntElementIterator, APInt>(
+          elems.begin(), counts, offsets, sizes, strides, &outValues);
+      newAttr = DenseElementsAttr::get(resultType, outValues);
+    } else if (auto elems = attr.dyn_cast<DenseFPElementsAttr>()) {
+      SmallVector<APFloat> outValues;
+      outValues.reserve(sourceType.getNumElements());
+      sliceElements<DenseElementsAttr::FloatElementIterator, APFloat>(
+          elems.begin(), counts, offsets, sizes, strides, &outValues);
+      newAttr = DenseElementsAttr::get(resultType, outValues);
+    }
+
+    if (newAttr) {
+      rewriter.replaceOpWithNewOp<arith::ConstantOp>(op, resultType, newAttr);
+      return success();
+    }
+
+    return failure();
+  }
+
+private:
+  /// This additionally controls whether the fold happens or not. Users can
+  /// impose their heuristics in the function.
+  ControlConstantExtractSliceFusionFn controlFn;
+};
+
 } // namespace
 
+void mlir::tensor::populateFoldConstantExtractSlicePatterns(
+    RewritePatternSet &patterns,
+    const ControlConstantExtractSliceFusionFn &controlFn) {
+  patterns.add<ConstantOpExtractSliceFolder>(patterns.getContext(), controlFn);
+}
+
 /// Return the canonical type of the result of an extract_slice op.
 struct SliceReturnTypeCanonicalizer {
   RankedTensorType operator()(ExtractSliceOp op,
                               ArrayRef<OpFoldResult> mixedOffsets,
                               ArrayRef<OpFoldResult> mixedSizes,
                               ArrayRef<OpFoldResult> mixedStrides) {
     return getCanonicalSliceResultType(op.getType().getRank(),
                                        op.getSourceType(), mixedOffsets,
static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) {
 
   auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; };
   if (insertOp && insertOp.source().getType() == extractOp.getType() &&
       insertOp.isSameAs(extractOp, isSame))
     return insertOp.source();
 
   return {};
 }
 
 OpFoldResult ExtractSliceOp::fold(ArrayRef<Attribute> operands) {
   if (auto splat = operands[0].dyn_cast_or_null<SplatElementsAttr>()) {
     auto resultType = result().getType().cast<ShapedType>();
     if (resultType.hasStaticShape())
       return splat.resizeSplat(resultType);
   }
   if (getSourceType() == getType() &&
       succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
     return this->source();
   if (Value slice = foldExtractAfterInsertSlice(*this))
     return slice;
+
   return OpFoldResult();
 }
 
 Value mlir::tensor::createCanonicalRankReducingExtractSliceOp(
     OpBuilder &b, Location loc, Value tensor, RankedTensorType targetType) {
   auto rankedTensorType = tensor.getType().cast<RankedTensorType>();
   unsigned rank = rankedTensorType.getRank();
   auto shape = rankedTensorType.getShape();
   SmallVector<OpFoldResult> offsets(rank, b.getIndexAttr(0));
   SmallVector<OpFoldResult> sizes;
   for (unsigned i = 0, e = rank; i < e; ++i) {
     OpFoldResult dim;
     if (rankedTensorType.isDynamicDim(i))
       dim = b.createOrFold<tensor::DimOp>(
           loc, tensor, b.create<arith::ConstantIndexOp>(loc, i));
     else
       dim = b.getIndexAttr(shape[i]);
     sizes.push_back(dim);
   }
   SmallVector<OpFoldResult> strides(rank, b.getIndexAttr(1));
   return b.createOrFold<tensor::ExtractSliceOp>(loc, targetType, tensor,
                                                 offsets, sizes, strides);
 }
 
 //===----------------------------------------------------------------------===//
 // InsertSliceOp
 //===----------------------------------------------------------------------===//
 
 // Build a InsertSliceOp with mixed static and dynamic entries.
 void InsertSliceOp::build(OpBuilder &b, OperationState &result, Value source,
                           Value dest, ArrayRef<OpFoldResult> offsets,
                           ArrayRef<OpFoldResult> sizes,
                           ArrayRef<OpFoldResult> strides,
                           ArrayRef<NamedAttribute> attrs) {
   SmallVector<int64_t> staticOffsets, staticSizes, staticStrides;
   SmallVector<Value> dynamicOffsets, dynamicSizes, dynamicStrides;
   dispatchIndexOpFoldResults(offsets, dynamicOffsets, staticOffsets,
                              ShapedType::kDynamicStrideOrOffset);
   dispatchIndexOpFoldResults(sizes, dynamicSizes, staticSizes,
                              ShapedType::kDynamicSize);
   dispatchIndexOpFoldResults(strides, dynamicStrides, staticStrides,
                              ShapedType::kDynamicStrideOrOffset);
   build(b, result, dest.getType(), source, dest, dynamicOffsets, dynamicSizes,
         dynamicStrides, b.getI64ArrayAttr(staticOffsets),
         b.getI64ArrayAttr(staticSizes), b.getI64ArrayAttr(staticStrides));
   result.addAttributes(attrs);
 }
 
 // Build a InsertSliceOp with dynamic entries.
 void InsertSliceOp::build(OpBuilder &b, OperationState &result, Value source,
                           Value dest, ValueRange offsets, ValueRange sizes,
                           ValueRange strides, ArrayRef<NamedAttribute> attrs) {
   SmallVector<OpFoldResult> offsetValues = llvm::to_vector<4>(
       llvm::map_range(offsets, [](Value v) -> OpFoldResult { return v; }));
   SmallVector<OpFoldResult> sizeValues = llvm::to_vector<4>(
       llvm::map_range(sizes, [](Value v) -> OpFoldResult { return v; }));
   SmallVector<OpFoldResult> strideValues = llvm::to_vector<4>(
       llvm::map_range(strides, [](Value v) -> OpFoldResult { return v; }));
   build(b, result, source, dest, offsetValues, sizeValues, strideValues);
 }
 
 static SliceVerificationResult
 verifyInsertSliceOp(ShapedType srcType, ShapedType dstType,
                     ArrayAttr staticOffsets, ArrayAttr staticSizes,
                     ArrayAttr staticStrides,
                     ShapedType *expectedType = nullptr) {
   // insert_slice is the inverse of extract_slice, use the same type inference.
   auto expected = ExtractSliceOp::inferRankReducedResultType(
                       srcType.getRank(), dstType.cast<RankedTensorType>(),
                       extractFromI64ArrayAttr(staticOffsets),
                       extractFromI64ArrayAttr(staticSizes),
                       extractFromI64ArrayAttr(staticStrides))
                       .cast<ShapedType>();
   if (expectedType)
     *expectedType = expected;
   return isRankReducedType(expected, srcType);
 }
 
 /// Verifier for InsertSliceOp.
 LogicalResult InsertSliceOp::verify() {
   ShapedType expectedType;

mlir/test/Dialect/Tensor/fold-constant-extract-slice.mlir

+// RUN: mlir-opt -split-input-file -test-tensor-transform-patterns=test-fold-constant-extract-slice %s | FileCheck %s
+
+// CHECK-LABEL: func @slice_constant
+//   CHECK-NOT:   tensor.extract_slice
+//       CHECK:   %[[CONST:.+]] = arith.constant dense<1.000000e+01> : tensor<1x1xf32>
+//       CHECK:   return %[[CONST]] :  tensor<1x1xf32>
+func @slice_constant(%arg0 : tensor<2x1xf32>) -> tensor<1x1xf32>
+{
+  %cst = arith.constant dense<[[10.0], [11.0]]> : tensor<2x1xf32>
+  %slice = tensor.extract_slice %cst[0, 0] [1, 1] [1, 1] : tensor<2x1xf32> to tensor<1x1xf32>
+  return %slice : tensor<1x1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @slice_constant_3x4
+//   CHECK-NOT:   tensor.extract_slice
+//       CHECK:   %[[CONST:.+]] = arith.constant dense<{{\[}}[1.000000e+01, 9.000000e+00], [1.100000e+01, 1.200000e+01]]> : tensor<2x2xf32>
+//       CHECK:   return %[[CONST]] :  tensor<2x2xf32>
+func @slice_constant_3x4(%arg0 : tensor<3x4xf32>) -> tensor<2x2xf32>
+{
+  %cst = arith.constant dense<[[10.0, 9.0, 8.0, 7.0], [11.0, 12.0, 13.0, 14.0], [1.0, 3.0, 5.0, 7.0]]> : tensor<3x4xf32>
+  %slice = tensor.extract_slice %cst[0, 0] [2, 2] [1, 1] : tensor<3x4xf32> to tensor<2x2xf32>
+  return %slice : tensor<2x2xf32>
+}
+
+// -----
+
+// CHECK-LABEL: func @slice_constant_3x4_offsets
+//   CHECK-NOT:   tensor.extract_slice
+//       CHECK:   %[[CONST:.+]] = arith.constant dense<{{\[}}[1.200000e+01, 1.300000e+01], [3.000000e+00, 5.000000e+00]]> : tensor<2x2xf32>
+//       CHECK:   return %[[CONST]] :  tensor<2x2xf32>
+func @slice_constant_3x4_offsets(%arg0 : tensor<3x4xf32>) -> tensor<2x2xf32>
+{
+  %cst = arith.constant dense<[[10.0, 9.0, 8.0, 7.0], [11.0, 12.0, 13.0, 14.0], [1.0, 3.0, 5.0, 7.0]]> : tensor<3x4xf32>
+  %slice = tensor.extract_slice %cst[1, 1] [2, 2] [1, 1] : tensor<3x4xf32> to tensor<2x2xf32>
+  return %slice : tensor<2x2xf32>
+}
+

mlir/test/lib/Dialect/Tensor/TestTensorTransforms.cpp

 //===- TestTensorTransforms.cpp - Test Tensor transformation patterns -----===//
 //
 // 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 file implements logic for testing Tensor transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Arithmetic/IR/Arithmetic.h"
 #include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tensor/Transforms/Transforms.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/GreedyPatternRewriteDriver.h"
 
 using namespace mlir;
 
 namespace {
 struct TestTensorTransforms
struct TestTensorTransforms
   }
 
   void runOnOperation() override;
 
   Option<bool> testSplitPaddingPatterns{
       *this, "test-split-padding-patterns",
       llvm::cl::desc("Test patterns to split tensor.pad ops"),
       llvm::cl::init(false)};
+
+  Option<bool> testFoldConstantExtractSlice{
+      *this, "test-fold-constant-extract-slice",
+      llvm::cl::desc("Test folding arith.constant and tensor.extract_slice"),
+      llvm::cl::init(false)};
 };
 } // namespace
 
 static void applySplitPaddingPatterns(FuncOp funcOp) {
   RewritePatternSet patterns(funcOp.getContext());
   tensor::populateSplitPaddingPatterns(patterns);
   (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
 }
 
+static void applyFoldConstantExtractSlicePatterns(FuncOp funcOp) {
+  RewritePatternSet patterns(funcOp.getContext());
+  tensor::ControlConstantExtractSliceFusionFn controlFn =
+      [](tensor::ExtractSliceOp op) {
+        if (!op.source().hasOneUse())
+          return false;
+
+        auto resultType = op.result().getType().cast<ShapedType>();
+        constexpr int64_t kConstantFoldingMaxNumElements = 1024;
+        if (resultType.getNumElements() > kConstantFoldingMaxNumElements)
+          return false;
+
+        return true;
+      };
+
+  tensor::populateFoldConstantExtractSlicePatterns(patterns, controlFn);
+  (void)applyPatternsAndFoldGreedily(funcOp, std::move(patterns));
+}
+
 void TestTensorTransforms::runOnOperation() {
   FuncOp func = getOperation();
   if (testSplitPaddingPatterns)
     applySplitPaddingPatterns(func);
+  if (testFoldConstantExtractSlice)
+    applyFoldConstantExtractSlicePatterns(func);
 }
 
 namespace mlir {
 namespace test {
 void registerTestTensorTransforms() {
   PassRegistration<TestTensorTransforms>();
 }
 } // namespace test
 } // namespace mlir