[mlir][Vector] Extend xfer drop unit dim patterns

[dcaballe]

Summary

This patch extends the transfer drop unit dim patterns to support cases where the vector shape should also be reduced
(e.g., transfer_read(memref<1x4x1xf32>, vector<1x4x1xf32>) -> transfer_read(memref<4xf32>, vector<4xf32>).

[hanchung]

Overall looks good to me, thanks! Just few nits.

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

287

I know that this is wrapped within anonymous namespace, so we don't need the static keyword. The style guide points out that we should use static keyword even if they are within anonymous namespace.

They reduce locality of reference: if you see a random function definition in a C++ file, it is easy to see if it is marked static, but seeing if it is in an anonymous namespace requires scanning a big chunk of the file.

https://llvm.org/docs/CodingStandards.html#anonymous-namespaces

292–297

ditto, add a static keyword.

346–347

style nit: use auto. Because ::get method already spells the type.

https://llvm.org/docs/CodingStandards.html#use-auto-type-deduction-to-make-code-more-readable

[pzread]

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

296

Can you add a test to check when we have the vector shape <1x1x1xf32>? I ran into a problem in that case since the pass reduces them into zero-dim vectors.

[dcaballe]

Address feedback

[dcaballe]

Landing. Thanks!

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

287

I think the doc means that we should move these cases out of the anonymous namespace and use static instead but not using both at the same time. Fixed!

[dcaballe]

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

296

Sure, let me try

[hanchung]

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

287

Good point.. You're right! (My point was echoing the "scanning a big chunk of the file" point. And yes, it's not saying that we should use it at the same time.)

[dcaballe]

Support and test 0-d

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

static Value rankReducingSubviewDroppingUnitDims(PatternRewriter &rewriter,
   if (canonicalizeStridedLayout(resultType) ==
       canonicalizeStridedLayout(inputType))
     return input;
   return rewriter.create<memref::SubViewOp>(
       loc, resultType, input, subViewOffsets, subViewSizes, subViewStrides);
 }
 
 /// Returns the number of dims that aren't unit dims.
-static int getReducedRank(ArrayRef<int64_t> shape) {
+int getReducedRank(ArrayRef<int64_t> shape) {
   return llvm::count_if(shape, [](int64_t dimSize) { return dimSize != 1; });
 }
 
+/// Returns a copy of `shape` without unit dims.
+SmallVector<int64_t> getReducedShape(ArrayRef<int64_t> shape) {
+  SmallVector<int64_t> reducedShape;
+  llvm::copy_if(shape, std::back_inserter(reducedShape),
+                [](int64_t dimSize) { return dimSize != 1; });
+  return reducedShape;
+}
+
 /// Returns true if all values are `arith.constant 0 : index`
 static bool isZero(Value v) {
   auto cst = v.getDefiningOp<arith::ConstantIndexOp>();
   return cst && cst.value() == 0;
 }
 
-/// Rewrites vector.transfer_read ops where the source has unit dims, by
-/// inserting a memref.subview dropping those unit dims.
+/// Rewrites `vector.transfer_read` ops where the source has unit dims, by
+/// inserting a memref.subview dropping those unit dims. The vector shapes are
+/// also reduced accordingly.
 class TransferReadDropUnitDimsPattern
     : public OpRewritePattern<vector::TransferReadOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferReadOp transferReadOp,
                                 PatternRewriter &rewriter) const override {
     auto loc = transferReadOp.getLoc();
     Value vector = transferReadOp.getVector();
     VectorType vectorType = cast<VectorType>(vector.getType());
     Value source = transferReadOp.getSource();
     MemRefType sourceType = dyn_cast<MemRefType>(source.getType());
     // TODO: support tensor types.
     if (!sourceType || !sourceType.hasStaticShape())
       return failure();
     if (sourceType.getNumElements() != vectorType.getNumElements())
       return failure();
     // TODO: generalize this pattern, relax the requirements here.
     if (transferReadOp.hasOutOfBoundsDim())
       return failure();
     if (!transferReadOp.getPermutationMap().isMinorIdentity())
       return failure();
+    // Check if the source shape can be further reduced.
     int reducedRank = getReducedRank(sourceType.getShape());
     if (reducedRank == sourceType.getRank())
-      return failure(); // The source shape can't be further reduced.
-    if (reducedRank != vectorType.getRank())
-      return failure(); // This pattern requires the vector shape to match the
-                        // reduced source shape.
+      return failure();
+    // Check if the reduced vector shape matches the reduced source shape.
+    // Otherwise, this case is not supported yet.
+    int vectorReducedRank = getReducedRank(vectorType.getShape());
+    if (reducedRank != vectorReducedRank)
+      return failure();
     if (llvm::any_of(transferReadOp.getIndices(),
                      [](Value v) { return !isZero(v); }))
       return failure();
     Value reducedShapeSource =
         rankReducingSubviewDroppingUnitDims(rewriter, loc, source);
     Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     SmallVector<Value> zeros(reducedRank, c0);
     auto identityMap = rewriter.getMultiDimIdentityMap(reducedRank);
-    rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
-        transferReadOp, vectorType, reducedShapeSource, zeros, identityMap);
+    VectorType reducedVectorType = VectorType::get(
+        getReducedShape(vectorType.getShape()), vectorType.getElementType());
+
+    auto newTransferReadOp = rewriter.create<vector::TransferReadOp>(
+        loc, reducedVectorType, reducedShapeSource, zeros, identityMap);
+    auto shapeCast = rewriter.createOrFold<vector::ShapeCastOp>(
+        loc, vectorType, newTransferReadOp);
+    rewriter.replaceOp(transferReadOp, shapeCast);
+
     return success();
   }
 };
 
-/// Rewrites vector.transfer_write ops where the "source" (i.e. destination) has
-/// unit dims, by inserting a memref.subview dropping those unit dims.
+/// Rewrites `vector.transfer_write` ops where the "source" (i.e. destination)
+/// has unit dims, by inserting a `memref.subview` dropping those unit dims. The
+/// vector shapes are also reduced accordingly.
 class TransferWriteDropUnitDimsPattern
     : public OpRewritePattern<vector::TransferWriteOp> {
   using OpRewritePattern::OpRewritePattern;
 
   LogicalResult matchAndRewrite(vector::TransferWriteOp transferWriteOp,
                                 PatternRewriter &rewriter) const override {
     auto loc = transferWriteOp.getLoc();
     Value vector = transferWriteOp.getVector();
     VectorType vectorType = cast<VectorType>(vector.getType());
     Value source = transferWriteOp.getSource();
     MemRefType sourceType = dyn_cast<MemRefType>(source.getType());
     // TODO: support tensor type.
     if (!sourceType || !sourceType.hasStaticShape())
       return failure();
     if (sourceType.getNumElements() != vectorType.getNumElements())
       return failure();
     // TODO: generalize this pattern, relax the requirements here.
     if (transferWriteOp.hasOutOfBoundsDim())
       return failure();
     if (!transferWriteOp.getPermutationMap().isMinorIdentity())
       return failure();
+    // Check if the destination shape can be further reduced.
     int reducedRank = getReducedRank(sourceType.getShape());
     if (reducedRank == sourceType.getRank())
-      return failure(); // The source shape can't be further reduced.
-    if (reducedRank != vectorType.getRank())
-      return failure(); // This pattern requires the vector shape to match the
-                        // reduced source shape.
+      return failure();
+    // Check if the reduced vector shape matches the reduced destination shape.
+    // Otherwise, this case is not supported yet.
+    int vectorReducedRank = getReducedRank(vectorType.getShape());
+    if (reducedRank != vectorReducedRank)
+      return failure();
     if (llvm::any_of(transferWriteOp.getIndices(),
                      [](Value v) { return !isZero(v); }))
       return failure();
     Value reducedShapeSource =
         rankReducingSubviewDroppingUnitDims(rewriter, loc, source);
     Value c0 = rewriter.create<arith::ConstantIndexOp>(loc, 0);
     SmallVector<Value> zeros(reducedRank, c0);
     auto identityMap = rewriter.getMultiDimIdentityMap(reducedRank);
+    VectorType reducedVectorType = VectorType::get(
+        getReducedShape(vectorType.getShape()), vectorType.getElementType());
+
+    auto shapeCast = rewriter.createOrFold<vector::ShapeCastOp>(
+        loc, reducedVectorType, vector);
     rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
-        transferWriteOp, vector, reducedShapeSource, zeros, identityMap);
+        transferWriteOp, shapeCast, reducedShapeSource, zeros, identityMap);
+
     return success();
   }
 };
 
 /// 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) {

mlir/test/Dialect/Vector/vector-transfer-drop-unit-dims-patterns.mlir

 }
 
 // CHECK-LABEL: func @transfer_read_rank_reducing
 //  CHECK-SAME:     %[[ARG:.+]]: memref<1x1x3x2xi8
 //       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 3, 2] [1, 1, 1, 1]
 //  CHECK-SAME:     memref<1x1x3x2xi8, {{.*}}> to memref<3x2xi8, {{.*}}>
 //       CHECK:   vector.transfer_read %[[SUBVIEW]]
 
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  transform.vector.apply_rank_reducing_subview_patterns %module_op
+      : (!pdl.operation) -> !pdl.operation
+}
+
+// -----
+
 func.func @transfer_write_rank_reducing(%arg : memref<1x1x3x2xi8, strided<[6, 6, 2, 1], offset: ?>>, %vec : vector<3x2xi8>) {
     %c0 = arith.constant 0 : index
     vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0] :
       vector<3x2xi8>, memref<1x1x3x2xi8, strided<[6, 6, 2, 1], offset: ?>>
     return
 }
 
 // CHECK-LABEL: func @transfer_write_rank_reducing
 //  CHECK-SAME:     %[[ARG:.+]]: memref<1x1x3x2xi8
 //       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0] [1, 1, 3, 2] [1, 1, 1, 1]
 //  CHECK-SAME:     memref<1x1x3x2xi8, {{.*}}> to memref<3x2xi8, {{.*}}>
 //       CHECK:   vector.transfer_write %{{.*}}, %[[SUBVIEW]]
 
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  transform.vector.apply_rank_reducing_subview_patterns %module_op
+      : (!pdl.operation) -> !pdl.operation
+}
+
+// -----
+
+func.func @transfer_read_and_vector_rank_reducing(
+      %arg : memref<1x1x3x2x1xf32>) -> vector<3x2x1xf32> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0.0 : f32
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0, %c0], %cst :
+      memref<1x1x3x2x1xf32>, vector<3x2x1xf32>
+    return %v : vector<3x2x1xf32>
+}
+
+// CHECK-LABEL: func @transfer_read_and_vector_rank_reducing
+//  CHECK-SAME:     %[[ARG:.+]]: memref<1x1x3x2x1xf32>
+//       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0, 0] [1, 1, 3, 2, 1] [1, 1, 1, 1, 1]
+//  CHECK-SAME:     memref<1x1x3x2x1xf32> to memref<3x2xf32>
+//       CHECK:   vector.transfer_read %[[SUBVIEW]]{{.*}} {in_bounds = [true, true]} : memref<3x2xf32>, vector<3x2xf32>
+
+transform.sequence failures(propagate) {
+^bb1(%module_op: !pdl.operation):
+  transform.vector.apply_rank_reducing_subview_patterns %module_op
+      : (!pdl.operation) -> !pdl.operation
+}
+
+// -----
+
+func.func @transfer_write_and_vector_rank_reducing(
+      %arg : memref<1x1x3x2x1xf32>,
+      %vec : vector<3x2x1xf32>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0, %c0] :
+      vector<3x2x1xf32>, memref<1x1x3x2x1xf32>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_and_vector_rank_reducing
+//  CHECK-SAME:     %[[ARG:.+]]: memref<1x1x3x2x1xf32>
+//       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[ARG]][0, 0, 0, 0, 0] [1, 1, 3, 2, 1] [1, 1, 1, 1, 1]
+//  CHECK-SAME:     memref<1x1x3x2x1xf32> to memref<3x2xf32>
+//       CHECK:   vector.transfer_write %{{.*}}, %[[SUBVIEW]]{{.*}} {in_bounds = [true, true]} : vector<3x2xf32>, memref<3x2xf32>
 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !pdl.operation):
   transform.vector.apply_rank_reducing_subview_patterns %module_op
       : (!pdl.operation) -> !pdl.operation
 }