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

[Diego Caballero]

================
@@ -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);
+}
+
+namespace {
+
+// Drops leading 1 dimensions from load-like memory operaitons. REmoves leading
+// unit dimensions from the result types and then broadcasts back in those 1s,
+// while also extracting (or shape_cast-ing) any leading unit dimensions on
+// the input operands.
+template <typename OpTy>
+struct CastAwayLoadLikeLeadingOneDim : public OpRewritePattern<OpTy> {
+  using OpRewritePattern<OpTy>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(OpTy op,
+                                PatternRewriter &rewriter) const override {
+    VectorType oldResultType = op.getVectorType();
+    VectorType newResultType = trimLeadingOneDims(oldResultType);
+    if (newResultType == oldResultType)
+      return failure();
+    int64_t nDropped = oldResultType.getRank() - newResultType.getRank();
+
+    Location loc = op.getLoc();
+    SmallVector<Value> newOperands;
+    newOperands.reserve(op->getNumOperands());
+    for (Value operand : op->getOperands()) {
+      if (isa<VectorType>(operand.getType())) {
+        newOperands.push_back(
+            dropLeadingOneDimsFromOperand(rewriter, loc, operand, nDropped));
+      } else {
+        newOperands.push_back(operand);
+      }
+    }
+
+    Operation *newOp =
+        rewriter.create(loc, op->getName().getIdentifier(), newOperands,
+                        TypeRange{newResultType}, op->getAttrs());
+    rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, oldResultType,
+                                                     newOp->getResult(0));
----------------
dcaballe wrote:

We are generating a `vector.shape_cast` to remove the unit dimension and a `vector.broadcast` to bring it back. Could we consistently use `vector.shape_cast`?

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