[mlir] Fold Tensor.extract_slice into a constant splat.

[okkwon]

Summary

Fold arith.extract_slice into arith.constant when the source is a constant
splat and the result type is statically shaped.

[mehdi_amini]

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

1254

The ArrayRef<Attribute> argument to the folder is providing you with the constant inputs, you're not supposed to match it.
So I would move this as the first folder before the two others (because it seems cheaper) and use the folder argument.

[okkwon]

Use the attribute operand instead of matching.

[mehdi_amini]

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

1246

Can you just pass in splat here instead of converting back to a generic attribute? There is overhead associated to this.
Otherwise we should add a method to be able to do: return splat.reshape(newShape);

[okkwon]

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

1246

Passing 'splat' causes an assertion in the get function because it expects an attribute of the element type, not a tensor type. Adding reshape sounds good to me. I will take a look. Thank you!

[okkwon]

Support a different shape for reshape for a splat type.

[okkwon]

There is already reshape(), but it didn't support a shape with a different number of element even for a splat type.
I relaxed the assert condition to allow a type with a different number of elements.

[mehdi_amini]

LGTM

[mravishankar]

mlir/include/mlir/IR/BuiltinAttributes.h

655 ↗

(On Diff #409456)

I think its clearer to keep reshape semantics to preserving the number of elements. Thats the general use of reshape through out MLIR.
We can add a separate slice method maybe to handle this case.

[okkwon]

Introduce resizeSplat instead of using reshape since reshape is not suppposed
to change the number of elements.

[okkwon]

mlir/include/mlir/IR/BuiltinAttributes.h

655 ↗

(On Diff #409456)

I introduced resizeSplat to create a new attribute instead of reshape.

[okkwon]

Merged into the LLVM main.

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

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>) {
+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 SplatElementsAttr::get(resultType,
+                                    splat.getSplatValue<Attribute>());
+  }
   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));

mlir/test/Dialect/Tensor/canonicalize.mlir

func @fold_extract_insert(%input : tensor<?x?x?xf32>, %slice: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<4x?x8xf32>) {
   %0 = tensor.insert_slice %slice into %input[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>
   %1 = tensor.extract_slice %0[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<?x?x?xf32> to tensor<4x?x8xf32>
   // CHECK: return %[[SLICE]]
   return %1 : tensor<4x?x8xf32>
 }
 
 // -----
 
+// CHECK-LABEL: func @fold_extract_constant_splat
+//   CHECK-NOT: tensor.extract_slice
+//       CHECK: arith.constant dense<42> : tensor<4x4xi32>
+func @fold_extract_constant_splat() -> (tensor<4x4xi32>) {
+  %cst = arith.constant dense<42> : tensor<1024x1024xi32>
+  %1 = tensor.extract_slice %cst[0,0] [4,4] [1, 1] : tensor<1024x1024xi32> to tensor<4x4xi32>
+  return %1 : tensor<4x4xi32>
+}
+
+// -----
+
 // CHECK-LABEL: func @fold_overlapping_insert
 //  CHECK-SAME: %[[INPUT:.+]]: tensor<?x?x?xf32>, %{{.+}}: tensor<4x?x8xf32>, %[[SLICE2:.+]]: tensor<4x?x8xf32>
 func @fold_overlapping_insert(%input : tensor<?x?x?xf32>, %slice1: tensor<4x?x8xf32>, %slice2: tensor<4x?x8xf32>, %i: index, %size: index) -> (tensor<?x?x?xf32>) {
   %c0 = arith.constant 0: index
   %c1 = arith.constant 1: index
   %0 = tensor.insert_slice %slice1 into %input[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>
   // CHECK: %[[INSERT:.+]] = tensor.insert_slice %[[SLICE2]] into %[[INPUT]]
   %1 = tensor.insert_slice %slice2 into %0[%c0, %i, 0] [4, %size, 8] [1, 1, %c1] : tensor<4x?x8xf32> into tensor<?x?x?xf32>