fixup! [mlir][Vector] Update patterns for flattening vector.xfer Ops

Final refactor requested by Nicolas

Author: banach-space

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

@@ -42,6 +42,39 @@ Value createOrFoldDimOp(OpBuilder &b, Location loc, Value source, int64_t dim);
 /// on a 2D slice. Otherwise, returns a failure.
 FailureOr<std::pair<int, int>> isTranspose2DSlice(vector::TransposeOp op);
 
+/// Return true if `vectorType` is a contiguous slice of `memrefType`.
+///
+/// Only the N = vectorType.getRank() trailing dims of `memrefType` are
+/// checked (the other dims are not relevant). Note that for `vectorType` to be
+/// a contiguous slice of `memrefType`, the trailing dims of the latter have
+/// to be contiguous - this is checked by looking at the corresponding strides.
+///
+/// There might be some restriction on the leading dim of `VectorType`:
+///
+/// Case 1. If all the trailing dims of `vectorType` match the trailing dims
+///         of `memrefType` then the leading dim of `vectorType` can be
+///         arbitrary.
+///
+///        Ex. 1.1 contiguous slice, perfect match
+///          vector<4x3x2xi32> from memref<5x4x3x2xi32>
+///        Ex. 1.2 contiguous slice, the leading dim does not match (2 != 4)
+///          vector<2x3x2xi32> from memref<5x4x3x2xi32>
+///
+/// Case 2. If an "internal" dim of `vectorType` does not match the
+///         corresponding trailing dim in `memrefType` then the remaining
+///         leading dims of `vectorType` have to be 1 (the first non-matching
+///         dim can be arbitrary).
+///
+///        Ex. 2.1 non-contiguous slice, 2 != 3 and the leading dim != <1>
+///          vector<2x2x2xi32> from memref<5x4x3x2xi32>
+///        Ex. 2.2  contiguous slice, 2 != 3 and the leading dim == <1>
+///          vector<1x2x2xi32> from memref<5x4x3x2xi32>
+///        Ex. 2.3. contiguous slice, 2 != 3 and the leading dims == <1x1>
+///          vector<1x1x2x2xi32> from memref<5x4x3x2xi32>
+///        Ex. 2.4. non-contiguous slice, 2 != 3 and the leading dims != <1x1>
+///         vector<2x1x2x2xi32> from memref<5x4x3x2xi32>)
+bool isContiguousSlice(MemRefType memrefType, VectorType vectorType);
+
 } // namespace vector
 
 /// Constructs a permutation map of invariant memref indices to vector

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

@@ -487,90 +487,6 @@ class TransferWriteDropUnitDimsPattern
 
 } // namespace
 
-/// Return true if `vectorType` is a contiguous slice of `memrefType`.
-///
-/// Compares `vectorType` against the trailing dimensions of `memrefType`
-/// to check whether `vectorType` is a contiguous slice of `memrefType`. This
-/// is implemented by iterating over the dims of `vectorType` and `memrefType`
-/// and comparing them starting from the inner-most/right-most dims.
-///
-/// Note that there might be some restriction on the leading dim of
-/// `VectorType`:
-///   1. if all the trailing dims of `vectorType` match the trailing dims
-///     of `memrefType` then the leading dim of `vectorType` can be arbitrary:
-///
-///       1.1 contiguous slice, perfect match
-///         vector<4x3x2xi32> from memref<5x4x3x2xi32>
-///       1.2 contiguous slice, all dims match except the leading dim: 2 != 4
-///         vector<2x3x2xi32> from memref<5x4x3x2xi32>
-///
-///   2. if an "internal" dim of `vectorType` does not match the corresponding
-///     trailing dim in `memrefType` then the remaining leading dims of
-///     `vectorType` have to be 1 (the first non-matching dim can be arbitrary):
-///
-///       2.1 non-contiguous slice, 2 != 3 and the leading dim != <1>
-///         vector<2x2x2xi32> from memref<5x4x3x2xi32>
-///       2.2  contiguous slice, 2 != 3 and the leading dim == <1>
-///         vector<1x2x2xi32> from memref<5x4x3x2xi32>
-///       2.3. contiguous slice, 2 != 3 and the leading dims == <1x1>
-///         vector<1x1x2x2xi32> from memref<5x4x3x2xi32>
-///       2.4. non-contiguous slice, 2 != 3 and the leading dims != <1x1>
-///         vector<2x1x2x2xi32> from memref<5x4x3x2xi32>)
-///
-/// In all cases `memrefType` has to be contiguous (this is checked by looking
-/// at strides).
-static bool isContiguousSlice(MemRefType memrefType, VectorType vectorType) {
-
-  // Get the shape of `vectorType`. The leading dim is treated separately.
-  ArrayRef<int64_t> targetShape = vectorType.getShape();
-  auto targetShapeTrailingDims = targetShape.drop_front(1);
-
-  // Get the strides of the memref.
-  int64_t offset;
-  SmallVector<int64_t> strides;
-  if (!succeeded(getStridesAndOffset(memrefType, strides, offset)))
-    return false;
-
-  // Non-unit stride in the trailing dimension means this memref is
-  // not contiguous.
-  if (strides.back() != 1)
-    return false;
-
-  // Do all but the leading dim of `vectorType` and `memrefType` match?
-  bool allTrailingDimsMatch = true;
-
-  // The trailing dimension of `memrefType` after collapsing/flattening the
-  // current dim. This will be a product of the leading dims, hence initialising
-  // to 1.
-  int64_t flatDim = 1;
-
-  // Iterate over all dim of `vectorType` (in reverse) excluding the leading dim
-  // and compare them against the trailing dims of `memrefType`.
-  strides.pop_back();
-  for (auto [targetDim, memrefDim, memrefStride] :
-       llvm::reverse(llvm::zip(targetShapeTrailingDims,
-                               memrefType.getShape().drop_front(1), strides))) {
-    flatDim *= memrefDim;
-    // If the memref stride does not match the flattened dim, then this is
-    // memref is not contiguous.
-    if (flatDim != memrefStride)
-      return false;
-
-    // If a non-matching dim was found previously, then the remaining dims of
-    // `VectorType` should be 1.
-    if (!allTrailingDimsMatch && (targetDim != 1))
-      return false;
-
-    allTrailingDimsMatch = (targetDim == memrefDim);
-  }
-
-  // If the trailing dims of `vectorType` and `memrefType` match, then this is a
-  // contiguous load. If there was a mismatch, then the internal dims have
-  // already been verified to be unit dims, but the leading dim still has to be
-  // checked.
-  return allTrailingDimsMatch ? true : (targetShape[0] == 1);
-}
-
 /// Creates a memref.collapse_shape collapsing all inner dimensions of the
 /// input starting at `firstDimToCollapse`.
 static Value collapseInnerDims(PatternRewriter &rewriter, mlir::Location loc,
@@ -630,7 +546,7 @@ class FlattenContiguousRowMajorTransferReadPattern
     if (vectorType.getRank() <= 1)
       // Already 0D/1D, nothing to do.
       return failure();
-    if (!isContiguousSlice(sourceType, vectorType))
+    if (!vector::isContiguousSlice(sourceType, vectorType))
       return failure();
     int64_t firstContiguousInnerDim =
         sourceType.getRank() - vectorType.getRank();
@@ -688,7 +604,7 @@ class FlattenContiguousRowMajorTransferWritePattern
     if (vectorType.getRank() <= 1)
       // Already 0D/1D, nothing to do.
       return failure();
-    if (!isContiguousSlice(sourceType, vectorType))
+    if (!vector::isContiguousSlice(sourceType, vectorType))
       return failure();
     int64_t firstContiguousInnerDim =
         sourceType.getRank() - vectorType.getRank();

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

@@ -249,3 +249,47 @@ bool matcher::operatesOnSuperVectorsOf(Operation &op,
   // between parallel, reduction and possibly other cases.
   return ratio.has_value();
 }
+
+bool vector::isContiguousSlice(MemRefType memrefType, VectorType vectorType) {
+  if (vectorType.isScalable())
+    return false;
+
+  ArrayRef<int64_t> vectorShape = vectorType.getShape();
+  auto vecRank = vectorType.getRank();
+
+  // Extract the trailing dims and strides of the input memref
+  auto memrefShape = memrefType.getShape().take_back(vecRank);
+  int64_t offset;
+  SmallVector<int64_t> stridesFull;
+  if (!succeeded(getStridesAndOffset(memrefType, stridesFull, offset)))
+    return false;
+  auto strides = ArrayRef<int64_t>(stridesFull).take_back(vecRank);
+
+  // Cond 1: A contiguous memref will always have a unit trailing stride.
+  if (strides.back() != 1)
+    return false;
+
+  // Cond 2: Strides of a contiguous memref have to match the flattened dims.
+  strides = strides.drop_back(1);
+  SmallVector<int64_t> flattenedDims;
+  for (size_t i = 1; i < memrefShape.size(); i++)
+    flattenedDims.push_back(mlir::computeProduct(memrefShape.take_back(i)));
+
+  if (!llvm::equal(strides, llvm::reverse(flattenedDims)))
+    return false;
+
+  // Cond 3: Compare the dims of `vectorType` against `memrefType` (in reverse).
+  // In the most basic case, all dims will match.
+  auto firstNonMatchingDim =
+      std::mismatch(vectorShape.rbegin(), vectorShape.rend(),
+                    memrefShape.rbegin(), memrefShape.rend());
+  if (firstNonMatchingDim.first == vectorShape.rend())
+    return true;
+
+  // One non-matching dim is still fine, however the remaining leading dims of
+  // `vectorType` need to be 1.
+  SmallVector<int64_t> leadingDims(++firstNonMatchingDim.first,
+                                   vectorShape.rend());
+
+  return llvm::all_of(leadingDims, [](auto x) { return x == 1; });
+}

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

@@ -41,7 +41,7 @@ func.func @transfer_read_dims_mismatch_contiguous(
 
 // -----
 
-func.func @transfer_read_dims_mismatch_contiguous(
+func.func @transfer_read_dims_mismatch_non_contiguous(
     %arg : memref<5x4x3x2xi8, strided<[24, 6, 2, 1], offset: ?>>) -> vector<2x1x2x2xi8> {
     %c0 = arith.constant 0 : index
     %cst = arith.constant 0 : i8
@@ -50,6 +50,7 @@ func.func @transfer_read_dims_mismatch_contiguous(
     return %v : vector<2x1x2x2xi8>
 }
 
+// CHECK-LABEL: func.func @transfer_read_dims_mismatch_non_contiguous
 // CHECK-NOT: memref.collapse_shape
 // CHECK-NOT: vector.shape_cast
 
@@ -100,6 +101,7 @@ func.func @transfer_write_dims_mismatch_non_contiguous(
     return
 }
 
+// CHECK-LABEL: func.func @transfer_write_dims_mismatch_non_contiguous
 // CHECK-NOT: memref.collapse_shape
 // CHECK-NOT: vector.shape_cast
 
@@ -110,6 +112,7 @@ func.func @transfer_write_0d(%arg : memref<i8>, %vec : vector<i8>) {
       return
 }
 
+// CHECK-LABEL: func.func @transfer_write_0d
 // CHECK-NOT: memref.collapse_shape
 // CHECK-NOT: vector.shape_cast
 
@@ -121,6 +124,7 @@ func.func @transfer_read_0d(%arg : memref<i8>) -> vector<i8> {
       return %0 : vector<i8>
 }
 
+// CHECK-LABEL: func.func @transfer_read_0d
 // CHECK-NOT: memref.collapse_shape
 // CHECK-NOT: vector.shape_cast