Patterns flattening vector transfers to 1D

[Benoit]

Summary

This is needed at the moment to get good codegen from 2d vector.transfer
ops that aim to compile to SIMD load/store instructions but that can
only do so if the whole 2d transfer shape is handled in one piece, in
particular taking advantage of the memref being contiguous rowmajor.

For instance, if the target architecture has 128bit SIMD then we would
expect that contiguous row-major transfers of <4x4xi8> map to one SIMD
load/store instruction each.

The current generic lowering of multi-dimensional vector.transfer ops
can't achieve that because it peels dimensions one by one, so a transfer
of <4x4xi8> becomes 4 transfers of <4xi8>.

The new patterns here are only enabled for now by
-test-vector-transfer-flatten-patterns.

[mravishankar]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3426

I think you can drop the dimsize != 1 condition. The strides[i] != productOfInnerMostSizes should still hold.

[nicolasvasilache]

Cool, thanks much for tackling this, glad to see that it seems to get you to reasonable assembly (inferring from your description?)

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3398

Can we be more specific here: collapseContiguousMemRefCollapseTo1D ?

3402

nit:trivial braces.

3412

This should be sliced a bit differently and better reusing API around BuiltinTypes.h::450.

Basically, getStridesAndOffset is the future-proof way to get the offset and strides.
You want to check the last stride is 1. You can have this as a special helper bool isContiguousMostMinorDimension() or bool isStrideOneMostMinor(int dim) (whichever you find most natural and reusable).

Then you need to determine if the whole type is contiguous.
For this you should add a helper that:

1. returns true empty or identity layout map (the weird trifecta I mentioned previously)
2. then returns false if not fully static
3. then perform a direct MemRefType comparison between a proper usage or getStridesAndOffset, makeStridedLinearLayoutMap, canonicalizeStridedLayout. There should be something similar already using similar logic that can be refactored.

Giving this helper a name and reusing it in multiple places will be a nice cleanup.

3415

Nit: we avoid trivial braces in LLVM, here and below

3423

nit: we use camelCase in LLVM, here and below.

3424

nit dimSize != 1

3447

nit: trivial braces here and below (lift the comment out the the condition)

3473

Nice!

3572

VectorTransforms.cpp is way too bloated and we should split it up (a bit like the standard dialect).
Please either start a new properly named Vector/XXXPatterns.cpp file to put these new patterns or find an existing one.

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

23

You could use the form "memref<4x3x2x1xi8, offset: ?, strides: [6, 2, 1, 1]>"

Due to weird biases, it is currently implemented with an underlying affine_map but the informatio may be clearer like this.
Unfortunately it will still print with affine_map for now.

[Benoit]

add rank reducing subview to drop unit dims

[Benoit]

Thanks for the review comments! I'll apply them on monday. For now I've just updated this diff with the new idea to use rank reducing subviews to drop unit dims, (thanks @ThomasRaoux for the suggestion), that removes my need for https://reviews.llvm.org/D114821 so I can drop it now.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransforms.cpp

3397

Please don't add more patterns to VectorTransforms.cpp, we need to split them out into better isolated logical units.
Either add a new .cpp file at the same level with a proper name (see other XXXPatternXXX.cpp files) or put them
in an already existing such file, depending on what is most appropriate.

3398

This should be its own pattern and return failure when it fails to apply.

3422

You could run this through clang-format. I locally have this in my .bashrc

function git-format-add-and-amend(){
  echo "git add $1 && git show --name-only | egrep \"*.(\.cpp|\.h)\" | xargs -i clang-format --style=file -i {}; git add $1; git commit --amend"
  git add $1 && git show --name-only | egrep "*.(\.cpp|\.h)" | xargs -i clang-format --style=file -i {}; git add $1; git commit --amend
}

3425

This should be its own pattern and return failure when it fails to apply.

3441

static bool isStaticShapeAndContiguousRowMajor(MemRefType memrefType) {

if (!memrefType.hasStaticShape())
  return false;
int64_t offset;
SmallVector<int64_t> strides;
LogicalResult res = getStridesAndOffset(memrefType, strides, offset);
if (failed(res)) 
  return false;

// You may want to improve the APIs here to minimize the code below to something 
// that is expected to be reusable by others.
AffineExpr expr = makeCanonicalStridedLayoutExpr(
  memrefTyp.getSizes(), memrefType.getContext());
MemRefType canonicalMemRefType = MemRefType::get(
  memrefTyp.getSizes(), AffineMap::infer({expr}));

int64_t canonicalOffset;
SmallVector<int64_t> canonicalStrides;
LogicalResult res = getStridesAndOffset(
  canonicalMemRefType, canonicalStrides, canonicalOffset);
if (failed(res)) 
  llvm_unreachable("Unexpected stride extraction error");

for (auto it : llvm::zip(strides, canonicalStrides))
  if (std::get<0>(it) != std::get<1>(it))
    return false;

return true;

}

[Benoit]

apply some review comments

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

213 ↗

(On Diff #393458)

you'll need comments on this function and everywhere below (ideally with a short IR example but not strictly necessary)

[Benoit]

split into 2 patterns

[Benoit]

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

23

Thanks for the tip. For consistency between the MLIR code and the CHECK's I will stick to the affine_map<...> form for now.

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

213 ↗

(On Diff #393479)

Nit, top-levle function and class comments take 3 slashes ///

214 ↗

(On Diff #393479)

nit:camelCase

226 ↗

(On Diff #393479)

use std::copy_if or one of the other stl transforms?

234 ↗

(On Diff #393479)

nice!

245 ↗

(On Diff #393479)

nit: camelCase Drop

250 ↗

(On Diff #393479)

This would only work for the static memref cases.
It feels like you should either early exit if the type is not fully static or use the OpFoldResult-based builders.

258 ↗

(On Diff #393479)

std::count_if or llvm::count_if IIRC

[nicolasvasilache]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

268 ↗

(On Diff #393479)

@gysity, does this relate to the example we just discussed?

[Benoit]

more review comments

[Benoit]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

226 ↗

(On Diff #393479)

I didn't find a simple way to do that, because of how we need to create reducedStrides, not just reducedShape. While it would be possible to let copy_if create reducedShape, I thought that it was more readable if both vectors were created in the same way. In particular, it makes it plain that they have the same length, and that their i-th elements correspond to the same i-th dim.

250 ↗

(On Diff #393479)

The caller already has such an early exit, so I put assertions here.

[gysit]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

268 ↗

(On Diff #393479)

I think @Benoit `s revision addresses a similar topic but at a different level of the stack. I was looking into generating good transfer ops without switching between tensors and vectors in between and making sure rank-reduction works well. @Benoit optimizes the lowering of what we generate higher up in the stack by flattening the vectors, AFAIU. Looking forward to understand the performance implications but I could imagine this helps quite a bit with in combination with hoisting. In particular, for convolutions where we work on very high-dimensional vectors.

[nicolasvasilache]

Thanks for pushing on this @Benoit !

I'd suggest slicing and dicing into smaller commits so we can better track if we ever need to bisect; some of the behavior is tricky to get right and the smaller the CLs + tests, the better we will be reviewing and revisiting in the future.
I think you can turn this into 4-5 commits that can then be more easily clicked.

Also, please don't forget about putting those outside of vectortransforms.cpp.

These are nice developments, sorry it is getting longer than what I think you initially signed for but OTOH you're doing things in the right and future-proof way, so I am grateful!

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

298 ↗

(On Diff #393479)

This is incorrect, the original transferReadOp may have some permutation (i.e. also be sure to insert a proper test).
You want some projection map and compose that with the transfer permutation.

@ThomasRaoux for off-EU-hours advice.

Ah nm, my apologies I see you have a permutation_map test above.
Could you just add a comment/TODO that would highlight this / provision for future work?

299 ↗

(On Diff #393479)

zeros is invalid as we discussed on discord.
You can't go around a project map and applying it to values here.

337 ↗

(On Diff #393479)

same as above re projection map and zeros / identity map.

342 ↗

(On Diff #393479)

This is generally useful and should go to BuiltinTypes.h with proper doc (/// prefix) plz.

352 ↗

(On Diff #393479)

This is generally useful and should go to BuiltinTypes.h with proper doc (/// prefix) plz.

366 ↗

(On Diff #393479)

This is generally useful and should go to BuiltinTypes.h with proper doc (/// prefix) plz.

Also, can we retire helper functions that were not good enough for your use case (if there are too many uses please ignore this last point).

380 ↗

(On Diff #393479)

This should be moved to a proper place in the memref dialect (maybe some utils file and maybe there is already something similar) ?

398 ↗

(On Diff #393479)

I'd use the name "contiguous" in the pattern name and def. in the description otherwise the doc by itself would have invalid assumptions.

431 ↗

(On Diff #393479)

same comment here re 0 and permutation map
this case may be a bit trickier though

441 ↗

(On Diff #393479)

same comments as the above pattern

476 ↗

(On Diff #393479)

same comments re 0 and map

[mravishankar]

mlir/lib/Dialect/Vector/VectorTransferOpTransforms.cpp

243 ↗

(On Diff #393499)

There might be a simpler way to do this. SubViewOp already has constructors to generate rank-reduced subviews. You can use the inferRankReducedSubview method [here[(https://github.com/llvm/llvm-project/blob/7f09aee0f6b4b00508d2cf86b0b1339c8d2ca2d1/mlir/include/mlir/Dialect/MemRef/IR/MemRefOps.td#L1488) and get the result type and use that.

[Benoit]

Apply Mahesh's suggestion to use inferRankReducedResultType.

[Benoit]

More review comments.

[Benoit]

Hi Nicolas, thanks for the kind supportive words here regarding the usefulness of these patterns/helpers.

I think I've addressed the "make it correct" part of your comments, in particular, allZeroConstantIndexValues now checks that the indices of the transfer ops are really all zeros. And I've applied the other "localized" comments that you and Mahesh had (thanks again for those! in particular, Mahesh's comment allowed dropping ~20 lines of code by making dropUnitDims trivial).

I haven't yet addressed your comments about splitting this into multiple commits and about contributing these helpers to core headers:

• Regarding splitting into multiple commits: note that these patterns are so far only enabled in by test-only flags in TestVectorTransforms.cpp, so they are not for now going to affect anything outside these tests. If I understand correctly, the place where granularity will affect how well people can bisect any issues, will be in how we eventually enable these patterns outside of these tests?
• Regarding sharing helpers into core headers: I wonder if this would be best done anyway as a second step after this, and maybe even delay a little further to wait until a second use case arises from someone else, to get more context before blessing a particular helper into a core header? If you feel that you already have enough context to make this call now, that may be a sign that you or someone else with this experience, not I, should be making this move :-)
• As you guessed, I am at this point looking for a time-economical way to wrap up this work :-)

[Benoit]

rebased

[nicolasvasilache]

Sliced a first independent commit and landed as 0aea49a7308322e6987c7b45e4e0d7ab15609e78 as we discussed offline to help you land this.

[nicolasvasilache]

Second part landed as aba437ceb2379f219935b98a10ca3c5081f0c8b7.

Note that I reduced the amount of tests as the combination did not bring much IMO, feel free to disagree and revive some of those.
Also, note that there are now 2 populate functions to get the behavior you wanted, both of the should be called in sequence.

mlir/include/mlir/Dialect/Vector/VectorOps.h

 /// Collect a set of leading one dimension removal patterns.
 ///
 /// These patterns insert vector.shape_cast to remove leading one dimensions
 /// to expose more canonical forms of read/write/insert/extract operations.
 /// With them, there are more chances that we can cancel out extract-insert
 /// pairs or forward write-read pairs.
 void populateCastAwayVectorLeadingOneDimPatterns(RewritePatternSet &patterns);
 
+void populateFlattenVectorTransferPatterns(RewritePatternSet &patterns);
+
 /// Collect a set of patterns that bubble up/down bitcast ops.
 ///
 /// These patterns move vector.bitcast ops to be before insert ops or after
 /// extract ops where suitable. With them, bitcast will happen on smaller
 /// vectors and there are more chances to share extract/insert ops.
 void populateBubbleVectorBitCastOpPatterns(RewritePatternSet &patterns);
 
 /// Collect a set of transfer read/write lowering patterns.

mlir/lib/Dialect/Vector/VectorTransforms.cpp

Value result = rewriter.create<vector::TransferReadOp>(
         loc, resultTargetVecType, rankedReducedView,
         readOp.indices().drop_back(dimsToDrop), permMap, readOp.padding(),
         inBounds);
     rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(readOp, targetType,
                                                      result);
     return success();
   }
 };
 
+static Value rankReducingSubviewDroppingUnitDims(PatternRewriter &rewriter, mlir::Location loc,
+                          Value input) {
+  MemRefType inputType = input.getType().cast<MemRefType>();
+  int64_t offset;
+  SmallVector<int64_t> strides;
+  LogicalResult res = getStridesAndOffset(inputType, strides, offset);
+  assert(succeeded(res));
+  (void) res;
+  ArrayRef<int64_t> shape = inputType.getShape();
+  SmallVector<int64_t> reducedShape;
+  SmallVector<int64_t> reducedStrides;
+  for (unsigned i = 0; i < inputType.getRank(); ++i) {
+    if (shape[i] != 1) {
+      reducedShape.push_back(shape[i]);
+      reducedStrides.push_back(strides[i]);
+    }
+  }
+  if (reducedShape.size() == shape.size())
+    return input;
+  AffineMap reducedStridesMap =
+      makeStridedLinearLayoutMap(reducedStrides, offset, inputType.getContext());
+  MemRefType resultType = MemRefType::Builder(inputType).setShape(reducedShape).setLayout(AffineMapAttr::get(reducedStridesMap));
+  SmallVector<int64_t> subviewOffsets(inputType.getRank(), 0);
+  SmallVector<int64_t> subviewStrides(inputType.getRank(), 1);
+  return rewriter.create<memref::SubViewOp>(loc, resultType, input, subviewOffsets, shape, subviewStrides);
+}
+
+static Value collapseTo1D(PatternRewriter &rewriter, mlir::Location loc,
+                          Value input) {
+  Value rankReducedInput = rankReducingSubviewDroppingUnitDims(rewriter, loc, input);
+  ShapedType rankReducedInputType = rankReducedInput.getType().cast<ShapedType>();
+  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});
+}
+
+// Helper determining if a memref is static-shape and contiguous-row-major
+// layout, still allowing an arbitrary offset (unlike some existing similar
+// functions).
+static bool isStaticShapeAndContiguousRowMajor(MemRefType memrefType) {
+  SmallVector<int64_t> strides;
+  int64_t offset;
+  if (!memrefType.hasStaticShape()) {
+    return false;
+  }
+  LogicalResult res = getStridesAndOffset(memrefType, strides, offset);
+  assert(succeeded(res));
+  (void) res;
+  int64_t productOfInnerMostSizes = 1;
+  for (int i = memrefType.getRank() - 1; i >= 0; --i) {
+    int64_t dimsize = memrefType.getDimSize(i);
+    // The dimsize!=1 condition here means that we ignore the strides of
+    // unit dims, as they don't make a practical difference.
+    if (dimsize != 1 && strides[i] != productOfInnerMostSizes) {
+      return false;
+    }
+    // This simple arithmetic is correct thanks to having ensured above that
+    // we have a static shape.
+    productOfInnerMostSizes *= dimsize;
+  }
+  return true;
+}
+
+class FlattenTransferReadPattern
+    : 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.vector();
+    VectorType vectorType = vector.getType().cast<VectorType>();
+    Value source = transferReadOp.source();
+    MemRefType sourceType = source.getType().cast<MemRefType>();
+    if (vectorType.getRank() == 1 && sourceType.getRank() == 1) {
+      // Already 1D, nothing to do.
+      return failure();
+    }
+    if (!isStaticShapeAndContiguousRowMajor(sourceType)) {
+      return failure();
+    }
+    if (sourceType.getNumElements() != vectorType.getNumElements()) {
+      return failure();
+    }
+    if (transferReadOp.hasOutOfBoundsDim()) {
+      return failure();
+    }
+    if (!transferReadOp.permutation_map().isMinorIdentity()) {
+      return failure();
+    }
+    if (transferReadOp.mask()) {
+      return failure();
+    }
+    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+    auto identityMap1D = rewriter.getMultiDimIdentityMap(1);
+    VectorType vectorType1d = VectorType::get({sourceType.getNumElements()},
+                                              sourceType.getElementType());
+    Value source1d = collapseTo1D(rewriter, loc, source);
+    Value read1d = rewriter.create<vector::TransferReadOp>(
+        loc, vectorType1d, source1d, ValueRange{c0}, identityMap1D);
+    rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
+        transferReadOp, vector.getType().cast<VectorType>(), read1d);
+    return success();
+  }
+};
+
+class FlattenTransferWritePattern
+    : public OpRewritePattern<vector::TransferWriteOp> {
+  using OpRewritePattern<vector::TransferWriteOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(vector::TransferWriteOp transferWriteOp,
+                                PatternRewriter &rewriter) const override {
+    auto loc = transferWriteOp.getLoc();
+    Value vector = transferWriteOp.vector();
+    VectorType vectorType = vector.getType().cast<VectorType>();
+    Value source = transferWriteOp.source();
+    MemRefType sourceType = source.getType().cast<MemRefType>();
+    if (vectorType.getRank() == 1 && sourceType.getRank() == 1) {
+      // Already 1D, nothing to do.
+      return failure();
+    }
+    if (!isStaticShapeAndContiguousRowMajor(sourceType)) {
+      return failure();
+    }
+    if (sourceType.getNumElements() != vectorType.getNumElements()) {
+      return failure();
+    }
+    if (transferWriteOp.hasOutOfBoundsDim()) {
+      return failure();
+    }
+    if (!transferWriteOp.permutation_map().isMinorIdentity()) {
+      return failure();
+    }
+    if (transferWriteOp.mask()) {
+      return failure();
+    }
+    Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
+    auto identityMap1D = rewriter.getMultiDimIdentityMap(1);
+    VectorType vectorType1d = VectorType::get({sourceType.getNumElements()},
+                                              sourceType.getElementType());
+    Value source1d = collapseTo1D(rewriter, loc, source);
+    Value vector1d =
+        rewriter.create<vector::ShapeCastOp>(loc, vectorType1d, vector);
+    rewriter.create<vector::TransferWriteOp>(loc, vector1d, source1d,
+                                             ValueRange{c0}, identityMap1D);
+    rewriter.eraseOp(transferWriteOp);
+    return success();
+  }
+};
+
 void mlir::vector::populateVectorMaskMaterializationPatterns(
     RewritePatternSet &patterns, bool indexOptimizations) {
   patterns.add<VectorCreateMaskOpConversion,
                MaterializeTransferMask<vector::TransferReadOp>,
                MaterializeTransferMask<vector::TransferWriteOp>>(
       patterns.getContext(), indexOptimizations);
 }
 
 void mlir::vector::populatePropagateVectorDistributionPatterns(
     RewritePatternSet &patterns) {
   patterns.add<PointwiseExtractPattern, ContractExtractPattern,
                TransferReadExtractPattern, TransferWriteInsertPattern>(
       patterns.getContext());
 }
 
 void mlir::vector::populateShapeCastFoldingPatterns(
     RewritePatternSet &patterns) {
   patterns.add<ShapeCastOpFolder>(patterns.getContext());
 }
 
+void mlir::vector::populateFlattenVectorTransferPatterns(
+    RewritePatternSet &patterns) {
+  patterns.add<FlattenTransferReadPattern, FlattenTransferWritePattern>(
+      patterns.getContext());
+  populateShapeCastFoldingPatterns(patterns);
+}
+
 void mlir::vector::populateBubbleVectorBitCastOpPatterns(
     RewritePatternSet &patterns) {
   patterns.add<BubbleDownVectorBitCastForExtract,
                BubbleDownBitCastForStridedSliceExtract,
                BubbleUpBitCastForStridedSliceInsert>(patterns.getContext());
 }
 
 void mlir::vector::populateVectorBroadcastLoweringPatterns(

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

+// RUN: mlir-opt %s -test-vector-transfer-flatten-patterns -split-input-file | FileCheck %s
+
+func @transfer_read_flattenable(%arg : memref<5x4x3x2xi8>) -> vector<5x4x3x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : memref<5x4x3x2xi8>, vector<5x4x3x2xi8>
+    return %v : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8>
+// CHECK:         %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]{{.}} : memref<5x4x3x2xi8> into memref<120xi8>
+// CHECK:         %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]
+// CHECK:         %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<120xi8> to vector<5x4x3x2xi8>
+// CHECK:         return %[[VEC2D]]
+
+// -----
+
+func @transfer_read_flattenable_with_offset(%arg : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 24 + d1 * 6 + d2 * 2 + d3 + s0)>>) -> vector<5x4x3x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 24 + d1 * 6 + d2 * 2 + d3 + s0)>>, vector<5x4x3x2xi8>
+    return %v : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable_with_offset
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8
+// CHECK:         %[[COLLAPSED:.+]] = memref.collapse_shape %[[ARG]] {{.}}[0, 1, 2, 3]
+// CHECK:         %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]
+// CHECK:         %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<120xi8> to vector<5x4x3x2xi8>
+// CHECK:         return %[[VEC2D]]
+
+// -----
+
+func @transfer_read_flattenable_with_offset_with_rank_reducing_subview(%arg : memref<1x1x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 6 + d1 * 6 + d2 * 2 + d3 + s0)>>) -> vector<3x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : memref<1x1x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 6 + d1 * 6 + d2 * 2 + d3 + s0)>>, vector<3x2xi8>
+    return %v : vector<3x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable_with_offset
+// CHECK-SAME:      %[[ARG:.+]]: memref<1x1x3x2xi8
+// CHECK:         %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 3, 2] [1, 1, 1, 1]
+// CHECK:         %[[COLLAPSED:.+]] = memref.collapse_shape %[[SUBVIEW]] {{.}}[0, 1]
+// CHECK:         %[[READ1D:.+]] = vector.transfer_read %[[COLLAPSED]]
+// CHECK:         %[[VEC2D:.+]] = vector.shape_cast %[[READ1D]] : vector<6xi8> to vector<3x2xi8>
+// CHECK:         return %[[VEC2D]]
+
+// -----
+
+func @transfer_read_flattenable_with_offset_with_rank_reducing_subview_and_no_collapse(%arg : memref<1x1x1x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 2 + d1 * 2 + d2 * 2 + d3 + s0)>>) -> vector<2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : memref<1x1x1x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 2 + d1 * 2 + d2 * 2 + d3 + s0)>>, vector<2xi8>
+    return %v : vector<2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_flattenable_with_offset
+// CHECK-SAME:      %[[ARG:.+]]: memref<1x1x1x2xi8
+// CHECK:         %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 1, 2] [1, 1, 1, 1]
+// CHECK:         %[[READ1D:.+]] = vector.transfer_read %[[SUBVIEW]]
+// CHECK:         return %[[READ1D]]
+
+// -----
+
+func @transfer_read_nonflattenable_out_of_bounds(%arg : memref<5x4x3x2xi8>, %i : index) -> vector<5x4x3x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%i, %c0, %c0, %c0], %cst {in_bounds = [false, true, true, true]} : memref<5x4x3x2xi8>, vector<5x4x3x2xi8>
+    return %v : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_nonflattenable_out_of_bounds
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8>,
+// CHECK-SAME:      %[[I:.+]]: index
+// CHECK:         %[[READ:.+]] = vector.transfer_read %[[ARG]][%[[I]]
+// CHECK:         return %[[READ]]
+
+// -----
+
+func @transfer_read_nonflattenable_non_contiguous(%arg : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 * 25 + d1 * 6 + d2 * 2 + d3)>>) -> vector<5x4x3x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 * 25 + d1 * 6 + d2 * 2 + d3)>>, vector<5x4x3x2xi8>
+    return %v : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_nonflattenable_non_contiguous
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8,
+// CHECK:         %[[READ:.+]] = vector.transfer_read %[[ARG]]
+// CHECK:         return %[[READ]]
+
+// -----
+
+func @transfer_read_nonflattenable_non_row_major(%arg : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 + d1 * 5 + d2 * 20 + d3 * 60)>>) -> vector<5x4x3x2xi8> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0 : i8
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0], %cst : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 + d1 * 5 + d2 * 20 + d3 * 60)>>, vector<5x4x3x2xi8>
+    return %v : vector<5x4x3x2xi8>
+}
+
+// CHECK-LABEL: func @transfer_read_nonflattenable_non_row_major
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8,
+// CHECK:         %[[READ:.+]] = vector.transfer_read %[[ARG]]
+// CHECK:         return %[[READ]]
+
+// -----
+
+func @transfer_write_flattenable(%arg : memref<5x4x3x2xi8>, %vec : vector<5x4x3x2xi8>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] : vector<5x4x3x2xi8>, memref<5x4x3x2xi8>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_flattenable
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8>,
+// CHECK-SAME:      %[[VEC:.+]]: 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 @transfer_write_flattenable_with_offset(%arg : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 24 + d1 * 6 + d2 * 2 + d3 + s0)>>, %vec : vector<5x4x3x2xi8>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] : vector<5x4x3x2xi8>, memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 24 + d1 * 6 + d2 * 2 + d3 + s0)>>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_flattenable_with_offset
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8, {{.+}}>,
+// CHECK-SAME:      %[[VEC:.+]]: 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 @transfer_write_flattenable_with_offset_with_rank_reducing_subview(%arg : memref<1x1x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 6 + d1 * 6 + d2 * 2 + d3 + s0)>>, %vec : vector<3x2xi8>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] : vector<3x2xi8>, memref<1x1x3x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 6 + d1 * 6 + d2 * 2 + d3 + s0)>>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_flattenable_with_offset_with_rank_reducing_subview
+// CHECK-SAME:      %[[ARG:.+]]: memref<1x1x3x2xi8, {{.+}}>,
+// CHECK-SAME:      %[[VEC:.+]]: vector<3x2xi8>
+// CHECK:         %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 3, 2] [1, 1, 1, 1]
+// CHECK-DAG:     %[[COLLAPSED:.+]] = memref.collapse_shape %[[SUBVIEW]] {{.}}[0, 1]{{.}} : memref<3x2xi8, {{.+}}> into memref<6xi8, {{.+}}>
+// CHECK-DAG:     %[[VEC1D:.+]] = vector.shape_cast %[[VEC]] : vector<3x2xi8> to vector<6xi8>
+// CHECK:         vector.transfer_write %[[VEC1D]], %[[COLLAPSED]]
+
+// -----
+
+func @transfer_write_flattenable_with_offset_with_rank_reducing_subview_and_no_collapse(%arg : memref<1x1x1x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 2 + d1 * 2 + d2 * 2 + d3 + s0)>>, %vec : vector<2xi8>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] : vector<2xi8>, memref<1x1x1x2xi8, affine_map<(d0, d1, d2, d3)[s0] -> (d0 * 2 + d1 * 2 + d2 * 2 + d3 + s0)>>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_flattenable_with_offset_with_rank_reducing_subview_and_no_collapse
+// CHECK-SAME:      %[[ARG:.+]]: memref<1x1x1x2xi8, {{.+}}>,
+// CHECK-SAME:      %[[VEC:.+]]: vector<2xi8>
+// CHECK:         %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 1, 2] [1, 1, 1, 1]
+// CHECK:         vector.transfer_write %[[VEC]], %[[SUBVIEW]]
+
+// -----
+
+func @transfer_write_nonflattenable_out_of_bounds(%arg : memref<5x4x3x2xi8>, %vec : vector<5x4x3x2xi8>, %i : index) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg [%i, %c0, %c0, %c0] {in_bounds = [false, true, true, true]} : vector<5x4x3x2xi8>, memref<5x4x3x2xi8>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_nonflattenable_out_of_bounds
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8>,
+// CHECK-SAME:      %[[VEC:.+]]: vector<5x4x3x2xi8>
+// CHECK-SAME:      %[[I:.+]]: index
+// CHECK:         vector.transfer_write %[[VEC]], %[[ARG]]
+
+// -----
+
+func @transfer_write_nonflattenable_non_contiguous(%arg : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 * 25 + d1 * 6 + d2 * 2 + d3)>>, %vec : vector<5x4x3x2xi8>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg[%c0, %c0, %c0, %c0] : vector<5x4x3x2xi8>, memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 * 25 + d1 * 6 + d2 * 2 + d3)>>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_nonflattenable_non_contiguous
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8,
+// CHECK-SAME:      %[[VEC:.+]]: vector<5x4x3x2xi8>
+// CHECK:         vector.transfer_write %[[VEC]], %[[ARG]]
+
+// -----
+
+func @transfer_write_nonflattenable_non_row_major(%arg : memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 + d1 * 4 + d2 * 12 + d3 * 24)>>, %vec : vector<5x4x3x2xi8>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg[%c0, %c0, %c0, %c0] : vector<5x4x3x2xi8>, memref<5x4x3x2xi8, affine_map<(d0, d1, d2, d3) -> (d0 + d1 * 4 + d2 * 12 + d3 * 24)>>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_nonflattenable_non_row_major
+// CHECK-SAME:      %[[ARG:.+]]: memref<5x4x3x2xi8,
+// CHECK-SAME:      %[[VEC:.+]]: vector<5x4x3x2xi8>
+// CHECK:         vector.transfer_write %[[VEC]], %[[ARG]]

mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp

struct TestVectorReduceToContractPatternsPatterns
   }
   void runOnFunction() override {
     RewritePatternSet patterns(&getContext());
     populateVectorReductionToContractPatterns(patterns);
     (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
   }
 };
 
+struct TestFlattenVectorTransferPatterns
+    : public PassWrapper<TestFlattenVectorTransferPatterns, FunctionPass> {
+  StringRef getArgument() const final {
+    return "test-vector-transfer-flatten-patterns";
+  }
+  StringRef getDescription() const final {
+    return "Test patterns to rewrite contiguous row-major N-dimensional "
+           "vector.transfer_{read,write} ops into 1D transfers";
+  }
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<memref::MemRefDialect>();
+  }
+  void runOnFunction() override {
+    RewritePatternSet patterns(&getContext());
+    populateFlattenVectorTransferPatterns(patterns);
+    (void)applyPatternsAndFoldGreedily(getFunction(), std::move(patterns));
+  }
+};
+
 } // end anonymous namespace
 
 namespace mlir {
 namespace test {
 void registerTestVectorLowerings() {
   PassRegistration<TestVectorToVectorLowering>();
 
   PassRegistration<TestVectorContractionLowering>();
void registerTestVectorLowerings() {
 
   PassRegistration<TestVectorTransferLoweringPatterns>();
 
   PassRegistration<TestVectorMultiReductionLoweringPatterns>();
 
   PassRegistration<TestVectorTransferCollapseInnerMostContiguousDims>();
 
   PassRegistration<TestVectorReduceToContractPatternsPatterns>();
+
+  PassRegistration<TestFlattenVectorTransferPatterns>();
 }
 } // namespace test
 } // namespace mlir