[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

288

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

293–298

ditto, add a static keyword.

349–358

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

297

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

288

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

297

Sure, let me try

[hanchung]

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

288

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

 private:
   RewriterBase &rewriter;
   bool isReachable(Operation *start, Operation *dest);
   DominanceInfo dominators;
   PostDominanceInfo postDominators;
   std::vector<Operation *> opToErase;
 };
 
+} // namespace
 /// Return true if there is a path from start operation to dest operation,
 /// otherwise return false. The operations have to be in the same region.
 bool TransferOptimization::isReachable(Operation *start, Operation *dest) {
   assert(start->getParentRegion() == dest->getParentRegion() &&
          "This function only works for ops i the same region");
   // Simple case where the start op dominate the destination.
   if (dominators.dominates(start, dest))
     return true;
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) {
   return llvm::count_if(shape, [](int64_t dimSize) { return dimSize != 1; });
 }
 
+/// Returns a copy of `shape` without unit dims.
+static 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.
+namespace {
+
+/// 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);
+    auto 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();
   }
 };
 
+} // namespace
+
 /// 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;
   if (!succeeded(getStridesAndOffset(memrefType, strides, offset)))
for (int64_t i = firstDimToCollapse; i < rank; ++i) {
     if (!cst || cst.value() != 0)
       return failure();
   }
   outIndices = indices;
   outIndices.resize(firstDimToCollapse + 1);
   return success();
 }
 
+namespace {
+
 /// 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::OpRewritePattern;
 
if (isa<MemRefType>(xferOp.getSource().getType())) {
           xferOp, scalar, xferOp.getSource(), xferOp.getIndices());
     } else {
       rewriter.replaceOpWithNewOp<tensor::InsertOp>(
           xferOp, scalar, xferOp.getSource(), xferOp.getIndices());
     }
     return success();
   }
 };
+
 } // namespace
 
 void mlir::vector::transferOpflowOpt(RewriterBase &rewriter,
                                      Operation *rootOp) {
   TransferOptimization opt(rewriter, rootOp);
   // Run store to load forwarding first since it can expose more dead store
   // opportunity.
   rootOp->walk([&](vector::TransferReadOp read) {

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: !transform.any_op):
+  transform.vector.apply_rank_reducing_subview_patterns %module_op
+      : (!transform.any_op) -> !transform.any_op
+}
+
+// -----
+
+func.func @transfer_read_and_vector_rank_reducing_to_0d(
+      %arg : memref<1x1x1x1x1xf32>) -> vector<1x1x1xf32> {
+    %c0 = arith.constant 0 : index
+    %cst = arith.constant 0.0 : f32
+    %v = vector.transfer_read %arg[%c0, %c0, %c0, %c0, %c0], %cst :
+      memref<1x1x1x1x1xf32>, vector<1x1x1xf32>
+    return %v : vector<1x1x1xf32>
+}
+
+// CHECK-LABEL: func @transfer_read_and_vector_rank_reducing_to_0d
+//  CHECK-SAME:     %[[MEMREF:.+]]: memref<1x1x1x1x1xf32>
+//       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[MEMREF]][0, 0, 0, 0, 0] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : memref<1x1x1x1x1xf32> to memref<f32>
+//       CHECK:   %[[READ:.+]] = vector.transfer_read %[[SUBVIEW]]{{.*}} : memref<f32>, vector<f32>
+//       CHECK:   vector.shape_cast %[[READ]] : vector<f32> to vector<1x1x1xf32>
+
+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_to_0d(
+      %arg : memref<1x1x1x1x1xf32>,
+      %vec : vector<1x1x1xf32>) {
+    %c0 = arith.constant 0 : index
+    vector.transfer_write %vec, %arg [%c0, %c0, %c0, %c0, %c0] :
+      vector<1x1x1xf32>, memref<1x1x1x1x1xf32>
+    return
+}
+
+// CHECK-LABEL: func @transfer_write_and_vector_rank_reducing_to_0d
+//  CHECK-SAME:     %[[MEMREF:.+]]: memref<1x1x1x1x1xf32>, %[[VECTOR:.+]]: vector<1x1x1xf32>
+//       CHECK:   %[[SUBVIEW:.+]] = memref.subview %[[MEMREF]][0, 0, 0, 0, 0] [1, 1, 1, 1, 1] [1, 1, 1, 1, 1] : memref<1x1x1x1x1xf32> to memref<f32>
+//       CHECK:   %[[SHCAST:.+]] = vector.shape_cast %[[VECTOR]] : vector<1x1x1xf32> to vector<f32>
+//       CHECK:   vector.transfer_write %[[SHCAST]], %[[SUBVIEW]]{{.*}} : vector<f32>, memref<f32>
 
 transform.sequence failures(propagate) {
 ^bb1(%module_op: !transform.any_op):
   transform.vector.apply_rank_reducing_subview_patterns %module_op
       : (!transform.any_op) -> !transform.any_op
 }