[Mlir-commits] [mlir] [MLIR][Vector] Generalize DropUnitDimFromElementwiseOps to non leading / trailing dimensions. (PR #98455)

[Han-Chung Wang]

================
@@ -1622,7 +1622,33 @@ struct ChainedReduction final : OpRewritePattern<vector::ReductionOp> {
   }
 };
 
-/// For vectors with either leading or trailing unit dim, replaces:
+// Helper function dropping unit non-scalable dimension from a VectorType.
+// Scalable unit dimensions are not dropped. Folding such dimensions would
+// require "shifting" the scalable flag onto some other fixed-width dim (e.g.
+// vector<[1]x4xf32> -> vector<[4]xf32>). This could be implemented in the
+// future.
+static VectorType dropNonScalableUnitDimFromType(VectorType inVecTy) {
+  auto inVecShape = inVecTy.getShape();
+  SmallVector<int64_t> newShape;
+  SmallVector<bool> newScalableDims;
+  for (auto [dim, isScalable] :
+       llvm::zip_equal(inVecShape, inVecTy.getScalableDims())) {
+    if (dim == 1 && !isScalable)
+      continue;
+
+    newShape.push_back(dim);
+    newScalableDims.push_back(isScalable);
+  }
+  // All dims have been dropped, we need to return a legal shape for VectorType.
----------------
hanhanW wrote:

We don't need to think about IREE here. We should generate valid IRs in upstream transformations. IMO, 0-D vector is weird. I can't really tell the distinction between 0-D vector and `vector<1xT>`. I'd suggest to return 1-D vector and document it in the function comment (i.e., l.1625-l.1629).

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