review comments

Author: Hardcode84

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

@@ -902,6 +902,8 @@ struct BreakDownVectorBitCast : public OpRewritePattern<vector::BitCastOp> {
 };
 
 /// Reorders elementwise(broadcast/splat) to broadcast(elementwise). Ex:
+///
+/// Example:
 /// ```
 /// %a = vector.broadcast %arg1 : index to vector<1x4xindex>
 /// %b = vector.broadcast %arg2 : index to vector<1x4xindex>
@@ -987,6 +989,8 @@ struct ReorderElementwiseOpsOnBroadcast final
 /// This may result in cleaner code when extracting a single value
 /// from multi-element vector and also to help canonicalize 1-element vectors to
 /// scalars.
+///
+/// Example:
 /// ```
 ///  %0 = arith.addf %arg0, %arg1 : vector<4xf32>
 ///  %1 = vector.extract %0[1] : f32 from vector<4xf32>
@@ -1044,6 +1048,8 @@ class ExtractOpFromElementwise final
 };
 
 /// Pattern to rewrite vector.extract(vector.load) -> vector/memref.load.
+///
+/// Example:
 /// ```
 ///  vector.load %arg0[%arg1] : memref<?xf32>, vector<4xf32>
 ///  vector.extract %0[1] : f32 from vector<4xf32>
@@ -1062,13 +1068,14 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
                                 PatternRewriter &rewriter) const override {
     auto loadOp = op.getVector().getDefiningOp<vector::LoadOp>();
     if (!loadOp)
-      return rewriter.notifyMatchFailure(op, "not a load op");
+      return rewriter.notifyMatchFailure(op, "expected a load op");
 
+    // Checking for single use so we won't duplicate load ops.
     if (!loadOp->hasOneUse())
       return rewriter.notifyMatchFailure(op, "expected single op use");
 
-    VectorType memVecType = loadOp.getVectorType();
-    if (memVecType.isScalable())
+    VectorType loadVecType = loadOp.getVectorType();
+    if (loadVecType.isScalable())
       return rewriter.notifyMatchFailure(op,
                                          "scalable vectors are not supported");
 
@@ -1077,7 +1084,7 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
       return rewriter.notifyMatchFailure(
           op, "memrefs of vectors are not supported");
 
-    int64_t rankOffset = memType.getRank() - memVecType.getRank();
+    int64_t rankOffset = memType.getRank() - loadVecType.getRank();
     if (rankOffset < 0)
       return rewriter.notifyMatchFailure(op, "unsupported ranks combination");
 
@@ -1089,6 +1096,9 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
     SmallVector<Value> indices = loadOp.getIndices();
     SmallVector<OpFoldResult> extractPos = op.getMixedPosition();
 
+    // There may be memory stores between the load and the extract op, so we
+    // need to make sure that the new load op is inserted at the same place as
+    // the original load op.
     OpBuilder::InsertionGuard g(rewriter);
     rewriter.setInsertionPoint(loadOp);
     Location loc = loadOp.getLoc();
@@ -1110,12 +1120,15 @@ class ExtractOpFromLoad final : public OpRewritePattern<vector::ExtractOp> {
     } else {
       rewriter.replaceOpWithNewOp<memref::LoadOp>(op, base, indices);
     }
+    // We checked for single use so we can safely erase the load op.
     rewriter.eraseOp(loadOp);
     return success();
   }
 };
 
 /// Pattern to rewrite vector.store(vector.splat) -> vector/memref.store.
+///
+/// Example:
 /// ```
 /// %0 = vector.splat %arg2 : vector<1xf32>
 /// vector.store %0, %arg0[%arg1] : memref<?xf32>, vector<1xf32>
@@ -1145,8 +1158,9 @@ class StoreFromSplat final : public OpRewritePattern<vector::StoreOp> {
 
     Operation *splat = op.getValueToStore().getDefiningOp();
     if (!isa_and_present<vector::BroadcastOp, vector::SplatOp>(splat))
-      return rewriter.notifyMatchFailure(op, "not a splat");
+      return rewriter.notifyMatchFailure(op, "neither a splat nor a broadcast");
 
+    // Checking for single use so we can remove splat.
     if (!splat->hasOneUse())
       return rewriter.notifyMatchFailure(op, "expected single op use");
 

mlir/test/Dialect/Vector/vector-sink.mlir

@@ -587,9 +587,9 @@ func.func @extract_load_vec_high_rank(%arg0: memref<?x?x?xf32>, %arg1: index, %a
   return %1 : vector<4xf32>
 }
 
-// CHECK-LABEL: @negative_load_scalar_from_vec_memref
+// CHECK-LABEL: @negative_extract_load_scalar_from_vec_memref
 //  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xvector<4xf32>>, %[[ARG1:.*]]: index)
-func.func @negative_load_scalar_from_vec_memref(%arg0: memref<?xvector<4xf32>>, %arg1: index) -> f32 {
+func.func @negative_extract_load_scalar_from_vec_memref(%arg0: memref<?xvector<4xf32>>, %arg1: index) -> f32 {
 // CHECK:   %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xvector<4xf32>>, vector<4xf32>
 // CHECK:   %[[EXT:.*]] = vector.extract %[[RES]][0] : f32 from vector<4xf32>
 // CHECK:   return %[[EXT]] : f32
@@ -609,9 +609,9 @@ func.func @negative_extract_load_no_single_use(%arg0: memref<?xf32>, %arg1: inde
   return %1, %0 : f32, vector<4xf32>
 }
 
-// CHECK-LABEL: @negative_load_scalable
+// CHECK-LABEL: @negative_extract_load_scalable
 //  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)
-func.func @negative_load_scalable(%arg0: memref<?xf32>, %arg1: index) -> f32 {
+func.func @negative_extract_load_scalable(%arg0: memref<?xf32>, %arg1: index) -> f32 {
 // CHECK:   %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<[1]xf32>
 // CHECK:   %[[EXT:.*]] = vector.extract %[[RES]][0] : f32 from vector<[1]xf32>
 // CHECK:   return %[[EXT]] : f32
@@ -620,17 +620,6 @@ func.func @negative_load_scalable(%arg0: memref<?xf32>, %arg1: index) -> f32 {
   return %1 : f32
 }
 
-// CHECK-LABEL: @negative_extract_load_unsupported_ranks
-//  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index)
-func.func @negative_extract_load_unsupported_ranks(%arg0: memref<?xf32>, %arg1: index) -> vector<4xf32> {
-// CHECK:   %[[RES:.*]] = vector.load %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<2x4xf32>
-// CHECK:   %[[EXT:.*]] = vector.extract %[[RES]][1] : vector<4xf32> from vector<2x4xf32>
-// CHECK:   return %[[EXT]] : vector<4xf32>
-  %0 = vector.load %arg0[%arg1] : memref<?xf32>, vector<2x4xf32>
-  %1 = vector.extract %0[1] : vector<4xf32> from vector<2x4xf32>
-  return %1 : vector<4xf32>
-}
-
 //-----------------------------------------------------------------------------
 // [Pattern: StoreFromSplat]
 //-----------------------------------------------------------------------------
@@ -653,9 +642,9 @@ func.func @store_broadcast(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {
   return
 }
 
-// CHECK-LABEL: @store_broadcast_1d_2d
+// CHECK-LABEL: @store_broadcast_1d_to_2d
 //  CHECK-SAME:   (%[[ARG0:.*]]: memref<?x?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: index, %[[ARG3:.*]]: vector<1xf32>)
-func.func @store_broadcast_1d_2d(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index, %arg3: vector<1xf32>) {
+func.func @store_broadcast_1d_to_2d(%arg0: memref<?x?xf32>, %arg1: index, %arg2: index, %arg3: vector<1xf32>) {
 // CHECK:   vector.store %[[ARG3]], %[[ARG0]][%[[ARG1]], %[[ARG2]]] : memref<?x?xf32>, vector<1xf32>
   %0 = vector.broadcast %arg3 : vector<1xf32> to vector<1x1xf32>
   vector.store %0, %arg0[%arg1, %arg2] : memref<?x?xf32>, vector<1x1xf32>
@@ -682,9 +671,9 @@ func.func @negative_store_vec_memref(%arg0: memref<?xvector<1xf32>>, %arg1: inde
   return
 }
 
-// CHECK-LABEL: @negative_store_non_1
+// CHECK-LABEL: @negative_store_more_than_one_element
 //  CHECK-SAME:   (%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: index, %[[ARG2:.*]]: f32)
-func.func @negative_store_non_1(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {
+func.func @negative_store_more_than_one_element(%arg0: memref<?xf32>, %arg1: index, %arg2: f32) {
 // CHECK:   %[[RES:.*]] = vector.splat %[[ARG2]] : vector<4xf32>
 // CHECK:   vector.store %[[RES]], %[[ARG0]][%[[ARG1]]] : memref<?xf32>, vector<4xf32>
   %0 = vector.splat %arg2 : vector<4xf32>