[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

1261

Drop trivial braces here.

1272

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

1304

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

1313

Please use m_Constant instead of hardcoding constant operations.

1317–1330

Can you format here?

1349

Drop trivial braces 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

1289

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.

1321

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.

1326

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

1289

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

1321

This is done. Thank you for your thoughtful comments!

1326

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

1245–1341

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

1245–1341

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.

1289

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

1246

When is the result not a ShapedType?

1250–1251

1250–1251

[okkwon]

Updated to use cast to get ShapedType.

[okkwon]

reuse resultType

[okkwon]

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

1245–1341

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

1246

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

1296

Nit : s/Sice/Since

1298

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

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

1279

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

1284

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

1303

nit: Spell out auto here.

1320

1322

Spell out auto here.

1348

Spell out auto here.

[mehdi_amini]

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

1284

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...

1298

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?

1321

This is unusual for the folder to look

[okkwon]

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

1284

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

1284

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

1284

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.

1288

IIUC, constants cannot have a dynamic shape.

1298

Replied above with more context on this.

1321

I am not sure I understand the comment here.

1330

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

1284

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

1284

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

1288

I believe this comment is outdated and the code shows a

[mravishankar]

Looks good to me.

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

1227

Nit: Combine this and next checks into one

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

1239

I am not sure we really need this case?

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

67 ↗

(On Diff #411432)

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

1227

nit: the logic should be ||. Updated.

1239

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

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

67 ↗

(On Diff #411432)

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

18–21

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

1179

Missing doc comment.

1181–1205

Code comments are completely missing.

1262–1274

Likewise.

1285

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

1262–1274

Please elaborate. I put a simple comment over newAttr.

[okkwon]

Addressed Uday's comments.

[bondhugula]

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

1266

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/lib/Dialect/Tensor/IR/TensorOps.cpp

 #include "mlir/Dialect/Complex/IR/Complex.h"
 #include "mlir/Dialect/StandardOps/Utils/Utils.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/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Matchers.h"
+#include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "llvm/ADT/STLExtras.h"
 
 using namespace mlir;
 using namespace mlir::tensor;
 
 /// Materialize a single constant operation from a given attribute value with
Value newSlice = rewriter.create<ExtractSliceOp>(
         sliceOp.static_sizes(), sliceOp.static_strides());
     rewriter.replaceOpWithNewOp<tensor::CastOp>(sliceOp, sliceOp.getType(),
                                                 newSlice);
     return success();
   }
 };
 } // namespace
 
 /// 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,
                                        mixedSizes, mixedStrides);
   }
 };
 
 /// A canonicalizer wrapper to replace ExtractSliceOps.
 struct SliceCanonicalizer {
   void operator()(PatternRewriter &rewriter, ExtractSliceOp op,
                   ExtractSliceOp newOp) {
     Value replacement = newOp.getResult();
     if (replacement.getType() != op.getType())
       replacement = rewriter.create<tensor::CastOp>(op.getLoc(), op.getType(),
                                                     replacement);
     rewriter.replaceOp(op, replacement);
   }
 };
 
 void ExtractSliceOp::getCanonicalizationPatterns(RewritePatternSet &results,
                                                  MLIRContext *context) {
   results.add<
       OpWithOffsetSizesAndStridesConstantArgumentFolder<
           ExtractSliceOp, SliceReturnTypeCanonicalizer, SliceCanonicalizer>,
       ExtractSliceOpCastFolder>(context);
 }
 
 //
 static LogicalResult
 foldIdentityOffsetSizeAndStrideOpInterface(OffsetSizeAndStrideOpInterface op,
                                            ShapedType shapedType) {
   OpBuilder b(op.getContext());
   for (OpFoldResult ofr : op.getMixedOffsets())
     if (getConstantIntValue(ofr) != static_cast<int64_t>(0))
       return failure();
   // Rank-reducing noops only need to inspect the leading dimensions: llvm::zip
   // is appropriate.
   auto shape = shapedType.getShape();
   for (auto it : llvm::zip(op.getMixedSizes(), shape))
     if (getConstantIntValue(std::get<0>(it)) != std::get<1>(it))
       return failure();
   for (OpFoldResult ofr : op.getMixedStrides())
     if (getConstantIntValue(ofr) != static_cast<int64_t>(1))
       return failure();
   return success();
 }
 
 /// If we have an ExtractSliceOp consuming an InsertSliceOp with the same slice,
 /// we can return the InsertSliceOp's source directly.
 // TODO: This only checks the immediate producer; extend to go up the
 // insert/extract chain if the slices are disjoint.
 static Value foldExtractAfterInsertSlice(ExtractSliceOp extractOp) {
   auto insertOp = extractOp.source().getDefiningOp<InsertSliceOp>();
 
   auto isSame = [](OpFoldResult a, OpFoldResult b) { return a == b; };
   if (insertOp && insertOp.source().getType() == extractOp.getType() &&
       insertOp.isSameAs(extractOp, isSame))
     return insertOp.source();
 
   return {};
 }
 
+template <typename I, typename E>
+static void sliceElements(I values, ArrayRef<int64_t> counts,
+                          ArrayRef<int64_t> offsets, ArrayRef<int64_t> sizes,
+                          ArrayRef<int64_t> strides,
+                          llvm::SmallVectorImpl<E> *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 (int i = offset; i < size; i += stride)
+      outValues->push_back(*(values + i));
+
+    return;
+  }
+
+  for (; offset < size; offset += stride) {
+    auto begin = values + offset * counts.front();
+    sliceElements<I, E>(begin, counts.drop_front(), offsets.drop_front(),
+                        sizes.drop_front(), strides.drop_front(), outValues);
+  }
+}
+
+template <typename IterTy, typename ElemTy>
+static Attribute foldExtractSliceAfterConstant(ExtractSliceOp op,
+                                               IterTy values) {
+  auto sourceType = op.source().getType().cast<ShapedType>();
+  if (!sourceType.hasStaticShape())
+    return {};
+
+  auto shape = sourceType.getShape();
+  int64_t count = sourceType.getNumElements();
+  if (count == 0) {
+    return DenseElementsAttr::get<ElemTy>(
+        op.result().getType().cast<ShapedType>(),
+        /*list=*/{});
+  }
+
+  // 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 {};
+  auto sizes = extractFromI64ArrayAttr(op.static_sizes());
+  if (llvm::is_contained(sizes, ShapedType::kDynamicSize))
+    return {};
+  auto strides = extractFromI64ArrayAttr(op.static_strides());
+  if (llvm::is_contained(strides, ShapedType::kDynamicStrideOrOffset))
+    return {};
+
+  // Compute the stride for each dimension.
+  SmallVector<int64_t, 6> counts;
+  counts.reserve(shape.size());
+  for (auto v : shape) {
+    count = count / v;
+    counts.push_back(count);
+  }
+
+  SmallVector<ElemTy> outValues;
+  outValues.reserve(sourceType.getNumElements());
+  sliceElements<IterTy, ElemTy>(values, counts, offsets, sizes, strides,
+                                &outValues);
+
+  return DenseElementsAttr::get(op.result().getType().cast<ShapedType>(),
+                                outValues);
+}
+
+static Attribute foldConstant(ExtractSliceOp op) {
+  DenseElementsAttr attr;
+  if (!matchPattern(op.source(), m_Constant(&attr)))
+    return {};
+  // TODO: support the splat case
+  if (!attr || attr.isSplat())
+    return {};
+
+  if (auto intElems = attr.dyn_cast<DenseIntElementsAttr>()) {
+    if (auto folded =
+            foldExtractSliceAfterConstant<DenseElementsAttr::IntElementIterator,
+                                          APInt>(op, intElems.begin())) {
+      return folded;
+    }
+  } else if (auto floatElems = attr.dyn_cast<DenseFPElementsAttr>()) {
+    if (auto folded = foldExtractSliceAfterConstant<
+            DenseElementsAttr::FloatElementIterator, APFloat>(
+            op, floatElems.begin())) {
+      return folded;
+    }
+  }
+
+  return {};
+}
+
 OpFoldResult ExtractSliceOp::fold(ArrayRef<Attribute>) {
   if (getSourceType() == getType() &&
       succeeded(foldIdentityOffsetSizeAndStrideOpInterface(*this, getType())))
     return this->source();
   if (Value slice = foldExtractAfterInsertSlice(*this))
     return slice;
+
+  if (auto slice = foldConstant(*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();

mlir/test/Dialect/Tensor/canonicalize.mlir

 //       CHECK:   return %[[ARG0]] :  tensor<4x6x16x32xi8>
 func @trivial_slice(%arg0 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {
   %0 = tensor.extract_slice %arg0[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> to tensor<4x6x16x32xi8>
   return %0 : tensor<4x6x16x32xi8>
 }
 
 // -----
 
+// 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 @trivial_insert_slice
 //  CHECK-SAME:   %[[ARG0:.[a-z0-9A-Z_]+]]: tensor<4x6x16x32xi8>
 //   CHECK-NOT:   tensor.extract_slice
 //       CHECK:   return %[[ARG0]] :  tensor<4x6x16x32xi8>
 func @trivial_insert_slice(%arg0 : tensor<4x6x16x32xi8>, %arg1 : tensor<4x6x16x32xi8>) -> tensor<4x6x16x32xi8> {
   %0 = tensor.insert_slice %arg0 into %arg1[0, 0, 0, 0] [4, 6, 16, 32] [1, 1, 1, 1] : tensor<4x6x16x32xi8> into tensor<4x6x16x32xi8>
   return %0 : tensor<4x6x16x32xi8>
 }