style fixes

Author: Hardcode84

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

@@ -1047,6 +1047,16 @@ struct ReorderElementwiseOpsOnBroadcast final
 /// This may result in more efficient code when we extracting a single value
 /// from multi-element vector and also to help canonicalize 1-element vectors to
 /// scalars.
+/// ```
+///  %0 = arith.addf %arg0, %arg1 : vector<4xf32>
+///  %1 = vector.extract %0[1] : f32 from vector<4xf32>
+/// ```
+/// Gets converted to:
+/// ```
+///  %0 = vector.extract %arg0[1] : f32 from vector<4xf32>
+///  %1 = vector.extract %arg1[1] : f32 from vector<4xf32>
+///  %2 = arith.addf %0, %1 : f32
+/// ```
 class ExtractOpFromElementwise final
     : public OpRewritePattern<vector::ExtractOp> {
 public:
@@ -1061,7 +1071,7 @@ class ExtractOpFromElementwise final
         eltwise->getNumResults() != 1 || !eltwise->hasOneUse())
       return failure();
 
-    // Arguments and result types must match.
+    // Arguments types must match.
     if (!llvm::all_equal(eltwise->getOperandTypes()))
       return failure();
 
@@ -1072,7 +1082,7 @@ class ExtractOpFromElementwise final
 
     IRMapping mapping;
     Location loc = eltwise->getLoc();
-    for (auto &&[i, arg] : llvm::enumerate(eltwise->getOperands())) {
+    for (auto arg : eltwise->getOperands()) {
       Value newArg =
           rewriter.create<vector::ExtractOp>(loc, arg, op.getMixedPosition());
       mapping.map(arg, newArg);

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

@@ -13,9 +13,9 @@ module attributes {transform.with_named_sequence} {
 }
 
 
-// CHECK-LABEL: @extract_elementwise
+// CHECK-LABEL: @extract_elementwise_scalar
 //  CHECK-SAME:   (%[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>)
-func.func @extract_elementwise(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f32 {
+func.func @extract_elementwise_scalar(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f32 {
 // CHECK:   %[[EXT0:.*]] = vector.extract %[[ARG0]][1] : f32 from vector<4xf32>
 // CHECK:   %[[EXT1:.*]] = vector.extract %[[ARG1]][1] : f32 from vector<4xf32>
 // CHECK:   %[[RES:.*]] = arith.addf %[[EXT0]], %[[EXT1]] : f32

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

@@ -426,9 +426,9 @@ func.func @transpose_elementwise_diff_map_scalable(%a : vector<[4]x6x3x2xf32>, %
 
 // -----
 
-// CHECK-LABEL: @extract_elementwise
+// CHECK-LABEL: @extract_elementwise_scalar
 //  CHECK-SAME:   (%[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>)
-func.func @extract_elementwise(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f32 {
+func.func @extract_elementwise_scalar(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f32 {
 // CHECK:   %[[EXT0:.*]] = vector.extract %[[ARG0]][1] : f32 from vector<4xf32>
 // CHECK:   %[[EXT1:.*]] = vector.extract %[[ARG1]][1] : f32 from vector<4xf32>
 // CHECK:   %[[RES:.*]] = arith.addf %[[EXT0]], %[[EXT1]] : f32
@@ -438,9 +438,9 @@ func.func @extract_elementwise(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> f3
   return %1 : f32
 }
 
-// CHECK-LABEL: @extract_vec_elementwise
+// CHECK-LABEL: @extract_elementwise_vec
 //  CHECK-SAME:   (%[[ARG0:.*]]: vector<2x4xf32>, %[[ARG1:.*]]: vector<2x4xf32>)
-func.func @extract_vec_elementwise(%arg0: vector<2x4xf32>, %arg1: vector<2x4xf32>) -> vector<4xf32> {
+func.func @extract_elementwise_vec(%arg0: vector<2x4xf32>, %arg1: vector<2x4xf32>) -> vector<4xf32> {
 // CHECK:   %[[EXT0:.*]] = vector.extract %[[ARG0]][1] : vector<4xf32> from vector<2x4xf32>
 // CHECK:   %[[EXT1:.*]] = vector.extract %[[ARG1]][1] : vector<4xf32> from vector<2x4xf32>
 // CHECK:   %[[RES:.*]] = arith.addf %[[EXT0]], %[[EXT1]] : vector<4xf32>
@@ -450,9 +450,9 @@ func.func @extract_vec_elementwise(%arg0: vector<2x4xf32>, %arg1: vector<2x4xf32
   return %1 : vector<4xf32>
 }
 
-// CHECK-LABEL: @extract_elementwise_use
+// CHECK-LABEL: @extract_elementwise_no_single_use
 //  CHECK-SAME:   (%[[ARG0:.*]]: vector<4xf32>, %[[ARG1:.*]]: vector<4xf32>)
-func.func @extract_elementwise_use(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {
+func.func @extract_elementwise_no_single_use(%arg0: vector<4xf32>, %arg1: vector<4xf32>) -> (f32, vector<4xf32>) {
 // Do not propagate extract, as elementwise has other uses.
 // CHECK:   %[[ELT:.*]] = arith.addf %[[ARG0]], %[[ARG1]] : vector<4xf32>
 // CHECK:   %[[EXT:.*]] = vector.extract %[[ELT]][1] : f32 from vector<4xf32>