[Mlir-commits] [mlir] [mlir][vector] Use `ShapeCastOp` in `castAwayContractionLeadingOneDim` (PR #226524)

[Andrzej Warzyński]

================
@@ -458,23 +503,25 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
       rewriter.getArrayAttr(newIteratorTypes), contractOp.getKind());
 
   if (maskingOp) {
-    auto newMask = vector::ExtractOp::create(rewriter, loc, maskingOp.getMask(),
-                                             splatZero(dropDim));
+    auto newMask = rewriter.createOrFold<ShapeCastOp>(
+        loc,
+        trimLeadingUnitDims(cast<VectorType>(maskingOp.getMask().getType()),
+                            /*trimOnlyOneDim=*/true, /*allowRank0=*/true),
+        maskingOp.getMask());
 
     newOp = mlir::vector::maskOperation(rewriter, newOp, newMask);
   }
 
-  return vector::BroadcastOp::create(rewriter, loc,
-                                     contractOp->getResultTypes()[0],
-                                     newOp->getResults()[0])
-      .getResult();
+  return restoerLeadingUnitDimViaShapeCastOrBcast(
+      rewriter, loc, newOp->getResult(0), contractOp->getResultTypes()[0]);
 }
 
 namespace {
 
 /// Turns vector.contract on vector with leading 1 dimensions into
-/// vector.extract followed by vector.contract on vector without leading
-/// 1 dimensions. Also performs transpose of lhs and rhs operands if required
+/// vector.shape_cast followed by vector.contract on vector without leading
+/// 1 dimensions.
+/// TODO (check): Also performs transpose of lhs and rhs operands if required
----------------
banach-space wrote:

Thanks for noticing, it's actually Option 1 - I forgot to update it. I will be a clarifying update shortly.

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