[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

1253

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

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

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

[okkwon]

Merged into the LLVM main.

mlir/include/mlir/IR/BuiltinAttributes.h

public:
   bool empty() const { return size() == 0; }
 
   //===--------------------------------------------------------------------===//
   // Mutation Utilities
   //===--------------------------------------------------------------------===//
 
   /// Return a new DenseElementsAttr that has the same data as the current
   /// attribute, but has been reshaped to 'newType'. The new type must have the
   /// same total number of elements as well as element type.
   DenseElementsAttr reshape(ShapedType newType);
 
   /// Return a new DenseElementsAttr that has the same data as the current
+  /// attribute, but with a different shape for a splat type. The new type must
+  /// have the same element type.
+  DenseElementsAttr resizeSplat(ShapedType newType);
+
+  /// Return a new DenseElementsAttr that has the same data as the current
   /// attribute, but has bitcast elements to 'newElType'. The new type must have
   /// the same bitwidth as the current element type.
   DenseElementsAttr bitcast(Type newElType);
 
   /// Generates a new DenseElementsAttr by mapping each int value to a new
   /// underlying APInt. The new values can represent either an integer or float.
   /// This underlying type must be an DenseIntElementsAttr.
   DenseElementsAttr mapValues(Type newElementType,

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

mlir/lib/IR/BuiltinAttributes.cpp

DenseElementsAttr DenseElementsAttr::reshape(ShapedType newType) {
 
   assert(newType.getElementType() == curType.getElementType() &&
          "expected the same element type");
   assert(newType.getNumElements() == curType.getNumElements() &&
          "expected the same number of elements");
   return DenseIntOrFPElementsAttr::getRaw(newType, getRawData(), isSplat());
 }
 
+DenseElementsAttr DenseElementsAttr::resizeSplat(ShapedType newType) {
+  assert(isSplat() && "expected a splat type");
+
+  ShapedType curType = getType();
+  if (curType == newType)
+    return *this;
+
+  assert(newType.getElementType() == curType.getElementType() &&
+         "expected the same element type");
+  return DenseIntOrFPElementsAttr::getRaw(newType, getRawData(), true);
+}
+
 /// Return a new DenseElementsAttr that has the same data as the current
 /// attribute, but has bitcast elements such that it is now 'newType'. The new
 /// type must have the same shape and element types of the same bitwidth as the
 /// current type.
 DenseElementsAttr DenseElementsAttr::bitcast(Type newElType) {
   ShapedType curType = getType();
   Type curElType = curType.getElementType();
   if (curElType == newElType)

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>