address comments

Author: c-rhodes

mlir/include/mlir/Dialect/Vector/Utils/VectorUtils.h

@@ -211,11 +211,6 @@ Value createReadOrMaskedRead(OpBuilder &builder, Location loc, Value source,
 ///      static sizes in `shape`.
 LogicalResult isValidMaskedInputVector(ArrayRef<int64_t> shape,
                                        ArrayRef<int64_t> inputVectorSizes);
-
-/// Returns true if the leading dim(s) of `type` are fixed and the trailing dim
-/// is scalable.
-bool isTrailingDimScalable(VectorType type);
-
 } // namespace vector
 
 /// Constructs a permutation map of invariant memref indices to vector

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

@@ -68,9 +68,8 @@ struct FlattenGather : OpRewritePattern<vector::GatherOp> {
 
     // Unrolling doesn't take vscale into account. Pattern is disabled for
     // vectors with leading scalable dim(s).
-    if (resultTy.isScalable() && !isTrailingDimScalable(resultTy))
-      return rewriter.notifyMatchFailure(
-          op, "vector type must be fixed-width or scalable in trailing dim");
+    if (resultTy.getScalableDims().front())
+      return rewriter.notifyMatchFailure(op, "cannot unroll scalable dim");
 
     Location loc = op.getLoc();
     Value indexVec = op.getIndexVec();

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

@@ -342,6 +342,11 @@ class ScalableShapeCastOpRewritePattern
     rewriter.replaceOp(op, result);
     return success();
   }
+
+  static bool isTrailingDimScalable(VectorType type) {
+    return type.getRank() >= 1 && type.getScalableDims().back() &&
+           !llvm::is_contained(type.getScalableDims().drop_back(), true);
+  }
 };
 
 } // namespace

mlir/lib/Dialect/Vector/Utils/VectorUtils.cpp

@@ -396,8 +396,3 @@ vector::isValidMaskedInputVector(ArrayRef<int64_t> shape,
   }
   return success();
 }
-
-bool vector::isTrailingDimScalable(VectorType type) {
-  return type.getRank() >= 1 && type.getScalableDims().back() &&
-         !llvm::is_contained(type.getScalableDims().drop_back(), true);
-}