maxjhandsome
/
llvm-project
forked from OSchip/llvm-project
1
0
Fork 0
Code Issues Pull Requests Packages Releases Wiki Activity

[mlir][vector] Add constant folder for vector.shuffle ops 

Reviewed By: ThomasRaoux

Differential Revision: https://reviews.llvm.org/D119032
This commit is contained in:
Lei Zhang 2022-02-04 16:32:42 -05:00
parent 3fd30958b7
commit 9dd4c2dcb6
3 changed files with 40 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
mlir/include/mlir/Dialect/Vector/IR/VectorOps.td
View File

@ -465,6 +465,7 @@ def Vector_ShuffleOp :
let builders = [
OpBuilder<(ins "Value":$v1, "Value":$v2, "ArrayRef<int64_t>")>
];
let hasFolder = 1;
let extraClassDeclaration = [{
static StringRef getMaskAttrName() { return "mask"; }
VectorType getV1VectorType() {

27
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
View File

@ -1803,6 +1803,33 @@ static ParseResult parseShuffleOp(OpAsmParser &parser, OperationState &result) {
return success();
}
OpFoldResult vector::ShuffleOp::fold(ArrayRef<Attribute> operands) {
Attribute lhs = operands.front(), rhs = operands.back();
if (!lhs || !rhs)
return {};
auto lhsType = lhs.getType().cast<VectorType>();
// Only support 1-D for now to avoid complicated n-D DenseElementsAttr
// manipulation.
if (lhsType.getRank() != 1)
return {};
int64_t lhsSize = lhsType.getDimSize(0);
SmallVector<Attribute> results;
auto lhsElements = lhs.cast<DenseElementsAttr>().getValues<Attribute>();
auto rhsElements = rhs.cast<DenseElementsAttr>().getValues<Attribute>();
for (const auto &index : this->mask().getAsValueRange<IntegerAttr>()) {
int64_t i = index.getZExtValue();
if (i >= lhsSize) {
results.push_back(rhsElements[i - lhsSize]);
} else {
results.push_back(lhsElements[i]);
}
}
return DenseElementsAttr::get(getVectorType(), results);
}
//===----------------------------------------------------------------------===//
// InsertElementOp
//===----------------------------------------------------------------------===//

12
mlir/test/Dialect/Vector/canonicalize.mlir
View File

@ -1254,3 +1254,15 @@ func @splat_fold() -> vector<4xf32> {
// CHECK-NEXT: [[V:%.*]] = arith.constant dense<1.000000e+00> : vector<4xf32>
// CHECK-NEXT: return [[V]] : vector<4xf32>
}
// -----
// CHECK-LABEL: func @shuffle_1d
// CHECK: %[[V:.+]] = arith.constant dense<[3, 2, 5, 1]> : vector<4xi32>
// CHECK: return %[[V]]
func @shuffle_1d() -> vector<4xi32> {
%v0 = arith.constant dense<[0, 1, 2]> : vector<3xi32>
%v1 = arith.constant dense<[3, 4, 5]> : vector<3xi32>
%shuffle = vector.shuffle %v0, %v1 [3, 2, 5, 1] : vector<3xi32>, vector<3xi32>
return %shuffle : vector<4xi32>
}