simplify and add one test

Author: newling

mlir/lib/Dialect/Vector/IR/VectorOps.cpp

@@ -6216,10 +6216,10 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
       bool prevNotOne = (inputIndex != 0 && inputShape[inputIndex - 1] != 1);
       bool groupEndFound = notOne || prevNotOne;
       if (groupEndFound) {
-        // Return failure if all not permutation destinations for indices in
+        int high = inputIndex + deltaRank;
+        // Return failure if not all permutation destinations for indices in
         // [low, high) are in [low, high), i.e. the permutation is not local to
         // the group.
-        int high = inputIndex + deltaRank;
         for (int i = low; i < high; ++i) {
           if (permutation[i] < low || permutation[i] >= high) {
             return rewriter.notifyMatchFailure(
@@ -6229,14 +6229,13 @@ class FoldTransposeBroadcast : public OpRewritePattern<vector::TransposeOp> {
       }
     }
 
-    // We don't need to check the final group [low, outputRank) because
-    // if it is not locally bound, there must be a preceding group that
-    // already failed the check (impossible to have just 1 non-locally
-    // bound group).
+    // We don't need to check the final group [low, outputRank) because if it is
+    // not locally bound, there must be a preceding group that already failed
+    // the check (impossible to have just 1 non-locally bound group).
 
     // The preceding logic also ensures that at this point, the output of the
-    // transpose is definitely broadcastable from the input shape, so we
-    // don't need to check vector::isBroadcastableTo now.
+    // transpose is definitely broadcastable from the input shape, so we don't
+    // need to check vector::isBroadcastableTo now.
 
     rewriter.replaceOpWithNewOp<vector::BroadcastOp>(
         transpose, transpose.getResultVectorType(), transpose.getVector());

mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir

@@ -2,9 +2,9 @@
 
 // This file contains some canonicalizations tests involving vector.transpose.
 
-// CHECK-LABEL: scalar_broadcast_transpose_to_broadcast
+// CHECK-LABEL: broadcast_transpose_scalar_to_broadcast
 //  CHECK-SAME:  %[[ARG:.*]]: i8) -> vector<2x3x4xi8> {
-func.func @scalar_broadcast_transpose_to_broadcast(%arg0 : i8) -> vector<2x3x4xi8> {
+func.func @broadcast_transpose_scalar_to_broadcast(%arg0 : i8) -> vector<2x3x4xi8> {
 //       CHECK:  %[[BC:.*]] = vector.broadcast %[[ARG]] : i8 to vector<2x3x4xi8>
   %0 = vector.broadcast %arg0 : i8 to vector<3x4x2xi8>
   %1 = vector.transpose %0, [2, 0, 1] : vector<3x4x2xi8> to vector<2x3x4xi8>
@@ -14,23 +14,23 @@ func.func @scalar_broadcast_transpose_to_broadcast(%arg0 : i8) -> vector<2x3x4xi
 
 // -----
 
-// CHECK-LABEL: ones_broadcast_transpose_to_broadcast
+// CHECK-LABEL: broadcast_transpose_ones_to_broadcast
 //  CHECK-SAME:  %[[ARG:.*]]: vector<1x1x1xi8>) -> vector<2x3x4xi8> {
 //       CHECK:  %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<1x1x1xi8> to vector<2x3x4xi8>
 //       CHECK:  return %[[RES]] : vector<2x3x4xi8>
-func.func @ones_broadcast_transpose_to_broadcast(%arg0 : vector<1x1x1xi8>) -> vector<2x3x4xi8> {
+func.func @broadcast_transpose_ones_to_broadcast(%arg0 : vector<1x1x1xi8>) -> vector<2x3x4xi8> {
   %0 = vector.broadcast %arg0 : vector<1x1x1xi8> to vector<3x4x2xi8>
   %1 = vector.transpose %0, [2, 0, 1] : vector<3x4x2xi8> to vector<2x3x4xi8>
   return %1 : vector<2x3x4xi8>
 }
 
 // -----
 
-// CHECK-LABEL: partial_ones_broadcast_transpose_to_broadcast
+// CHECK-LABEL: broadcast_transpose_partial_ones_to_broadcast
 //  CHECK-SAME:  %[[ARG:.*]]: vector<1xi8>) -> vector<8x1xi8> {
 //       CHECK:  %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<1xi8> to vector<8x1xi8>
 //       CHECK:  return %[[RES]] : vector<8x1xi8>
-func.func @partial_ones_broadcast_transpose_to_broadcast(%arg0 : vector<1xi8>) -> vector<8x1xi8> {
+func.func @broadcast_transpose_partial_ones_to_broadcast(%arg0 : vector<1xi8>) -> vector<8x1xi8> {
   %0 = vector.broadcast %arg0 : vector<1xi8> to vector<1x8xi8>
   %1 = vector.transpose %0, [1, 0] : vector<1x8xi8> to vector<8x1xi8>
   return %1 : vector<8x1xi8>
@@ -50,6 +50,18 @@ func.func @broadcast_transpose_mixed_example(%arg0 : vector<4x1x1x7xi8>) -> vect
 
 // -----
 
+// CHECK-LABEL: broadcast_transpose_final_group
+//  CHECK-SAME:  %[[ARG:.*]]: vector<4x7x1x1xi8>) -> vector<4x7x2x3xi8> {
+//       CHECK:  %[[RES:.*]] = vector.broadcast %[[ARG]] : vector<4x7x1x1xi8> to vector<4x7x2x3xi8>
+//       CHECK:  return %[[RES]] : vector<4x7x2x3xi8>
+func.func @broadcast_transpose_final_group(%arg0 : vector<4x7x1x1xi8>) -> vector<4x7x2x3xi8> {
+  %0 = vector.broadcast %arg0 : vector<4x7x1x1xi8> to vector<4x7x3x2xi8>
+  %1 = vector.transpose %0, [0, 1, 3, 2] : vector<4x7x3x2xi8> to vector<4x7x2x3xi8>
+  return %1 : vector<4x7x2x3xi8>
+}
+
+// -----
+
 // CHECK-LABEL: negative_broadcast_transpose_square
 //  CHECK-SAME:  %[[ARG:.*]]:
 //       CHECK:  %[[BCT:.*]] = vector.broadcast %[[ARG]]
@@ -94,7 +106,7 @@ func.func @negative_broadcast_transpose_102(%arg0 : vector<3x1x3xi8>) -> vector<
 //       CHECK:  %[[BCT:.*]] = vector.broadcast %[[ARG]]
 //       CHECK:  %[[TRP:.*]] = vector.transpose %[[BCT]], [0, 2, 1]
 //       CHECK:  return %[[TRP]] : vector<3x3x3xi8>
-func.func @neagtive_broadcast_transpose_021(%arg0 : vector<3x1x3xi8>) -> vector<3x3x3xi8> {
+func.func @negative_broadcast_transpose_021(%arg0 : vector<3x1x3xi8>) -> vector<3x3x3xi8> {
   %0 = vector.broadcast %arg0 : vector<3x1x3xi8> to vector<3x3x3xi8>
   %1 = vector.transpose %0, [0, 2, 1] : vector<3x3x3xi8> to vector<3x3x3xi8>
   return %1 : vector<3x3x3xi8>