[Mlir-commits] [mlir] [mlir][vector] Use `ShapeCastOp` in `castAwayContractionLeadingOneDim` (PR #226524)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Sep 29 08:56:20 PDT 2026
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/226524
>From 552101aac10cdaeb1e2c2e6d775499c4a050b775 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 25 Sep 2026 16:02:50 +0000
Subject: [PATCH 1/2] [mlir][vector] Use `ShapeCastOp` in
`castAwayContractionLeadingOneDim`
Use `ShapeCastOp`, instead of `ExtractOp` and `BroadcastOp`, to drop the
leading unit dimension in `castAwayContractionLeadingOneDim`. This PR
continues the work started in:
* #208752, #210902, #211777, #219499, #219494
Essentially, this implements the approach already agreed on in the
previous PRs. See also #196206, which originally implemented this change
and was subsequently reverted in #199546.
`TransposeOp`s are left intact in light of the ongoing discussion about
whether folding `TransposeOp` into `ShapeCastOp` is valid, see e.g.
https://github.com/llvm/llvm-project/pull/219611.
Note that in the tests, `TransposeOp`s are nevertheless replaced with
`ShapeCastOp`s. This is due to the existing `TransposeOp` folders.
---
.../Transforms/VectorDropLeadUnitDim.cpp | 89 ++++++++++++++-----
.../vector-dropleadunitdim-transforms.mlir | 77 +++++++---------
2 files changed, 101 insertions(+), 65 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index 2d87493f9e070..f302b21494998 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -22,9 +22,11 @@
using namespace mlir;
using namespace mlir::vector;
-// Trims leading one dimensions from `oldType` and returns the result type.
-// Returns `vector<1xT>` if `oldType` only has one element.
-static VectorType trimLeadingUnitDims(VectorType oldType) {
+// Trims leading one dimensions (fixed-width) from `oldType` and returns the
+// result type. Returns `vector<1xT>` if `oldType` only has one element.
+static VectorType trimLeadingUnitDims(VectorType oldType,
+ bool trimOnlyOneDim = false,
+ bool allowRank0 = false) {
ArrayRef<int64_t> oldShape = oldType.getShape();
ArrayRef<int64_t> newShape = oldShape;
@@ -35,10 +37,13 @@ static VectorType trimLeadingUnitDims(VectorType oldType) {
!newScalableDims.front()) {
newShape = newShape.drop_front(1);
newScalableDims = newScalableDims.drop_front(1);
+
+ if (trimOnlyOneDim)
+ break;
}
// Make sure we have at least 1 dimension per vector type requirements.
- if (newShape.empty()) {
+ if (newShape.empty() && !allowRank0) {
newShape = oldShape.take_back();
newScalableDims = oldType.getScalableDims().take_back();
}
@@ -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,
+ Location loc,
+ mlir::Value oldVal,
+ mlir::Type newTy) {
+ if (!isa<VectorType>(oldVal.getType())) {
+ return rewriter.createOrFold<BroadcastOp>(loc, newTy, oldVal);
+ }
+
+ return rewriter.createOrFold<ShapeCastOp>(loc, newTy, oldVal);
+}
+
FailureOr<Value>
mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
MaskingOpInterface maskingOp,
@@ -337,11 +377,8 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
return failure();
if (oldAccType.getRank() < 1)
return failure();
- if (oldAccType.getShape()[0] != 1)
+ if (oldAccType.getShape()[0] != 1 || oldAccType.getScalableDims()[0])
return failure();
- // currently we support only dropping one dim but the pattern can be applied
- // greedily to drop more.
- int64_t dropDim = 1;
auto oldIndexingMaps = contractOp.getIndexingMapsArray();
SmallVector<AffineMap> newIndexingMaps;
@@ -349,6 +386,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
auto oldIteratorTypes = contractOp.getIteratorTypes();
SmallVector<Attribute> newIteratorTypes;
+ // 0-th dim from the accuumlator
int64_t dimToDrop = oldIndexingMaps[2].getDimPosition(0);
if (!isParallelIterator(oldIteratorTypes[dimToDrop]))
@@ -418,6 +456,13 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
map = AffineMap::get(map.getNumDims(), 0, transposeResults,
contractOp.getContext());
if (transposeNonOuterUnitDims) {
+ // TODO: While the discussion on the validity of folding
+ // TransposeOp into ShapeCastOp continues, see e.g.
+ // * https://github.com/llvm/llvm-project/pull/219611,
+ // keep the explicit TransposeOp here. Note that existing TransposeOp
+ // folders already turn it into a ShapeCastOp, as demonstrated by the
+ // tests. Revisit this and consider inserting ShapeCastOp directly
+ // once the discussion progresses.
operands[it.index()] = rewriter.createOrFold<vector::TransposeOp>(
loc, operands[it.index()], perm);
}
@@ -443,11 +488,11 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
contractOp.getContext()));
// Extract if its a valid extraction, otherwise use the operand
// without extraction.
- newOperands.push_back(validExtract
- ? vector::ExtractOp::create(rewriter, loc,
- operands[it.index()],
- splatZero(dropDim))
- : operands[it.index()]);
+ auto oldVal = operands[it.index()];
+ newOperands.push_back(
+ validExtract
+ ? dropLeadingUnitDimViaShapeCastOrExtract(rewriter, loc, oldVal)
+ : oldVal);
}
// Depending on whether this vector.contract is masked, the replacing Op
@@ -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
/// prior to extract.
struct CastAwayContractionLeadingOneDim
: public MaskableOpRewritePattern<vector::ContractionOp> {
diff --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index 5deef2aee9a26..46894b63fff9c 100644
--- a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
@@ -5,13 +5,13 @@
// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK-LABEL: cast_away_contraction_leading_one_dims
-// CHECK-NEXT: %[[R0:.+]] = vector.extract %{{.*}}[0] : vector<16x8xf32> from vector<1x16x8xf32>
-// CHECK-NEXT: %[[R1:.+]] = vector.extract %{{.*}}[0] : vector<8x16xf32> from vector<1x8x16xf32>
-// CHECK-NEXT: %[[R2:.+]] = vector.extract %{{.*}}[0] : vector<16x16xf32> from vector<1x16x16xf32>
+// CHECK-NEXT: %[[R0:.+]] = vector.shape_cast %{{.*}} : vector<1x16x8xf32> to vector<16x8xf32>
+// CHECK-NEXT: %[[R1:.+]] = vector.shape_cast %{{.*}} : vector<1x8x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[R2:.+]] = vector.shape_cast %{{.*}} : vector<1x16x16xf32> to vector<16x16xf32>
// CHECK-NEXT: %[[R3:.+]] = vector.contract {indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
// CHECK-SAME: %[[R0]], %[[R1]], %[[R2]] : vector<16x8xf32>, vector<8x16xf32> into vector<16x16xf32>
-// CHECK-NEXT: %[[R4:.+]] = vector.broadcast %[[R3]] : vector<16x16xf32> to vector<1x16x16xf32>
+// CHECK-NEXT: %[[R4:.+]] = vector.shape_cast %[[R3]] : vector<16x16xf32> to vector<1x16x16xf32>
// CHECK-NEXT: return %[[R4]] : vector<1x16x16xf32>
#contraction_accesses0 = [
@@ -36,14 +36,14 @@ func.func @cast_away_contraction_leading_one_dims(%arg0: vector<1x16x8xf32>, %ar
// CHECK-LABEL: func.func @cast_away_contraction_leading_one_dim_under_const_mask
// CHECK: %[[MASK:.*]] = vector.constant_mask [15, 15, 8] : vector<16x16x8xi1>
-// CHECK: %[[R0:.*]] = vector.extract %{{.*}}[0] : vector<16x8xf32> from vector<1x16x8xf32>
-// CHECK: %[[R1:.*]] = vector.extract %{{.*}}[0] : vector<8x16xf32> from vector<1x8x16xf32>
-// CHECK: %[[R2:.*]] = vector.extract %{{.*}}[0] : vector<16x16xf32> from vector<1x16x16xf32>
+// CHECK: %[[R0:.*]] = vector.shape_cast %{{.*}} : vector<1x16x8xf32> to vector<16x8xf32>
+// CHECK: %[[R1:.*]] = vector.shape_cast %{{.*}} : vector<1x8x16xf32> to vector<8x16xf32>
+// CHECK: %[[R2:.*]] = vector.shape_cast %{{.*}} : vector<1x16x16xf32> to vector<16x16xf32>
// CHECK: %[[CONTRACT:.*]] = vector.mask %[[MASK]] {
// CHECK-SAME: vector.contract {indexing_maps = [#[[$MAP_0]], #[[$MAP_1]], #[[$MAP_2]]], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
// CHECK-SAME: %[[R0]], %[[R1]], %[[R2]] : vector<16x8xf32>, vector<8x16xf32> into vector<16x16xf32>
// CHECK-SAME: } : vector<16x16x8xi1> -> vector<16x16xf32>
-// CHECK: %[[RES:.*]] = vector.broadcast %[[CONTRACT]] : vector<16x16xf32> to vector<1x16x16xf32>
+// CHECK: %[[RES:.*]] = vector.shape_cast %[[CONTRACT]] : vector<16x16xf32> to vector<1x16x16xf32>
// CHECK: return %[[RES]] : vector<1x16x16xf32>
#contraction_accesses0 = [
@@ -70,15 +70,15 @@ func.func @cast_away_contraction_leading_one_dim_under_const_mask(%arg0: vector<
// CHECK-DAG: #[[$MAP2:.+]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK-LABEL: func.func @cast_away_contraction_leading_one_dim_under_mask
-// CHECK: %[[R0:.*]] = vector.extract %{{.*}} : vector<16x8xf32> from vector<1x16x8xf32>
-// CHECK: %[[R1:.*]] = vector.extract %{{.*}} : vector<8x16xf32> from vector<1x8x16xf32>
-// CHECK: %[[R2:.*]] = vector.extract %{{.*}} : vector<16x16xf32> from vector<1x16x16xf32>
-// CHECK: %[[M:.*]] = vector.extract %{{.*}} : vector<16x16x8xi1> from vector<1x16x16x8xi1>
+// CHECK: %[[R0:.*]] = vector.shape_cast %{{.*}} : vector<1x16x8xf32> to vector<16x8xf32>
+// CHECK: %[[R1:.*]] = vector.shape_cast %{{.*}} : vector<1x8x16xf32> to vector<8x16xf32>
+// CHECK: %[[R2:.*]] = vector.shape_cast %{{.*}} : vector<1x16x16xf32> to vector<16x16xf32>
+// CHECK: %[[M:.*]] = vector.shape_cast %{{.*}} : vector<1x16x16x8xi1> to vector<16x16x8xi1>
// CHECK: %[[CONTRACT:.*]] = vector.mask %[[M]] {
// CHECK-SAME: vector.contract {indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]], iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
// CHECK-SAME: %[[R0]], %[[R1]], %[[R2]] : vector<16x8xf32>, vector<8x16xf32> into vector<16x16xf32>
// CHECK-SAME: } : vector<16x16x8xi1> -> vector<16x16xf32>
-// CHECK-NEXT: %[[RES:.*]] = vector.broadcast %[[CONTRACT]] : vector<16x16xf32> to vector<1x16x16xf32>
+// CHECK-NEXT: %[[RES:.*]] = vector.shape_cast %[[CONTRACT]] : vector<16x16xf32> to vector<1x16x16xf32>
// CHECK-NEXT: return %[[RES]] : vector<1x16x16xf32>
#contraction_accesses0 = [
@@ -109,14 +109,13 @@ func.func @cast_away_contraction_leading_one_dim_under_mask(
// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1) -> (d0)>
// CHECK-LABEL: cast_away_contraction_leading_one_dims_transposeneeded
-// CHECK-NEXT: %[[R0:.+]] = vector.extract %{{.*}}[0] : vector<8x16xf32> from vector<1x8x16xf32>
-// CHECK-NEXT: %[[R1:.+]] = vector.extract %{{.*}}[0, 0] : vector<8xf32> from vector<1x1x8xf32>
-// CHECK-NEXT: %[[R2:.+]] = vector.extract %{{.*}}[0, 0] : vector<16xf32> from vector<1x1x16xf32>
+// CHECK-NEXT: %[[R0:.+]] = vector.shape_cast %{{.*}} : vector<1x8x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[R1:.+]] = vector.shape_cast %{{.*}} : vector<1x1x8xf32> to vector<8xf32>
+// CHECK-NEXT: %[[R2:.+]] = vector.shape_cast %{{.*}} : vector<1x1x16xf32> to vector<16xf32>
// CHECK-NEXT: %[[R3:.+]] = vector.contract {indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]],
// CHECK-SAME: iterator_types = ["parallel", "reduction"], kind = #vector.kind<mul>}
// CHECK-SAME: %[[R1]], %[[R0]], %[[R2]] : vector<8xf32>, vector<8x16xf32> into vector<16xf32>
-// CHECK-NEXT: %[[R4:.+]] = vector.broadcast %[[R3]] : vector<16xf32> to vector<1x16xf32>
-// CHECK-NEXT: %[[R5:.+]] = vector.broadcast %[[R4]] : vector<1x16xf32> to vector<1x1x16xf32>
+// CHECK-NEXT: %[[R5:.+]] = vector.shape_cast %[[R3]] : vector<16xf32> to vector<1x1x16xf32>
// CHECK-NEXT: return %[[R5]] : vector<1x1x16xf32>
#contraction_accesses1 = [
@@ -141,15 +140,13 @@ func.func @cast_away_contraction_leading_one_dims_transposeneeded(%arg0: vector<
// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK-LABEL: cast_away_contraction_leading_one_dims_transposeneeded2
-// CHECK-NEXT: %[[R0:.+]] = vector.transpose %{{.*}}[1, 0, 2] : vector<8x1x16xf32> to vector<1x8x16xf32>
-// CHECK-NEXT: %[[R1:.+]] = vector.extract %[[R0]][0] : vector<8x16xf32> from vector<1x8x16xf32>
-// CHECK-NEXT: %[[R2:.+]] = vector.transpose %{{.*}}[2, 0, 1] : vector<2x8x1xf32> to vector<1x2x8xf32>
-// CHECK-NEXT: %[[R3:.+]] = vector.extract %[[R2]][0] : vector<2x8xf32> from vector<1x2x8xf32>
-// CHECK-NEXT: %[[R4:.+]] = vector.extract %{{.*}}[0] : vector<2x16xf32> from vector<1x2x16xf32>
+// CHECK-NEXT: %[[LHS:.*]] = vector.shape_cast %{{.*}} : vector<8x1x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[RHS:.*]] = vector.shape_cast %{{.*}} : vector<2x8x1xf32> to vector<2x8xf32>
+// CHECK-NEXT: %[[ACC:.*]] = vector.shape_cast %{{.*}} : vector<1x2x16xf32> to vector<2x16xf32>
// CHECK-NEXT: %[[R5:.+]] = vector.contract {indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
-// CHECK-SAME: %[[R1]], %[[R3]], %[[R4]] : vector<8x16xf32>, vector<2x8xf32> into vector<2x16xf32>
-// CHECK-NEXT: %[[R6:.+]] = vector.broadcast %[[R5]] : vector<2x16xf32> to vector<1x2x16xf32>
+// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]] : vector<8x16xf32>, vector<2x8xf32> into vector<2x16xf32>
+// CHECK-NEXT: %[[R6:.+]] = vector.shape_cast %[[R5]] : vector<2x16xf32> to vector<1x2x16xf32>
// CHECK-NEXT: return %[[R6]] : vector<1x2x16xf32>
#contraction_accesses2 = [
@@ -175,18 +172,13 @@ func.func @cast_away_contraction_leading_one_dims_transposeneeded2(%arg0: vector
// CHECK-LABEL: cast_away_contraction_leading_one_dims_nonleadingunitdim_rank4
-// CHECK-NEXT: %[[R0:.+]] = vector.extract %{{.*}}[0] : vector<8x1x16xf32> from vector<1x8x1x16xf32>
-// CHECK-NEXT: %[[R1:.+]] = vector.extract %{{.*}}[0] : vector<2x8x1xf32> from vector<1x2x8x1xf32>
-// CHECK-NEXT: %[[R2:.+]] = vector.transpose %[[R0]], [1, 0, 2] : vector<8x1x16xf32> to vector<1x8x16xf32>
-// CHECK-NEXT: %[[R3:.+]] = vector.extract %[[R2]][0] : vector<8x16xf32> from vector<1x8x16xf32>
-// CHECK-NEXT: %[[R4:.+]] = vector.transpose %[[R1]], [2, 0, 1] : vector<2x8x1xf32> to vector<1x2x8xf32>
-// CHECK-NEXT: %[[R5:.+]] = vector.extract %[[R4]][0] : vector<2x8xf32> from vector<1x2x8xf32>
-// CHECK-NEXT: %[[R6:.+]] = vector.extract %{{.*}}[0, 0] : vector<2x16xf32> from vector<1x1x2x16xf32>
+// CHECK-NEXT: %[[LHS:.*]] = vector.shape_cast %{{.*}} : vector<1x8x1x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[RHS:.*]] = vector.shape_cast %{{.*}} : vector<1x2x8x1xf32> to vector<2x8xf32>
+// CHECK-NEXT: %[[ACC:.*]] = vector.shape_cast %{{.*}} : vector<1x1x2x16xf32> to vector<2x16xf32>
// CHECK-NEXT: %[[R7:.+]] = vector.contract {indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
-// CHECK-SAME: %[[R3]], %[[R5]], %[[R6]] : vector<8x16xf32>, vector<2x8xf32> into vector<2x16xf32>
-// CHECK-NEXT: %[[R8:.+]] = vector.broadcast %[[R7]] : vector<2x16xf32> to vector<1x2x16xf32>
-// CHECK-NEXT: %[[R9:.+]] = vector.broadcast %[[R8]] : vector<1x2x16xf32> to vector<1x1x2x16xf32>
+// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]] : vector<8x16xf32>, vector<2x8xf32> into vector<2x16xf32>
+// CHECK-NEXT: %[[R9:.+]] = vector.shape_cast %[[R7]] : vector<2x16xf32> to vector<1x1x2x16xf32>
// CHECK-NEXT: return %[[R9]] : vector<1x1x2x16xf32>
#contraction_accesses2 = [
@@ -211,16 +203,13 @@ func.func @cast_away_contraction_leading_one_dims_nonleadingunitdim_rank4(%arg0:
// CHECK-DAG: #[[$map2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
// CHECK-LABEL: cast_away_contraction_leading_one_dims_nonleadingunitdim_rank4_acctranspose
-// CHECK-NEXT: %[[R0:.+]] = vector.transpose %{{.*}}, [2, 0, 1, 3] : vector<1x8x1x16xf32> to vector<1x1x8x16xf32>
-// CHECK-NEXT: %[[R1:.+]] = vector.transpose %{{.*}}, [3, 0, 1, 2] : vector<1x2x8x1xf32> to vector<1x1x2x8xf32>
-// CHECK-NEXT: %[[R2:.+]] = vector.extract %[[R0]][0, 0] : vector<8x16xf32> from vector<1x1x8x16xf32>
-// CHECK-NEXT: %[[R3:.+]] = vector.extract %[[R1]][0, 0] : vector<2x8xf32> from vector<1x1x2x8xf32>
-// CHECK-NEXT: %[[R4:.+]] = vector.extract %{{.*}}[0, 0] : vector<2x16xf32> from vector<1x1x2x16xf32>
+// CHECK-NEXT: %[[LHS:.*]] = vector.shape_cast %arg0 : vector<1x8x1x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[RHS:.*]] = vector.shape_cast %arg1 : vector<1x2x8x1xf32> to vector<2x8xf32>
+// CHECK-NEXT: %[[ACC:.*]] = vector.shape_cast %arg2 : vector<1x1x2x16xf32> to vector<2x16xf32>
// CHECK-NEXT: %[[R5:.+]] = vector.contract {indexing_maps = [#[[$map0]], #[[$map1]], #[[$map2]]],
// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"], kind = #vector.kind<add>}
-// CHECK-SAME: %[[R2]], %[[R3]], %[[R4]] : vector<8x16xf32>, vector<2x8xf32> into vector<2x16xf32>
-// CHECK-NEXT: %[[R6:.+]] = vector.broadcast %[[R5]] : vector<2x16xf32> to vector<1x2x16xf32>
-// CHECK-NEXT: %[[R7:.+]] = vector.broadcast %[[R6]] : vector<1x2x16xf32> to vector<1x1x2x16xf32>
+// CHECK-SAME: %[[LHS]], %[[RHS]], %[[ACC]] : vector<8x16xf32>, vector<2x8xf32> into vector<2x16xf32>
+// CHECK-NEXT: %[[R7:.+]] = vector.shape_cast %[[R5]] : vector<2x16xf32> to vector<1x1x2x16xf32>
// CHECK-NEXT: return %[[R7]] : vector<1x1x2x16xf32>
#contraction_accesses3 = [
@@ -256,7 +245,7 @@ func.func @cast_away_contraction_does_not_transpose_leading_unit_dims(%lhs: vect
// CHECK-DAG: #[[$map_dp1:.*]] = affine_map<(d0) -> ()>
// CHECK-LABEL: cast_away_contraction_leading_one_dims_to_dot_product
-// CHECK-NEXT: %[[R0:.+]] = vector.extract %{{.*}}[0] : vector<64xf32> from vector<1x64xf32>
+// CHECK-NEXT: %[[R0:.+]] = vector.shape_cast %{{.*}} : vector<1x64xf32> to vector<64xf32>
// CHECK-NEXT: %[[R1:.+]] = vector.extract %{{.*}}[0] : f32 from vector<1xf32>
// CHECK-NEXT: %[[R2:.+]] = vector.contract {indexing_maps = [#[[$map_dp0]], #[[$map_dp0]], #[[$map_dp1]]],
// CHECK-SAME: iterator_types = ["reduction"], kind = #vector.kind<add>}
>From ed9c8185680b38070f31659d8bdd547471413db4 Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 29 Sep 2026 15:56:00 +0000
Subject: [PATCH 2/2] Address comments from Fede
---
.../Vector/Transforms/VectorDropLeadUnitDim.cpp | 14 +++++++++-----
1 file changed, 9 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index f302b21494998..f84e7450fc0ad 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -357,7 +357,7 @@ static Value dropLeadingUnitDimViaShapeCastOrExtract(RewriterBase &rewriter,
// 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,
+static Value restoreLeadingUnitDimViaShapeCastOrBcast(RewriterBase &rewriter,
Location loc,
mlir::Value oldVal,
mlir::Type newTy) {
@@ -386,7 +386,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
auto oldIteratorTypes = contractOp.getIteratorTypes();
SmallVector<Attribute> newIteratorTypes;
- // 0-th dim from the accuumlator
+ // 0-th dim from the accumulator
int64_t dimToDrop = oldIndexingMaps[2].getDimPosition(0);
if (!isParallelIterator(oldIteratorTypes[dimToDrop]))
@@ -520,9 +520,13 @@ namespace {
/// Turns vector.contract on vector with leading 1 dimensions into
/// 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
-/// prior to extract.
+/// 1 dimensions. Also performs transpose of lhs and rhs operands if required.
+///
+/// TODO: While the discussion on the validity of folding TransposeOp into
+/// ShapeCastOp continues, see e.g.
+/// * https://github.com/llvm/llvm-project/pull/219611,
+/// keep the explicit TransposeOp here. Once the discussion settles, revisit and
+/// consider replacing TransposeOp with ShapeCastOp.
struct CastAwayContractionLeadingOneDim
: public MaskableOpRewritePattern<vector::ContractionOp> {
using MaskableOpRewritePattern::MaskableOpRewritePattern;
More information about the Mlir-commits
mailing list