[Mlir-commits] [mlir] [mlir][vector] Migrate drop-lead-unit-dim to shape_cast (PR #196206)

[Andrzej Warzyński]

================
@@ -45,14 +48,111 @@ static VectorType trimLeadingOneDims(VectorType oldType) {
   return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
 }
 
-/// Return a smallVector of size `rank` containing all zeros.
-static SmallVector<int64_t> splatZero(int64_t rank) {
-  return SmallVector<int64_t>(rank, 0);
+static bool isNonScalableUnitDim(VectorType type, int64_t dim) {
+  assert(dim >= 0 && dim < type.getRank() &&
+         "expected a valid vector dimension");
+  return type.getShape()[dim] == 1 && !type.getScalableDims()[dim];
 }
+
+/// Returns true if the first `k` dimensions of `type` are non-scalable unit
+/// dimensions.
+static bool leadingDimsAreUnit(VectorType type, int64_t k) {
+  assert(k >= 0 && k <= type.getRank() &&
+         "expected a valid leading dimension count");
+  return llvm::all_of(llvm::seq<int64_t>(0, k), [&](int64_t dim) {
+    return isNonScalableUnitDim(type, dim);
+  });
+}
+
+static bool leadingDimsAreUnitAfterPermutation(VectorType type,
+                                               ArrayRef<int64_t> permutation,
+                                               int64_t k) {
+  assert(k >= 0 && k <= static_cast<int64_t>(permutation.size()) &&
+         "expected a valid leading dimension count");
+  return llvm::all_of(permutation.take_front(k), [&](int64_t dim) {
+    return isNonScalableUnitDim(type, dim);
+  });
+}
+
+/// Shape-casts `operand` to the vector type obtained by dropping dimension
+/// `dim`, which must be non-scalable and unit-sized.
+static Value dropUnitDim(OpBuilder &b, Location loc, Value operand,
+                         int64_t dim) {
+  auto oldType = cast<VectorType>(operand.getType());
+  assert(isNonScalableUnitDim(oldType, dim) &&
+         "expected a non-scalable unit dim to drop");
+  int64_t rank = oldType.getRank();
+  assert(rank > 1 && "cannot shape_cast to a 0-D vector");
+
+  SmallVector<int64_t> newShape;
+  SmallVector<bool> newScalableDims;
+  newShape.reserve(rank - 1);
+  newScalableDims.reserve(rank - 1);
+  for (auto [i, size, scalable] :
+       llvm::enumerate(oldType.getShape(), oldType.getScalableDims())) {
+    if (static_cast<int64_t>(i) == dim)
+      continue;
+    newShape.push_back(size);
+    newScalableDims.push_back(scalable);
+  }
+
+  return b.createOrFold<vector::ShapeCastOp>(
+      loc, VectorType::get(newShape, oldType.getElementType(), newScalableDims),
+      operand);
+}
+
+/// Shape-casts `operand` to the vector type obtained by dropping the first
+/// `k` non-scalable unit dimensions. Callers must ensure at least one vector
+/// dimension remains after the drop.
+static Value dropLeadingUnitDims(OpBuilder &b, Location loc, Value operand,
+                                 int64_t k) {
+  auto oldType = cast<VectorType>(operand.getType());
+  assert(leadingDimsAreUnit(oldType, k) &&
+         "expected non-scalable leading unit dims to drop");
+  assert(k < oldType.getRank() &&
+         "shape_cast cannot drop all vector dimensions");
+  VectorType newType = VectorType::get(oldType.getShape().drop_front(k),
+                                       oldType.getElementType(),
+                                       oldType.getScalableDims().drop_front(k));
+  return b.createOrFold<vector::ShapeCastOp>(loc, newType, operand);
+}
+
+/// Returns the vector type obtained by applying `permutation` to `type`.
+static VectorType permuteVectorType(VectorType type,
+                                    ArrayRef<int64_t> permutation) {
+  assert(static_cast<int64_t>(permutation.size()) == type.getRank() &&
+         "expected a permutation matching the operand rank");
+  SmallVector<int64_t> permutedShape =
+      applyPermutation(type.getShape(), permutation);
+  SmallVector<bool> permutedScalableDims =
+      applyPermutation(type.getScalableDims(), permutation);
+  return VectorType::get(permutedShape, type.getElementType(),
+                         permutedScalableDims);
+}
+
+/// Like `dropLeadingUnitDims` except that if all dimensions would be dropped,
+/// the single element inside that vector is extracted and returned.
+static Value dropLeadingUnitDims0DIsScalar(OpBuilder &b, Location loc,
+                                           Value operand, int64_t k) {
+  auto oldType = cast<VectorType>(operand.getType());
+  assert(leadingDimsAreUnit(oldType, k) &&
+         "expected non-scalable leading unit dims to drop");
+
+  if (k == oldType.getRank()) {
+    SmallVector<int64_t> zeros(k, static_cast<int64_t>(0));
+    return vector::ExtractOp::create(b, loc, operand, zeros);
+  }
+
+  VectorType newType = VectorType::get(oldType.getShape().drop_front(k),
+                                       oldType.getElementType(),
+                                       oldType.getScalableDims().drop_front(k));
+  return vector::ShapeCastOp::create(b, loc, newType, operand);
----------------
banach-space wrote:

`createOrFold`?

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