[Mlir-commits] [mlir] [MLIR][Vector]Generalize DropUnitDimFromElementwiseOps (PR #92934)

[Han-Chung Wang]

================
@@ -1607,7 +1607,23 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
   }
 };
 
-/// For vectors with either leading or trailing unit dim, replaces:
+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++;
+    }
+  }
----------------
hanhanW wrote:

This is O(N^2) because `dropDim` could use `vector.erase`. I'd prefer having a variable to store the shape, and construct the vector type at the end.

https://github.com/llvm/llvm-project/pull/92934