Skip to content

Commit 88610b7

Browse files
authored
[mlir][vector] Disable transpose -> shuffle lowering for scalable vectors (#79979)
vector.shuffle is not supported for scalable vectors (outside of splats)
1 parent 95ef8e3 commit 88610b7

File tree

2 files changed

+25
-0
lines changed

2 files changed

+25
-0
lines changed

mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -434,6 +434,10 @@ class TransposeOp2DToShuffleLowering
434434
return rewriter.notifyMatchFailure(
435435
op, "not using vector shuffle based lowering");
436436

437+
if (op.getSourceVectorType().isScalable())
438+
return rewriter.notifyMatchFailure(
439+
op, "vector shuffle lowering not supported for scalable vectors");
440+
437441
auto srcGtOneDims = isTranspose2DSlice(op);
438442
if (failed(srcGtOneDims))
439443
return rewriter.notifyMatchFailure(

mlir/test/Dialect/Vector/vector-transpose-lowering.mlir

Lines changed: 21 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -849,3 +849,24 @@ module attributes {transform.with_named_sequence} {
849849
transform.yield
850850
}
851851
}
852+
853+
// -----
854+
855+
// Scalable transposes should not be lowered to vector.shuffle.
856+
857+
// CHECK-LABEL: func @transpose_nx8x2xf32
858+
func.func @transpose_nx8x2xf32(%arg0: vector<[8]x2xf32>) -> vector<2x[8]xf32> {
859+
// CHECK-NOT: vector.shuffle
860+
// CHECK: vector.transpose %{{.*}} : vector<[8]x2xf32> to vector<2x[8]xf32>
861+
%0 = vector.transpose %arg0, [1, 0] : vector<[8]x2xf32> to vector<2x[8]xf32>
862+
return %0 : vector<2x[8]xf32>
863+
}
864+
865+
module attributes {transform.with_named_sequence} {
866+
transform.named_sequence @__transform_main(%func_op: !transform.op<"func.func"> {transform.readonly}) {
867+
transform.apply_patterns to %func_op {
868+
transform.apply_patterns.vector.lower_transpose lowering_strategy = "shuffle_1d"
869+
} : !transform.op<"func.func">
870+
transform.yield
871+
}
872+
}

0 commit comments

Comments
 (0)