[mlir][vector] Fix double rank reducing folding bug

[antiagainst]

Summary

In https://reviews.llvm.org/D133883, we changed the
FoldExtractSliceIntoTransferRead pattern from requiring
full identity map to minor identity map. This effectively
allows rank reducing vector.transfer_read ops. However,
the logic for checking tensor.extract_slice rank reducing
still looks at the vector rank, which now could be smaller
than the tensor.extract_slice's output tensor rank.
It ends up we can have incorrect index cacluation after
folding due to this double rank reducing behavior.

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

 void TransferReadOp::getEffects(
     SmallVectorImpl<SideEffects::EffectInstance<MemoryEffects::Effect>>
         &effects) {
   if (getShapedType().isa<MemRefType>())
     effects.emplace_back(MemoryEffects::Read::get(), getSource(),
                          SideEffects::DefaultResource::get());
 }
 
+/// Returns true if all rank reduced in the given `extractOp` happen in leading
+/// dimensions earlier than last `trailingRank` dimensions.
+static bool areAllRankReducedLeadingDim(tensor::ExtractSliceOp extractOp,
+                                        unsigned trailingRank) {
+  // If no ranks are reduced at all, it's a degenerated case; always true.
+  if (extractOp.getSourceType().getRank() == extractOp.getType().getRank())
+    return true;
+
+  RankedTensorType inferredType = extractOp.inferResultType(
+      extractOp.getSourceType(), extractOp.getMixedOffsets(),
+      extractOp.getMixedSizes(), extractOp.getMixedStrides());
+  return extractOp.getType().getShape().take_back(trailingRank) ==
+         inferredType.getShape().take_back(trailingRank);
+}
+
 namespace {
 /// Fold transfer_reads of a tensor.extract_slice op. E.g.:
 ///
 /// ```
 /// %0 = tensor.extract_slice %t[%a, %b] [%c, %d] [1, 1]
 ///     : tensor<?x?xf32> to tensor<?x?xf32>
 /// %1 = vector.transfer_read %0[%e, %f], %cst {in_bounds = [true, true]}
 ///     : tensor<?x?xf32>, vector<4x5xf32>
LogicalResult matchAndRewrite(TransferReadOp xferOp,
     //
     // Cannot fold into:
     // ```
     //    %0 = vector.transfer_read %t[0,0,0], %cst :
     //      tensor<2x1x4xf32>, vector<2x4xf32>
     // ```
     // For this, check the trailing `vectorRank` dims of the extract_slice
     // result tensor match the trailing dims of the inferred result tensor.
+    if (!areAllRankReducedLeadingDim(extractOp, extractOp.getType().getRank()))
+      return failure();
+
     int64_t rankReduced =
         extractOp.getSourceType().getRank() - extractOp.getType().getRank();
-    int64_t vectorRank = xferOp.getVectorType().getRank();
-    RankedTensorType inferredDestTensorType =
-        tensor::ExtractSliceOp::inferResultType(
-            extractOp.getSourceType(), extractOp.getMixedOffsets(),
-            extractOp.getMixedSizes(), extractOp.getMixedStrides());
-    auto actualDestTensorShape = extractOp.getType().getShape();
-    if (rankReduced > 0 &&
-        actualDestTensorShape.take_back(vectorRank) !=
-            inferredDestTensorType.getShape().take_back(vectorRank))
-      return failure();
 
     SmallVector<Value> newIndices;
     // In case this is a rank-reducing ExtractSliceOp, copy rank-reduced
     // indices first.
     for (int64_t i = 0; i < rankReduced; ++i) {
       OpFoldResult offset = extractOp.getMixedOffsets()[i];
       newIndices.push_back(getValueOrCreateConstantIndexOp(
           rewriter, extractOp.getLoc(), offset));

mlir/test/Dialect/Vector/canonicalize.mlir

 //  CHECK-NEXT:   %[[SPLAT:.*]] = vector.splat %[[ARG]] : vector<2x4x3xi32>
 //  CHECK-NEXT:   return %[[SPLAT]] : vector<2x4x3xi32>
 func.func @insert_splat(%x : i32) -> vector<2x4x3xi32> {
   %v0 = vector.splat %x : vector<4x3xi32>
   %v1 = vector.splat %x : vector<2x4x3xi32>
   %insert = vector.insert %v0, %v1[0] : vector<4x3xi32> into vector<2x4x3xi32>
   return %insert : vector<2x4x3xi32>
 }
+
+// -----
+
+// CHECK-LABEL: func.func @transfer_read_from_rank_reducing_extract_slice
+//       CHECK:   tensor.extract_slice
+//       CHECK:   vector.transfer_read
+func.func @transfer_read_from_rank_reducing_extract_slice(%src: tensor<1x8x8x8xf32>, %i1: index, %i2: index, %i3: index, %i4: index) -> vector<4xf32> {
+  %c0 = arith.constant 0 : index
+  %f0 = arith.constant 0.000000e+00 : f32
+  %0 = tensor.extract_slice %src[0, %i1, %i2, %i3] [1, 4, 1, 4] [1, 1, 1, 1] : tensor<1x8x8x8xf32> to tensor<1x4x4xf32>
+  %1 = vector.transfer_read %0[%c0, %i4, %c0], %f0 {in_bounds = [true]} : tensor<1x4x4xf32>, vector<4xf32>
+  return %1 : vector<4xf32>
+}