[Mlir-commits] [mlir] [mlir][vector] Extend `combineContractAndBroadcast` to accept `vector.shape_cast` (PR #208752)
Andrzej Warzyński
llvmlistbot at llvm.org
Tue Jul 14 02:10:56 PDT 2026
https://github.com/banach-space updated https://github.com/llvm/llvm-project/pull/208752
>From 643775665538881e17a0ddb2337ae7f170a30bcb Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 10 Jul 2026 14:56:23 +0000
Subject: [PATCH 1/2] [mlir][vector] Extend `combineContractAndBroadcast` to
accept `vector.shape_cast`
`combineContractAndBroadcast` is a pattern that matches `vector.contract`
operations whose inputs are defined by `vector.broadcast`, e.g.:
```mlir
%0 = vector.broadcast %lhs : vector<8x4xi32> to vector<1x8x4xi32>
%1 = vector.broadcast %rhs : vector<8x4xi32> to vector<1x8x4xi32>
%result = vector.contract {
indexing_maps = [#map0, #map1, #map2],
iterator_types = ["reduction", "parallel", "parallel", "reduction"],
kind = #vector.kind<add>
} %0, %1, %acc : vector<1x8x4xi32>, vector<1x8x4xi32> into vector<8x8xi32>
```
The pattern folds away the reduction over a unit dimension introduced by
the defining operations, producing a simplified `vector.contract`:
```mlir
vector.contract {
indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
iterator_types = ["parallel", "parallel", "reduction"]
} %0, %1, %acc : vector<8x4xi32>, vector<8x4xi32> into vector<8x8xi32>
```
This change extends the pattern to also accept broadcast-like
`vector.shape_cast` operations:
```mlir
%0 = vector.shape_cast %lhs : vector<8x4xi32> to vector<1x8x4xi32>
%1 = vector.shape_cast %rhs : vector<8x4xi32> to vector<1x8x4xi32>
%result = vector.contract {
indexing_maps = [#map0, #map1, #map2],
iterator_types = ["reduction", "parallel", "parallel", "reduction"],
kind = #vector.kind<add>
} %0, %1, %acc : vector<1x8x4xi32>, vector<1x8x4xi32> into vector<8x8xi32>
```
Note that `mlir::vector::isBroadcastableTo` was only moved to make it
available to the newly introduced `ShapeCastOp` helper,
`ShapeCastOp::isBroadcastLike()`.
---
.../mlir/Dialect/Vector/IR/VectorOps.h | 2 +-
.../mlir/Dialect/Vector/IR/VectorOps.td | 3 +
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 121 ++++++++++--------
.../Vector/Transforms/VectorTransforms.cpp | 36 ++++--
.../Vector/vector-reduce-to-contract.mlir | 103 ++++++++++++++-
5 files changed, 194 insertions(+), 71 deletions(-)
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
index b3a0653b90765..1a7d37d989f5a 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.h
@@ -67,7 +67,7 @@ enum class ConstantMaskKind { AllFalse = 0, AllTrue };
/// arguments.
void buildTerminatedBody(OpBuilder &builder, Location loc);
-/// Return whether `srcType` can be broadcast to `dstVectorType` under the
+/// Models whether `srcType` can be broadcast to `dstVectorType` under the
/// semantics of the `vector.broadcast` op.
enum class BroadcastableToResult {
Success = 0,
diff --git a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
index a519e6d89a472..f2a8ce2951106 100644
--- a/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
+++ b/mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
@@ -2484,6 +2484,9 @@ def Vector_ShapeCastOp :
VectorType getResultVectorType() {
return ::llvm::cast<VectorType>(getResult().getType());
}
+ // Return true if this Op is effectively a vector.broadcast (i.e. the input
+ // and output shapes satisfy the vector.broadcast constraints).
+ bool isBroadcastLike();
}];
let assemblyFormat = "$source attr-dict `:` type($source) `to` type($result)";
let hasFolder = 1;
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index e0cb9923b1dd4..fb535cdae461f 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -429,6 +429,63 @@ static Attribute convertNumericAttr(Attribute attr, Type expectedType) {
return attr;
}
+/// Return whether `srcType` can be broadcast to `dstVectorType` under the
+/// semantics of the `vector.broadcast` op.
+BroadcastableToResult mlir::vector::isBroadcastableTo(
+ Type srcType, VectorType dstVectorType,
+ std::pair<VectorDim, VectorDim> *mismatchingDims) {
+ // Broadcast scalar to vector of the same element type.
+ if (isa<VectorElementTypeInterface>(srcType) && dstVectorType &&
+ srcType == getElementTypeOrSelf(dstVectorType))
+ return BroadcastableToResult::Success;
+ // From now on, only vectors broadcast.
+ VectorType srcVectorType = llvm::dyn_cast<VectorType>(srcType);
+ if (!srcVectorType)
+ return BroadcastableToResult::SourceTypeNotAVector;
+
+ int64_t srcRank = srcVectorType.getRank();
+ int64_t dstRank = dstVectorType.getRank();
+ if (srcRank > dstRank)
+ return BroadcastableToResult::SourceRankHigher;
+ // Source has an exact match or singleton value for all trailing dimensions
+ // (all leading dimensions are simply duplicated).
+ int64_t lead = dstRank - srcRank;
+ for (int64_t dimIdx = 0; dimIdx < srcRank; ++dimIdx) {
+ // Have mismatching dims (in the sense of vector.broadcast semantics) been
+ // encountered?
+ bool foundMismatchingDims = false;
+
+ // Check fixed-width dims.
+ int64_t srcDim = srcVectorType.getDimSize(dimIdx);
+ int64_t dstDim = dstVectorType.getDimSize(lead + dimIdx);
+ if (srcDim != 1 && srcDim != dstDim)
+ foundMismatchingDims = true;
+
+ // Check scalable flags.
+ bool srcDimScalableFlag = srcVectorType.getScalableDims()[dimIdx];
+ bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + dimIdx];
+ if ((srcDim == 1 && srcDimScalableFlag && dstDim != 1) ||
+ // 1 -> [N] is fine, everything else should be rejected when mixing
+ // fixed-width and scalable dims
+ (srcDimScalableFlag != dstDimScalableFlag &&
+ (srcDim != 1 || srcDimScalableFlag)))
+ foundMismatchingDims = true;
+
+ if (foundMismatchingDims) {
+ if (mismatchingDims != nullptr) {
+ mismatchingDims->first.dim = srcDim;
+ mismatchingDims->first.isScalable = srcDimScalableFlag;
+
+ mismatchingDims->second.dim = dstDim;
+ mismatchingDims->second.isScalable = dstDimScalableFlag;
+ }
+ return BroadcastableToResult::DimensionMismatch;
+ }
+ }
+
+ return BroadcastableToResult::Success;
+}
+
//===----------------------------------------------------------------------===//
// CombiningKindAttr
//===----------------------------------------------------------------------===//
@@ -3101,61 +3158,6 @@ Value BroadcastOp::createOrFoldBroadcastOp(
return res;
}
-BroadcastableToResult mlir::vector::isBroadcastableTo(
- Type srcType, VectorType dstVectorType,
- std::pair<VectorDim, VectorDim> *mismatchingDims) {
- // Broadcast scalar to vector of the same element type.
- if (isa<VectorElementTypeInterface>(srcType) && dstVectorType &&
- srcType == getElementTypeOrSelf(dstVectorType))
- return BroadcastableToResult::Success;
- // From now on, only vectors broadcast.
- VectorType srcVectorType = llvm::dyn_cast<VectorType>(srcType);
- if (!srcVectorType)
- return BroadcastableToResult::SourceTypeNotAVector;
-
- int64_t srcRank = srcVectorType.getRank();
- int64_t dstRank = dstVectorType.getRank();
- if (srcRank > dstRank)
- return BroadcastableToResult::SourceRankHigher;
- // Source has an exact match or singleton value for all trailing dimensions
- // (all leading dimensions are simply duplicated).
- int64_t lead = dstRank - srcRank;
- for (int64_t dimIdx = 0; dimIdx < srcRank; ++dimIdx) {
- // Have mismatching dims (in the sense of vector.broadcast semantics) been
- // encountered?
- bool foundMismatchingDims = false;
-
- // Check fixed-width dims.
- int64_t srcDim = srcVectorType.getDimSize(dimIdx);
- int64_t dstDim = dstVectorType.getDimSize(lead + dimIdx);
- if (srcDim != 1 && srcDim != dstDim)
- foundMismatchingDims = true;
-
- // Check scalable flags.
- bool srcDimScalableFlag = srcVectorType.getScalableDims()[dimIdx];
- bool dstDimScalableFlag = dstVectorType.getScalableDims()[lead + dimIdx];
- if ((srcDim == 1 && srcDimScalableFlag && dstDim != 1) ||
- // 1 -> [N] is fine, everything else should be rejected when mixing
- // fixed-width and scalable dims
- (srcDimScalableFlag != dstDimScalableFlag &&
- (srcDim != 1 || srcDimScalableFlag)))
- foundMismatchingDims = true;
-
- if (foundMismatchingDims) {
- if (mismatchingDims != nullptr) {
- mismatchingDims->first.dim = srcDim;
- mismatchingDims->first.isScalable = srcDimScalableFlag;
-
- mismatchingDims->second.dim = dstDim;
- mismatchingDims->second.isScalable = dstDimScalableFlag;
- }
- return BroadcastableToResult::DimensionMismatch;
- }
- }
-
- return BroadcastableToResult::Success;
-}
-
LogicalResult BroadcastOp::verify() {
std::pair<VectorDim, VectorDim> mismatchingDims;
BroadcastableToResult res = isBroadcastableTo(
@@ -6757,6 +6759,15 @@ LogicalResult ShapeCastOp::verify() {
return success();
}
+bool ShapeCastOp::isBroadcastLike() {
+ auto srcType = getSourceVectorType();
+ auto resType = getResultVectorType();
+
+ std::pair<VectorDim, VectorDim> mismatchingDims;
+ return isBroadcastableTo(srcType, resType, &mismatchingDims) ==
+ BroadcastableToResult::Success;
+}
+
/// Return true if `transpose` does not permute a pair of non-unit dims.
/// By `order preserving` we mean that the flattened versions of the input and
/// output vectors are (numerically) identical. In other words `transpose` is
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 752610efc6992..22e7742128ff2 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -282,21 +282,32 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
bool changed = false;
for (Value *operand : {&lhs, &rhs}) {
AffineMap &map = maps[index++];
+
+ // Make sure that the source
+ auto sc = operand->getDefiningOp<vector::ShapeCastOp>();
auto broadcast = operand->getDefiningOp<vector::BroadcastOp>();
- if (!broadcast)
+ if (!broadcast && !sc)
continue;
+
+ if (sc && !sc.isBroadcastLike())
+ return rewriter.notifyMatchFailure(
+ contractOp, "Operand defined via vector.shape_cast that has "
+ "non-broadcast semantics");
+
+ VectorType srcType = (sc) ? sc.getSourceVectorType()
+ : dyn_cast<VectorType>(broadcast.getSourceType());
+ VectorType resType =
+ (sc) ? sc.getResultVectorType() : broadcast.getResultVectorType();
+
// contractionOp can only take vector as operands.
- auto srcType = dyn_cast<VectorType>(broadcast.getSourceType());
- if (!srcType ||
- srcType.getRank() == broadcast.getResultVectorType().getRank())
+ // auto srcType = dyn_cast<VectorType>(broadcast.getSourceVectorType());
+ if (!srcType || srcType.getRank() >= resType.getRank())
continue;
- int64_t rankDiff =
- broadcast.getResultVectorType().getRank() - srcType.getRank();
+ int64_t rankDiff = resType.getRank() - srcType.getRank();
bool innerDimBroadcast = false;
SmallVector<AffineExpr> originalDims;
for (const auto &dim : llvm::enumerate(srcType.getShape())) {
- if (dim.value() !=
- broadcast.getResultVectorType().getDimSize(rankDiff + dim.index())) {
+ if (dim.value() != resType.getDimSize(rankDiff + dim.index())) {
innerDimBroadcast = true;
break;
}
@@ -311,7 +322,7 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
// of non-unit size.
bool nonUnitDimReductionBroadcast = false;
for (int64_t i = 0; i < rankDiff; ++i) {
- if (broadcast.getResultVectorType().getDimSize(i) != 1 &&
+ if (resType.getDimSize(i) != 1 &&
isReductionIterator(contractOp.getIteratorTypes()
.getValue()[map.getDimPosition(i)])) {
nonUnitDimReductionBroadcast = true;
@@ -321,11 +332,10 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
if (nonUnitDimReductionBroadcast)
continue;
- AffineMap broadcastMap =
- AffineMap::get(broadcast.getResultVectorType().getRank(), 0,
- originalDims, contractOp.getContext());
+ AffineMap broadcastMap = AffineMap::get(resType.getRank(), 0, originalDims,
+ contractOp.getContext());
map = broadcastMap.compose(map);
- *operand = broadcast.getSource();
+ *operand = broadcast ? broadcast.getSource() : sc.getSource();
changed = true;
}
diff --git a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
index 3b51e6b1e1b6f..8e53a815f70f8 100644
--- a/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
+++ b/mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir
@@ -202,7 +202,37 @@ func.func @contract_broadcast_unit_dim_reduction(%arg0 : vector<8x4xi32>, %arg1
// -----
-// Same as above, but with a mask.
+// Same as above (`@contract_broadcast_unit_dim_reduction`), but
+// `vector.broadcast` is replaced with `vector.shape_cast`.
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: contract_shape_cast_unit_dim_reduction
+// CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<8x4xi32>, %[[ARG2:.+]]: vector<8x8xi32>)
+// CHECK: vector.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x4xi32>, vector<8x4xi32> into vector<8x8xi32>
+func.func @contract_shape_cast_unit_dim_reduction(%arg0 : vector<8x4xi32>, %arg1 : vector<8x4xi32>, %arg2 : vector<8x8xi32>) -> vector<8x8xi32> {
+ %0 = vector.shape_cast %arg0 : vector<8x4xi32> to vector<1x8x4xi32>
+ %1 = vector.shape_cast %arg1 : vector<8x4xi32> to vector<1x8x4xi32>
+ %result = vector.contract {
+ indexing_maps = [#map0, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>
+ } %0, %1, %arg2 : vector<1x8x4xi32>, vector<1x8x4xi32> into vector<8x8xi32>
+ return %result : vector<8x8xi32>
+}
+
+// -----
+
+// Same as above (`@contract_broadcast_unit_dim_reduction`), but with a mask.
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
@@ -235,7 +265,41 @@ func.func @contract_broadcast_unit_dim_reduction_masked(%arg0 : vector<8x4xi32>,
// -----
-// Same as above, but with a scalable dim.
+// Same as above (`@contract_broadcast_unit_dim_reduction_masked`), but with
+// `vector.shape_cast` instead of `vector.broadcast`.
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: contract_shape_cast_unit_dim_reduction_masked
+// CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<8x4xi32>, %[[ARG2:.+]]: vector<8x8xi32>, %[[MASK:.+]]: vector<1x8x8x4xi1>)
+// CHECK: %[[MASK_SC:.*]] = vector.shape_cast %[[MASK]] : vector<1x8x8x4xi1> to vector<8x8x4xi1>
+// CHECK: %[[R:.*]] = vector.mask %[[MASK_SC]] {
+// CHECK-SAME: vector.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x4xi32>, vector<8x4xi32> into vector<8x8xi32>
+func.func @contract_shape_cast_unit_dim_reduction_masked(%arg0 : vector<8x4xi32>, %arg1 : vector<8x4xi32>, %arg2 : vector<8x8xi32>, %mask: vector<1x8x8x4xi1>) -> vector<8x8xi32> {
+ %0 = vector.shape_cast %arg0 : vector<8x4xi32> to vector<1x8x4xi32>
+ %1 = vector.shape_cast %arg1 : vector<8x4xi32> to vector<1x8x4xi32>
+ %result = vector.mask %mask {
+ vector.contract {
+ indexing_maps = [#map0, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>
+ } %0, %1, %arg2 : vector<1x8x4xi32>, vector<1x8x4xi32> into vector<8x8xi32>
+ } : vector<1x8x8x4xi1> -> vector<8x8xi32>
+ return %result : vector<8x8xi32>
+}
+
+// -----
+
+// Same as above (`@contract_broadcast_unit_dim_reduction_masked`), but with a scalable dim.
#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
@@ -267,6 +331,41 @@ func.func @contract_broadcast_unit_dim_reduction_masked_scalable(%arg0 : vector<
}
// -----
+
+// Same as above (`@contract_broadcast_unit_dim_reduction_masked_scalable`), but with
+// vector.shape_cast instead of vector.broadcast.
+
+#map0 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d3)>
+#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d3)>
+#map2 = affine_map<(d0, d1, d2, d3) -> (d1, d2)>
+
+// CHECK-DAG: #[[$MAP0:.*]] = affine_map<(d0, d1, d2) -> (d0, d2)>
+// CHECK-DAG: #[[$MAP1:.*]] = affine_map<(d0, d1, d2) -> (d1, d2)>
+// CHECK-DAG: #[[$MAP2:.*]] = affine_map<(d0, d1, d2) -> (d0, d1)>
+
+// CHECK-LABEL: contract_shape_cast_unit_dim_reduction_masked_scalable
+// CHECK-SAME: (%[[ARG0:.+]]: vector<8x4xi32>, %[[ARG1:.+]]: vector<[8]x4xi32>, %[[ARG2:.+]]: vector<8x[8]xi32>, %[[MASK:.+]]: vector<1x8x[8]x4xi1>)
+// CHECK: %[[MASK_SC:.*]] = vector.shape_cast %[[MASK]] : vector<1x8x[8]x4xi1> to vector<8x[8]x4xi1>
+// CHECK: %[[R:.*]] = vector.mask %[[MASK_SC]] {
+// CHECK-SAME: vector.contract
+// CHECK-SAME: indexing_maps = [#[[$MAP0]], #[[$MAP1]], #[[$MAP2]]]
+// CHECK-SAME: iterator_types = ["parallel", "parallel", "reduction"]
+// CHECK-SAME: %[[ARG0]], %[[ARG1]], %[[ARG2]] : vector<8x4xi32>, vector<[8]x4xi32> into vector<8x[8]xi32>
+func.func @contract_shape_cast_unit_dim_reduction_masked_scalable(%arg0 : vector<8x4xi32>, %arg1 : vector<[8]x4xi32>, %arg2 : vector<8x[8]xi32>, %mask: vector<1x8x[8]x4xi1>) -> vector<8x[8]xi32> {
+ %0 = vector.shape_cast %arg0 : vector<8x4xi32> to vector<1x8x4xi32>
+ %1 = vector.shape_cast %arg1 : vector<[8]x4xi32> to vector<1x[8]x4xi32>
+ %result = vector.mask %mask {
+ vector.contract {
+ indexing_maps = [#map0, #map1, #map2],
+ iterator_types = ["reduction", "parallel", "parallel", "reduction"],
+ kind = #vector.kind<add>
+ } %0, %1, %arg2 : vector<1x8x4xi32>, vector<1x[8]x4xi32> into vector<8x[8]xi32>
+ } : vector<1x8x[8]x4xi1> -> vector<8x[8]xi32>
+ return %result : vector<8x[8]xi32>
+}
+
+// -----
+
// Test that CombineContractBroadcast will not combine a broadcast that creates
// a non-unit dim that is consumed by a reduction iterator.
// Moreover, the affine_map's are permuting the position of that reduction
>From 5a9232f71f80268e93cfd7790a2229b839c43fae Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Tue, 14 Jul 2026 09:09:30 +0000
Subject: [PATCH 2/2] Address PR comments
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 28 +++++++++++++++++--
.../Vector/Transforms/VectorTransforms.cpp | 10 ++++---
2 files changed, 32 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index fb535cdae461f..dd1a3895ad669 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -6759,13 +6759,37 @@ LogicalResult ShapeCastOp::verify() {
return success();
}
+/// Check whether this ShapeCastOp is effectively a BroadcastOp.
+///
+/// The only case in which this method can return `true` is when the underlying
+/// op merely adds leading unit dimensions, e.g.:
+/// %res = vector.shape_cast %src : vector<8x4xi32> to vector<1x8x4xi32>
+///
bool ShapeCastOp::isBroadcastLike() {
auto srcType = getSourceVectorType();
auto resType = getResultVectorType();
+ // Is srcType broadcastable to resType?
std::pair<VectorDim, VectorDim> mismatchingDims;
- return isBroadcastableTo(srcType, resType, &mismatchingDims) ==
- BroadcastableToResult::Success;
+ if (isBroadcastableTo(srcType, resType, &mismatchingDims) !=
+ BroadcastableToResult::Success)
+ return false;
+
+ // Do ranks mismatch?
+ //
+ // The only case where ranks match and this ShapeCastOp is also a broadcast,
+ // is when it's effectively a NOp, but that's an uninteresting edge case.
+ size_t rankDiff = resType.getRank() - srcType.getRank();
+ if (rankDiff == 0)
+ return false;
+
+ // Are all newly added leading dims unit?
+ if (!llvm::all_of(resType.getShape().take_front(rankDiff),
+ [](int64_t dim) { return dim == 1; }))
+ return false;
+
+ // Do all trailing dims match?
+ return resType.getShape().take_back(srcType.getRank()) == srcType.getShape();
}
/// Return true if `transpose` does not permute a pair of non-unit dims.
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
index 22e7742128ff2..ad35465179e94 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp
@@ -283,7 +283,8 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
for (Value *operand : {&lhs, &rhs}) {
AffineMap &map = maps[index++];
- // Make sure that the source
+ // Accept operands defined by vector.broadcast and broadcast-like
+ // vector.shape_cast.
auto sc = operand->getDefiningOp<vector::ShapeCastOp>();
auto broadcast = operand->getDefiningOp<vector::BroadcastOp>();
if (!broadcast && !sc)
@@ -294,10 +295,11 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
contractOp, "Operand defined via vector.shape_cast that has "
"non-broadcast semantics");
- VectorType srcType = (sc) ? sc.getSourceVectorType()
- : dyn_cast<VectorType>(broadcast.getSourceType());
+ // Get the source and the result types.
+ VectorType srcType = sc ? sc.getSourceVectorType()
+ : dyn_cast<VectorType>(broadcast.getSourceType());
VectorType resType =
- (sc) ? sc.getResultVectorType() : broadcast.getResultVectorType();
+ sc ? sc.getResultVectorType() : broadcast.getResultVectorType();
// contractionOp can only take vector as operands.
// auto srcType = dyn_cast<VectorType>(broadcast.getSourceVectorType());
More information about the Mlir-commits
mailing list