Generalize the vector transfer flattening patterns (dyn shapes).

[Benoit]

Summary

This generalizes this pattern to support dynamic shapes, and the source memref is no longer flattened all the way to 1D. Only a sufficient number of innermost dimensions are collapsed, to match the number of elements of the vectorType (which is still being flattened to 1D).

Not flattening the outer dimensions of the memrefType is key to simple, correct handling of the indices for these outer dims: we can just keep those as-is, whereas if we kept the flattening of all those outer dims, we would have to generate arithmetic to compute the correct flattened index.

This pattern was also conflating another thing into it: a helper was calling rankReducingSubviewDroppingUnitDims, which was duplicating the work of the other patterns Transfer{Read,Write}DropUnitDimsPattern. This is dropped. That didn't seem to break anything, possibly because whichever places needed that were also adding those patterns anyway. The reason why I specifically needed to drop that now was that conflating that was making the logic here really messy.

[Benoit]

fix comment.

[Benoit]

actually return 0 for out-of-bounds case.

[Benoit]

Fixes issues found in IREE integration; limit generality: not handling nonzero indices after all.

[Benoit]

clang-format

[Benoit]

correct handling of outer indices. No longer flattening the outer dims.

[Benoit]

drop no-longer-needed hunk

[Benoit]

comments

[ThomasRaoux]

LGTM once the last minor comments are addressed

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

397

Nit: remove braces

489

Should be cast<> if you don’t expect null result (and don’t check for it)

[Benoit]

review comments

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

 }
 
 /// Creates a rank-reducing memref.subview op that drops unit dims from its
 /// input. Or just returns the input if it was already without unit dims.
 static Value rankReducingSubviewDroppingUnitDims(PatternRewriter &rewriter,
                                                  mlir::Location loc,
                                                  Value input) {
   MemRefType inputType = input.getType().cast<MemRefType>();
-  assert(inputType.hasStaticShape());
   SmallVector<int64_t> subViewOffsets(inputType.getRank(), 0);
   SmallVector<int64_t> subViewStrides(inputType.getRank(), 1);
   ArrayRef<int64_t> subViewSizes = inputType.getShape();
   MemRefType resultType =
       dropUnitDims(inputType, subViewOffsets, subViewSizes, subViewStrides);
   if (canonicalizeStridedLayout(resultType) ==
       canonicalizeStridedLayout(inputType))
     return input;
if (rankReducedInputType.getRank() == 1)
     return rankReducedInput;
   ReassociationIndices indices;
   for (int i = 0; i < rankReducedInputType.getRank(); ++i)
     indices.push_back(i);
   return rewriter.create<memref::CollapseShapeOp>(
       loc, rankReducedInput, std::array<ReassociationIndices, 1>{indices});
 }
 
+/// Returns how many elements are statically known to be contiguous in the
+/// layout of this `memrefType`, at the position given by `indices`.
+/// For example, for `memref<8xf32>` and `indices=[%c2]` where
+/// `%c2 = arith.constant 2 : index`, the return value should be 6 = 8 - 2.
+/// Returns 0 in error cases.
+static int64_t
+getContiguousStaticRowMajorElementsAtPosition(MemRefType memrefType,
+                                              ValueRange indices) {
+  auto shape = memrefType.getShape();
+  if (shape.size() != indices.size())
+    return 0;
+  SmallVector<int64_t> strides;
+  int64_t offset;
+  if (failed(getStridesAndOffset(memrefType, strides, offset)))
+    return 0;
+  // contiguousSize will be the number of contiguous elements in memrefType in
+  // the dimensions traversed by the for loop below.
+  int64_t contiguousSize = 1;
+  // indicesOffset will be the offset of the position given by `indices` in the
+  // dimensions traversed by the for loop below.
+  int64_t indicesOffset = 0;
+  for (int i = static_cast<int>(strides.size()) - 1; i >= 0; --i) {
+    if (shape[i] == ShapedType::kDynamicSize)
+      break;
+    if (strides[i] != contiguousSize)
+      break;
+    arith::ConstantIndexOp cstIndex =
+        indices[i].getDefiningOp<arith::ConstantIndexOp>();
+    if (!cstIndex)
+      break;
+    int64_t indexValue = cstIndex.value();
+    if (indexValue < 0 || indexValue >= shape[i])
+      break;
+    indicesOffset += indexValue * contiguousSize;
+    contiguousSize *= shape[i];
+  }
+  return contiguousSize - indicesOffset;
+}
+
 /// Rewrites contiguous row-major vector.transfer_read ops by inserting
 /// memref.collapse_shape on the source so that the resulting
 /// vector.transfer_read has a 1D source. Requires the source shape to be
 /// already reduced i.e. without unit dims.
 class FlattenContiguousRowMajorTransferReadPattern
     : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern<vector::TransferReadOp>::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferReadOp transferReadOp,
                                 PatternRewriter &rewriter) const override {
     auto loc = transferReadOp.getLoc();
     Value vector = transferReadOp.getVector();
     VectorType vectorType = vector.getType().cast<VectorType>();
     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();
-    if (!isStaticShapeAndContiguousRowMajor(sourceType))
-      return failure();
-    if (getReducedRank(sourceType.getShape()) != sourceType.getRank())
-      // This pattern requires the source to already be rank-reduced.
-      return failure();
-    if (sourceType.getNumElements() != vectorType.getNumElements())
+    if (getContiguousStaticRowMajorElementsAtPosition(
+            sourceType, transferReadOp.getIndices()) <
+        vectorType.getNumElements())
       return failure();
     // TODO: generalize this pattern, relax the requirements here.
     if (transferReadOp.hasOutOfBoundsDim())
       return failure();
     if (!transferReadOp.getPermutationMap().isMinorIdentity())
       return failure();
     if (transferReadOp.getMask())
       return failure();
-    if (llvm::any_of(transferReadOp.getIndices(),
-                     [](Value v) { return !isZero(v); }))
-      return failure();
     Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     auto identityMap1D = rewriter.getMultiDimIdentityMap(1);
-    VectorType vectorType1d = VectorType::get({sourceType.getNumElements()},
-                                              sourceType.getElementType());
+    VectorType vectorType1d = VectorType::get({vectorType.getNumElements()},
+                                              vectorType.getElementType());
     Value source1d =
         collapseContiguousRowMajorMemRefTo1D(rewriter, loc, source);
-    Value read1d = rewriter.create<vector::TransferReadOp>(
+    vector::TransferReadOp read1d = rewriter.create<vector::TransferReadOp>(
         loc, vectorType1d, source1d, ValueRange{c0}, identityMap1D);
+    read1d.setInBoundsAttr(rewriter.getBoolArrayAttr({true}));
     rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
         transferReadOp, vector.getType().cast<VectorType>(), read1d);
     return success();
   }
 };
 
 /// Rewrites contiguous row-major vector.transfer_write ops by inserting
 /// memref.collapse_shape on the source so that the resulting
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();
-    if (!isStaticShapeAndContiguousRowMajor(sourceType))
-      return failure();
-    if (getReducedRank(sourceType.getShape()) != sourceType.getRank())
-      // This pattern requires the source to already be rank-reduced.
-      return failure();
-    if (sourceType.getNumElements() != vectorType.getNumElements())
+    if (getContiguousStaticRowMajorElementsAtPosition(
+            sourceType, transferWriteOp.getIndices()) <
+        vectorType.getNumElements())
       return failure();
     // TODO: generalize this pattern, relax the requirements here.
     if (transferWriteOp.hasOutOfBoundsDim())
       return failure();
     if (!transferWriteOp.getPermutationMap().isMinorIdentity())
       return failure();
     if (transferWriteOp.getMask())
       return failure();
-    if (llvm::any_of(transferWriteOp.getIndices(),
-                     [](Value v) { return !isZero(v); }))
-      return failure();
     Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     auto identityMap1D = rewriter.getMultiDimIdentityMap(1);
-    VectorType vectorType1d = VectorType::get({sourceType.getNumElements()},
-                                              sourceType.getElementType());
+    VectorType vectorType1d = VectorType::get({vectorType.getNumElements()},
+                                              vectorType.getElementType());
     Value source1d =
         collapseContiguousRowMajorMemRefTo1D(rewriter, loc, source);
     Value vector1d =
         rewriter.create<vector::ShapeCastOp>(loc, vectorType1d, vector);
-    rewriter.create<vector::TransferWriteOp>(loc, vector1d, source1d,
-                                             ValueRange{c0}, identityMap1D);
+    vector::TransferWriteOp write1d = rewriter.create<vector::TransferWriteOp>(
+        loc, vector1d, source1d, ValueRange{c0}, identityMap1D);
+    write1d.setInBoundsAttr(rewriter.getBoolArrayAttr({true}));
     rewriter.eraseOp(transferWriteOp);
     return success();
   }
 };
 
 } // namespace
 
 void mlir::vector::transferOpflowOpt(Operation *rootOp) {
   TransferOptimization opt(rootOp);

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

func.func @transfer_read_0d(%arg : memref<i8>) -> vector<i8> {
       return %0 : vector<i8>
 }
 
 // CHECK-LABEL: func @transfer_read_0d
 // CHECK-SAME:       %[[ARG:.+]]: memref<i8>
 // CHECK:            %[[CST:.+]] = arith.constant 0 : i8
 // CHECK:            %[[READ:.+]] = vector.transfer_read %[[ARG]][], %[[CST]] : memref<i8>
 // CHECK:            return %[[READ]]
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3)[s0, s1] -> (d0 * s1 + s0 + d1 * 32 + d2 * 4 + d3)>
+
+func.func @transfer_read_flattenable_with_dynamic_dims_and_indices(%arg0 : memref<?x?x8x4xi8, #map0>, %arg1 : index, %arg2 : index) -> vector<8x4xi8> {
+    %c0_i8 = arith.constant 0 : i8
+    %c0 = arith.constant 0 : index
+    %result = vector.transfer_read %arg0[%arg1, %arg2, %c0, %c0], %c0_i8 {in_bounds = [true, true]} : memref<?x?x8x4xi8, #map0>, vector<8x4xi8>
+    return %result : vector<8x4xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable_with_dynamic_dims_and_indices
+// CHECK-SAME:  %[[ARG0:.+]]: memref<?x?x8x4xi8
+// CHECK:       %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG0]]
+// CHECK:       %[[VEC1D:.+]] = vector.transfer_read %[[COLLAPSED]]
+// CHECK-SAME:  {in_bounds = [true]}
+// CHECK:       %[[VEC2D:.+]] = vector.shape_cast %[[VEC1D]] : vector<32xi8> to vector<8x4xi8>
+// CHECK:       return %[[VEC2D]] : vector<8x4xi8>
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3)[s0, s1] -> (d0 * s1 + s0 + d1 * 32 + d2 * 4 + d3)>
+
+func.func @transfer_read_flattenable_with_dynamic_dims_and_nonzero_static_indices(%arg0 : memref<?x?x8x4xi8, #map0>, %arg1 : index, %arg2 : index) -> vector<5x4xi8> {
+    %c0_i8 = arith.constant 0 : i8
+    %c0 = arith.constant 0 : index
+    %c3 = arith.constant 3 : index
+    %result = vector.transfer_read %arg0[%arg1, %arg2, %c3, %c0], %c0_i8 {in_bounds = [true, true]} : memref<?x?x8x4xi8, #map0>, vector<5x4xi8>
+    return %result : vector<5x4xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable_with_dynamic_dims_and_nonzero_static_indices
+// CHECK-SAME:  %[[ARG0:.+]]: memref<?x?x8x4xi8
+// CHECK:       %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG0]]
+// CHECK:       %[[VEC1D:.+]] = vector.transfer_read %[[COLLAPSED]]
+// CHECK-SAME:  {in_bounds = [true]}
+// CHECK:       %[[VEC2D:.+]] = vector.shape_cast %[[VEC1D]] : vector<20xi8> to vector<5x4xi8>
+// CHECK:       return %[[VEC2D]] : vector<5x4xi8>
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3)[s0, s1] -> (d0 * s1 + s0 + d1 * 32 + d2 * 4 + d3)>
+
+func.func @transfer_read_nonflattenable_due_to_too_large_nonzero_static_indices(%arg0 : memref<?x?x8x4xi8, #map0>, %arg1 : index, %arg2 : index) -> vector<5x4xi8> {
+    %c0_i8 = arith.constant 0 : i8
+    %c0 = arith.constant 0 : index
+    %c3 = arith.constant 4 : index
+    %result = vector.transfer_read %arg0[%arg1, %arg2, %c3, %c0], %c0_i8 {in_bounds = [true, true]} : memref<?x?x8x4xi8, #map0>, vector<5x4xi8>
+    return %result : vector<5x4xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_nonflattenable_due_to_too_large_nonzero_static_indices
+// CHECK-SAME:  %[[ARG0:.+]]: memref<?x?x8x4xi8
+// CHECK:       %[[VEC2D:.+]] = vector.transfer_read %[[ARG0]]
+// CHECK:       return %[[VEC2D]] : vector<5x4xi8>
+
+// -----
+
+#map0 = affine_map<(d0, d1, d2, d3)[s0, s1] -> (d0 * s1 + s0 + d1 * 32 + d2 * 4 + d3)>
+
+func.func @transfer_write_flattenable_with_dynamic_dims_and_indices(%vec : vector<8x4xi8>, %dst : memref<?x?x8x4xi8, #map0>, %arg1 : index, %arg2 : index) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %dst[%arg1, %arg2, %c0, %c0] {in_bounds = [true, true]} : vector<8x4xi8>, memref<?x?x8x4xi8, #map0>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_flattenable_with_dynamic_dims_and_indices
+// CHECK-SAME:  %[[VEC2D:.+]]: vector<8x4xi8>
+// CHECK-SAME:  %[[DST:.+]]: memref<?x?x8x4xi8
+// CHECK:       %[[COLLAPSED:.+]] = memref.collapse_shape %[[DST]]
+// CHECK:       %[[VEC1D:.+]] = vector.shape_cast %[[VEC2D]] : vector<8x4xi8> to vector<32xi8>
+// CHECK:       vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
+// CHECK-SAME:  {in_bounds = [true]}