[mlir][vector] Extend combineContractAndBroadcast to accept vector.shape_cast (#208752)
`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()`.
This commit is contained in:
Andrzej Warzyński
GitHub
parent
709b7d1bd4
commit
de5cef3127
5 changed files with 221 additions and 72 deletions
|
|
@ -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,
|
||||
|
|
|
|||
|
|
@ -2496,6 +2496,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;
|
||||
|
|
|
|||
|
|
@ -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(
|
||||
|
|
@ -6787,6 +6789,39 @@ 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;
|
||||
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.
|
||||
/// By `order preserving` we mean that the flattened versions of the input and
|
||||
/// output vectors are (numerically) identical. In other words `transpose` is
|
||||
|
|
|
|||
|
|
@ -244,7 +244,7 @@ struct CombineContractResultTranspose final
|
|||
}
|
||||
};
|
||||
|
||||
/// Merge BroadcastOp into ContractionOp user.
|
||||
/// Merge BroadcastOp (and broadcast-like ShapeCastOp) into ContractionOp user.
|
||||
/// Ex:
|
||||
/// ```
|
||||
/// %0 = vector.broadcast %arg0 : vector<32x16xf32> to vector<8x32x16xf32>
|
||||
|
|
@ -282,21 +282,34 @@ FailureOr<Value> combineContractAndBroadcast(vector::ContractionOp contractOp,
|
|||
bool changed = false;
|
||||
for (Value *operand : {&lhs, &rhs}) {
|
||||
AffineMap &map = maps[index++];
|
||||
|
||||
// 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)
|
||||
if (!broadcast && !sc)
|
||||
continue;
|
||||
|
||||
if (sc && !sc.isBroadcastLike())
|
||||
return rewriter.notifyMatchFailure(
|
||||
contractOp, "Operand defined via vector.shape_cast that has "
|
||||
"non-broadcast semantics");
|
||||
|
||||
// Get the source and the result types.
|
||||
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 +324,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 +334,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;
|
||||
}
|
||||
|
||||
|
|
|
|||
|
|
@ -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
|
||||
|
|
|
|||
Loading…
Reference in a new issue