[mlir][vector] Add folder for ExtractStridedSliceOp

[ThomasRaoux]

Summary

Add folder for the case where ExtractStridedSliceOp source comes from a chain of InsertStridedSliceOp. Also add a folder for the trivial case where the ExtractStridedSliceOp is a no-op.

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

1637

document function value

1640

give explicit integer decl rather than auto, here and below?

perhaps it is allowed by style, but you switch to explicit types below for int

1665

This is very subjective, so feel free to ignore but I always find

start <= offset && offset < end

easier to read :-)

1689

this really belongs to the else part right? this seems a bit strangely placed

[ThomasRaoux]

Address review comments.

[aartbik]

mlir/lib/Dialect/Vector/VectorOps.cpp

1651

not: bigger -> greater than (more mathy)

[ThomasRaoux]

Thanks Aart!

[nicolasvasilache]

LGTM, minor post-commit comments (just address and land if appropriate).
Thanks Thomas !

mlir/lib/Dialect/Vector/VectorOps.cpp

1633

s/o/of

1659

Put a TODO here for a more powerful analysis?

mlir/test/Dialect/Vector/canonicalize.mlir

170

Make this: %0 = ... offsets = [0, 1] to make the test a bit more interesting ?
The resulting offset should then be vector.extract_strided_slice ... offsets = [0, 0]

mlir/include/mlir/Dialect/Vector/VectorOps.td

def Vector_ExtractStridedSliceOp :
   let extraClassDeclaration = [{
     static StringRef getOffsetsAttrName() { return "offsets"; }
     static StringRef getSizesAttrName() { return "sizes"; }
     static StringRef getStridesAttrName() { return "strides"; }
     VectorType getVectorType(){ return vector().getType().cast<VectorType>(); }
     void getOffsets(SmallVectorImpl<int64_t> &results);
   }];
   let hasCanonicalizer = 1;
+  let hasFolder = 1;
   let assemblyFormat = "$vector attr-dict `:` type($vector) `to` type(results)";
 }
 
 def Vector_TransferReadOp :
   Vector_Op<"transfer_read", [
       DeclareOpInterfaceMethods<VectorTransferOpInterface>,
       DeclareOpInterfaceMethods<VectorUnrollOpInterface, ["getShapeForUnroll"]>
     ]>,

mlir/lib/Dialect/Vector/VectorOps.cpp

static LogicalResult verify(ExtractStridedSliceOp op) {
   if (op.getResult().getType() != resultType) {
     op.emitOpError("expected result type to be ") << resultType;
     return failure();
   }
 
   return success();
 }
 
+// When the source of ExtractStrided comes from a chain of InsertStrided ops try
+// to use the source o the InsertStrided ops if we can detect that the extracted
+// vector is a subset of one of the vector inserted.
+static LogicalResult
+foldExtractStridedOpFromInsertChain(ExtractStridedSliceOp op) {
+  // Helper to extract integer out of ArrayAttr.
+  auto getElement = [](ArrayAttr array, int idx) {
+    return array[idx].cast<IntegerAttr>().getInt();
+  };
+  ArrayAttr extractOffsets = op.offsets();
+  ArrayAttr extractStrides = op.strides();
+  ArrayAttr extractSizes = op.sizes();
+  auto insertOp = op.vector().getDefiningOp<InsertStridedSliceOp>();
+  while (insertOp) {
+    if (op.getVectorType().getRank() !=
+        insertOp.getSourceVectorType().getRank())
+      return failure();
+    ArrayAttr insertOffsets = insertOp.offsets();
+    ArrayAttr insertStrides = insertOp.strides();
+    // If the rank of extract is greater than the rank of insert, we are likely
+    // extracting a partial chunk of the vector inserted.
+    if (extractOffsets.size() > insertOffsets.size())
+      return failure();
+    bool patialoverlap = false;
+    bool disjoint = false;
+    SmallVector<int64_t, 4> offsetDiffs;
+    for (unsigned dim = 0, e = extractOffsets.size(); dim < e; ++dim) {
+      if (getElement(extractStrides, dim) != getElement(insertStrides, dim))
+        return failure();
+      int64_t start = getElement(insertOffsets, dim);
+      int64_t end = start + insertOp.getSourceVectorType().getDimSize(dim);
+      int64_t offset = getElement(extractOffsets, dim);
+      int64_t size = getElement(extractSizes, dim);
+      // Check if the start of the extract offset is in the interval inserted.
+      if (start <= offset && offset < end) {
+        // If the extract interval overlaps but is not fully included we may
+        // have a partial overlap that will prevent any folding.
+        if (offset + size > end)
+          patialoverlap = true;
+        offsetDiffs.push_back(offset - start);
+        continue;
+      }
+      disjoint = true;
+      break;
+    }
+    // The extract element chunk is a subset of the insert element.
+    if (!disjoint && !patialoverlap) {
+      op.setOperand(insertOp.source());
+      // OpBuilder is only used as a helper to build an I64ArrayAttr.
+      OpBuilder b(op.getContext());
+      op.setAttr(ExtractStridedSliceOp::getOffsetsAttrName(),
+                 b.getI64ArrayAttr(offsetDiffs));
+      return success();
+    }
+    // If the chunk extracted is disjoint from the chunk inserted, keep looking
+    // in the insert chain.
+    if (disjoint)
+      insertOp = insertOp.dest().getDefiningOp<InsertStridedSliceOp>();
+    else {
+      // The extracted vector partially overlap the inserted vector, we cannot
+      // fold.
+      return failure();
+    }
+  }
+  return failure();
+}
+
+OpFoldResult ExtractStridedSliceOp::fold(ArrayRef<Attribute> operands) {
+  if (getVectorType() == getResult().getType())
+    return vector();
+  if (succeeded(foldExtractStridedOpFromInsertChain(*this)))
+    return getResult();
+  return {};
+}
+
 void ExtractStridedSliceOp::getOffsets(SmallVectorImpl<int64_t> &results) {
   populateFromInt64AttrArray(offsets(), results);
 }
 
 namespace {
 
 // Pattern to rewrite a ExtractStridedSliceOp(ConstantMaskOp) -> ConstantMaskOp.
 class StridedSliceConstantMaskFolder final

mlir/test/Dialect/Vector/canonicalize.mlir

%1 = vector.extract_strided_slice %0
     {offsets = [1, 1], sizes = [2, 1], strides = [1, 1]}
       : vector<4x3xi1> to vector<2x1xi1>
   // CHECK: vector.constant_mask [1, 1] : vector<2x1xi1>
   return %1 : vector<2x1xi1>
 }
 
 // -----
 
+// CHECK-LABEL: extract_strided_fold
+//  CHECK-SAME: (%[[ARG:.*]]: vector<4x3xi1>)
+//  CHECK-NEXT:   return %[[ARG]] : vector<4x3xi1>
+func @extract_strided_fold(%arg : vector<4x3xi1>) -> (vector<4x3xi1>) {
+  %0 = vector.extract_strided_slice %arg
+    {offsets = [0, 0], sizes = [4, 3], strides = [1, 1]}
+      : vector<4x3xi1> to vector<4x3xi1>
+  return %0 : vector<4x3xi1>
+}
+
+// -----
+
+// CHECK-LABEL: extract_strided_fold_insert
+//  CHECK-SAME: (%[[ARG:.*]]: vector<4x4xf32>
+//  CHECK-NEXT:   return %[[ARG]] : vector<4x4xf32>
+func @extract_strided_fold_insert(%a: vector<4x4xf32>, %b: vector<8x16xf32>)
+  -> (vector<4x4xf32>) {
+  %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]}
+    : vector<4x4xf32> into vector<8x16xf32>
+  %1 = vector.extract_strided_slice %0
+    {offsets = [2, 2], sizes = [4, 4], strides = [1, 1]}
+      : vector<8x16xf32> to vector<4x4xf32>
+  return %1 : vector<4x4xf32>
+}
+
+// -----
+
+// Case where the vector inserted is a subset of the vector extracted.
+// CHECK-LABEL: extract_strided_fold_insert
+//  CHECK-SAME: (%[[ARG0:.*]]: vector<6x4xf32>
+//  CHECK-NEXT:   %[[EXT:.*]] = vector.extract_strided_slice %[[ARG0]]
+//  CHECK-SAME:     {offsets = [0, 0], sizes = [4, 4], strides = [1, 1]}
+//  CHECK-SAME:       : vector<6x4xf32> to vector<4x4xf32>
+//  CHECK-NEXT:   return %[[EXT]] : vector<4x4xf32>
+func @extract_strided_fold_insert(%a: vector<6x4xf32>, %b: vector<8x16xf32>)
+  -> (vector<4x4xf32>) {
+  %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]}
+    : vector<6x4xf32> into vector<8x16xf32>
+  %1 = vector.extract_strided_slice %0
+    {offsets = [2, 2], sizes = [4, 4], strides = [1, 1]}
+      : vector<8x16xf32> to vector<4x4xf32>
+  return %1 : vector<4x4xf32>
+}
+
+// -----
+
+// Negative test where the extract is not a subset of the element inserted.
+// CHECK-LABEL: extract_strided_fold_negative
+//  CHECK-SAME: (%[[ARG0:.*]]: vector<4x4xf32>, %[[ARG1:.*]]: vector<8x16xf32>
+//       CHECK:   %[[INS:.*]] = vector.insert_strided_slice %[[ARG0]], %[[ARG1]]
+//  CHECK-SAME:     {offsets = [2, 2], strides = [1, 1]}
+//  CHECK-SAME:       : vector<4x4xf32> into vector<8x16xf32>
+//       CHECK:   %[[EXT:.*]] = vector.extract_strided_slice %[[INS]]
+//  CHECK-SAME:     {offsets = [2, 2], sizes = [6, 4], strides = [1, 1]}
+//  CHECK-SAME:       : vector<8x16xf32> to vector<6x4xf32>
+//  CHECK-NEXT:   return %[[EXT]] : vector<6x4xf32>
+func @extract_strided_fold_negative(%a: vector<4x4xf32>, %b: vector<8x16xf32>)
+  -> (vector<6x4xf32>) {
+  %0 = vector.insert_strided_slice %a, %b {offsets = [2, 2], strides = [1, 1]}
+    : vector<4x4xf32> into vector<8x16xf32>
+  %1 = vector.extract_strided_slice %0
+    {offsets = [2, 2], sizes = [6, 4], strides = [1, 1]}
+      : vector<8x16xf32> to vector<6x4xf32>
+  return %1 : vector<6x4xf32>
+}
+
+// -----
+
+// Case where we need to go through 2 level of insert element.
+// CHECK-LABEL: extract_strided_fold_insert
+//  CHECK-SAME: (%[[ARG0:.*]]: vector<2x4xf32>, %[[ARG1:.*]]: vector<1x4xf32>,
+//  CHECK-NEXT:   %[[EXT:.*]] = vector.extract_strided_slice %[[ARG1]]
+//  CHECK-SAME:     {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]}
+//  CHECK-SAME:       : vector<1x4xf32> to vector<1x1xf32>
+//  CHECK-NEXT:   return %[[EXT]] : vector<1x1xf32>
+func @extract_strided_fold_insert(%a: vector<2x4xf32>, %b: vector<1x4xf32>,
+                                  %c : vector<1x4xf32>) -> (vector<1x1xf32>) {
+  %0 = vector.insert_strided_slice %b, %a {offsets = [0, 0], strides = [1, 1]}
+    : vector<1x4xf32> into vector<2x4xf32>
+  %1 = vector.insert_strided_slice %c, %0 {offsets = [1, 0], strides = [1, 1]}
+    : vector<1x4xf32> into vector<2x4xf32>
+  %2 = vector.extract_strided_slice %1
+      {offsets = [0, 1], sizes = [1, 1], strides = [1, 1]}
+        : vector<2x4xf32> to vector<1x1xf32>
+  return %2 : vector<1x1xf32>
+}
+
+// -----
+
 // CHECK-LABEL: transpose_1D_identity
 // CHECK-SAME: ([[ARG:%.*]]: vector<4xf32>)
 func @transpose_1D_identity(%arg : vector<4xf32>) -> vector<4xf32> {
   // CHECK-NOT: transpose
   %0 = vector.transpose %arg, [0] : vector<4xf32> to vector<4xf32>
   // CHECK-NEXT: return [[ARG]]
   return %0 : vector<4xf32>
 }