[mlir][vector] Fix bug in transfer op flattening

[ThomasRaoux]

Summary

The logic to figure out if a transfer op can be flattened wasn't
considering the shape being loaded therefore it was incorrectly assuming
some transfer ops were reading contiguous data.

[Benoit]

Sorry about that! I realize my mistake now seeing the new lit testcase.

mlir/lib/Dialect/Vector/Transforms/VectorTransferOpTransforms.cpp

LogicalResult matchAndRewrite(vector::TransferWriteOp transferWriteOp,
     SmallVector<Value> zeros(reducedRank, c0);
     auto identityMap = rewriter.getMultiDimIdentityMap(reducedRank);
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
         transferWriteOp, vector, reducedShapeSource, zeros, identityMap);
     return success();
   }
 };
 
-/// Returns the position of the first inner dimension that has contiguous layout
-/// with at least `requiredContiguousSize` contiguous elements.
-/// When such a dimension is found, the return value satisfies:
-///   0 <= return_value <= memrefType.getRank() - 1.
-/// When no such dimension is found, the return value is memrefType.getRank().
-static int64_t getContiguousInnerDim(MemRefType memrefType,
-                                     int64_t requiredContiguousSize) {
+/// Return true if the memref type has its inner dimension matching the given
+/// shape. Otherwise return false.
+static int64_t hasMatchingInnerContigousShape(MemRefType memrefType,
+                                              ArrayRef<int64_t> targetShape) {
   auto shape = memrefType.getShape();
   SmallVector<int64_t> strides;
   int64_t offset;
-  int64_t innerDim = shape.size();
-  if (succeeded(getStridesAndOffset(memrefType, strides, offset))) {
-    int64_t innerSize = 1;
-    while (true) {
-      if (innerDim == 0)
-        break;
-      const int64_t nextDim = innerDim - 1;
-      if (shape[nextDim] == ShapedType::kDynamicSize)
-        break;
-      if (strides[nextDim] != innerSize)
-        break;
-      innerSize *= shape[nextDim];
-      innerDim = nextDim;
-      if (innerSize >= requiredContiguousSize)
-        break;
-    }
+  if (!succeeded(getStridesAndOffset(memrefType, strides, offset)))
+    return false;
+  if (strides.back() != 1)
+    return false;
+  strides.pop_back();
+  int64_t flatDim = 1;
+  for (auto [targetDim, memrefDim, memrefStride] :
+       llvm::reverse(llvm::zip(targetShape, shape, strides))) {
+    flatDim *= memrefDim;
+    if (flatDim != memrefStride || targetDim != memrefDim)
+      return false;
   }
-  return innerDim;
+  return true;
 }
 
 /// Creates a memref.collapse_shape collapsing all inner dimensions of the
 /// input starting at `firstDimToCollapse`.
 static Value collapseInnerDims(PatternRewriter &rewriter, mlir::Location loc,
                                Value input, int64_t firstDimToCollapse) {
   ShapedType inputType = input.getType().cast<ShapedType>();
   if (inputType.getRank() == 1)
LogicalResult matchAndRewrite(vector::TransferReadOp transferReadOp,
     Value source = transferReadOp.getSource();
     MemRefType sourceType = source.getType().dyn_cast<MemRefType>();
     // Contiguity check is valid on tensors only.
     if (!sourceType)
       return failure();
     if (vectorType.getRank() <= 1)
       // Already 0D/1D, nothing to do.
       return failure();
-    int64_t firstContiguousInnerDim =
-        getContiguousInnerDim(sourceType, vectorType.getNumElements());
-    if (firstContiguousInnerDim >= sourceType.getRank() - 1)
+    if (!hasMatchingInnerContigousShape(
+            sourceType,
+            vectorType.getShape().take_back(vectorType.getRank() - 1)))
       return failure();
+    int64_t firstContiguousInnerDim =
+        sourceType.getRank() - vectorType.getRank();
     // TODO: generalize this pattern, relax the requirements here.
     if (transferReadOp.hasOutOfBoundsDim())
       return failure();
     if (!transferReadOp.getPermutationMap().isMinorIdentity())
       return failure();
     if (transferReadOp.getMask())
       return failure();
     SmallVector<Value> collapsedIndices;
LogicalResult matchAndRewrite(vector::TransferWriteOp transferWriteOp,
     Value source = transferWriteOp.getSource();
     MemRefType sourceType = source.getType().dyn_cast<MemRefType>();
     // Contiguity check is valid on tensors only.
     if (!sourceType)
       return failure();
     if (vectorType.getRank() <= 1)
       // Already 0D/1D, nothing to do.
       return failure();
-    int64_t firstContiguousInnerDim =
-        getContiguousInnerDim(sourceType, vectorType.getNumElements());
-    if (firstContiguousInnerDim >= sourceType.getRank() - 1)
+    if (!hasMatchingInnerContigousShape(
+            sourceType,
+            vectorType.getShape().take_back(vectorType.getRank() - 1)))
       return failure();
+    int64_t firstContiguousInnerDim =
+        sourceType.getRank() - vectorType.getRank();
     // TODO: generalize this pattern, relax the requirements here.
     if (transferWriteOp.hasOutOfBoundsDim())
       return failure();
     if (!transferWriteOp.getPermutationMap().isMinorIdentity())
       return failure();
     if (transferWriteOp.getMask())
       return failure();
     SmallVector<Value> collapsedIndices;

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

 // RUN: mlir-opt %s -test-vector-transfer-flatten-patterns -split-input-file | FileCheck %s
 
 func.func @transfer_read_flattenable_with_offset(
       %arg : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<5x4x3x2xi8> {
     %c0 = arith.constant 0 : index
     %cst = arith.constant 0 : i8
     %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : 
       memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<5x4x3x2xi8>
     return %v : vector<5x4x3x2xi8>
 }
 
 // CHECK-LABEL: func @transfer_read_flattenable_with_offset
 // CHECK-SAME:      %[[ARG:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8
 // CHECK:         %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]
-// C-HECK:         %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]
-// C-HECK:         %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<120xi8> to vector<5x4x3x2xi8>
-// C-HECK:         return %[[VEC2D]]
+// CHECK:         %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]
+// CHECK:         %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<120xi8> to vector<5x4x3x2xi8>
+// CHECK:         return %[[VEC2D]]
 
 // -----
 
 func.func @transfer_write_flattenable_with_offset(
       %arg : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, %vec : vector<5x4x3x2xi8>) {
     %c0 = arith.constant 0 : index
     vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] : 
       vector<5x4x3x2xi8>, memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>
     return
 }
 
-// C-HECK-LABEL: func @transfer_write_flattenable_with_offset
-// C-HECK-SAME:      %[[ARG:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8
-// C-HECK-SAME:      %[[VEC:[0-9a-zA-Z]+]]: vector<5x4x3x2xi8>
-// C-HECK-DAG:     %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]{{.}} : memref<5x4x3x2xi8, {{.+}}> into memref<120xi8, {{.+}}>
-// C-HECK-DAG:     %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<5x4x3x2xi8> to vector<120xi8>
-// C-HECK:         vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
+// CHECK-LABEL: func @transfer_write_flattenable_with_offset
+// CHECK-SAME:      %[[ARG:[0-9a-zA-Z]+]]: memref<5x4x3x2xi8
+// CHECK-SAME:      %[[VEC:[0-9a-zA-Z]+]]: vector<5x4x3x2xi8>
+// CHECK-DAG:     %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]{{.}} : memref<5x4x3x2xi8, {{.+}}> into memref<120xi8, {{.+}}>
+// CHECK-DAG:     %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<5x4x3x2xi8> to vector<120xi8>
+// CHECK:         vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
 
 // -----
 
 func.func @transfer_write_0d(%arg : memref<i8>, %vec : vector<i8>) {
       vector.transfer_write %vec, %arg[] : vector<i8>, memref<i8>
       return
 }
 

 // CHECK:       %[[C0:.+]] = arith.constant 0 : index
 // CHECK:       %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG1]] {{\[}}[0], [1], [2, 3]{{\]}}
 // CHECK-SAME:    : memref<?x?x8x4xi8, {{.+}}> into memref<?x?x32xi8, {{.+}}>
 // CHECK:       %[[VEC1D:.+]] = vector.shape_cast %[[ARG0]] : vector<8x4xi8> to vector<32xi8>
 // CHECK:       vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
 // CHECK-SAME:    [%[[ARG2]], %[[ARG3]], %[[C0]]]
 // CHECK-SAME:    {in_bounds = [true]}
 // CHECK-SAME:    : vector<32xi8>, memref<?x?x32xi8, {{.+}}>
+
+// -----
+
+func.func @transfer_read_flattenable_negative(
+      %arg : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<2x2x2x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : 
+      memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>, vector<2x2x2x2xi8>
+    return %v : vector<2x2x2x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable_negative
+//       CHECK:   vector.transfer_read {{.*}} vector<2x2x2x2xi8>
+
+// -----
+
+func.func @transfer_read_flattenable_negative2(
+      %arg : memref<5x4x3x2xi8, strided<[24, 8, 2, 1], offset: ?>>) -> vector<5x4x3x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : 
+      memref<5x4x3x2xi8, strided<[24, 8, 2, 1], offset: ?>>, vector<5x4x3x2xi8>
+    return %v : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable_negative2
+//       CHECK:   vector.transfer_read {{.*}} vector<5x4x3x2xi8>