[Mlir-commits] [mlir] [mlir][vector] Migrate drop-lead-unit-dim to shape_cast (PR #196206)
Krzysztof Drewniak
llvmlistbot at llvm.org
Fri May 8 13:52:42 PDT 2026
https://github.com/krzysz00 updated https://github.com/llvm/llvm-project/pull/196206
>From 5b2ae1b3f3e48f6629d4f6a3ea36b98e069ab99a Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Wed, 6 May 2026 19:53:50 +0000
Subject: [PATCH 1/8] [mlir][vector] Migrate drop-lead-unit-dim to shape_cast
Post-merge discussion on #195686 led to the conclusion that we should
change the behavior of drealeadunitdim to use shape_cast instead of
extracts and broadcasts since those are now the canonical form of such
unit-dimension striping. This commit implements that change.
The one exception is that vector contractions where the accumulator is
reduced to a scalar still use extract/broadcast.
AI: Codex 5.5 did most of the work on this one.
---
.../Transforms/VectorDropLeadUnitDim.cpp | 260 +++++++------
.../vector-dropleadunitdim-transforms.mlir | 341 ++++++++++--------
.../Dialect/Vector/vector-transforms.mlir | 4 +-
3 files changed, 347 insertions(+), 258 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index 26a702ef0f512..931a87ff83e9c 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -22,6 +22,11 @@
using namespace mlir;
using namespace mlir::vector;
+/// Return a smallVector of size `rank` containing all zeros.
+static SmallVector<int64_t> splatZero(int64_t rank) {
+ return SmallVector<int64_t>(rank, 0);
+}
+
// Trims leading one dimensions from `oldType` and returns the result type.
// Returns `vector<1xT>` if `oldType` only has one element.
static VectorType trimLeadingOneDims(VectorType oldType) {
@@ -45,14 +50,82 @@ static VectorType trimLeadingOneDims(VectorType oldType) {
return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
}
-/// Return a smallVector of size `rank` containing all zeros.
-static SmallVector<int64_t> splatZero(int64_t rank) {
- return SmallVector<int64_t>(rank, 0);
+/// Returns `value` if it already has `newType`, otherwise inserts a
+/// vector.shape_cast to `newType`.
+static Value shapeCastVector(OpBuilder &b, Location loc, Value value,
+ VectorType newType) {
+ if (value.getType() == newType)
+ return value;
+ return vector::ShapeCastOp::create(b, loc, newType, value);
+}
+
+static bool hasNonScalableUnitLeadingDims(VectorType type, int64_t dropCount) {
+ if (dropCount < 0 || dropCount > type.getRank())
+ return false;
+ ArrayRef<int64_t> leadingShape = type.getShape().take_front(dropCount);
+ ArrayRef<bool> leadingScalable = type.getScalableDims().take_front(dropCount);
+ return llvm::all_of(leadingShape, [](int64_t dim) { return dim == 1; }) &&
+ llvm::none_of(leadingScalable,
+ [](bool scalable) { return scalable; });
+}
+
+static bool isNonScalableUnitDim(VectorType type, int64_t dim) {
+ return dim >= 0 && dim < type.getRank() && type.getShape()[dim] == 1 &&
+ !type.getScalableDims()[dim];
+}
+
+/// Shape-casts `operand` to the vector type obtained by dropping the first
+/// `dropCount` dimensions. Callers must ensure at least one vector dimension
+/// remains after the drop.
+static Value shapeCastDroppingLeadingDims(OpBuilder &b, Location loc,
+ Value operand, int64_t dropCount) {
+ auto oldType = cast<VectorType>(operand.getType());
+ assert(dropCount < oldType.getRank() &&
+ "shape_cast cannot drop all vector dimensions");
+ VectorType newType = VectorType::get(
+ oldType.getShape().drop_front(dropCount), oldType.getElementType(),
+ oldType.getScalableDims().drop_front(dropCount));
+ return shapeCastVector(b, loc, operand, newType);
+}
+
+static Value shapeCastDroppingDim(OpBuilder &b, Location loc, Value operand,
+ int64_t dim) {
+ auto oldType = cast<VectorType>(operand.getType());
+ assert(isNonScalableUnitDim(oldType, dim) &&
+ "expected a non-scalable unit dim to drop");
+
+ SmallVector<int64_t> newShape;
+ SmallVector<bool> newScalableDims;
+ for (int64_t i = 0, e = oldType.getRank(); i < e; ++i) {
+ if (i == dim)
+ continue;
+ newShape.push_back(oldType.getShape()[i]);
+ newScalableDims.push_back(oldType.getScalableDims()[i]);
+ }
+
+ return shapeCastVector(
+ b, loc, operand,
+ VectorType::get(newShape, oldType.getElementType(), newScalableDims));
+}
+
+static Value dropLeadingDimsForContraction(OpBuilder &b, Location loc,
+ Value operand, int64_t dropCount) {
+ auto oldType = cast<VectorType>(operand.getType());
+ assert(hasNonScalableUnitLeadingDims(oldType, dropCount) &&
+ "expected non-scalable leading unit dims to drop");
+
+ // vector.contract rejects 0-D vector accumulators/results. When every vector
+ // dimension is dropped, use the scalar path that vector.contract accepts.
+ if (dropCount == oldType.getRank())
+ return vector::ExtractOp::create(b, loc, operand, splatZero(dropCount));
+
+ return shapeCastDroppingLeadingDims(b, loc, operand, dropCount);
}
+
namespace {
// Casts away leading one dimensions in vector.extract_strided_slice's vector
-// input by inserting vector.broadcast.
+// input by inserting vector.shape_cast.
struct CastAwayExtractStridedSliceLeadingOneDim
: public OpRewritePattern<vector::ExtractStridedSliceOp> {
using Base::Base;
@@ -78,8 +151,8 @@ struct CastAwayExtractStridedSliceLeadingOneDim
Location loc = extractOp.getLoc();
- Value newSrcVector = vector::ExtractOp::create(
- rewriter, loc, extractOp.getSource(), splatZero(dropCount));
+ Value newSrcVector =
+ shapeCastVector(rewriter, loc, extractOp.getSource(), newSrcType);
// The offsets/sizes/strides attribute can have a less number of elements
// than the input vector's rank: it is meant for the leading dimensions.
@@ -94,7 +167,7 @@ struct CastAwayExtractStridedSliceLeadingOneDim
rewriter, loc, newDstType, newSrcVector, newOffsets, newSizes,
newStrides);
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(extractOp, oldDstType,
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(extractOp, oldDstType,
newExtractOp);
return success();
@@ -102,7 +175,7 @@ struct CastAwayExtractStridedSliceLeadingOneDim
};
// Casts away leading one dimensions in vector.insert_strided_slice's vector
-// inputs by inserting vector.broadcast.
+// inputs by inserting vector.shape_cast.
struct CastAwayInsertStridedSliceLeadingOneDim
: public OpRewritePattern<vector::InsertStridedSliceOp> {
using Base::Base;
@@ -122,10 +195,10 @@ struct CastAwayInsertStridedSliceLeadingOneDim
// Trim leading one dimensions from both operands.
Location loc = insertOp.getLoc();
- Value newSrcVector = vector::ExtractOp::create(
- rewriter, loc, insertOp.getValueToStore(), splatZero(srcDropCount));
- Value newDstVector = vector::ExtractOp::create(
- rewriter, loc, insertOp.getDest(), splatZero(dstDropCount));
+ Value newSrcVector = shapeCastVector(rewriter, loc,
+ insertOp.getValueToStore(), newSrcType);
+ Value newDstVector =
+ shapeCastVector(rewriter, loc, insertOp.getDest(), newDstType);
auto newOffsets = rewriter.getArrayAttr(
insertOp.getOffsets().getValue().take_back(newDstType.getRank()));
@@ -136,7 +209,7 @@ struct CastAwayInsertStridedSliceLeadingOneDim
rewriter, loc, newDstType, newSrcVector, newDstVector, newOffsets,
newStrides);
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(insertOp, oldDstType,
newInsertOp);
return success();
@@ -144,7 +217,7 @@ struct CastAwayInsertStridedSliceLeadingOneDim
};
// Casts away leading one dimensions in vector.insert's vector inputs by
-// inserting vector.broadcast.
+// inserting vector.shape_cast.
struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
using Base::Base;
@@ -171,12 +244,11 @@ struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
Location loc = insertOp.getLoc();
Value newSrcVector = insertOp.getValueToStore();
- if (oldSrcRank != 0) {
- newSrcVector = vector::ExtractOp::create(
- rewriter, loc, insertOp.getValueToStore(), splatZero(srcDropCount));
- }
- Value newDstVector = vector::ExtractOp::create(
- rewriter, loc, insertOp.getDest(), splatZero(dstDropCount));
+ if (oldSrcRank != 0)
+ newSrcVector = shapeCastVector(rewriter, loc, insertOp.getValueToStore(),
+ cast<VectorType>(newSrcType));
+ Value newDstVector =
+ shapeCastVector(rewriter, loc, insertOp.getDest(), newDstType);
// New position rank needs to be computed in two steps: (1) if destination
// type has leading unit dims, we also trim the position array accordingly,
@@ -193,7 +265,7 @@ struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
auto newInsertOp = vector::InsertOp::create(rewriter, loc, newSrcVector,
newDstVector, newPosition);
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(insertOp, oldDstType,
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(insertOp, oldDstType,
newInsertOp);
return success();
@@ -201,20 +273,10 @@ struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
};
static Value dropUnitDimsFromMask(OpBuilder &b, Location loc, Value mask,
- VectorType newType, AffineMap newMap,
- VectorType oldMaskType) {
+ VectorType newType, AffineMap newMap) {
// Infer the type of the new mask from the new map.
VectorType newMaskType = inferTransferOpMaskType(newType, newMap);
-
- // If the new mask is broadcastable to the old result type, we can safely
- // use a `vector.extract` to get the new mask. Otherwise the best we can
- // do is shape cast.
- if (vector::isBroadcastableTo(newMaskType, oldMaskType) ==
- BroadcastableToResult::Success) {
- int64_t dropDim = oldMaskType.getRank() - newMaskType.getRank();
- return vector::ExtractOp::create(b, loc, mask, splatZero(dropDim));
- }
- return vector::ShapeCastOp::create(b, loc, newMaskType, mask);
+ return shapeCastVector(b, loc, mask, newMaskType);
}
// Turns vector.transfer_read on vector with leading 1 dimensions into
@@ -256,16 +318,14 @@ struct CastAwayTransferReadLeadingOneDim
read.getInBoundsAttr().getValue().take_back(newType.getRank()));
Value mask = Value();
- if (read.getMask()) {
- VectorType maskType = read.getMaskType();
+ if (read.getMask())
mask = dropUnitDimsFromMask(rewriter, read.getLoc(), read.getMask(),
- newType, newMap, maskType);
- }
+ newType, newMap);
auto newRead = vector::TransferReadOp::create(
rewriter, read.getLoc(), newType, read.getBase(), read.getIndices(),
AffineMapAttr::get(newMap), read.getPadding(), mask, inBoundsAttr);
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(read, oldType, newRead);
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(read, oldType, newRead);
return success();
}
@@ -295,8 +355,6 @@ struct CastAwayTransferWriteLeadingOneDim
VectorType newType = trimLeadingOneDims(oldType);
if (newType == oldType)
return failure();
- int64_t dropDim = oldType.getRank() - newType.getRank();
-
AffineMap oldMap = write.getPermutationMap();
ArrayRef<AffineExpr> newResults =
oldMap.getResults().take_back(newType.getRank());
@@ -309,13 +367,12 @@ struct CastAwayTransferWriteLeadingOneDim
inBoundsAttr = rewriter.getArrayAttr(
write.getInBoundsAttr().getValue().take_back(newType.getRank()));
- auto newVector = vector::ExtractOp::create(
- rewriter, write.getLoc(), write.getVector(), splatZero(dropDim));
+ auto newVector =
+ shapeCastVector(rewriter, write.getLoc(), write.getVector(), newType);
if (write.getMask()) {
- VectorType maskType = write.getMaskType();
Value newMask = dropUnitDimsFromMask(
- rewriter, write.getLoc(), write.getMask(), newType, newMap, maskType);
+ rewriter, write.getLoc(), write.getMask(), newType, newMap);
rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
write, newVector, write.getBase(), write.getIndices(),
AffineMapAttr::get(newMap), newMask, inBoundsAttr);
@@ -340,7 +397,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
return failure();
if (oldAccType.getRank() < 1)
return failure();
- if (oldAccType.getShape()[0] != 1)
+ if (!isNonScalableUnitDim(oldAccType, 0))
return failure();
// currently we support only dropping one dim but the pattern can be applied
// greedily to drop more.
@@ -372,8 +429,8 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
for (const auto &it : llvm::enumerate(oldIndexingMaps)) {
// Check if the dim to be dropped exists as a leading dim in the operand
- // if it does then we use vector.extract to drop it.
- bool validExtract = false;
+ // if it does then we use vector.shape_cast to drop it.
+ bool needsShapeCast = false;
SmallVector<AffineExpr> results;
auto map = it.value();
int64_t orginalZeroDim = it.value().getDimPosition(0);
@@ -415,8 +472,8 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
}
}
- // Do the transpose now if needed so that we can drop the
- // correct dim using extract later.
+ // Do the transpose now if needed so that we can drop the correct dim
+ // with shape_cast later.
if (transposeNeeded) {
map = AffineMap::get(map.getNumDims(), 0, transposeResults,
contractOp.getContext());
@@ -428,10 +485,9 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
}
// We have taken care to have the dim to be dropped be
// the leading dim. If its still not leading that means it
- // does not exist in this operand and hence we do not need
- // an extract.
+ // does not exist in this operand and hence we do not need a shape_cast.
if (map.getDimPosition(0) == dimToDrop)
- validExtract = true;
+ needsShapeCast = true;
for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
int64_t currDim = map.getDimPosition(i);
@@ -444,13 +500,16 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
}
newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
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()]);
+ if (needsShapeCast) {
+ auto operandType = cast<VectorType>(operands[it.index()].getType());
+ if (operandType.getRank() < dropDim ||
+ !hasNonScalableUnitLeadingDims(operandType, dropDim))
+ return failure();
+ newOperands.push_back(dropLeadingDimsForContraction(
+ rewriter, loc, operands[it.index()], dropDim));
+ } else {
+ newOperands.push_back(operands[it.index()]);
+ }
}
// Depending on whether this vector.contract is masked, the replacing Op
@@ -461,12 +520,22 @@ 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 oldMaskType = cast<VectorType>(maskingOp.getMask().getType());
+ if (oldMaskType.getRank() <= 1 ||
+ !isNonScalableUnitDim(oldMaskType, dimToDrop))
+ return failure();
+ Value newMask =
+ shapeCastDroppingDim(rewriter, loc, maskingOp.getMask(), dimToDrop);
newOp = mlir::vector::maskOperation(rewriter, newOp, newMask);
}
+ if (isa<VectorType>(newOp->getResults()[0].getType()))
+ return vector::ShapeCastOp::create(rewriter, loc,
+ contractOp->getResultTypes()[0],
+ newOp->getResults()[0])
+ .getResult();
+
return vector::BroadcastOp::create(rewriter, loc,
contractOp->getResultTypes()[0],
newOp->getResults()[0])
@@ -476,9 +545,9 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
namespace {
/// Turns vector.contract on vector with leading 1 dimensions into
-/// vector.extract followed by vector.contract on vector without leading
+/// vector.shape_cast followed by vector.contract on vector without leading
/// 1 dimensions. Also performs transpose of lhs and rhs operands if required
-/// prior to extract.
+/// prior to the shape cast.
struct CastAwayContractionLeadingOneDim
: public MaskableOpRewritePattern<vector::ContractionOp> {
using MaskableOpRewritePattern::MaskableOpRewritePattern;
@@ -493,14 +562,15 @@ struct CastAwayContractionLeadingOneDim
/// Looks at elementwise operations on vectors with at least one leading
/// dimension equal 1, e.g. vector<1x[4]x1xf32> (but not vector<2x[4]x1xf32>),
-/// and cast aways the leading one dimensions (_plural_) and then broadcasts
-/// the results.
+/// and casts away the leading one dimensions (_plural_) with shape_cast.
///
/// Example before:
/// %1 = arith.mulf %arg0, %arg1 : vector<1x4x1xf32>
/// Example after:
-/// %2 = arith.mulf %0, %1 : vector<4x1xf32>
-/// %3 = vector.broadcast %2 : vector<4x1xf32> to vector<1x4x1xf32>
+/// %2 = vector.shape_cast %arg0 : vector<1x4x1xf32> to vector<4x1xf32>
+/// %3 = vector.shape_cast %arg1 : vector<1x4x1xf32> to vector<4x1xf32>
+/// %4 = arith.mulf %2, %3 : vector<4x1xf32>
+/// %5 = vector.shape_cast %4 : vector<4x1xf32> to vector<1x4x1xf32>
///
/// Does support scalable vectors.
class CastAwayElementwiseLeadingOneDim : public RewritePattern {
@@ -519,52 +589,29 @@ class CastAwayElementwiseLeadingOneDim : public RewritePattern {
VectorType newVecType = trimLeadingOneDims(vecType);
if (newVecType == vecType)
return failure();
- int64_t dropDim = vecType.getRank() - newVecType.getRank();
SmallVector<Value, 4> newOperands;
for (Value operand : op->getOperands()) {
- if (auto opVecType = dyn_cast<VectorType>(operand.getType())) {
- newOperands.push_back(vector::ExtractOp::create(
- rewriter, op->getLoc(), operand, splatZero(dropDim)));
- } else {
+ if (auto opVecType = dyn_cast<VectorType>(operand.getType()))
+ newOperands.push_back(shapeCastVector(
+ rewriter, op->getLoc(), operand, trimLeadingOneDims(opVecType)));
+ else
newOperands.push_back(operand);
- }
}
Operation *newOp =
rewriter.create(op->getLoc(), op->getName().getIdentifier(),
newOperands, newVecType, op->getAttrs());
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, vecType,
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, vecType,
newOp->getResult(0));
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.
+// Drops leading unit dimensions from load-like memory operations by
+// shape_casting each vector operand and shape_casting the result back to the
+// original type.
template <typename OpTy>
struct CastAwayLoadLikeLeadingOneDim : public OpRewritePattern<OpTy> {
using OpRewritePattern<OpTy>::OpRewritePattern;
@@ -583,7 +630,7 @@ struct CastAwayLoadLikeLeadingOneDim : public OpRewritePattern<OpTy> {
for (Value operand : op->getOperands()) {
if (isa<VectorType>(operand.getType())) {
newOperands.push_back(
- dropLeadingOneDimsFromOperand(rewriter, loc, operand, nDropped));
+ shapeCastDroppingLeadingDims(rewriter, loc, operand, nDropped));
} else {
newOperands.push_back(operand);
}
@@ -592,15 +639,14 @@ struct CastAwayLoadLikeLeadingOneDim : public OpRewritePattern<OpTy> {
Operation *newOp =
rewriter.create(loc, op->getName().getIdentifier(), newOperands,
TypeRange{newResultType}, op->getAttrs());
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(op, oldResultType,
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(op, oldResultType,
newOp->getResult(0));
return success();
}
};
-// Drops leading 1 dimensions from store-like memory ops. Extracts or
-// `shape_cast`s away those leading unit dimensions and leaves any scalar
-// operands alone.
+// Drops leading unit dimensions from store-like memory operations by
+// shape_casting each vector operand and leaving any scalar operands alone.
template <typename OpTy>
struct CastAwayStoreLikeLeadingOneDim : public OpRewritePattern<OpTy> {
using OpRewritePattern<OpTy>::OpRewritePattern;
@@ -619,7 +665,7 @@ struct CastAwayStoreLikeLeadingOneDim : public OpRewritePattern<OpTy> {
for (Value operand : op->getOperands()) {
if (isa<VectorType>(operand.getType())) {
newOperands.push_back(
- dropLeadingOneDimsFromOperand(rewriter, loc, operand, nDropped));
+ shapeCastDroppingLeadingDims(rewriter, loc, operand, nDropped));
} else {
newOperands.push_back(operand);
}
@@ -633,8 +679,8 @@ struct CastAwayStoreLikeLeadingOneDim : public OpRewritePattern<OpTy> {
}
};
-// Drops leading 1 dimensions from vector.constant_mask and inserts a
-// vector.broadcast back to the original shape.
+// Drops leading 1 dimensions from vector.constant_mask and shape_casts back to
+// the original shape.
struct CastAwayConstantMaskLeadingOneDim
: public OpRewritePattern<vector::ConstantMaskOp> {
using Base::Base;
@@ -659,7 +705,7 @@ struct CastAwayConstantMaskLeadingOneDim
auto newMask = vector::ConstantMaskOp::create(rewriter, mask.getLoc(),
newType, newDimSizes);
- rewriter.replaceOpWithNewOp<vector::BroadcastOp>(mask, oldType, newMask);
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(mask, oldType, newMask);
return success();
}
};
diff --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index bf01c8a8589d9..e5b1cc07319e3 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,15 +109,14 @@ 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: return %[[R5]] : vector<1x1x16xf32>
+// CHECK-NEXT: %[[R4:.+]] = vector.shape_cast %[[R3]] : vector<16xf32> to vector<1x1x16xf32>
+// CHECK-NEXT: return %[[R4]] : vector<1x1x16xf32>
#contraction_accesses1 = [
affine_map<(l, i, j, k) -> (i, l, k)>,
@@ -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: %[[R1:.+]] = vector.shape_cast %{{.*}} : vector<8x1x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[R3:.+]] = vector.shape_cast %{{.*}} : vector<2x8x1xf32> to vector<2x8xf32>
+// CHECK-NEXT: %[[R4:.+]] = 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-NEXT: %[[R6:.+]] = vector.shape_cast %[[R5]] : vector<2x16xf32> to vector<1x2x16xf32>
// CHECK-NEXT: return %[[R6]] : vector<1x2x16xf32>
#contraction_accesses2 = [
@@ -175,19 +172,14 @@ 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: %[[R3:.+]] = vector.shape_cast %{{.*}} : vector<1x8x1x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[R5:.+]] = vector.shape_cast %{{.*}} : vector<1x2x8x1xf32> to vector<2x8xf32>
+// CHECK-NEXT: %[[R6:.+]] = 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-NEXT: return %[[R9]] : vector<1x1x2x16xf32>
+// CHECK-NEXT: %[[R8:.+]] = vector.shape_cast %[[R7]] : vector<2x16xf32> to vector<1x1x2x16xf32>
+// CHECK-NEXT: return %[[R8]] : vector<1x1x2x16xf32>
#contraction_accesses2 = [
affine_map<(m, l, i, j, k) -> (m, k, l, j)>,
@@ -211,17 +203,14 @@ 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: %[[R2:.+]] = vector.shape_cast %{{.*}} : vector<1x8x1x16xf32> to vector<8x16xf32>
+// CHECK-NEXT: %[[R3:.+]] = vector.shape_cast %{{.*}} : vector<1x2x8x1xf32> to vector<2x8xf32>
+// CHECK-NEXT: %[[R4:.+]] = vector.shape_cast %{{.*}} : 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-NEXT: return %[[R7]] : vector<1x1x2x16xf32>
+// CHECK-NEXT: %[[R6:.+]] = vector.shape_cast %[[R5]] : vector<2x16xf32> to vector<1x1x2x16xf32>
+// CHECK-NEXT: return %[[R6]] : vector<1x1x2x16xf32>
#contraction_accesses3 = [
affine_map<(m, l, i, j, k) -> (m, k, l, j)>,
@@ -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>}
@@ -269,45 +258,82 @@ func.func @cast_away_contraction_leading_one_dims_to_dot_product(%arg0: vector<6
return %0 : vector<1xf32>
}
+// -----
+
+// CHECK-DAG: #[[$DOT_MAP:.*]] = affine_map<(d0) -> (d0)>
+// CHECK-DAG: #[[$SCALAR_MAP:.*]] = affine_map<(d0) -> ()>
+
+// CHECK-LABEL: cast_away_masked_contraction_with_rank1_acc
+// CHECK-NEXT: %[[RHS:.+]] = vector.shape_cast %{{.*}} : vector<1x64xf32> to vector<64xf32>
+// CHECK-NEXT: %[[ACC:.+]] = vector.extract %{{.*}}[0] : f32 from vector<1xf32>
+// CHECK-NEXT: %[[MASK:.+]] = vector.shape_cast %{{.*}} : vector<64x1xi1> to vector<64xi1>
+// CHECK-NEXT: %[[DOT:.+]] = vector.mask %[[MASK]] {
+// CHECK-SAME: vector.contract {indexing_maps = [#[[$DOT_MAP]], #[[$DOT_MAP]], #[[$SCALAR_MAP]]], iterator_types = ["reduction"], kind = #vector.kind<add>}
+// CHECK-SAME: %{{.*}}, %[[RHS]], %[[ACC]] : vector<64xf32>, vector<64xf32> into f32
+// CHECK-SAME: } : vector<64xi1> -> f32
+// CHECK-NEXT: %[[RES:.+]] = vector.broadcast %[[DOT]] : f32 to vector<1xf32>
+// CHECK-NEXT: return %[[RES]] : vector<1xf32>
+
+func.func @cast_away_masked_contraction_with_rank1_acc(%arg0: vector<64xf32>, %arg1: vector<1x64xf32>, %arg2: vector<1xf32>, %mask: vector<64x1xi1>) -> vector<1xf32> {
+ %0 = vector.mask %mask {
+ vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<64xf32>, vector<1x64xf32> into vector<1xf32>
+ } : vector<64x1xi1> -> vector<1xf32>
+ return %0 : vector<1xf32>
+}
+
+// -----
+
+// CHECK-LABEL: do_not_cast_away_contraction_with_scalable_rank1_acc
+// CHECK-NOT: vector.shape_cast
+// CHECK-NOT: vector.extract
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: vector.contract
+// CHECK-NEXT: return
+
+func.func @do_not_cast_away_contraction_with_scalable_rank1_acc(%arg0: vector<64xf32>, %arg1: vector<[1]x64xf32>, %arg2: vector<[1]xf32>) -> vector<[1]xf32> {
+ %0 = vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<64xf32>, vector<[1]x64xf32> into vector<[1]xf32>
+ return %0 : vector<[1]xf32>
+}
+
// -----
// CHECK-LABEL: func @cast_away_extract_strided_slice_leading_one_dims
func.func @cast_away_extract_strided_slice_leading_one_dims(%arg0: vector<1x8x8xf16>) -> vector<1x1x8xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<8x8xf16> from vector<1x8x8xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x8x8xf16> to vector<8x8xf16>
// CHECK: %[[EXTRACT:.+]] = vector.extract_strided_slice %[[SRC]] {offsets = [4], sizes = [1], strides = [1]} : vector<8x8xf16> to vector<1x8xf16>
%0 = vector.extract_strided_slice %arg0 {offsets = [0, 4], sizes = [1, 1], strides = [1, 1]} : vector<1x8x8xf16> to vector<1x1x8xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[EXTRACT]] : vector<1x8xf16> to vector<1x1x8xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[EXTRACT]] : vector<1x8xf16> to vector<1x1x8xf16>
// CHECK: return %[[RET]]
return %0: vector<1x1x8xf16>
}
// CHECK-LABEL: func @cast_away_extract_strided_slice_leading_one_dims_scalable
func.func @cast_away_extract_strided_slice_leading_one_dims_scalable(%arg0: vector<1x8x[8]xf16>) -> vector<1x1x[8]xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<8x[8]xf16> from vector<1x8x[8]xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x8x[8]xf16> to vector<8x[8]xf16>
// CHECK: %[[EXTRACT:.+]] = vector.extract_strided_slice %[[SRC]] {offsets = [4], sizes = [1], strides = [1]} : vector<8x[8]xf16> to vector<1x[8]xf16>
%0 = vector.extract_strided_slice %arg0 {offsets = [0, 4], sizes = [1, 1], strides = [1, 1]} : vector<1x8x[8]xf16> to vector<1x1x[8]xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[EXTRACT]] : vector<1x[8]xf16> to vector<1x1x[8]xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[EXTRACT]] : vector<1x[8]xf16> to vector<1x1x[8]xf16>
// CHECK: return %[[RET]]
return %0: vector<1x1x[8]xf16>
}
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims
func.func @cast_away_insert_strided_slice_leading_one_dims(%arg0: vector<1x8xf16>, %arg1: vector<1x8x8xf16>) -> vector<1x8x8xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<8xf16> from vector<1x8xf16>
- // CHECK: %[[DST:.+]] = vector.extract %{{.*}}[0] : vector<8x8xf16> from vector<1x8x8xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x8xf16> to vector<8xf16>
+ // CHECK: %[[DST:.+]] = vector.shape_cast %{{.*}} : vector<1x8x8xf16> to vector<8x8xf16>
// CHECK: %[[INSERT:.+]] = vector.insert_strided_slice %[[SRC]], %[[DST]] {offsets = [0, 0], strides = [1]} : vector<8xf16> into vector<8x8xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x8xf16> into vector<1x8x8xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[INSERT]] : vector<8x8xf16> to vector<1x8x8xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[INSERT]] : vector<8x8xf16> to vector<1x8x8xf16>
// CHECK: return %[[RET]]
return %0: vector<1x8x8xf16>
}
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims_scalable
func.func @cast_away_insert_strided_slice_leading_one_dims_scalable(%arg0: vector<1x[8]xf16>, %arg1: vector<1x8x[8]xf16>) -> vector<1x8x[8]xf16> {
- // CHECK: %[[SRC:.+]] = vector.extract %{{.*}}[0] : vector<[8]xf16> from vector<1x[8]xf16>
- // CHECK: %[[DST:.+]] = vector.extract %{{.*}}[0] : vector<8x[8]xf16> from vector<1x8x[8]xf16>
+ // CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x[8]xf16> to vector<[8]xf16>
+ // CHECK: %[[DST:.+]] = vector.shape_cast %{{.*}} : vector<1x8x[8]xf16> to vector<8x[8]xf16>
// CHECK: %[[INSERT:.+]] = vector.insert_strided_slice %[[SRC]], %[[DST]] {offsets = [0, 0], strides = [1]} : vector<[8]xf16> into vector<8x[8]xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x[8]xf16> into vector<1x8x[8]xf16>
- // CHECK: %[[RET:.+]] = vector.broadcast %[[INSERT]] : vector<8x[8]xf16> to vector<1x8x[8]xf16>
+ // CHECK: %[[RET:.+]] = vector.shape_cast %[[INSERT]] : vector<8x[8]xf16> to vector<1x8x[8]xf16>
// CHECK: return %[[RET]]
return %0: vector<1x8x[8]xf16>
}
@@ -315,8 +341,7 @@ func.func @cast_away_insert_strided_slice_leading_one_dims_scalable(%arg0: vecto
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims_one_element
// CHECK-SAME: %[[ARG0:.+]]: vector<1x1xf16>, %{{.+}}: vector<1x1x1xf16>
func.func @cast_away_insert_strided_slice_leading_one_dims_one_element(%arg0: vector<1x1xf16>, %arg1: vector<1x1x1xf16>) -> vector<1x1x1xf16> {
- // CHECK: %[[EXT:.+]] = vector.extract %{{.*}}[0] : vector<1xf16> from vector<1x1xf16>
- // CHECK: %[[B:.+]] = vector.broadcast %[[EXT]] : vector<1xf16> to vector<1x1x1xf16>
+ // CHECK: %[[B:.+]] = vector.shape_cast %{{.*}} : vector<1x1xf16> to vector<1x1x1xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x1xf16> into vector<1x1x1xf16>
// CHECK: return %[[B]]
return %0: vector<1x1x1xf16>
@@ -325,8 +350,7 @@ func.func @cast_away_insert_strided_slice_leading_one_dims_one_element(%arg0: ve
// CHECK-LABEL: func @cast_away_insert_strided_slice_leading_one_dims_one_element_scalable
// CHECK-SAME: %[[ARG0:.+]]: vector<1x[1]xf16>, %{{.+}}: vector<1x1x[1]xf16>
func.func @cast_away_insert_strided_slice_leading_one_dims_one_element_scalable(%arg0: vector<1x[1]xf16>, %arg1: vector<1x1x[1]xf16>) -> vector<1x1x[1]xf16> {
- // CHECK: %[[EXT:.+]] = vector.extract %{{.*}}[0] : vector<[1]xf16> from vector<1x[1]xf16>
- // CHECK: %[[B:.+]] = vector.broadcast %[[EXT]] : vector<[1]xf16> to vector<1x1x[1]xf16>
+ // CHECK: %[[B:.+]] = vector.shape_cast %{{.*}} : vector<1x[1]xf16> to vector<1x1x[1]xf16>
%0 = vector.insert_strided_slice %arg0, %arg1 {offsets = [0, 0, 0], strides = [1, 1]} : vector<1x[1]xf16> into vector<1x1x[1]xf16>
// CHECK: return %[[B]]
return %0: vector<1x1x[1]xf16>
@@ -339,7 +363,7 @@ func.func @cast_away_transfer_read_leading_one_dims(%arg0: memref<1x4x8x16xf16>)
// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f16
%f0 = arith.constant 0. : f16
// CHECK: %[[READ:.+]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]], %[[F0]] {in_bounds = [true]} : memref<1x4x8x16xf16>, vector<4xf16>
- // CHECK: %[[CAST:.+]] = vector.broadcast %[[READ]] : vector<4xf16> to vector<1x4xf16>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %[[READ]] : vector<4xf16> to vector<1x4xf16>
%0 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %f0 {in_bounds = [true, true]} : memref<1x4x8x16xf16>, vector<1x4xf16>
// CHECK: return %[[CAST]]
return %0: vector<1x4xf16>
@@ -351,9 +375,9 @@ func.func @cast_away_masked_transfer_read_leading_one_dims(%arg0: memref<1x4x8x1
%c0 = arith.constant 0 : index
// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f16
%f0 = arith.constant 0. : f16
- // CHECK: %[[MASK_CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
+ // CHECK: %[[MASK_CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
// CHECK: %[[READ:.+]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]], %[[F0]], %[[MASK_CAST]] {in_bounds = [true]} : memref<1x4x8x16xf16>, vector<4xf16>
- // CHECK: %[[CAST:.+]] = vector.broadcast %[[READ]] : vector<4xf16> to vector<1x4xf16>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %[[READ]] : vector<4xf16> to vector<1x4xf16>
%0 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %f0, %arg1 {in_bounds = [true, true]} : memref<1x4x8x16xf16>, vector<1x4xf16>
// CHECK: return %[[CAST]]
return %0: vector<1x4xf16>
@@ -363,7 +387,7 @@ func.func @cast_away_masked_transfer_read_leading_one_dims(%arg0: memref<1x4x8x1
func.func @cast_away_transfer_read_leading_one_dims_one_element(%arg0: memref<1x1x1x1xf16>) -> vector<1x1xf16> {
%c0 = arith.constant 0 : index
%f0 = arith.constant 0. : f16
- // CHECK: vector.broadcast %{{.+}} : vector<1xf16> to vector<1x1xf16>
+ // CHECK: vector.shape_cast %{{.+}} : vector<1xf16> to vector<1x1xf16>
%0 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %f0 {in_bounds = [true, true]} : memref<1x1x1x1xf16>, vector<1x1xf16>
return %0: vector<1x1xf16>
}
@@ -380,7 +404,7 @@ func.func @cast_away_nontrivial_map_masked_transfer_read(%arg0: memref<1x4x8xf16
// CHECK: %[[MASK_CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4x1xi1> to vector<4xi1>
// CHECK: %[[READ:.+]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %[[F0]], %[[MASK_CAST]] {in_bounds = [true]
// CHECK-SAME: permutation_map = #[[$MAP]]} : memref<1x4x8xf16>, vector<4xf16>
- // CHECK: %[[CAST:.+]] = vector.broadcast %[[READ]] : vector<4xf16> to vector<1x1x4xf16>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %[[READ]] : vector<4xf16> to vector<1x1x4xf16>
%0 = vector.transfer_read %arg0[%c0, %c0, %c0], %f0, %arg1 {in_bounds = [true, true, true],
permutation_map = affine_map<(d0, d1, d2) -> (d0, d2, d1)>} : memref<1x4x8xf16>, vector<1x1x4xf16>
// CHECK: return %[[CAST]]
@@ -391,7 +415,7 @@ func.func @cast_away_nontrivial_map_masked_transfer_read(%arg0: memref<1x4x8xf16
// CHECK-LABEL: func @not_insert_cast_fo4_transfer_read_under_mask
// CHECK: %[[MASK:.+]] = vector.constant_mask
-// CHECK: %[[CASTED_MASK:.+]] = vector.broadcast %[[MASK]]
+// CHECK: %[[CASTED_MASK:.+]] = vector.shape_cast %[[MASK]]
// CHECK: %[[RET:.+]] = vector.mask %[[CASTED_MASK]] {
// CHECK-SAME: vector.transfer_read {{.*}} : memref<1x1x4xf16>, vector<1x4xf16> }
// CHECK: return %[[RET]] : vector<1x4xf16>
@@ -411,7 +435,7 @@ func.func @not_insert_cast_fo4_transfer_read_under_mask(%arg0: memref<1x1x4xf16>
func.func @cast_away_transfer_write_leading_one_dims(%arg0: memref<1x4x8x16xf16>, %arg1: vector<1x4xf16>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- // CHECK: %[[CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xf16> from vector<1x4xf16>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf16> to vector<4xf16>
// CHECK: vector.transfer_write %[[CAST]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true]} : vector<4xf16>, memref<1x4x8x16xf16>
vector.transfer_write %arg1, %arg0[%c0, %c0, %c0, %c0] {in_bounds = [true, true]} : vector<1x4xf16>, memref<1x4x8x16xf16>
@@ -422,8 +446,8 @@ func.func @cast_away_transfer_write_leading_one_dims(%arg0: memref<1x4x8x16xf16>
func.func @cast_away_masked_transfer_write_leading_one_dims(%arg0: memref<1x4x8x16xf16>, %arg1: vector<1x4xf16>, %arg2: vector<1x4xi1>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- // CHECK: %[[CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xf16> from vector<1x4xf16>
- // CHECK: %[[MASK_CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf16> to vector<4xf16>
+ // CHECK: %[[MASK_CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
// CHECK: vector.transfer_write %[[CAST]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]], %[[MASK_CAST]] {in_bounds = [true]} : vector<4xf16>, memref<1x4x8x16xf16>
vector.transfer_write %arg1, %arg0[%c0, %c0, %c0, %c0], %arg2 {in_bounds = [true, true]} : vector<1x4xf16>, memref<1x4x8x16xf16>
@@ -433,7 +457,7 @@ func.func @cast_away_masked_transfer_write_leading_one_dims(%arg0: memref<1x4x8x
// CHECK-LABEL: func @cast_away_transfer_write_leading_one_dims_one_element
func.func @cast_away_transfer_write_leading_one_dims_one_element(%arg0: memref<1x1x1x1xf16>, %arg1: vector<1x1xf16>) {
%c0 = arith.constant 0 : index
- // CHECK: vector.extract %{{.+}}[0] : vector<1xf16> from vector<1x1xf16>
+ // CHECK: vector.shape_cast %{{.+}} : vector<1x1xf16> to vector<1xf16>
vector.transfer_write %arg1, %arg0[%c0, %c0, %c0, %c0] {in_bounds = [true, true]} : vector<1x1xf16>, memref<1x1x1x1xf16>
return
}
@@ -442,7 +466,7 @@ func.func @cast_away_transfer_write_leading_one_dims_one_element(%arg0: memref<1
// CHECK-LABEL: func @not_insert_cast_for_transfer_write_under_mask
// CHECK: %[[MASK:.+]] = vector.constant_mask
-// CHECK: %[[CASTED_MASK:.+]] = vector.broadcast %[[MASK]]
+// CHECK: %[[CASTED_MASK:.+]] = vector.shape_cast %[[MASK]]
// CHECK: vector.mask %[[CASTED_MASK]] {
// CHECK-SAME: vector.transfer_write {{.*}} : vector<1x4xf16>, memref<1x1x4xf16> }
// CHECK: return
@@ -462,7 +486,7 @@ func.func @not_insert_cast_for_transfer_write_under_mask(%arg0: memref<1x1x4xf16
func.func @cast_away_nontrivial_map_masked_transfer_write(%arg0: memref<1x4x8xf16>, %arg1: vector<1x1x4xf16>, %arg2: vector<1x4x1xi1>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- // CHECK: %[[CAST:.+]] = vector.extract %{{.*}}[0, 0] : vector<4xf16> from vector<1x1x4xf16>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x1x4xf16> to vector<4xf16>
// CHECK: %[[MASK_CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4x1xi1> to vector<4xi1>
// CHECK: vector.transfer_write %[[CAST]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %[[MASK_CAST]] {in_bounds = [true]
// CHECK-SAME: permutation_map = #[[$MAP]]} : vector<4xf16>, memref<1x4x8xf16>
@@ -479,25 +503,25 @@ func.func @cast_away_elementwise_leading_one_dims(
%arg0: vector<1x1x8xf32>, %arg1: f32, %arg2: vector<1x4xf32>,
%arg3: vector<1x4xf32>, %arg4: i1) ->
(vector<1x1x8xf32>, vector<1x4xi1>, vector<1x4xf32>, vector<1x4xf32>) {
- // CHECK: vector.extract %{{.*}}[0, 0] : vector<8xf32> from vector<1x1x8xf32>
- // CHECK: vector.extract %{{.*}}[0, 0] : vector<8xf32> from vector<1x1x8xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x1x8xf32> to vector<8xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x1x8xf32> to vector<8xf32>
// CHECK: arith.addf %{{.*}}, %{{.*}} : vector<8xf32>
- // CHECK: vector.broadcast %{{.*}} : vector<8xf32> to vector<1x1x8xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<8xf32> to vector<1x1x8xf32>
%0 = arith.addf %arg0, %arg0 : vector<1x1x8xf32>
- // CHECK: vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
- // CHECK: vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: arith.cmpf ogt, %{{.*}}, %{{.*}} : vector<4xf32>
- // CHECK: vector.broadcast %{{.*}} : vector<4xi1> to vector<1x4xi1>
+ // CHECK: vector.shape_cast %{{.*}} : vector<4xi1> to vector<1x4xi1>
%1 = arith.cmpf ogt, %arg2, %arg3 : vector<1x4xf32>
- // CHECK: vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
- // CHECK: vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: select %{{.*}}, %{{.*}}, %{{.*}} : vector<4xi1>, vector<4xf32>
- // CHECK: vector.broadcast %{{.*}} : vector<4xf32> to vector<1x4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<4xf32> to vector<1x4xf32>
%2 = arith.select %1, %arg3, %arg2 : vector<1x4xi1>, vector<1x4xf32>
- // CHECK: vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
- // CHECK: vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: select %arg4, %12, %{{.*}} : vector<4xf32>
- // CHECK: vector.broadcast %{{.*}} : vector<4xf32> to vector<1x4xf32>
+ // CHECK: vector.shape_cast %{{.*}} : vector<4xf32> to vector<1x4xf32>
%3 = arith.select %arg4, %arg3, %arg2 : vector<1x4xf32>
return %0, %1, %2, %3: vector<1x1x8xf32>, vector<1x4xi1>, vector<1x4xf32>, vector<1x4xf32>
}
@@ -506,10 +530,10 @@ func.func @cast_away_elementwise_leading_one_dims(
// CHECK-LABEL: func @cast_away_insert_leading_one_dims_scalar
// CHECK-SAME: (%[[S:.+]]: f32, %[[V:.+]]: vector<1x1x4xf32>)
-// CHECK: %[[EXTRACT:.+]] = vector.extract %[[V]][0, 0] : vector<4xf32> from vector<1x1x4xf32>
-// CHECK: %[[INSERT:.+]] = vector.insert %[[S]], %[[EXTRACT]] [0] : f32 into vector<4xf32>
-// CHECK: %[[BCAST:.+]] = vector.broadcast %[[INSERT]] : vector<4xf32> to vector<1x1x4xf32>
-// CHECK: return %[[BCAST]]
+// CHECK: %[[DST_CAST:.+]] = vector.shape_cast %[[V]] : vector<1x1x4xf32> to vector<4xf32>
+// CHECK: %[[INSERT:.+]] = vector.insert %[[S]], %[[DST_CAST]] [0] : f32 into vector<4xf32>
+// CHECK: %[[RESULT_CAST:.+]] = vector.shape_cast %[[INSERT]] : vector<4xf32> to vector<1x1x4xf32>
+// CHECK: return %[[RESULT_CAST]]
func.func @cast_away_insert_leading_one_dims_scalar(%s: f32, %v: vector<1x1x4xf32>) -> vector<1x1x4xf32> {
%0 = vector.insert %s, %v [0, 0, 0] : f32 into vector<1x1x4xf32>
return %0: vector<1x1x4xf32>
@@ -521,10 +545,10 @@ func.func @cast_away_insert_leading_one_dims_scalar(%s: f32, %v: vector<1x1x4xf3
// CHECK-SAME: %[[S:.*]]: f32,
// CHECK-SAME: %[[V:.*]]: vector<1x1x[4]xf32>) -> vector<1x1x[4]xf32> {
func.func @cast_away_insert_leading_one_dims_scalar_scalable(%s: f32, %v: vector<1x1x[4]xf32>) -> vector<1x1x[4]xf32> {
-// CHECK: %[[EXTRACT:.*]] = vector.extract %[[V]][0, 0] : vector<[4]xf32> from vector<1x1x[4]xf32>
-// CHECK: %[[INSERT:.*]] = vector.insert %[[S]], %[[EXTRACT]] [0] : f32 into vector<[4]xf32>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[INSERT]] : vector<[4]xf32> to vector<1x1x[4]xf32>
-// CHECK: return %[[BCAST]] : vector<1x1x[4]xf32>
+// CHECK: %[[DST_CAST:.*]] = vector.shape_cast %[[V]] : vector<1x1x[4]xf32> to vector<[4]xf32>
+// CHECK: %[[INSERT:.*]] = vector.insert %[[S]], %[[DST_CAST]] [0] : f32 into vector<[4]xf32>
+// CHECK: %[[RESULT_CAST:.*]] = vector.shape_cast %[[INSERT]] : vector<[4]xf32> to vector<1x1x[4]xf32>
+// CHECK: return %[[RESULT_CAST]] : vector<1x1x[4]xf32>
%0 = vector.insert %s, %v [0, 0, 0] : f32 into vector<1x1x[4]xf32>
return %0: vector<1x1x[4]xf32>
}
@@ -535,10 +559,10 @@ func.func @cast_away_insert_leading_one_dims_scalar_scalable(%s: f32, %v: vector
// CHECK-SAME: %[[S:.*]]: f32,
// CHECK-SAME: %[[V:.*]]: vector<1x[1]x4xf32>) -> vector<1x[1]x4xf32> {
func.func @cast_away_insert_leading_one_dims_scalar_skip_scalable_dim(%s: f32, %v: vector<1x[1]x4xf32>) -> vector<1x[1]x4xf32> {
-// CHECK: %[[EXTRACT:.*]] = vector.extract %[[V]][0] : vector<[1]x4xf32> from vector<1x[1]x4xf32>
-// CHECK: %[[INSERT:.*]] = vector.insert %[[S]], %[[EXTRACT]] [0, 0] : f32 into vector<[1]x4xf32>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[INSERT]] : vector<[1]x4xf32> to vector<1x[1]x4xf32>
-// CHECK: return %[[BCAST]] : vector<1x[1]x4xf32>
+// CHECK: %[[DST_CAST:.*]] = vector.shape_cast %[[V]] : vector<1x[1]x4xf32> to vector<[1]x4xf32>
+// CHECK: %[[INSERT:.*]] = vector.insert %[[S]], %[[DST_CAST]] [0, 0] : f32 into vector<[1]x4xf32>
+// CHECK: %[[RESULT_CAST:.*]] = vector.shape_cast %[[INSERT]] : vector<[1]x4xf32> to vector<1x[1]x4xf32>
+// CHECK: return %[[RESULT_CAST]] : vector<1x[1]x4xf32>
%0 = vector.insert %s, %v [0, 0, 0] : f32 into vector<1x[1]x4xf32>
return %0: vector<1x[1]x4xf32>
}
@@ -547,8 +571,8 @@ func.func @cast_away_insert_leading_one_dims_scalar_skip_scalable_dim(%s: f32, %
// CHECK-LABEL: func @cast_away_insert_leading_one_dims_rank1
// CHECK-SAME: (%[[S:.+]]: vector<4xf32>, %[[V:.+]]: vector<1x1x4xf32>)
-// CHECK: %[[BCAST:.+]] = vector.broadcast %[[S]] : vector<4xf32> to vector<1x1x4xf32>
-// CHECK: return %[[BCAST]]
+// CHECK: %[[RESULT_CAST:.+]] = vector.shape_cast %[[S]] : vector<4xf32> to vector<1x1x4xf32>
+// CHECK: return %[[RESULT_CAST]]
func.func @cast_away_insert_leading_one_dims_rank1(%s: vector<4xf32>, %v: vector<1x1x4xf32>) -> vector<1x1x4xf32> {
%0 = vector.insert %s, %v [0, 0] : vector<4xf32> into vector<1x1x4xf32>
return %0: vector<1x1x4xf32>
@@ -559,8 +583,8 @@ func.func @cast_away_insert_leading_one_dims_rank1(%s: vector<4xf32>, %v: vector
// CHECK-LABEL: func.func @cast_away_insert_leading_one_dims_rank1_scalable(
// CHECK-SAME: %[[S:.*]]: vector<[4]xf32>,
// CHECK-SAME: %[[V:.*]]: vector<1x1x[4]xf32>) -> vector<1x1x[4]xf32> {
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[S]] : vector<[4]xf32> to vector<1x1x[4]xf32>
-// CHECK: return %[[BCAST]] : vector<1x1x[4]xf32>
+// CHECK: %[[RESULT_CAST:.*]] = vector.shape_cast %[[S]] : vector<[4]xf32> to vector<1x1x[4]xf32>
+// CHECK: return %[[RESULT_CAST]] : vector<1x1x[4]xf32>
func.func @cast_away_insert_leading_one_dims_rank1_scalable(%s: vector<[4]xf32>, %v: vector<1x1x[4]xf32>) -> vector<1x1x[4]xf32> {
%0 = vector.insert %s, %v [0, 0] : vector<[4]xf32> into vector<1x1x[4]xf32>
return %0: vector<1x1x[4]xf32>
@@ -570,9 +594,8 @@ func.func @cast_away_insert_leading_one_dims_rank1_scalable(%s: vector<[4]xf32>,
// CHECK-LABEL: func @cast_away_insert_leading_one_dims_rank2
// CHECK-SAME: (%[[S:.+]]: vector<1x4xf32>, %[[V:.+]]: vector<1x1x4xf32>)
-// CHECK: %[[EXTRACT:.+]] = vector.extract %[[S]][0] : vector<4xf32> from vector<1x4xf32>
-// CHECK: %[[BCAST:.+]] = vector.broadcast %[[EXTRACT]] : vector<4xf32> to vector<1x1x4xf32>
-// CHECK: return %[[BCAST]]
+// CHECK: %[[SRC_CAST:.+]] = vector.shape_cast %[[S]] : vector<1x4xf32> to vector<1x1x4xf32>
+// CHECK: return %[[SRC_CAST]]
func.func @cast_away_insert_leading_one_dims_rank2(%s: vector<1x4xf32>, %v: vector<1x1x4xf32>) -> vector<1x1x4xf32> {
%0 = vector.insert %s, %v [0] : vector<1x4xf32> into vector<1x1x4xf32>
return %0: vector<1x1x4xf32>
@@ -583,9 +606,8 @@ func.func @cast_away_insert_leading_one_dims_rank2(%s: vector<1x4xf32>, %v: vect
// CHECK-LABEL: func.func @cast_away_insert_leading_one_dims_rank2_scalable(
// CHECK-SAME: %[[S:.*]]: vector<1x[4]xf32>,
// CHECK-SAME: %[[V:.*]]: vector<1x1x[4]xf32>) -> vector<1x1x[4]xf32> {
-// CHECK: %[[EXTRACT:.*]] = vector.extract %[[S]][0] : vector<[4]xf32> from vector<1x[4]xf32>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[EXTRACT]] : vector<[4]xf32> to vector<1x1x[4]xf32>
-// CHECK: return %[[BCAST]] : vector<1x1x[4]xf32>
+// CHECK: %[[SRC_CAST:.*]] = vector.shape_cast %[[S]] : vector<1x[4]xf32> to vector<1x1x[4]xf32>
+// CHECK: return %[[SRC_CAST]] : vector<1x1x[4]xf32>
func.func @cast_away_insert_leading_one_dims_rank2_scalable(%s: vector<1x[4]xf32>, %v: vector<1x1x[4]xf32>) -> vector<1x1x[4]xf32> {
%0 = vector.insert %s, %v [0] : vector<1x[4]xf32> into vector<1x1x[4]xf32>
return %0: vector<1x1x[4]xf32>
@@ -595,11 +617,11 @@ func.func @cast_away_insert_leading_one_dims_rank2_scalable(%s: vector<1x[4]xf32
// CHECK-LABEL: func @cast_away_insert_leading_one_dims_rank2_one_dest
// CHECK-SAME: (%[[S:.+]]: vector<1x4xf32>, %[[V:.+]]: vector<1x2x1x4xf32>)
-// CHECK: %[[EXTRACTS:.+]] = vector.extract %[[S]][0] : vector<4xf32> from vector<1x4xf32>
-// CHECK: %[[EXTRACTV:.+]] = vector.extract %[[V]][0] : vector<2x1x4xf32> from vector<1x2x1x4xf32>
-// CHECK: %[[INSERT:.+]] = vector.insert %[[EXTRACTS]], %[[EXTRACTV]] [1, 0] : vector<4xf32> into vector<2x1x4xf32>
-// CHECK: %[[BCAST:.+]] = vector.broadcast %[[INSERT]] : vector<2x1x4xf32> to vector<1x2x1x4xf32>
-// CHECK: return %[[BCAST]]
+// CHECK: %[[SRC_CAST:.+]] = vector.shape_cast %[[S]] : vector<1x4xf32> to vector<4xf32>
+// CHECK: %[[DST_CAST:.+]] = vector.shape_cast %[[V]] : vector<1x2x1x4xf32> to vector<2x1x4xf32>
+// CHECK: %[[INSERT:.+]] = vector.insert %[[SRC_CAST]], %[[DST_CAST]] [1, 0] : vector<4xf32> into vector<2x1x4xf32>
+// CHECK: %[[RESULT_CAST:.+]] = vector.shape_cast %[[INSERT]] : vector<2x1x4xf32> to vector<1x2x1x4xf32>
+// CHECK: return %[[RESULT_CAST]]
func.func @cast_away_insert_leading_one_dims_rank2_one_dest(%s: vector<1x4xf32>, %v: vector<1x2x1x4xf32>) -> vector<1x2x1x4xf32> {
%0 = vector.insert %s, %v [0, 1] : vector<1x4xf32> into vector<1x2x1x4xf32>
return %0: vector<1x2x1x4xf32>
@@ -610,11 +632,11 @@ func.func @cast_away_insert_leading_one_dims_rank2_one_dest(%s: vector<1x4xf32>,
// CHECK-LABEL: func.func @cast_away_insert_leading_one_dims_rank2_one_dest_scalable(
// CHECK-SAME: %[[S:.*]]: vector<1x[4]xf32>,
// CHECK-SAME: %[[V:.*]]: vector<1x2x1x[4]xf32>) -> vector<1x2x1x[4]xf32> {
-// CHECK: %[[EXTRACTS:.*]] = vector.extract %[[S]][0] : vector<[4]xf32> from vector<1x[4]xf32>
-// CHECK: %[[EXTRACTV:.*]] = vector.extract %[[V]][0] : vector<2x1x[4]xf32> from vector<1x2x1x[4]xf32>
-// CHECK: %[[INSERT:.*]] = vector.insert %[[EXTRACTS]], %[[EXTRACTV]] [1, 0] : vector<[4]xf32> into vector<2x1x[4]xf32>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[INSERT]] : vector<2x1x[4]xf32> to vector<1x2x1x[4]xf32>
-// CHECK: return %[[BCAST]] : vector<1x2x1x[4]xf32>
+// CHECK: %[[SRC_CAST:.*]] = vector.shape_cast %[[S]] : vector<1x[4]xf32> to vector<[4]xf32>
+// CHECK: %[[DST_CAST:.*]] = vector.shape_cast %[[V]] : vector<1x2x1x[4]xf32> to vector<2x1x[4]xf32>
+// CHECK: %[[INSERT:.*]] = vector.insert %[[SRC_CAST]], %[[DST_CAST]] [1, 0] : vector<[4]xf32> into vector<2x1x[4]xf32>
+// CHECK: %[[RESULT_CAST:.*]] = vector.shape_cast %[[INSERT]] : vector<2x1x[4]xf32> to vector<1x2x1x[4]xf32>
+// CHECK: return %[[RESULT_CAST]] : vector<1x2x1x[4]xf32>
func.func @cast_away_insert_leading_one_dims_rank2_one_dest_scalable(%s: vector<1x[4]xf32>, %v: vector<1x2x1x[4]xf32>) -> vector<1x2x1x[4]xf32> {
%0 = vector.insert %s, %v [0, 1] : vector<1x[4]xf32> into vector<1x2x1x[4]xf32>
return %0: vector<1x2x1x[4]xf32>
@@ -624,8 +646,8 @@ func.func @cast_away_insert_leading_one_dims_rank2_one_dest_scalable(%s: vector<
// CHECK-LABEL: func @cast_away_insert_leading_one_dims_non_one_dest
// CHECK-SAME: (%[[S:.+]]: vector<1x4xf32>, %[[V:.+]]: vector<8x1x4xf32>)
-// CHECK: %[[EXTRACT:.+]] = vector.extract %[[S]][0] : vector<4xf32> from vector<1x4xf32>
-// CHECK: %[[INSERT:.+]] = vector.insert %[[EXTRACT]], %[[V]] [5, 0] : vector<4xf32> into vector<8x1x4xf32>
+// CHECK: %[[SRC_CAST:.+]] = vector.shape_cast %[[S]] : vector<1x4xf32> to vector<4xf32>
+// CHECK: %[[INSERT:.+]] = vector.insert %[[SRC_CAST]], %[[V]] [5, 0] : vector<4xf32> into vector<8x1x4xf32>
// CHECK: return %[[INSERT]]
func.func @cast_away_insert_leading_one_dims_non_one_dest(%s: vector<1x4xf32>, %v: vector<8x1x4xf32>) -> vector<8x1x4xf32> {
%0 = vector.insert %s, %v [5] : vector<1x4xf32> into vector<8x1x4xf32>
@@ -637,8 +659,8 @@ func.func @cast_away_insert_leading_one_dims_non_one_dest(%s: vector<1x4xf32>, %
// CHECK-LABEL: func.func @cast_away_insert_leading_one_dims_non_one_dest_scalable(
// CHECK-SAME: %[[S:.*]]: vector<1x[4]xf32>,
// CHECK-SAME: %[[V:.*]]: vector<8x1x[4]xf32>) -> vector<8x1x[4]xf32> {
-// CHECK: %[[EXTRACT:.*]] = vector.extract %[[S]][0] : vector<[4]xf32> from vector<1x[4]xf32>
-// CHECK: %[[INSERT:.*]] = vector.insert %[[EXTRACT]], %[[V]] [5, 0] : vector<[4]xf32> into vector<8x1x[4]xf32>
+// CHECK: %[[SRC_CAST:.*]] = vector.shape_cast %[[S]] : vector<1x[4]xf32> to vector<[4]xf32>
+// CHECK: %[[INSERT:.*]] = vector.insert %[[SRC_CAST]], %[[V]] [5, 0] : vector<[4]xf32> into vector<8x1x[4]xf32>
// CHECK: return %[[INSERT]] : vector<8x1x[4]xf32>
func.func @cast_away_insert_leading_one_dims_non_one_dest_scalable(%s: vector<1x[4]xf32>, %v: vector<8x1x[4]xf32>) -> vector<8x1x[4]xf32> {
%0 = vector.insert %s, %v [5] : vector<1x[4]xf32> into vector<8x1x[4]xf32>
@@ -649,11 +671,11 @@ func.func @cast_away_insert_leading_one_dims_non_one_dest_scalable(%s: vector<1x
// CHECK-LABEL: func @cast_away_insert_leading_one_dims_one_two_dest
// CHECK-SAME: (%[[S:.+]]: vector<1x8xi1>, %[[V:.+]]: vector<1x1x8x1x8xi1>)
-// CHECK: %[[EXTRACTS:.+]] = vector.extract %[[S]][0] : vector<8xi1> from vector<1x8xi1>
-// CHECK: %[[EXTRACTV:.+]] = vector.extract %[[V]][0, 0] : vector<8x1x8xi1> from vector<1x1x8x1x8xi1>
-// CHECK: %[[INSERT:.+]] = vector.insert %[[EXTRACTS]], %[[EXTRACTV]] [7, 0] : vector<8xi1> into vector<8x1x8xi1>
-// CHECK: %[[BCAST:.+]] = vector.broadcast %[[INSERT]] : vector<8x1x8xi1> to vector<1x1x8x1x8xi1>
-// CHECK: return %[[BCAST]]
+// CHECK: %[[SRC_CAST:.+]] = vector.shape_cast %[[S]] : vector<1x8xi1> to vector<8xi1>
+// CHECK: %[[DST_CAST:.+]] = vector.shape_cast %[[V]] : vector<1x1x8x1x8xi1> to vector<8x1x8xi1>
+// CHECK: %[[INSERT:.+]] = vector.insert %[[SRC_CAST]], %[[DST_CAST]] [7, 0] : vector<8xi1> into vector<8x1x8xi1>
+// CHECK: %[[RESULT_CAST:.+]] = vector.shape_cast %[[INSERT]] : vector<8x1x8xi1> to vector<1x1x8x1x8xi1>
+// CHECK: return %[[RESULT_CAST]]
func.func @cast_away_insert_leading_one_dims_one_two_dest(%s: vector<1x8xi1>, %v: vector<1x1x8x1x8xi1>) -> vector<1x1x8x1x8xi1> {
%0 = vector.insert %s, %v [0, 0, 7] : vector<1x8xi1> into vector<1x1x8x1x8xi1>
return %0: vector<1x1x8x1x8xi1>
@@ -664,11 +686,11 @@ func.func @cast_away_insert_leading_one_dims_one_two_dest(%s: vector<1x8xi1>, %v
// CHECK-LABEL: func.func @cast_away_insert_leading_one_dims_one_two_dest_scalable(
// CHECK-SAME: %[[S:.*]]: vector<1x[8]xi1>,
// CHECK-SAME: %[[V:.*]]: vector<1x1x8x1x[8]xi1>) -> vector<1x1x8x1x[8]xi1> {
-// CHECK: %[[EXTRACTS:.*]] = vector.extract %[[S]][0] : vector<[8]xi1> from vector<1x[8]xi1>
-// CHECK: %[[EXTRACTV:.*]] = vector.extract %[[V]][0, 0] : vector<8x1x[8]xi1> from vector<1x1x8x1x[8]xi1>
-// CHECK: %[[INSERT:.*]] = vector.insert %[[EXTRACTS]], %[[EXTRACTV]] [7, 0] : vector<[8]xi1> into vector<8x1x[8]xi1>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[INSERT]] : vector<8x1x[8]xi1> to vector<1x1x8x1x[8]xi1>
-// CHECK: return %[[BCAST]] : vector<1x1x8x1x[8]xi1>
+// CHECK: %[[SRC_CAST:.*]] = vector.shape_cast %[[S]] : vector<1x[8]xi1> to vector<[8]xi1>
+// CHECK: %[[DST_CAST:.*]] = vector.shape_cast %[[V]] : vector<1x1x8x1x[8]xi1> to vector<8x1x[8]xi1>
+// CHECK: %[[INSERT:.*]] = vector.insert %[[SRC_CAST]], %[[DST_CAST]] [7, 0] : vector<[8]xi1> into vector<8x1x[8]xi1>
+// CHECK: %[[RESULT_CAST:.*]] = vector.shape_cast %[[INSERT]] : vector<8x1x[8]xi1> to vector<1x1x8x1x[8]xi1>
+// CHECK: return %[[RESULT_CAST]] : vector<1x1x8x1x[8]xi1>
func.func @cast_away_insert_leading_one_dims_one_two_dest_scalable(%s: vector<1x[8]xi1>, %v: vector<1x1x8x1x[8]xi1>) -> vector<1x1x8x1x[8]xi1> {
%0 = vector.insert %s, %v [0, 0, 7] : vector<1x[8]xi1> into vector<1x1x8x1x[8]xi1>
return %0: vector<1x1x8x1x[8]xi1>
@@ -678,8 +700,8 @@ func.func @cast_away_insert_leading_one_dims_one_two_dest_scalable(%s: vector<1x
// CHECK-LABEL: func.func @cast_away_constant_mask() -> vector<1x1x8x2x1xi1> {
// CHECK: %[[MASK:.*]] = vector.constant_mask [6, 1, 1] : vector<8x2x1xi1>
-// CHECK: %[[BCAST:.*]] = vector.broadcast %[[MASK]] : vector<8x2x1xi1> to vector<1x1x8x2x1xi1>
-// CHECK: return %[[BCAST]] : vector<1x1x8x2x1xi1>
+// CHECK: %[[MASK_CAST:.*]] = vector.shape_cast %[[MASK]] : vector<8x2x1xi1> to vector<1x1x8x2x1xi1>
+// CHECK: return %[[MASK_CAST]] : vector<1x1x8x2x1xi1>
func.func @cast_away_constant_mask() -> vector<1x1x8x2x1xi1> {
%0 = vector.constant_mask [1, 1, 6, 1, 1] : vector<1x1x8x2x1xi1>
return %0: vector<1x1x8x2x1xi1>
@@ -698,7 +720,7 @@ func.func @drop_unit_dims_scalar_cond_select(%cond: i1, %arg0: vector<1x16xi1>,
// CHECK-LABEL: func.func @cast_away_load_leading_one_dims
// CHECK: %[[L:.+]] = vector.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<8x16xf32>, vector<4xf32>
-// CHECK: %[[B:.+]] = vector.broadcast %[[L]] : vector<4xf32> to vector<1x4xf32>
+// CHECK: %[[B:.+]] = vector.shape_cast %[[L]] : vector<4xf32> to vector<1x4xf32>
// CHECK: return %[[B]] : vector<1x4xf32>
func.func @cast_away_load_leading_one_dims(%base: memref<8x16xf32>, %i: index, %j: index) -> vector<1x4xf32> {
%0 = vector.load %base[%i, %j] : memref<8x16xf32>, vector<1x4xf32>
@@ -707,11 +729,22 @@ func.func @cast_away_load_leading_one_dims(%base: memref<8x16xf32>, %i: index, %
// -----
+// CHECK-LABEL: func.func @cast_away_load_leading_one_dims_scalable
+// CHECK: %[[L:.+]] = vector.load %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>, vector<[4]xf32>
+// CHECK: %[[B:.+]] = vector.shape_cast %[[L]] : vector<[4]xf32> to vector<1x[4]xf32>
+// CHECK: return %[[B]] : vector<1x[4]xf32>
+func.func @cast_away_load_leading_one_dims_scalable(%base: memref<?x?xf32>, %i: index, %j: index) -> vector<1x[4]xf32> {
+ %0 = vector.load %base[%i, %j] : memref<?x?xf32>, vector<1x[4]xf32>
+ return %0 : vector<1x[4]xf32>
+}
+
+// -----
+
// CHECK-LABEL: func.func @cast_away_maskedload_leading_one_dims
-// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
-// CHECK: %[[P:.+]] = vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+// CHECK: %[[M:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
+// CHECK: %[[P:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: %[[L:.+]] = vector.maskedload %{{.*}}[%{{.*}}], %[[M]], %[[P]] : memref<16xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
-// CHECK: %[[B:.+]] = vector.broadcast %[[L]] : vector<4xf32> to vector<1x4xf32>
+// CHECK: %[[B:.+]] = vector.shape_cast %[[L]] : vector<4xf32> to vector<1x4xf32>
// CHECK: return %[[B]] : vector<1x4xf32>
func.func @cast_away_maskedload_leading_one_dims(%base: memref<16xf32>, %i: index, %mask: vector<1x4xi1>, %pass: vector<1x4xf32>) -> vector<1x4xf32> {
%0 = vector.maskedload %base[%i], %mask, %pass : memref<16xf32>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>
@@ -721,10 +754,10 @@ func.func @cast_away_maskedload_leading_one_dims(%base: memref<16xf32>, %i: inde
// -----
// CHECK-LABEL: func.func @cast_away_expandload_leading_one_dims
-// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
-// CHECK: %[[P:.+]] = vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+// CHECK: %[[M:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
+// CHECK: %[[P:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: %[[L:.+]] = vector.expandload %{{.*}}[%{{.*}}], %[[M]], %[[P]] : memref<16xf32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
-// CHECK: %[[B:.+]] = vector.broadcast %[[L]] : vector<4xf32> to vector<1x4xf32>
+// CHECK: %[[B:.+]] = vector.shape_cast %[[L]] : vector<4xf32> to vector<1x4xf32>
// CHECK: return %[[B]] : vector<1x4xf32>
func.func @cast_away_expandload_leading_one_dims(%base: memref<16xf32>, %i: index, %mask: vector<1x4xi1>, %pass: vector<1x4xf32>) -> vector<1x4xf32> {
%0 = vector.expandload %base[%i], %mask, %pass : memref<16xf32>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>
@@ -734,11 +767,11 @@ func.func @cast_away_expandload_leading_one_dims(%base: memref<16xf32>, %i: inde
// -----
// CHECK-LABEL: func.func @cast_away_gather_leading_one_dims
-// CHECK: %[[I:.+]] = vector.extract %{{.*}}[0] : vector<4xi32> from vector<1x4xi32>
-// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
-// CHECK: %[[P:.+]] = vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+// CHECK: %[[I:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi32> to vector<4xi32>
+// CHECK: %[[M:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
+// CHECK: %[[P:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: %[[G:.+]] = vector.gather %{{.*}}[%{{.*}}] [%[[I]]], %[[M]], %[[P]] : memref<16xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32> into vector<4xf32>
-// CHECK: %[[B:.+]] = vector.broadcast %[[G]] : vector<4xf32> to vector<1x4xf32>
+// CHECK: %[[B:.+]] = vector.shape_cast %[[G]] : vector<4xf32> to vector<1x4xf32>
// CHECK: return %[[B]] : vector<1x4xf32>
func.func @cast_away_gather_leading_one_dims(%base: memref<16xf32>, %i: index, %idx: vector<1x4xi32>, %mask: vector<1x4xi1>, %pass: vector<1x4xf32>) -> vector<1x4xf32> {
%0 = vector.gather %base[%i] [%idx], %mask, %pass : memref<16xf32>, vector<1x4xi32>, vector<1x4xi1>, vector<1x4xf32> into vector<1x4xf32>
@@ -748,7 +781,7 @@ func.func @cast_away_gather_leading_one_dims(%base: memref<16xf32>, %i: index, %
// -----
// CHECK-LABEL: func.func @cast_away_store_leading_one_dims
-// CHECK: %[[V:.+]] = vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+// CHECK: %[[V:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: vector.store %[[V]], %{{.*}}[%{{.*}}, %{{.*}}] : memref<8x16xf32>, vector<4xf32>
func.func @cast_away_store_leading_one_dims(%val: vector<1x4xf32>, %base: memref<8x16xf32>, %i: index, %j: index) {
vector.store %val, %base[%i, %j] : memref<8x16xf32>, vector<1x4xf32>
@@ -757,9 +790,19 @@ func.func @cast_away_store_leading_one_dims(%val: vector<1x4xf32>, %base: memref
// -----
+// CHECK-LABEL: func.func @cast_away_store_leading_one_dims_scalable
+// CHECK: %[[V:.+]] = vector.shape_cast %{{.*}} : vector<1x[4]xf32> to vector<[4]xf32>
+// CHECK: vector.store %[[V]], %{{.*}}[%{{.*}}, %{{.*}}] : memref<?x?xf32>, vector<[4]xf32>
+func.func @cast_away_store_leading_one_dims_scalable(%val: vector<1x[4]xf32>, %base: memref<?x?xf32>, %i: index, %j: index) {
+ vector.store %val, %base[%i, %j] : memref<?x?xf32>, vector<1x[4]xf32>
+ return
+}
+
+// -----
+
// CHECK-LABEL: func.func @cast_away_maskedstore_leading_one_dims
-// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
-// CHECK: %[[V:.+]] = vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+// CHECK: %[[M:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
+// CHECK: %[[V:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: vector.maskedstore %{{.*}}[%{{.*}}], %[[M]], %[[V]] : memref<16xf32>, vector<4xi1>, vector<4xf32>
func.func @cast_away_maskedstore_leading_one_dims(%base: memref<16xf32>, %i: index, %mask: vector<1x4xi1>, %val: vector<1x4xf32>) {
vector.maskedstore %base[%i], %mask, %val : memref<16xf32>, vector<1x4xi1>, vector<1x4xf32>
@@ -769,8 +812,8 @@ func.func @cast_away_maskedstore_leading_one_dims(%base: memref<16xf32>, %i: ind
// -----
// CHECK-LABEL: func.func @cast_away_compressstore_leading_one_dims
-// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
-// CHECK: %[[V:.+]] = vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+// CHECK: %[[M:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
+// CHECK: %[[V:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: vector.compressstore %{{.*}}[%{{.*}}], %[[M]], %[[V]] : memref<16xf32>, vector<4xi1>, vector<4xf32>
func.func @cast_away_compressstore_leading_one_dims(%base: memref<16xf32>, %i: index, %mask: vector<1x4xi1>, %val: vector<1x4xf32>) {
vector.compressstore %base[%i], %mask, %val : memref<16xf32>, vector<1x4xi1>, vector<1x4xf32>
@@ -780,9 +823,9 @@ func.func @cast_away_compressstore_leading_one_dims(%base: memref<16xf32>, %i: i
// -----
// CHECK-LABEL: func.func @cast_away_scatter_leading_one_dims
-// CHECK: %[[I:.+]] = vector.extract %{{.*}}[0] : vector<4xi32> from vector<1x4xi32>
-// CHECK: %[[M:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
-// CHECK: %[[V:.+]] = vector.extract %{{.*}}[0] : vector<4xf32> from vector<1x4xf32>
+// CHECK: %[[I:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi32> to vector<4xi32>
+// CHECK: %[[M:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
+// CHECK: %[[V:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf32> to vector<4xf32>
// CHECK: vector.scatter %{{.*}}[%{{.*}}] [%[[I]]], %[[M]], %[[V]] : memref<16xf32>, vector<4xi32>, vector<4xi1>, vector<4xf32>
func.func @cast_away_scatter_leading_one_dims(%base: memref<16xf32>, %i: index, %idx: vector<1x4xi32>, %mask: vector<1x4xi1>, %val: vector<1x4xf32>) {
vector.scatter %base[%i] [%idx], %mask, %val : memref<16xf32>, vector<1x4xi32>, vector<1x4xi1>, vector<1x4xf32>
diff --git a/mlir/test/Dialect/Vector/vector-transforms.mlir b/mlir/test/Dialect/Vector/vector-transforms.mlir
index de12a87253a67..37ca4358d97bd 100644
--- a/mlir/test/Dialect/Vector/vector-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-transforms.mlir
@@ -36,7 +36,7 @@ func.func @no_change(%arg0: vector<2x[4]x1xf32>, %arg1: vector<2x[4]x1xf32>) ->
// CHECK-LABEL: func.func @cast_away_leading_one_dim(
// CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<4x1xf32>
-// CHECK: vector.broadcast %[[MUL]] : vector<4x1xf32> to vector<1x4x1xf32>
+// CHECK: vector.shape_cast %[[MUL]] : vector<4x1xf32> to vector<1x4x1xf32>
func.func @cast_away_leading_one_dim(%arg0: vector<1x4x1xf32>, %arg1: vector<1x4x1xf32>) -> vector<1x4x1xf32> {
%1 = arith.mulf %arg0, %arg1 : vector<1x4x1xf32>
return %1: vector<1x4x1xf32>
@@ -44,7 +44,7 @@ func.func @cast_away_leading_one_dim(%arg0: vector<1x4x1xf32>, %arg1: vector<1x4
// CHECK-LABEL: func.func @cast_away_leading_one_dim_scalable(
// CHECK: %[[MUL:.*]] = arith.mulf %{{.*}}, %{{.*}} : vector<[4]x1xf32>
-// CHECK: vector.broadcast %[[MUL]] : vector<[4]x1xf32> to vector<1x[4]x1xf32>
+// CHECK: vector.shape_cast %[[MUL]] : vector<[4]x1xf32> to vector<1x[4]x1xf32>
func.func @cast_away_leading_one_dim_scalable(%arg0: vector<1x[4]x1xf32>, %arg1: vector<1x[4]x1xf32>) -> vector<1x[4]x1xf32> {
%1 = arith.mulf %arg0, %arg1 : vector<1x[4]x1xf32>
return %1: vector<1x[4]x1xf32>
>From c3e4ac9247ffc5ec9fa86e05108bc0ecd3aec4a0 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Wed, 6 May 2026 23:18:18 +0000
Subject: [PATCH 2/8] Clang-format
---
.../Vector/Transforms/VectorDropLeadUnitDim.cpp | 15 +++++++--------
1 file changed, 7 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index 931a87ff83e9c..c6cdca156c9f4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -65,8 +65,7 @@ static bool hasNonScalableUnitLeadingDims(VectorType type, int64_t dropCount) {
ArrayRef<int64_t> leadingShape = type.getShape().take_front(dropCount);
ArrayRef<bool> leadingScalable = type.getScalableDims().take_front(dropCount);
return llvm::all_of(leadingShape, [](int64_t dim) { return dim == 1; }) &&
- llvm::none_of(leadingScalable,
- [](bool scalable) { return scalable; });
+ llvm::none_of(leadingScalable, [](bool scalable) { return scalable; });
}
static bool isNonScalableUnitDim(VectorType type, int64_t dim) {
@@ -195,8 +194,8 @@ struct CastAwayInsertStridedSliceLeadingOneDim
// Trim leading one dimensions from both operands.
Location loc = insertOp.getLoc();
- Value newSrcVector = shapeCastVector(rewriter, loc,
- insertOp.getValueToStore(), newSrcType);
+ Value newSrcVector =
+ shapeCastVector(rewriter, loc, insertOp.getValueToStore(), newSrcType);
Value newDstVector =
shapeCastVector(rewriter, loc, insertOp.getDest(), newDstType);
@@ -371,8 +370,8 @@ struct CastAwayTransferWriteLeadingOneDim
shapeCastVector(rewriter, write.getLoc(), write.getVector(), newType);
if (write.getMask()) {
- Value newMask = dropUnitDimsFromMask(
- rewriter, write.getLoc(), write.getMask(), newType, newMap);
+ Value newMask = dropUnitDimsFromMask(rewriter, write.getLoc(),
+ write.getMask(), newType, newMap);
rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
write, newVector, write.getBase(), write.getIndices(),
AffineMapAttr::get(newMap), newMask, inBoundsAttr);
@@ -592,8 +591,8 @@ class CastAwayElementwiseLeadingOneDim : public RewritePattern {
SmallVector<Value, 4> newOperands;
for (Value operand : op->getOperands()) {
if (auto opVecType = dyn_cast<VectorType>(operand.getType()))
- newOperands.push_back(shapeCastVector(
- rewriter, op->getLoc(), operand, trimLeadingOneDims(opVecType)));
+ newOperands.push_back(shapeCastVector(rewriter, op->getLoc(), operand,
+ trimLeadingOneDims(opVecType)));
else
newOperands.push_back(operand);
}
>From f8ce0504ab2b6311c427f67bf36284057b822525 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Fri, 8 May 2026 00:27:11 +0000
Subject: [PATCH 3/8] Review comments (dropped unneeded bounds checks, fixed
modify-before-failing)
---
.../Transforms/VectorDropLeadUnitDim.cpp | 126 +++++++++++-------
.../vector-dropleadunitdim-transforms.mlir | 15 +++
2 files changed, 90 insertions(+), 51 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index c6cdca156c9f4..33efc5a1e0546 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -7,7 +7,9 @@
//===----------------------------------------------------------------------===//
#include <numeric>
+#include <utility>
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
@@ -15,6 +17,7 @@
#include "mlir/Dialect/Vector/Utils/VectorUtils.h"
#include "mlir/IR/Builders.h"
#include "mlir/IR/TypeUtilities.h"
+#include "llvm/ADT/Repeated.h"
#include "llvm/ADT/STLExtras.h"
#define DEBUG_TYPE "vector-drop-unit-dim"
@@ -22,11 +25,6 @@
using namespace mlir;
using namespace mlir::vector;
-/// Return a smallVector of size `rank` containing all zeros.
-static SmallVector<int64_t> splatZero(int64_t rank) {
- return SmallVector<int64_t>(rank, 0);
-}
-
// Trims leading one dimensions from `oldType` and returns the result type.
// Returns `vector<1xT>` if `oldType` only has one element.
static VectorType trimLeadingOneDims(VectorType oldType) {
@@ -60,8 +58,8 @@ static Value shapeCastVector(OpBuilder &b, Location loc, Value value,
}
static bool hasNonScalableUnitLeadingDims(VectorType type, int64_t dropCount) {
- if (dropCount < 0 || dropCount > type.getRank())
- return false;
+ assert(dropCount >= 0 && dropCount <= type.getRank() &&
+ "expected a valid leading dimension count");
ArrayRef<int64_t> leadingShape = type.getShape().take_front(dropCount);
ArrayRef<bool> leadingScalable = type.getScalableDims().take_front(dropCount);
return llvm::all_of(leadingShape, [](int64_t dim) { return dim == 1; }) &&
@@ -69,8 +67,14 @@ static bool hasNonScalableUnitLeadingDims(VectorType type, int64_t dropCount) {
}
static bool isNonScalableUnitDim(VectorType type, int64_t dim) {
- return dim >= 0 && dim < type.getRank() && type.getShape()[dim] == 1 &&
- !type.getScalableDims()[dim];
+ return type.getShape()[dim] == 1 && !type.getScalableDims()[dim];
+}
+
+static VectorType transposeVectorType(VectorType type,
+ ArrayRef<int64_t> permutation) {
+ return VectorType::get(applyPermutation(type.getShape(), permutation),
+ type.getElementType(),
+ applyPermutation(type.getScalableDims(), permutation));
}
/// Shape-casts `operand` to the vector type obtained by dropping the first
@@ -115,8 +119,10 @@ static Value dropLeadingDimsForContraction(OpBuilder &b, Location loc,
// vector.contract rejects 0-D vector accumulators/results. When every vector
// dimension is dropped, use the scalar path that vector.contract accepts.
- if (dropCount == oldType.getRank())
- return vector::ExtractOp::create(b, loc, operand, splatZero(dropCount));
+ if (dropCount == oldType.getRank()) {
+ llvm::Repeated<int64_t> zeros(static_cast<size_t>(dropCount), 0);
+ return vector::ExtractOp::create(b, loc, operand, llvm::to_vector(zeros));
+ }
return shapeCastDroppingLeadingDims(b, loc, operand, dropCount);
}
@@ -387,6 +393,16 @@ struct CastAwayTransferWriteLeadingOneDim
} // namespace
+namespace {
+struct ContractionOperandDropInfo {
+ AffineMap map;
+ bool needsShapeCast = false;
+ bool transposeNeeded = false;
+ bool transposeNonOuterUnitDims = false;
+ SmallVector<int64_t> permutation;
+};
+} // namespace
+
FailureOr<Value>
mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
MaskingOpInterface maskingOp,
@@ -423,35 +439,34 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
SmallVector<Value> operands = {contractOp.getLhs(), contractOp.getRhs(),
contractOp.getAcc()};
+ SmallVector<ContractionOperandDropInfo> operandDropInfo;
SmallVector<Value> newOperands;
auto loc = contractOp.getLoc();
for (const auto &it : llvm::enumerate(oldIndexingMaps)) {
// Check if the dim to be dropped exists as a leading dim in the operand
// if it does then we use vector.shape_cast to drop it.
- bool needsShapeCast = false;
+ ContractionOperandDropInfo dropInfo;
SmallVector<AffineExpr> results;
- auto map = it.value();
- int64_t orginalZeroDim = it.value().getDimPosition(0);
- if (orginalZeroDim != dimToDrop) {
+ dropInfo.map = it.value();
+ int64_t originalZeroDim = it.value().getDimPosition(0);
+ if (originalZeroDim != dimToDrop) {
// There are two reasons to be in this path, 1. We need to
// transpose the operand to make the dim to be dropped
// leading. 2. The dim to be dropped does not exist and in
// that case we dont want to add a unit transpose but we must
// check all the indices to make sure this is the case.
- bool transposeNeeded = false;
- SmallVector<int64_t> perm;
SmallVector<AffineExpr> transposeResults;
- for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
- int64_t currDim = map.getDimPosition(i);
+ for (int64_t i = 0, e = dropInfo.map.getNumResults(); i < e; ++i) {
+ int64_t currDim = dropInfo.map.getDimPosition(i);
if (currDim == dimToDrop) {
- transposeNeeded = true;
- perm.insert(perm.begin(), i);
+ dropInfo.transposeNeeded = true;
+ dropInfo.permutation.insert(dropInfo.permutation.begin(), i);
auto targetExpr = rewriter.getAffineDimExpr(currDim);
transposeResults.insert(transposeResults.begin(), targetExpr);
} else {
- perm.push_back(i);
+ dropInfo.permutation.push_back(i);
auto targetExpr = rewriter.getAffineDimExpr(currDim);
transposeResults.push_back(targetExpr);
}
@@ -460,36 +475,31 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
// Checks if only the outer, unit dimensions (of size 1) are permuted.
// Such transposes do not materially effect the underlying vector and can
// be omitted. EG: perm [1, 0, 2] applied to vector<1x1x8xi32>
- bool transposeNonOuterUnitDims = false;
auto operandShape = cast<ShapedType>(operands[it.index()].getType());
- for (auto [index, dim] :
- llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
+ for (auto [index, dim] : llvm::enumerate(
+ ArrayRef<int64_t>(dropInfo.permutation).drop_back(1))) {
if (dim != static_cast<int64_t>(index) &&
operandShape.getDimSize(index) != 1) {
- transposeNonOuterUnitDims = true;
+ dropInfo.transposeNonOuterUnitDims = true;
break;
}
}
// Do the transpose now if needed so that we can drop the correct dim
// with shape_cast later.
- if (transposeNeeded) {
- map = AffineMap::get(map.getNumDims(), 0, transposeResults,
- contractOp.getContext());
- if (transposeNonOuterUnitDims) {
- operands[it.index()] = rewriter.createOrFold<vector::TransposeOp>(
- loc, operands[it.index()], perm);
- }
- }
+ if (dropInfo.transposeNeeded)
+ dropInfo.map =
+ AffineMap::get(dropInfo.map.getNumDims(), 0, transposeResults,
+ contractOp.getContext());
}
// We have taken care to have the dim to be dropped be
// the leading dim. If its still not leading that means it
// does not exist in this operand and hence we do not need a shape_cast.
- if (map.getDimPosition(0) == dimToDrop)
- needsShapeCast = true;
+ if (dropInfo.map.getDimPosition(0) == dimToDrop)
+ dropInfo.needsShapeCast = true;
- for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
- int64_t currDim = map.getDimPosition(i);
+ for (int64_t i = 0, e = dropInfo.map.getNumResults(); i < e; ++i) {
+ int64_t currDim = dropInfo.map.getDimPosition(i);
if (currDim == dimToDrop)
// This is the dim we are dropping.
continue;
@@ -497,18 +507,36 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
currDim < dimToDrop ? currDim : currDim - 1);
results.push_back(targetExpr);
}
- newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
- contractOp.getContext()));
- if (needsShapeCast) {
+ newIndexingMaps.push_back(AffineMap::get(dropInfo.map.getNumDims() - 1, 0,
+ results, contractOp.getContext()));
+ if (dropInfo.needsShapeCast) {
auto operandType = cast<VectorType>(operands[it.index()].getType());
- if (operandType.getRank() < dropDim ||
- !hasNonScalableUnitLeadingDims(operandType, dropDim))
+ VectorType typeToDrop =
+ dropInfo.transposeNeeded && dropInfo.transposeNonOuterUnitDims
+ ? transposeVectorType(operandType, dropInfo.permutation)
+ : operandType;
+ if (!hasNonScalableUnitLeadingDims(typeToDrop, dropDim))
return failure();
- newOperands.push_back(dropLeadingDimsForContraction(
- rewriter, loc, operands[it.index()], dropDim));
- } else {
- newOperands.push_back(operands[it.index()]);
}
+ operandDropInfo.push_back(std::move(dropInfo));
+ }
+
+ if (maskingOp) {
+ auto oldMaskType = cast<VectorType>(maskingOp.getMask().getType());
+ if (oldMaskType.getRank() <= 1 ||
+ !isNonScalableUnitDim(oldMaskType, dimToDrop))
+ return failure();
+ }
+
+ for (const auto &it : llvm::enumerate(operandDropInfo)) {
+ Value operand = operands[it.index()];
+ const ContractionOperandDropInfo &dropInfo = it.value();
+ if (dropInfo.transposeNeeded && dropInfo.transposeNonOuterUnitDims)
+ operand = rewriter.createOrFold<vector::TransposeOp>(
+ loc, operand, dropInfo.permutation);
+ if (dropInfo.needsShapeCast)
+ operand = dropLeadingDimsForContraction(rewriter, loc, operand, dropDim);
+ newOperands.push_back(operand);
}
// Depending on whether this vector.contract is masked, the replacing Op
@@ -519,10 +547,6 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
rewriter.getArrayAttr(newIteratorTypes), contractOp.getKind());
if (maskingOp) {
- auto oldMaskType = cast<VectorType>(maskingOp.getMask().getType());
- if (oldMaskType.getRank() <= 1 ||
- !isNonScalableUnitDim(oldMaskType, dimToDrop))
- return failure();
Value newMask =
shapeCastDroppingDim(rewriter, loc, maskingOp.getMask(), dimToDrop);
diff --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index e5b1cc07319e3..ab2f566ff0cb0 100644
--- a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
@@ -296,6 +296,21 @@ func.func @do_not_cast_away_contraction_with_scalable_rank1_acc(%arg0: vector<64
}
// -----
+
+// CHECK-LABEL: do_not_cast_away_contraction_with_scalable_operand_dim
+// CHECK-NOT: vector.shape_cast
+// CHECK-NOT: vector.extract
+// CHECK-NOT: vector.broadcast
+// CHECK-NEXT: vector.contract
+// CHECK-NEXT: return
+
+func.func @do_not_cast_away_contraction_with_scalable_operand_dim(%arg0: vector<64xf32>, %arg1: vector<[1]x64xf32>, %arg2: vector<1xf32>) -> vector<1xf32> {
+ %0 = vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<64xf32>, vector<[1]x64xf32> into vector<1xf32>
+ return %0 : vector<1xf32>
+}
+
+// -----
+
// CHECK-LABEL: func @cast_away_extract_strided_slice_leading_one_dims
func.func @cast_away_extract_strided_slice_leading_one_dims(%arg0: vector<1x8x8xf16>) -> vector<1x1x8xf16> {
// CHECK: %[[SRC:.+]] = vector.shape_cast %{{.*}} : vector<1x8x8xf16> to vector<8x8xf16>
>From e2d5d377ec04929d647551c8c164437493b729e5 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Fri, 8 May 2026 18:29:15 +0000
Subject: [PATCH 4/8] Human changes
---
.../Transforms/VectorDropLeadUnitDim.cpp | 262 +++++++++---------
.../vector-dropleadunitdim-transforms.mlir | 8 +-
2 files changed, 132 insertions(+), 138 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index 33efc5a1e0546..9ba53c978d36d 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -7,9 +7,7 @@
//===----------------------------------------------------------------------===//
#include <numeric>
-#include <utility>
-#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
@@ -48,83 +46,89 @@ static VectorType trimLeadingOneDims(VectorType oldType) {
return VectorType::get(newShape, oldType.getElementType(), newScalableDims);
}
-/// Returns `value` if it already has `newType`, otherwise inserts a
-/// vector.shape_cast to `newType`.
-static Value shapeCastVector(OpBuilder &b, Location loc, Value value,
- VectorType newType) {
- if (value.getType() == newType)
- return value;
- return vector::ShapeCastOp::create(b, loc, newType, value);
-}
-
-static bool hasNonScalableUnitLeadingDims(VectorType type, int64_t dropCount) {
- assert(dropCount >= 0 && dropCount <= type.getRank() &&
- "expected a valid leading dimension count");
- ArrayRef<int64_t> leadingShape = type.getShape().take_front(dropCount);
- ArrayRef<bool> leadingScalable = type.getScalableDims().take_front(dropCount);
- return llvm::all_of(leadingShape, [](int64_t dim) { return dim == 1; }) &&
- llvm::none_of(leadingScalable, [](bool scalable) { return scalable; });
-}
-
static bool isNonScalableUnitDim(VectorType type, int64_t dim) {
+ assert(dim >= 0 && dim < type.getRank() &&
+ "expected a valid vector dimension");
return type.getShape()[dim] == 1 && !type.getScalableDims()[dim];
}
-static VectorType transposeVectorType(VectorType type,
- ArrayRef<int64_t> permutation) {
- return VectorType::get(applyPermutation(type.getShape(), permutation),
- type.getElementType(),
- applyPermutation(type.getScalableDims(), permutation));
+/// Returns true if the first `k` dimensions of `type` are non-scalable unit
+/// dimensions.
+static bool leadingDimsAreUnit(VectorType type, int64_t k) {
+ assert(k >= 0 && k <= type.getRank() &&
+ "expected a valid leading dimension count");
+ return llvm::all_of(llvm::seq<int64_t>(0, k), [&](int64_t dim) {
+ return isNonScalableUnitDim(type, dim);
+ });
}
-/// Shape-casts `operand` to the vector type obtained by dropping the first
-/// `dropCount` dimensions. Callers must ensure at least one vector dimension
-/// remains after the drop.
-static Value shapeCastDroppingLeadingDims(OpBuilder &b, Location loc,
- Value operand, int64_t dropCount) {
- auto oldType = cast<VectorType>(operand.getType());
- assert(dropCount < oldType.getRank() &&
- "shape_cast cannot drop all vector dimensions");
- VectorType newType = VectorType::get(
- oldType.getShape().drop_front(dropCount), oldType.getElementType(),
- oldType.getScalableDims().drop_front(dropCount));
- return shapeCastVector(b, loc, operand, newType);
+static bool leadingDimsAreUnitAfterPermutation(VectorType type,
+ ArrayRef<int64_t> permutation,
+ int64_t k) {
+ assert(k >= 0 && k <= static_cast<int64_t>(permutation.size()) &&
+ "expected a valid leading dimension count");
+ return llvm::all_of(permutation.take_front(k), [&](int64_t dim) {
+ return isNonScalableUnitDim(type, dim);
+ });
}
-static Value shapeCastDroppingDim(OpBuilder &b, Location loc, Value operand,
- int64_t dim) {
+/// Shape-casts `operand` to the vector type obtained by dropping dimension
+/// `dim`, which must be non-scalable and unit-sized.
+static Value dropUnitDim(OpBuilder &b, Location loc, Value operand,
+ int64_t dim) {
auto oldType = cast<VectorType>(operand.getType());
assert(isNonScalableUnitDim(oldType, dim) &&
"expected a non-scalable unit dim to drop");
+ int64_t rank = oldType.getRank();
+ assert(rank > 1 && "cannot shape_cast to a 0-D vector");
SmallVector<int64_t> newShape;
SmallVector<bool> newScalableDims;
- for (int64_t i = 0, e = oldType.getRank(); i < e; ++i) {
- if (i == dim)
+ newShape.reserve(rank - 1);
+ newScalableDims.reserve(rank - 1);
+ for (auto [i, size, scalable] :
+ llvm::enumerate(oldType.getShape(), oldType.getScalableDims())) {
+ if (static_cast<int64_t>(i) == dim)
continue;
- newShape.push_back(oldType.getShape()[i]);
- newScalableDims.push_back(oldType.getScalableDims()[i]);
+ newShape.push_back(size);
+ newScalableDims.push_back(scalable);
}
- return shapeCastVector(
- b, loc, operand,
- VectorType::get(newShape, oldType.getElementType(), newScalableDims));
+ return b.createOrFold<vector::ShapeCastOp>(
+ loc, VectorType::get(newShape, oldType.getElementType(), newScalableDims),
+ operand);
}
-static Value dropLeadingDimsForContraction(OpBuilder &b, Location loc,
- Value operand, int64_t dropCount) {
+/// Shape-casts `operand` to the vector type obtained by dropping the first
+/// `k` non-scalable unit dimensions. Callers must ensure at least one vector
+/// dimension remains after the drop.
+static Value dropLeadingUnitDims(OpBuilder &b, Location loc, Value operand,
+ int64_t k) {
+ auto oldType = cast<VectorType>(operand.getType());
+ assert(leadingDimsAreUnit(oldType, k) &&
+ "expected non-scalable leading unit dims to drop");
+ assert(k < oldType.getRank() &&
+ "shape_cast cannot drop all vector dimensions");
+ VectorType newType = VectorType::get(oldType.getShape().drop_front(k),
+ oldType.getElementType(),
+ oldType.getScalableDims().drop_front(k));
+ return b.createOrFold<vector::ShapeCastOp>(loc, newType, operand);
+}
+
+/// Like `dropLeadingUnitDims` except that if all dimensions would be dropped,
+/// the single element inside that vector is extracted and returned.
+static Value dropLeadingUnitDims0DIsScalar(OpBuilder &b, Location loc,
+ Value operand, int64_t k) {
auto oldType = cast<VectorType>(operand.getType());
- assert(hasNonScalableUnitLeadingDims(oldType, dropCount) &&
+ assert(leadingDimsAreUnit(oldType, k) &&
"expected non-scalable leading unit dims to drop");
- // vector.contract rejects 0-D vector accumulators/results. When every vector
- // dimension is dropped, use the scalar path that vector.contract accepts.
- if (dropCount == oldType.getRank()) {
- llvm::Repeated<int64_t> zeros(static_cast<size_t>(dropCount), 0);
+ if (k == oldType.getRank()) {
+ llvm::Repeated<int64_t> zeros(static_cast<size_t>(k), 0);
return vector::ExtractOp::create(b, loc, operand, llvm::to_vector(zeros));
}
- return shapeCastDroppingLeadingDims(b, loc, operand, dropCount);
+ return dropLeadingUnitDims(b, loc, operand, k);
}
namespace {
@@ -156,8 +160,8 @@ struct CastAwayExtractStridedSliceLeadingOneDim
Location loc = extractOp.getLoc();
- Value newSrcVector =
- shapeCastVector(rewriter, loc, extractOp.getSource(), newSrcType);
+ Value newSrcVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc, newSrcType, extractOp.getSource());
// The offsets/sizes/strides attribute can have a less number of elements
// than the input vector's rank: it is meant for the leading dimensions.
@@ -200,10 +204,10 @@ struct CastAwayInsertStridedSliceLeadingOneDim
// Trim leading one dimensions from both operands.
Location loc = insertOp.getLoc();
- Value newSrcVector =
- shapeCastVector(rewriter, loc, insertOp.getValueToStore(), newSrcType);
- Value newDstVector =
- shapeCastVector(rewriter, loc, insertOp.getDest(), newDstType);
+ Value newSrcVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc, newSrcType, insertOp.getValueToStore());
+ Value newDstVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc, newDstType, insertOp.getDest());
auto newOffsets = rewriter.getArrayAttr(
insertOp.getOffsets().getValue().take_back(newDstType.getRank()));
@@ -250,10 +254,10 @@ struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
Value newSrcVector = insertOp.getValueToStore();
if (oldSrcRank != 0)
- newSrcVector = shapeCastVector(rewriter, loc, insertOp.getValueToStore(),
- cast<VectorType>(newSrcType));
- Value newDstVector =
- shapeCastVector(rewriter, loc, insertOp.getDest(), newDstType);
+ newSrcVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc, cast<VectorType>(newSrcType), insertOp.getValueToStore());
+ Value newDstVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc, newDstType, insertOp.getDest());
// New position rank needs to be computed in two steps: (1) if destination
// type has leading unit dims, we also trim the position array accordingly,
@@ -281,7 +285,7 @@ static Value dropUnitDimsFromMask(OpBuilder &b, Location loc, Value mask,
VectorType newType, AffineMap newMap) {
// Infer the type of the new mask from the new map.
VectorType newMaskType = inferTransferOpMaskType(newType, newMap);
- return shapeCastVector(b, loc, mask, newMaskType);
+ return b.createOrFold<vector::ShapeCastOp>(loc, newMaskType, mask);
}
// Turns vector.transfer_read on vector with leading 1 dimensions into
@@ -372,8 +376,8 @@ struct CastAwayTransferWriteLeadingOneDim
inBoundsAttr = rewriter.getArrayAttr(
write.getInBoundsAttr().getValue().take_back(newType.getRank()));
- auto newVector =
- shapeCastVector(rewriter, write.getLoc(), write.getVector(), newType);
+ auto newVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ write.getLoc(), newType, write.getVector());
if (write.getMask()) {
Value newMask = dropUnitDimsFromMask(rewriter, write.getLoc(),
@@ -393,16 +397,6 @@ struct CastAwayTransferWriteLeadingOneDim
} // namespace
-namespace {
-struct ContractionOperandDropInfo {
- AffineMap map;
- bool needsShapeCast = false;
- bool transposeNeeded = false;
- bool transposeNonOuterUnitDims = false;
- SmallVector<int64_t> permutation;
-};
-} // namespace
-
FailureOr<Value>
mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
MaskingOpInterface maskingOp,
@@ -439,16 +433,22 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
SmallVector<Value> operands = {contractOp.getLhs(), contractOp.getRhs(),
contractOp.getAcc()};
- SmallVector<ContractionOperandDropInfo> operandDropInfo;
SmallVector<Value> newOperands;
auto loc = contractOp.getLoc();
+ if (maskingOp) {
+ auto oldMaskType = cast<VectorType>(maskingOp.getMask().getType());
+ if (oldMaskType.getRank() <= 1 || dimToDrop >= oldMaskType.getRank() ||
+ !isNonScalableUnitDim(oldMaskType, dimToDrop))
+ return failure();
+ }
+
for (const auto &it : llvm::enumerate(oldIndexingMaps)) {
// Check if the dim to be dropped exists as a leading dim in the operand
// if it does then we use vector.shape_cast to drop it.
- ContractionOperandDropInfo dropInfo;
+ bool needsDrop = false;
SmallVector<AffineExpr> results;
- dropInfo.map = it.value();
+ auto map = it.value();
int64_t originalZeroDim = it.value().getDimPosition(0);
if (originalZeroDim != dimToDrop) {
// There are two reasons to be in this path, 1. We need to
@@ -456,17 +456,19 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
// leading. 2. The dim to be dropped does not exist and in
// that case we dont want to add a unit transpose but we must
// check all the indices to make sure this is the case.
+ bool transposeNeeded = false;
+ SmallVector<int64_t> perm;
SmallVector<AffineExpr> transposeResults;
- for (int64_t i = 0, e = dropInfo.map.getNumResults(); i < e; ++i) {
- int64_t currDim = dropInfo.map.getDimPosition(i);
+ for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
+ int64_t currDim = map.getDimPosition(i);
if (currDim == dimToDrop) {
- dropInfo.transposeNeeded = true;
- dropInfo.permutation.insert(dropInfo.permutation.begin(), i);
+ transposeNeeded = true;
+ perm.insert(perm.begin(), i);
auto targetExpr = rewriter.getAffineDimExpr(currDim);
transposeResults.insert(transposeResults.begin(), targetExpr);
} else {
- dropInfo.permutation.push_back(i);
+ perm.push_back(i);
auto targetExpr = rewriter.getAffineDimExpr(currDim);
transposeResults.push_back(targetExpr);
}
@@ -475,31 +477,48 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
// Checks if only the outer, unit dimensions (of size 1) are permuted.
// Such transposes do not materially effect the underlying vector and can
// be omitted. EG: perm [1, 0, 2] applied to vector<1x1x8xi32>
- auto operandShape = cast<ShapedType>(operands[it.index()].getType());
- for (auto [index, dim] : llvm::enumerate(
- ArrayRef<int64_t>(dropInfo.permutation).drop_back(1))) {
+ bool transposeNonOuterUnitDims = false;
+ auto operandType = cast<VectorType>(operands[it.index()].getType());
+ for (auto [index, dim] :
+ llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
if (dim != static_cast<int64_t>(index) &&
- operandShape.getDimSize(index) != 1) {
- dropInfo.transposeNonOuterUnitDims = true;
+ !isNonScalableUnitDim(operandType, index)) {
+ transposeNonOuterUnitDims = true;
break;
}
}
// Do the transpose now if needed so that we can drop the correct dim
// with shape_cast later.
- if (dropInfo.transposeNeeded)
- dropInfo.map =
- AffineMap::get(dropInfo.map.getNumDims(), 0, transposeResults,
- contractOp.getContext());
+ if (transposeNeeded) {
+ map = AffineMap::get(map.getNumDims(), 0, transposeResults,
+ contractOp.getContext());
+ if (map.getDimPosition(0) == dimToDrop) {
+ bool leadingDimsCanBeDropped =
+ transposeNonOuterUnitDims
+ ? leadingDimsAreUnitAfterPermutation(operandType, perm,
+ dropDim)
+ : leadingDimsAreUnit(operandType, dropDim);
+ if (!leadingDimsCanBeDropped)
+ return failure();
+ }
+ if (transposeNonOuterUnitDims)
+ operands[it.index()] = rewriter.createOrFold<vector::TransposeOp>(
+ loc, operands[it.index()], perm);
+ }
}
// We have taken care to have the dim to be dropped be
// the leading dim. If its still not leading that means it
// does not exist in this operand and hence we do not need a shape_cast.
- if (dropInfo.map.getDimPosition(0) == dimToDrop)
- dropInfo.needsShapeCast = true;
+ if (map.getDimPosition(0) == dimToDrop)
+ needsDrop = true;
+ if (needsDrop && originalZeroDim == dimToDrop &&
+ !leadingDimsAreUnit(cast<VectorType>(operands[it.index()].getType()),
+ dropDim))
+ return failure();
- for (int64_t i = 0, e = dropInfo.map.getNumResults(); i < e; ++i) {
- int64_t currDim = dropInfo.map.getDimPosition(i);
+ for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
+ int64_t currDim = map.getDimPosition(i);
if (currDim == dimToDrop)
// This is the dim we are dropping.
continue;
@@ -507,36 +526,12 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
currDim < dimToDrop ? currDim : currDim - 1);
results.push_back(targetExpr);
}
- newIndexingMaps.push_back(AffineMap::get(dropInfo.map.getNumDims() - 1, 0,
- results, contractOp.getContext()));
- if (dropInfo.needsShapeCast) {
- auto operandType = cast<VectorType>(operands[it.index()].getType());
- VectorType typeToDrop =
- dropInfo.transposeNeeded && dropInfo.transposeNonOuterUnitDims
- ? transposeVectorType(operandType, dropInfo.permutation)
- : operandType;
- if (!hasNonScalableUnitLeadingDims(typeToDrop, dropDim))
- return failure();
- }
- operandDropInfo.push_back(std::move(dropInfo));
- }
-
- if (maskingOp) {
- auto oldMaskType = cast<VectorType>(maskingOp.getMask().getType());
- if (oldMaskType.getRank() <= 1 ||
- !isNonScalableUnitDim(oldMaskType, dimToDrop))
- return failure();
- }
-
- for (const auto &it : llvm::enumerate(operandDropInfo)) {
- Value operand = operands[it.index()];
- const ContractionOperandDropInfo &dropInfo = it.value();
- if (dropInfo.transposeNeeded && dropInfo.transposeNonOuterUnitDims)
- operand = rewriter.createOrFold<vector::TransposeOp>(
- loc, operand, dropInfo.permutation);
- if (dropInfo.needsShapeCast)
- operand = dropLeadingDimsForContraction(rewriter, loc, operand, dropDim);
- newOperands.push_back(operand);
+ newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
+ contractOp.getContext()));
+ newOperands.push_back(
+ needsDrop ? dropLeadingUnitDims0DIsScalar(rewriter, loc,
+ operands[it.index()], dropDim)
+ : operands[it.index()]);
}
// Depending on whether this vector.contract is masked, the replacing Op
@@ -547,8 +542,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
rewriter.getArrayAttr(newIteratorTypes), contractOp.getKind());
if (maskingOp) {
- Value newMask =
- shapeCastDroppingDim(rewriter, loc, maskingOp.getMask(), dimToDrop);
+ Value newMask = dropUnitDim(rewriter, loc, maskingOp.getMask(), dimToDrop);
newOp = mlir::vector::maskOperation(rewriter, newOp, newMask);
}
@@ -615,8 +609,8 @@ class CastAwayElementwiseLeadingOneDim : public RewritePattern {
SmallVector<Value, 4> newOperands;
for (Value operand : op->getOperands()) {
if (auto opVecType = dyn_cast<VectorType>(operand.getType()))
- newOperands.push_back(shapeCastVector(rewriter, op->getLoc(), operand,
- trimLeadingOneDims(opVecType)));
+ newOperands.push_back(rewriter.createOrFold<vector::ShapeCastOp>(
+ op->getLoc(), trimLeadingOneDims(opVecType), operand));
else
newOperands.push_back(operand);
}
@@ -653,7 +647,7 @@ struct CastAwayLoadLikeLeadingOneDim : public OpRewritePattern<OpTy> {
for (Value operand : op->getOperands()) {
if (isa<VectorType>(operand.getType())) {
newOperands.push_back(
- shapeCastDroppingLeadingDims(rewriter, loc, operand, nDropped));
+ dropLeadingUnitDims(rewriter, loc, operand, nDropped));
} else {
newOperands.push_back(operand);
}
@@ -688,7 +682,7 @@ struct CastAwayStoreLikeLeadingOneDim : public OpRewritePattern<OpTy> {
for (Value operand : op->getOperands()) {
if (isa<VectorType>(operand.getType())) {
newOperands.push_back(
- shapeCastDroppingLeadingDims(rewriter, loc, operand, nDropped));
+ dropLeadingUnitDims(rewriter, loc, operand, nDropped));
} else {
newOperands.push_back(operand);
}
diff --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index ab2f566ff0cb0..81b68baadca55 100644
--- a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
@@ -283,28 +283,28 @@ func.func @cast_away_masked_contraction_with_rank1_acc(%arg0: vector<64xf32>, %a
// -----
-// CHECK-LABEL: do_not_cast_away_contraction_with_scalable_rank1_acc
+// CHECK-LABEL: negative_cast_away_contraction_with_scalable_rank1_acc
// CHECK-NOT: vector.shape_cast
// CHECK-NOT: vector.extract
// CHECK-NOT: vector.broadcast
// CHECK-NEXT: vector.contract
// CHECK-NEXT: return
-func.func @do_not_cast_away_contraction_with_scalable_rank1_acc(%arg0: vector<64xf32>, %arg1: vector<[1]x64xf32>, %arg2: vector<[1]xf32>) -> vector<[1]xf32> {
+func.func @negative_cast_away_contraction_with_scalable_rank1_acc(%arg0: vector<64xf32>, %arg1: vector<[1]x64xf32>, %arg2: vector<[1]xf32>) -> vector<[1]xf32> {
%0 = vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<64xf32>, vector<[1]x64xf32> into vector<[1]xf32>
return %0 : vector<[1]xf32>
}
// -----
-// CHECK-LABEL: do_not_cast_away_contraction_with_scalable_operand_dim
+// CHECK-LABEL: negative_cast_away_contraction_with_scalable_operand_dim
// CHECK-NOT: vector.shape_cast
// CHECK-NOT: vector.extract
// CHECK-NOT: vector.broadcast
// CHECK-NEXT: vector.contract
// CHECK-NEXT: return
-func.func @do_not_cast_away_contraction_with_scalable_operand_dim(%arg0: vector<64xf32>, %arg1: vector<[1]x64xf32>, %arg2: vector<1xf32>) -> vector<1xf32> {
+func.func @negative_cast_away_contraction_with_scalable_operand_dim(%arg0: vector<64xf32>, %arg1: vector<[1]x64xf32>, %arg2: vector<1xf32>) -> vector<1xf32> {
%0 = vector.contract {indexing_maps = [affine_map<(d0, d1) -> (d0)>, affine_map<(d0, d1) -> (d1, d0)>, affine_map<(d0, d1) -> (d1)>], iterator_types = ["reduction", "parallel"], kind = #vector.kind<add>} %arg0, %arg1, %arg2 : vector<64xf32>, vector<[1]x64xf32> into vector<1xf32>
return %0 : vector<1xf32>
}
>From bf0cbc5272fff7d5f158b409cea3f0b064570405 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Fri, 8 May 2026 18:38:31 +0000
Subject: [PATCH 5/8] Name change
---
.../Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp | 8 ++++----
1 file changed, 4 insertions(+), 4 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index 9ba53c978d36d..11561012aaff9 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -446,7 +446,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
for (const auto &it : llvm::enumerate(oldIndexingMaps)) {
// Check if the dim to be dropped exists as a leading dim in the operand
// if it does then we use vector.shape_cast to drop it.
- bool needsDrop = false;
+ bool needsCast = false;
SmallVector<AffineExpr> results;
auto map = it.value();
int64_t originalZeroDim = it.value().getDimPosition(0);
@@ -511,8 +511,8 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
// the leading dim. If its still not leading that means it
// does not exist in this operand and hence we do not need a shape_cast.
if (map.getDimPosition(0) == dimToDrop)
- needsDrop = true;
- if (needsDrop && originalZeroDim == dimToDrop &&
+ needsCast = true;
+ if (needsCast && originalZeroDim == dimToDrop &&
!leadingDimsAreUnit(cast<VectorType>(operands[it.index()].getType()),
dropDim))
return failure();
@@ -529,7 +529,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
contractOp.getContext()));
newOperands.push_back(
- needsDrop ? dropLeadingUnitDims0DIsScalar(rewriter, loc,
+ needsCast ? dropLeadingUnitDims0DIsScalar(rewriter, loc,
operands[it.index()], dropDim)
: operands[it.index()]);
}
>From 7348f5b6a270acd29e8916866f7ff985c9f84ead Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Fri, 8 May 2026 20:21:26 +0000
Subject: [PATCH 6/8] Simplify things since we don't need transposes
---
.../Transforms/VectorDropLeadUnitDim.cpp | 138 ++++++++++--------
1 file changed, 77 insertions(+), 61 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index 11561012aaff9..c408cc3e2a745 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -7,7 +7,9 @@
//===----------------------------------------------------------------------===//
#include <numeric>
+#include <utility>
+#include "mlir/Dialect/Utils/IndexingUtils.h"
#include "mlir/Dialect/Utils/StructuredOpsUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/Dialect/Vector/Transforms/VectorRewritePatterns.h"
@@ -115,6 +117,19 @@ static Value dropLeadingUnitDims(OpBuilder &b, Location loc, Value operand,
return b.createOrFold<vector::ShapeCastOp>(loc, newType, operand);
}
+/// Returns the vector type obtained by applying `permutation` to `type`.
+static VectorType permuteVectorType(VectorType type,
+ ArrayRef<int64_t> permutation) {
+ assert(static_cast<int64_t>(permutation.size()) == type.getRank() &&
+ "expected a permutation matching the operand rank");
+ SmallVector<int64_t> permutedShape =
+ applyPermutation(type.getShape(), permutation);
+ SmallVector<bool> permutedScalableDims =
+ applyPermutation(type.getScalableDims(), permutation);
+ return VectorType::get(permutedShape, type.getElementType(),
+ permutedScalableDims);
+}
+
/// Like `dropLeadingUnitDims` except that if all dimensions would be dropped,
/// the single element inside that vector is extracted and returned.
static Value dropLeadingUnitDims0DIsScalar(OpBuilder &b, Location loc,
@@ -128,7 +143,10 @@ static Value dropLeadingUnitDims0DIsScalar(OpBuilder &b, Location loc,
return vector::ExtractOp::create(b, loc, operand, llvm::to_vector(zeros));
}
- return dropLeadingUnitDims(b, loc, operand, k);
+ VectorType newType = VectorType::get(oldType.getShape().drop_front(k),
+ oldType.getElementType(),
+ oldType.getScalableDims().drop_front(k));
+ return vector::ShapeCastOp::create(b, loc, newType, operand);
}
namespace {
@@ -397,6 +415,15 @@ struct CastAwayTransferWriteLeadingOneDim
} // namespace
+namespace {
+struct VectorContractOperandCastPlan {
+ AffineMap map;
+ SmallVector<int64_t> permutation;
+ bool dropLeadingUnitDim = false;
+ bool permuteOperand = false;
+};
+} // namespace
+
FailureOr<Value>
mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
MaskingOpInterface maskingOp,
@@ -433,6 +460,7 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
SmallVector<Value> operands = {contractOp.getLhs(), contractOp.getRhs(),
contractOp.getAcc()};
+ SmallVector<VectorContractOperandCastPlan> operandCastPlans;
SmallVector<Value> newOperands;
auto loc = contractOp.getLoc();
@@ -446,79 +474,56 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
for (const auto &it : llvm::enumerate(oldIndexingMaps)) {
// Check if the dim to be dropped exists as a leading dim in the operand
// if it does then we use vector.shape_cast to drop it.
- bool needsCast = false;
+ VectorContractOperandCastPlan plan;
SmallVector<AffineExpr> results;
- auto map = it.value();
- int64_t originalZeroDim = it.value().getDimPosition(0);
+ plan.map = it.value();
+ int64_t originalZeroDim = plan.map.getDimPosition(0);
if (originalZeroDim != dimToDrop) {
// There are two reasons to be in this path, 1. We need to
- // transpose the operand to make the dim to be dropped
+ // permute the operand type to make the dim to be dropped
// leading. 2. The dim to be dropped does not exist and in
- // that case we dont want to add a unit transpose but we must
+ // that case we dont want to add a unit permutation but we must
// check all the indices to make sure this is the case.
- bool transposeNeeded = false;
- SmallVector<int64_t> perm;
- SmallVector<AffineExpr> transposeResults;
+ SmallVector<AffineExpr> permutedResults;
- for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
- int64_t currDim = map.getDimPosition(i);
+ for (int64_t i = 0, e = plan.map.getNumResults(); i < e; ++i) {
+ int64_t currDim = plan.map.getDimPosition(i);
if (currDim == dimToDrop) {
- transposeNeeded = true;
- perm.insert(perm.begin(), i);
+ plan.permuteOperand = true;
+ plan.permutation.insert(plan.permutation.begin(), i);
auto targetExpr = rewriter.getAffineDimExpr(currDim);
- transposeResults.insert(transposeResults.begin(), targetExpr);
+ permutedResults.insert(permutedResults.begin(), targetExpr);
} else {
- perm.push_back(i);
+ plan.permutation.push_back(i);
auto targetExpr = rewriter.getAffineDimExpr(currDim);
- transposeResults.push_back(targetExpr);
+ permutedResults.push_back(targetExpr);
}
}
- // Checks if only the outer, unit dimensions (of size 1) are permuted.
- // Such transposes do not materially effect the underlying vector and can
- // be omitted. EG: perm [1, 0, 2] applied to vector<1x1x8xi32>
- bool transposeNonOuterUnitDims = false;
- auto operandType = cast<VectorType>(operands[it.index()].getType());
- for (auto [index, dim] :
- llvm::enumerate(ArrayRef<int64_t>(perm).drop_back(1))) {
- if (dim != static_cast<int64_t>(index) &&
- !isNonScalableUnitDim(operandType, index)) {
- transposeNonOuterUnitDims = true;
- break;
- }
- }
-
- // Do the transpose now if needed so that we can drop the correct dim
- // with shape_cast later.
- if (transposeNeeded) {
- map = AffineMap::get(map.getNumDims(), 0, transposeResults,
- contractOp.getContext());
- if (map.getDimPosition(0) == dimToDrop) {
- bool leadingDimsCanBeDropped =
- transposeNonOuterUnitDims
- ? leadingDimsAreUnitAfterPermutation(operandType, perm,
- dropDim)
- : leadingDimsAreUnit(operandType, dropDim);
- if (!leadingDimsCanBeDropped)
+ // Update the map now so that the later shape_cast drops the correct dim.
+ if (plan.permuteOperand) {
+ plan.map = AffineMap::get(plan.map.getNumDims(), 0, permutedResults,
+ contractOp.getContext());
+ if (plan.map.getDimPosition(0) == dimToDrop) {
+ auto operandType = cast<VectorType>(operands[it.index()].getType());
+ if (!leadingDimsAreUnitAfterPermutation(operandType, plan.permutation,
+ dropDim))
return failure();
}
- if (transposeNonOuterUnitDims)
- operands[it.index()] = rewriter.createOrFold<vector::TransposeOp>(
- loc, operands[it.index()], perm);
}
}
// We have taken care to have the dim to be dropped be
// the leading dim. If its still not leading that means it
// does not exist in this operand and hence we do not need a shape_cast.
- if (map.getDimPosition(0) == dimToDrop)
- needsCast = true;
- if (needsCast && originalZeroDim == dimToDrop &&
+ if (plan.map.getDimPosition(0) == dimToDrop)
+ plan.dropLeadingUnitDim = true;
+ if (plan.dropLeadingUnitDim && originalZeroDim == dimToDrop &&
!leadingDimsAreUnit(cast<VectorType>(operands[it.index()].getType()),
dropDim))
return failure();
- for (int64_t i = 0, e = map.getNumResults(); i < e; ++i) {
- int64_t currDim = map.getDimPosition(i);
+ for (int64_t i = 0, e = plan.map.getNumResults(); i < e; ++i) {
+ int64_t currDim = plan.map.getDimPosition(i);
if (currDim == dimToDrop)
// This is the dim we are dropping.
continue;
@@ -526,12 +531,23 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
currDim < dimToDrop ? currDim : currDim - 1);
results.push_back(targetExpr);
}
- newIndexingMaps.push_back(AffineMap::get(map.getNumDims() - 1, 0, results,
- contractOp.getContext()));
- newOperands.push_back(
- needsCast ? dropLeadingUnitDims0DIsScalar(rewriter, loc,
- operands[it.index()], dropDim)
- : operands[it.index()]);
+ newIndexingMaps.push_back(AffineMap::get(plan.map.getNumDims() - 1, 0,
+ results, contractOp.getContext()));
+ operandCastPlans.push_back(std::move(plan));
+ }
+
+ for (auto [plan, operand] : llvm::zip_equal(operandCastPlans, operands)) {
+ Value newOperand = operand;
+ if (plan.permuteOperand)
+ newOperand = rewriter.createOrFold<vector::ShapeCastOp>(
+ loc,
+ permuteVectorType(cast<VectorType>(newOperand.getType()),
+ plan.permutation),
+ newOperand);
+ if (plan.dropLeadingUnitDim)
+ newOperand =
+ dropLeadingUnitDims0DIsScalar(rewriter, loc, newOperand, dropDim);
+ newOperands.push_back(newOperand);
}
// Depending on whether this vector.contract is masked, the replacing Op
@@ -547,13 +563,13 @@ mlir::vector::castAwayContractionLeadingOneDim(vector::ContractionOp contractOp,
newOp = mlir::vector::maskOperation(rewriter, newOp, newMask);
}
- if (isa<VectorType>(newOp->getResults()[0].getType()))
- return vector::ShapeCastOp::create(rewriter, loc,
+ if (!isa<VectorType>(newOp->getResults()[0].getType()))
+ return vector::BroadcastOp::create(rewriter, loc,
contractOp->getResultTypes()[0],
newOp->getResults()[0])
.getResult();
- return vector::BroadcastOp::create(rewriter, loc,
+ return vector::ShapeCastOp::create(rewriter, loc,
contractOp->getResultTypes()[0],
newOp->getResults()[0])
.getResult();
@@ -563,8 +579,8 @@ namespace {
/// Turns vector.contract on vector with leading 1 dimensions into
/// vector.shape_cast followed by vector.contract on vector without leading
-/// 1 dimensions. Also performs transpose of lhs and rhs operands if required
-/// prior to the shape cast.
+/// 1 dimensions. Non-leading unit dimensions are dropped via direct
+/// shape_casts.
struct CastAwayContractionLeadingOneDim
: public MaskableOpRewritePattern<vector::ContractionOp> {
using MaskableOpRewritePattern::MaskableOpRewritePattern;
>From ab5f5a7d9b28f6940e37bf8ecc90c890d97ed927 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Fri, 8 May 2026 20:38:25 +0000
Subject: [PATCH 7/8] Stray to_vector
---
mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index c408cc3e2a745..a1ecd1086fdbe 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -140,7 +140,7 @@ static Value dropLeadingUnitDims0DIsScalar(OpBuilder &b, Location loc,
if (k == oldType.getRank()) {
llvm::Repeated<int64_t> zeros(static_cast<size_t>(k), 0);
- return vector::ExtractOp::create(b, loc, operand, llvm::to_vector(zeros));
+ return vector::ExtractOp::create(b, loc, operand, zeros);
}
VectorType newType = VectorType::get(oldType.getShape().drop_front(k),
>From 7717cf468fb424ce2f425c2042a94e1109e3b563 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Fri, 8 May 2026 20:52:17 +0000
Subject: [PATCH 8/8] Ok, actually, we need OpFoldResult range
---
mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index a1ecd1086fdbe..8be2fa071a9d7 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -139,7 +139,7 @@ static Value dropLeadingUnitDims0DIsScalar(OpBuilder &b, Location loc,
"expected non-scalable leading unit dims to drop");
if (k == oldType.getRank()) {
- llvm::Repeated<int64_t> zeros(static_cast<size_t>(k), 0);
+ SmallVector<int64_t> zeros(k, static_cast<int64_t>(0));
return vector::ExtractOp::create(b, loc, operand, zeros);
}
More information about the Mlir-commits
mailing list