[Mlir-commits] [mlir] [mlir][vector] Extend `combineContractAndBroadcast` to accept `vector.shape_cast` (PR #208752)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Fri Jul 10 08:20:31 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Andrzej Warzyński (banach-space)
<details>
<summary>Changes</summary>
`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()`.
---
Full diff: https://github.com/llvm/llvm-project/pull/208752.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.h (+1-1)
- (modified) mlir/include/mlir/Dialect/Vector/IR/VectorOps.td (+3)
- (modified) mlir/lib/Dialect/Vector/IR/VectorOps.cpp (+66-55)
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorTransforms.cpp (+23-13)
- (modified) mlir/test/Dialect/Vector/vector-reduce-to-contract.mlir (+101-2)
``````````diff
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
``````````
</details>
https://github.com/llvm/llvm-project/pull/208752
More information about the Mlir-commits
mailing list