Skip to content

[mlir][vector] Additional transpose folding #138347

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 4 commits into from
May 14, 2025
Merged

[mlir][vector] Additional transpose folding #138347

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
4 commits
Select commit Hold shift + click to select a range
e5ca0d8
catch additional foldable case
newling May 2, 2025
1f21d4d
Update test with transpose that is now folded
newling May 5, 2025
2a2af7c
better grouping of tests
newling May 12, 2025
d92bf2c
test grouping polish
newling May 14, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
32 changes: 16 additions & 16 deletions mlir/lib/Dialect/Vector/IR/VectorOps.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -5573,13 +5573,11 @@ LogicalResult ShapeCastOp::verify() {
return success();
}

namespace {

/// 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
/// effectively a shape cast.
bool isOrderPreserving(TransposeOp transpose) {
static bool isOrderPreserving(TransposeOp transpose) {
ArrayRef<int64_t> permutation = transpose.getPermutation();
VectorType sourceType = transpose.getSourceVectorType();
ArrayRef<int64_t> inShape = sourceType.getShape();
Expand All @@ -5599,8 +5597,6 @@ bool isOrderPreserving(TransposeOp transpose) {
return true;
}

} // namespace

OpFoldResult ShapeCastOp::fold(FoldAdaptor adaptor) {

VectorType resultType = getType();
Expand Down Expand Up @@ -5997,18 +5993,22 @@ OpFoldResult vector::TransposeOp::fold(FoldAdaptor adaptor) {
if (llvm::dyn_cast_if_present<ub::PoisonAttr>(adaptor.getVector()))
return ub::PoisonAttr::get(getContext());

// Eliminate identity transpose ops. This happens when the dimensions of the
// input vector remain in their original order after the transpose operation.
ArrayRef<int64_t> perm = getPermutation();

// Check if the permutation of the dimensions contains sequential values:
// {0, 1, 2, ...}.
for (int64_t i = 0, e = perm.size(); i < e; i++) {
if (perm[i] != i)
return {};
}
// Eliminate identity transposes, and more generally any transposes that
// preserves the shape without permuting elements.
//
// Examples of what to fold:
// %0 = vector.transpose %arg, [0, 1] : vector<1x1xi8> to vector<1x1xi8>
// %0 = vector.transpose %arg, [0, 1] : vector<2x2xi8> to vector<2x2xi8>
// %0 = vector.transpose %arg, [1, 0] : vector<1x1xi8> to vector<1x1xi8>
//
// Example of what NOT to fold:
// %0 = vector.transpose %arg, [1, 0] : vector<2x2xi8> to vector<2x2xi8>
//
if (getSourceVectorType() == getResultVectorType() &&
isOrderPreserving(*this))
return getVector();

return getVector();
return {};
}

LogicalResult vector::TransposeOp::verify() {
Expand Down
22 changes: 0 additions & 22 deletions mlir/test/Dialect/Vector/canonicalize.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -450,28 +450,6 @@ func.func @extract_strided_fold_insert(%a: vector<2x8xf32>, %b: vector<1x4xf32>,

// -----

// CHECK-LABEL: transpose_1D_identity
// CHECK-SAME: ([[ARG:%.*]]: vector<4xf32>)
func.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.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.func @transpose_3D_identity(%arg : vector<4x3x2xf32>) -> vector<4x3x2xf32> {
Expand Down
50 changes: 49 additions & 1 deletion mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
View file
Open in desktop
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

How about adding a negative test?

Original file line number Diff line number Diff line change
@@ -1,6 +1,10 @@
// RUN: mlir-opt %s -canonicalize="test-convergence" -split-input-file -allow-unregistered-dialect | FileCheck %s

// This file contains some canonicalizations tests involving vector.transpose.
// This file contains some tests of canonicalizations and foldings involving vector.transpose.

// +---------------------------------------------------------------------------
// Tests of FoldTransposeBroadcast: transpose(broadcast) -> broadcast
// +---------------------------------------------------------------------------

// CHECK-LABEL: func @transpose_scalar_broadcast1
// CHECK-SAME: (%[[ARG:.+]]: vector<1xf32>)
Expand Down Expand Up @@ -248,3 +252,47 @@ func.func @negative_transpose_of_shape_cast(%arg : vector<6xi8>) -> vector<2x3xi
%1 = vector.transpose %0, [1, 0] : vector<3x2xi8> to vector<2x3xi8>
return %1 : vector<2x3xi8>
}

// -----

// +-----------------------------------
// Tests of TransposeOp::fold
// +-----------------------------------

// CHECK-LABEL: transpose_1D_identity
// CHECK-SAME: [[ARG:%.*]]: vector<4xf32>
// CHECK-NEXT: return [[ARG]]
func.func @transpose_1D_identity(%arg : vector<4xf32>) -> vector<4xf32> {
%0 = vector.transpose %arg, [0] : vector<4xf32> to vector<4xf32>
return %0 : vector<4xf32>
}

// -----

// CHECK-LABEL: transpose_2D_identity
// CHECK-SAME: [[ARG:%.*]]: vector<4x3xf32>
// CHECK-NEXT: return [[ARG]]
func.func @transpose_2D_identity(%arg : vector<4x3xf32>) -> vector<4x3xf32> {
%0 = vector.transpose %arg, [0, 1] : vector<4x3xf32> to vector<4x3xf32>
return %0 : vector<4x3xf32>
}

// -----

// CHECK-LABEL: transpose_shape_and_order_preserving
// CHECK-SAME: [[ARG:%.*]]: vector<6x1x1x4xi8>
// CHECK-NEXT: return [[ARG]]
func.func @transpose_shape_and_order_preserving(%arg : vector<6x1x1x4xi8>) -> vector<6x1x1x4xi8> {
%0 = vector.transpose %arg, [0, 2, 1, 3] : vector<6x1x1x4xi8> to vector<6x1x1x4xi8>
return %0 : vector<6x1x1x4xi8>
}

// -----

// CHECK-LABEL: negative_transpose_fold
// CHECK: [[TRANSP:%.*]] = vector.transpose
// CHECK: return [[TRANSP]]
func.func @negative_transpose_fold(%arg : vector<2x2xi8>) -> vector<2x2xi8> {
%0 = vector.transpose %arg, [1, 0] : vector<2x2xi8> to vector<2x2xi8>
return %0 : vector<2x2xi8>
}
16 changes: 9 additions & 7 deletions mlir/test/Dialect/Vector/vector-transpose-lowering.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,13 +65,15 @@ func.func @transpose102_8x1x8xf32(%arg0: vector<8x1x8xf32>) -> vector<1x8x8xf32>
return %0 : vector<1x8x8xf32>
}

// CHECK-LABEL: func @transpose1023_1x1x8x8xf32(
func.func @transpose1023_1x1x8x8xf32(%arg0: vector<1x1x8x8xf32>) -> vector<1x1x8x8xf32> {
// Note the single 2-D extract/insert pair since 2 and 3 are not transposed!
// CHECK: vector.extract {{.*}}[0, 0] : vector<8x8xf32> from vector<1x1x8x8xf32>
// CHECK-NEXT: vector.insert {{.*}} [0, 0] : vector<8x8xf32> into vector<1x1x8x8xf32>
%0 = vector.transpose %arg0, [1, 0, 2, 3] : vector<1x1x8x8xf32> to vector<1x1x8x8xf32>
return %0 : vector<1x1x8x8xf32>
// CHECK-LABEL: func @transpose1023_2x1x8x4xf32(
func.func @transpose1023_2x1x8x4xf32(%arg0: vector<2x1x8x4xf32>) -> vector<1x2x8x4xf32> {
// Note the 2-D extract/insert pair since dimensions 2 and 3 are not transposed!
// CHECK: vector.extract {{.*}}[0, 0] : vector<8x4xf32> from vector<2x1x8x4xf32>
// CHECK-NEXT: vector.insert {{.*}} [0, 0] : vector<8x4xf32> into vector<1x2x8x4xf32>
// CHECK-NEXT: vector.extract {{.*}}[1, 0] : vector<8x4xf32> from vector<2x1x8x4xf32>
// CHECK-NEXT: vector.insert {{.*}} [0, 1] : vector<8x4xf32> into vector<1x2x8x4xf32>
Comment on lines +73 to +74
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

IIUC, these were missing before?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I've changed the type of the (result) vector from <1x1x..> to <1x2x..>

Why? Because now the <1x1x...> case gets folded, so we end up with return %arg0. Based on the original comment in the test

"Note the single 2-D extract/insert pair since 2 and 3 are not transposed!"

I assume changing the sizes of dims 0/1 retains the original goal of the test.

%0 = vector.transpose %arg0, [1, 0, 2, 3] : vector<2x1x8x4xf32> to vector<1x2x8x4xf32>
return %0 : vector<1x2x8x4xf32>
}

/// Scalable dim should not be unrolled.
Expand Down