Skip to content

[mlir][vector][nfc] Update comments in vector-transpose.mlir #139699

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 2 commits into from
May 14, 2025
Merged

[mlir][vector][nfc] Update comments in vector-transpose.mlir #139699

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
ff87b87
[mlir][vector][nfc] Update comments in vector-transpose.mlir
banach-space May 13, 2025
c36107e
fixup! [mlir][vector][nfc] Update comments in vector-transpose.mlir
banach-space 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
43 changes: 22 additions & 21 deletions mlir/test/Dialect/Vector/canonicalize/vector-transpose.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -137,20 +137,22 @@ func.func @negative_broadcast_transpose_021(%arg0 : vector<3x1x3xi8>) -> vector<
return %1 : vector<3x3x3xi8>
}


// -----

// Test of FoldTransposeShapeCast
/// +--------------------------------------------------------------------------
/// Tests of ShapeCastOp::fold: shape_cast(transpose) -> shape_cast
/// +--------------------------------------------------------------------------

// In this test, the permutation maps the non-unit dimensions (1 and 2) as follows:
// 1 -> 0
// 2 -> 4
// Because 0 < 4, this permutation is order preserving and effectively a shape_cast.
// CHECK-LABEL: @transpose_shape_cast
// CHECK-LABEL: @shape_cast_of_transpose
// CHECK-SAME: %[[ARG:.*]]: vector<1x4x4x1x1xi8>) -> vector<4x4xi8> {
// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :
// CHECK-SAME: vector<1x4x4x1x1xi8> to vector<4x4xi8>
// CHECK: return %[[SHAPE_CAST]] : vector<4x4xi8>
func.func @transpose_shape_cast(%arg : vector<1x4x4x1x1xi8>) -> vector<4x4xi8> {
func.func @shape_cast_of_transpose(%arg : vector<1x4x4x1x1xi8>) -> vector<4x4xi8> {
%0 = vector.transpose %arg, [1, 0, 3, 4, 2]
: vector<1x4x4x1x1xi8> to vector<4x1x1x1x4xi8>
%1 = vector.shape_cast %0 : vector<4x1x1x1x4xi8> to vector<4x4xi8>
Expand All @@ -159,18 +161,17 @@ func.func @transpose_shape_cast(%arg : vector<1x4x4x1x1xi8>) -> vector<4x4xi8> {

// -----

// Test of FoldTransposeShapeCast
// In this test, the mapping of non-unit dimensions (1 and 2) is as follows:
// 1 -> 2
// 2 -> 1
// As this is not increasing (2 > 1), this transpose is not order
// preserving and cannot be treated as a shape_cast.
// CHECK-LABEL: @negative_transpose_shape_cast
// CHECK-LABEL: @negative_shape_cast_of_transpose
// CHECK-SAME: %[[ARG:.*]]: vector<1x4x4x1xi8>) -> vector<4x4xi8> {
// CHECK: %[[TRANSPOSE:.*]] = vector.transpose %[[ARG]]
// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[TRANSPOSE]]
// CHECK: return %[[SHAPE_CAST]] : vector<4x4xi8>
func.func @negative_transpose_shape_cast(%arg : vector<1x4x4x1xi8>) -> vector<4x4xi8> {
func.func @negative_shape_cast_of_transpose(%arg : vector<1x4x4x1xi8>) -> vector<4x4xi8> {
%0 = vector.transpose %arg, [0, 2, 1, 3]
: vector<1x4x4x1xi8> to vector<1x4x4x1xi8>
%1 = vector.shape_cast %0 : vector<1x4x4x1xi8> to vector<4x4xi8>
Expand All @@ -179,42 +180,43 @@ func.func @negative_transpose_shape_cast(%arg : vector<1x4x4x1xi8>) -> vector<4x

// -----

// Test of FoldTransposeShapeCast
// Currently the conversion shape_cast(transpose) -> shape_cast is disabled for
// scalable vectors because of bad interaction with ConvertIllegalShapeCastOpsToTransposes
// CHECK-LABEL: @negative_transpose_shape_cast_scalable
// CHECK-LABEL: @negative_shape_cast_of_transpose_scalable
// CHECK: vector.transpose
// CHECK: vector.shape_cast
func.func @negative_transpose_shape_cast_scalable(%arg : vector<[4]x1xi8>) -> vector<[4]xi8> {
func.func @negative_shape_cast_of_transpose_scalable(%arg : vector<[4]x1xi8>) -> vector<[4]xi8> {
%0 = vector.transpose %arg, [1, 0] : vector<[4]x1xi8> to vector<1x[4]xi8>
%1 = vector.shape_cast %0 : vector<1x[4]xi8> to vector<[4]xi8>
return %1 : vector<[4]xi8>
}

// -----

// Test of shape_cast folding.
/// +--------------------------------------------------------------------------
/// Tests of FoldTransposeShapeCast: transpose(shape_cast) -> shape_cast
/// +--------------------------------------------------------------------------

// The conversion transpose(shape_cast) -> shape_cast is not disabled for scalable
// vectors.
// CHECK-LABEL: @shape_cast_transpose_scalable
// CHECK-LABEL: @transpose_of_shape_cast_scalable
// CHECK: vector.shape_cast
// CHECK-SAME: vector<[4]xi8> to vector<[4]x1xi8>
func.func @shape_cast_transpose_scalable(%arg : vector<[4]xi8>) -> vector<[4]x1xi8> {
func.func @transpose_of_shape_cast_scalable(%arg : vector<[4]xi8>) -> vector<[4]x1xi8> {
%0 = vector.shape_cast %arg : vector<[4]xi8> to vector<1x[4]xi8>
%1 = vector.transpose %0, [1, 0] : vector<1x[4]xi8> to vector<[4]x1xi8>
return %1 : vector<[4]x1xi8>
}

// -----

// Test of shape_cast folding.
// A transpose that is 'order preserving' can be treated like a shape_cast.
// CHECK-LABEL: @shape_cast_transpose
// CHECK-LABEL: @transpose_of_shape_cast
// CHECK-SAME: %[[ARG:.*]]: vector<2x3x1x1xi8>) -> vector<6x1x1xi8> {
// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :
// CHECK-SAME: vector<2x3x1x1xi8> to vector<6x1x1xi8>
// CHECK: return %[[SHAPE_CAST]] : vector<6x1x1xi8>
func.func @shape_cast_transpose(%arg : vector<2x3x1x1xi8>) -> vector<6x1x1xi8> {
func.func @transpose_of_shape_cast(%arg : vector<2x3x1x1xi8>) -> vector<6x1x1xi8> {
%0 = vector.shape_cast %arg : vector<2x3x1x1xi8> to vector<6x1x1xi8>
%1 = vector.transpose %0, [0, 2, 1]
: vector<6x1x1xi8> to vector<6x1x1xi8>
Expand All @@ -223,12 +225,11 @@ func.func @shape_cast_transpose(%arg : vector<2x3x1x1xi8>) -> vector<6x1x1xi8>

// -----

// Test of shape_cast folding.
// Scalable dimensions should be treated as non-unit dimensions.
// CHECK-LABEL: @shape_cast_transpose_scalable
// CHECK-LABEL: @transpose_of_shape_cast_scalable
// CHECK: vector.shape_cast
// CHECK: vector.transpose
func.func @shape_cast_transpose_scalable_unit(%arg : vector<[1]x4x1xi8>) -> vector<4x[1]xi8> {
func.func @transpose_of_shape_cast_scalable_unit(%arg : vector<[1]x4x1xi8>) -> vector<4x[1]xi8> {
%0 = vector.shape_cast %arg : vector<[1]x4x1xi8> to vector<[1]x4xi8>
%1 = vector.transpose %0, [1, 0] : vector<[1]x4xi8> to vector<4x[1]xi8>
return %1 : vector<4x[1]xi8>
Expand All @@ -237,12 +238,12 @@ func.func @shape_cast_transpose_scalable_unit(%arg : vector<[1]x4x1xi8>) -> vect
// -----

// Test of shape_cast (not) folding.
// CHECK-LABEL: @negative_shape_cast_transpose
// CHECK-LABEL: @negative_transpose_of_shape_cast
// CHECK-SAME: %[[ARG:.*]]: vector<6xi8>) -> vector<2x3xi8> {
// CHECK: %[[SHAPE_CAST:.*]] = vector.shape_cast %[[ARG]] :
// CHECK: %[[TRANSPOSE:.*]] = vector.transpose %[[SHAPE_CAST]]
// CHECK: return %[[TRANSPOSE]] : vector<2x3xi8>
func.func @negative_shape_cast_transpose(%arg : vector<6xi8>) -> vector<2x3xi8> {
func.func @negative_transpose_of_shape_cast(%arg : vector<6xi8>) -> vector<2x3xi8> {
%0 = vector.shape_cast %arg : vector<6xi8> to vector<3x2xi8>
%1 = vector.transpose %0, [1, 0] : vector<3x2xi8> to vector<2x3xi8>
return %1 : vector<2x3xi8>
Expand Down
Loading