[Mlir-commits] [mlir] [mlir][Vector] Add load, store, etc. to dropleadunitdim (PR #195686)

[Andrzej Warzyński]

================
@@ -537,6 +537,101 @@ class CastAwayElementwiseLeadingOneDim : public RewritePattern {
     return success();
   }
 };
+} // namespace
+
+// Drops `dropDim` leading dimensions from `operand` using vector.extract when
+// those dims are all non-scalable units (the cheap, structural rewrite); falls
+// back to vector.shape_cast otherwise.
+static Value dropLeadingOneDimsFromOperand(OpBuilder &b, Location loc,
+                                           Value operand, int64_t nDropped) {
+  auto oldType = cast<VectorType>(operand.getType());
+  ArrayRef<int64_t> leadingShape = oldType.getShape().take_front(nDropped);
+  ArrayRef<bool> leadingScalable =
+      oldType.getScalableDims().take_front(nDropped);
+  bool extractable =
+      llvm::all_of(leadingShape, [](int64_t d) { return d == 1; }) &&
+      llvm::none_of(leadingScalable, [](bool s) { return s; });
+  if (extractable)
+    return vector::ExtractOp::create(b, loc, operand, splatZero(nDropped));
+  VectorType newType = VectorType::get(
+      oldType.getShape().drop_front(nDropped), oldType.getElementType(),
+      oldType.getScalableDims().drop_front(nDropped));
+  return vector::ShapeCastOp::create(b, loc, newType, operand);
----------------
banach-space wrote:

I would much prefer if this was simplified and returned `vector.shape_cast` unconditionally. 

Also, could you add tests for scalable vectors, thanks!

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