Fix naming and test

Author: nujaa

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

@@ -1607,20 +1607,20 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
   }
 };
 
-FailureOr<VectorType> dropNonScalableUnitDimType(VectorType VT) {
-  int removed = 0;
-  auto shape = VT.getShape();
-  auto builder = VectorType::Builder(VT);
-  for (unsigned i = 0; i < shape.size(); i++) {
-    if (shape[i] == 1 && !VT.getScalableDims()[i]) {
-      builder.dropDim(i - removed);
-      removed++;
+FailureOr<VectorType> dropNonScalableUnitDimType(VectorType inVecTy) {
+  int numUnitDimsDropped = 0;
+  auto inVecShape = inVecTy.getShape();
+  auto newVecBuilder = VectorType::Builder(inVecTy);
+  for (unsigned i = 0; i < inVecShape.size(); i++) {
+    if (inVecShape[i] == 1 && !inVecTy.getScalableDims()[i]) {
+      newVecBuilder.dropDim(i - numUnitDimsDropped);
+      numUnitDimsDropped++;
     }
   }
 
-  if (removed == 0)
+  if (numUnitDimsDropped == 0)
     return failure();
-  return VectorType(builder);
+  return VectorType(newVecBuilder);
 }
 
 /// For vectors with at least an unit dim, replaces:

mlir/test/Dialect/Vector/vector-transfer-flatten.mlir

@@ -509,7 +509,7 @@ func.func @fold_unit_inner_dim(%arg0 : vector<8x1x3xf128>,
 
 // CHECK-LABEL: func.func @fold_unit_inner_dim(
 // CHECK-SAME:    %[[VAL_0:.*]]: vector<8x1x3xf128>,
-// CHECK-SAME:    %[[VAL_1:.*]]: vector<1x8x3xf128>) -> vector<8xxf128> {
+// CHECK-SAME:    %[[VAL_1:.*]]: vector<1x8x3xf128>) -> vector<8x3xf128> {
 // CHECK:         %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1x3xf128> to vector<8x3xf128>
 // CHECK:         %[[VAL_3:.*]] = vector.shape_cast %[[VAL_1]] : vector<1x8x3xf128> to vector<8x3xf128>
 // CHECK:         %[[VAL_4:.*]] = arith.mulf %[[VAL_2]], %[[VAL_3]] : vector<8x3xf128>
@@ -521,8 +521,8 @@ func.func @fold_unit_inner_dim_scalable(%arg0 : vector<8x1x[1]x3xf128>,
                               %arg1 : vector<1x8x[1]x3xf128>) -> vector<8x[1]x3xf128> {
    %sc_arg1 = vector.shape_cast %arg1 : vector<1x8x[1]x3xf128> to vector<8x1x[1]x3xf128>
    %mul = arith.mulf %arg0, %sc_arg1 : vector<8x1x[1]x3xf128>
-   %res = vector.shape_cast %add : vector<8x1x[1]x3xf128> to vector<8x[1]x3xf128>
-   return %mul : vector<8x[1]x3xf128>
+   %res = vector.shape_cast %mul : vector<8x1x[1]x3xf128> to vector<8x[1]x3xf128>
+   return %res : vector<8x[1]x3xf128>
 }
 
 // CHECK-LABEL: func.func @fold_unit_inner_dim_scalable(