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

[Federico Bruzzone]

================
@@ -328,6 +333,41 @@ struct CastAwayTransferWriteLeadingOneDim
 
 } // namespace
 
+// Takes `oldVal` and "drops" the leading unit dim with either ShapeCastOp or
+// ExtractOp. The latter is used for rank-1 vectors to make sure that a scalar
+// (as opposed to rank-0 vector) is generated. This is a requirement of e.g.
+// ContractOp.
+static Value dropLeadingUnitDimViaShapeCastOrExtract(RewriterBase &rewriter,
+                                                     Location loc,
+                                                     mlir::Value oldVal) {
+  auto oldValTy = cast<VectorType>(oldVal.getType());
+  if (oldValTy.getRank() == 1) {
+    return rewriter.createOrFold<ExtractOp>(loc, oldVal, 0);
+  }
+
+  return rewriter.createOrFold<ShapeCastOp>(
+      loc,
+      trimLeadingUnitDims(oldValTy,
+                          /*trimOnlyOneDim=*/true,
+                          /*allowRank0=*/true),
+      oldVal);
+}
+
+// Takes `oldVal` and "adds" leading unit dim with either ShapeCastOp or
+// BroadcastOp. The latter is used for scalar inputs as ShapeCastOp cannot
+// "broadcast" from a scalar. Scalars are used (instead of rank-0 vectors) as
+// ContractOp operands.
+static Value restoerLeadingUnitDimViaShapeCastOrBcast(RewriterBase &rewriter,
----------------
FedericoBruzzone wrote:

[nit]

```suggestion
static Value restoreLeadingUnitDimViaShapeCastOrBcast(RewriterBase &rewriter,
```

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