[VectorOps] Implement a simple folder for identity vector.transpose operations.

Differential Revision: https://reviews.llvm.org/D77088
This commit is contained in:
Alex Grosul 2020-03-31 13:43:32 -07:00
parent 43e576593e
commit 855e738be2
3 changed files with 92 additions and 0 deletions

View File

@ -1315,10 +1315,12 @@ def Vector_TransposeOp :
VectorType getResultType() {
return result().getType().cast<VectorType>();
}
void getTransp(SmallVectorImpl<int64_t> &results);
}];
let assemblyFormat = [{
$vector `,` $transp attr-dict `:` type($vector) `to` type($result)
}];
let hasFolder = 1;
}
def Vector_TupleGetOp :

View File

@ -1524,6 +1524,23 @@ static LogicalResult verify(TupleOp op) { return success(); }
// TransposeOp
//===----------------------------------------------------------------------===//
// Eliminates transpose operations, which produce values identical to their
// input values. This happens when the dimensions of the input vector remain in
// their original order after the transpose operation.
OpFoldResult TransposeOp::fold(ArrayRef<Attribute> operands) {
SmallVector<int64_t, 4> transp;
getTransp(transp);
// Check if the permutation of the dimensions contains sequential values:
// {0, 1, 2, ...}.
for (int64_t i = 0, e = transp.size(); i < e; i++) {
if (transp[i] != i)
return {};
}
return vector();
}
static LogicalResult verify(TransposeOp op) {
VectorType vectorType = op.getVectorType();
VectorType resultType = op.getResultType();
@ -1549,6 +1566,10 @@ static LogicalResult verify(TransposeOp op) {
return success();
}
void TransposeOp::getTransp(SmallVectorImpl<int64_t> &results) {
populateFromInt64AttrArray(transp(), results);
}
//===----------------------------------------------------------------------===//
// TupleGetOp
//===----------------------------------------------------------------------===//

View File

@ -87,3 +87,72 @@ func @strided_slice_of_constant_mask() -> (vector<2x1xi1>) {
// CHECK: vector.constant_mask [1, 1] : vector<2x1xi1>
return %1 : vector<2x1xi1>
}
// -----
// CHECK-LABEL: transpose_1D_identity
// CHECK-SAME: ([[ARG:%.*]]: vector<4xf32>)
func @transpose_1D_identity(%arg : vector<4xf32>) -> vector<4xf32> {
// CHECK-NOT: transpose
%0 = vector.transpose %arg, [0] : vector<4xf32> to vector<4xf32>
// CHECK-NEXT: return [[ARG]]
return %0 : vector<4xf32>
}
// -----
// CHECK-LABEL: transpose_2D_identity
// CHECK-SAME: ([[ARG:%.*]]: vector<4x3xf32>)
func @transpose_2D_identity(%arg : vector<4x3xf32>) -> vector<4x3xf32> {
// CHECK-NOT: transpose
%0 = vector.transpose %arg, [0, 1] : vector<4x3xf32> to vector<4x3xf32>
// CHECK-NEXT: return [[ARG]]
return %0 : vector<4x3xf32>
}
// -----
// CHECK-LABEL: transpose_3D_identity
// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>)
func @transpose_3D_identity(%arg : vector<4x3x2xf32>) -> vector<4x3x2xf32> {
// CHECK-NOT: transpose
%0 = vector.transpose %arg, [0, 1, 2] : vector<4x3x2xf32> to vector<4x3x2xf32>
// CHECK-NEXT: return [[ARG]]
return %0 : vector<4x3x2xf32>
}
// -----
// CHECK-LABEL: transpose_2D_sequence
// CHECK-SAME: ([[ARG:%.*]]: vector<4x3xf32>)
func @transpose_2D_sequence(%arg : vector<4x3xf32>) -> vector<3x4xf32> {
// CHECK-NOT: transpose
%0 = vector.transpose %arg, [0, 1] : vector<4x3xf32> to vector<4x3xf32>
// CHECK: [[T1:%.*]] = vector.transpose [[ARG]], [1, 0]
%1 = vector.transpose %0, [1, 0] : vector<4x3xf32> to vector<3x4xf32>
// CHECK-NOT: transpose
%2 = vector.transpose %1, [0, 1] : vector<3x4xf32> to vector<3x4xf32>
// CHECK: [[ADD:%.*]] = addf [[T1]], [[T1]]
%4 = addf %1, %2 : vector<3x4xf32>
// CHECK-NEXT: return [[ADD]]
return %4 : vector<3x4xf32>
}
// -----
// CHECK-LABEL: transpose_3D_sequence
// CHECK-SAME: ([[ARG:%.*]]: vector<4x3x2xf32>)
func @transpose_3D_sequence(%arg : vector<4x3x2xf32>) -> vector<2x3x4xf32> {
// CHECK: [[T0:%.*]] = vector.transpose [[ARG]], [1, 2, 0]
%0 = vector.transpose %arg, [1, 2, 0] : vector<4x3x2xf32> to vector<3x2x4xf32>
// CHECK-NOT: transpose
%1 = vector.transpose %0, [0, 1, 2] : vector<3x2x4xf32> to vector<3x2x4xf32>
// CHECK: [[T2:%.*]] = vector.transpose [[T0]], [1, 0, 2]
%2 = vector.transpose %1, [1, 0, 2] : vector<3x2x4xf32> to vector<2x3x4xf32>
// CHECK: [[ADD:%.*]] = addf [[T2]], [[T2]]
%3 = addf %2, %2 : vector<2x3x4xf32>
// CHECK-NOT: transpose
%4 = vector.transpose %3, [0, 1, 2] : vector<2x3x4xf32> to vector<2x3x4xf32>
// CHECK-NEXT: return [[ADD]]
return %4 : vector<2x3x4xf32>
}