Fixup : comment hook

Author: nujaa

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

@@ -1607,7 +1607,10 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
   }
 };
 
-VectorType dropNonScalableUnitDimType(VectorType inVecTy) {
+/// Drop unit dimension of the input VectorType. Scalable dimensions cannot be
+/// folded as we do not want to merge them through a shape_cast and create ill
+/// shaped scalable sizes.
+static VectorType dropNonScalableUnitDimFromType(VectorType inVecTy) {
   auto newVecBuilder = VectorType::Builder(inVecTy);
   auto inVecShape = inVecTy.getShape();
   SmallVector<int64_t> newShape;
@@ -1676,7 +1679,7 @@ struct DropUnitDimFromElementwiseOps final
     auto loc = op->getLoc();
     for (auto operand : op->getOperands()) {
       auto opVectorType = cast<VectorType>(operand.getType());
-      auto newVType = dropNonScalableUnitDimType(opVectorType);
+      auto newVType = dropNonScalableUnitDimFromType(opVectorType);
       if (newVType == opVectorType)
         return rewriter.notifyMatchFailure(op, "No unit dimension to remove.");
 
@@ -1685,7 +1688,7 @@ struct DropUnitDimFromElementwiseOps final
     }
 
     VectorType newResultVectorType =
-        dropNonScalableUnitDimType(resultVectorType);
+        dropNonScalableUnitDimFromType(resultVectorType);
     // Create an updated elementwise Op without unit dim
     Operation *elementwiseOp =
         rewriter.create(loc, op->getName().getIdentifier(), newOperands,

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

@@ -499,15 +499,15 @@ func.func @fold_unit_dims_entirely(%arg0 : vector<8xi32>,
 
 // -----
 
-func.func @fold_unit_inner_dim(%arg0 : vector<8x1x3xf128>,
+func.func @fold_inner_unit_dim(%arg0 : vector<8x1x3xf128>,
                               %arg1 : vector<1x8x3xf128>) -> vector<8x3xf128> {
    %sc_arg1 = vector.shape_cast %arg1 : vector<1x8x3xf128> to vector<8x1x3xf128>
    %mul = arith.mulf %arg0, %sc_arg1 : vector<8x1x3xf128>
    %res = vector.shape_cast %mul : vector<8x1x3xf128> to vector<8x3xf128>
    return %res : vector<8x3xf128>
 }
 
-// CHECK-LABEL: func.func @fold_unit_inner_dim(
+// CHECK-LABEL: func.func @fold_inner_unit_dim(
 // CHECK-SAME:    %[[VAL_0:.*]]: vector<8x1x3xf128>,
 // CHECK-SAME:    %[[VAL_1:.*]]: vector<1x8x3xf128>) -> vector<8x3xf128> {
 // CHECK:         %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1x3xf128> to vector<8x3xf128>
@@ -517,15 +517,15 @@ func.func @fold_unit_inner_dim(%arg0 : vector<8x1x3xf128>,
 
 // -----
 
-func.func @fold_unit_inner_dim_scalable(%arg0 : vector<8x1x[1]x3xf128>,
+func.func @fold_inner_unit_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 %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(
+// CHECK-LABEL: func.func @fold_inner_unit_dim_scalable(
 // CHECK-SAME:    %[[VAL_0:.*]]: vector<8x1x[1]x3xf128>,
 // CHECK-SAME:    %[[VAL_1:.*]]: vector<1x8x[1]x3xf128>) -> vector<8x[1]x3xf128> {
 // CHECK:         %[[VAL_2:.*]] = vector.shape_cast %[[VAL_0]] : vector<8x1x[1]x3xf128> to vector<8x[1]x3xf128>