[Mlir-commits] [mlir] 88610b7 - [mlir][vector] Disable transpose -> shuffle lowering for scalable vectors (#79979)

[llvmlistbot at llvm.org]

Author: Benjamin Maxwell
Date: 2024-01-31T09:21:44Z
New Revision: 88610b79510983d7838e4f97411d5a2ac1f8fee8

URL: https://github.com/llvm/llvm-project/commit/88610b79510983d7838e4f97411d5a2ac1f8fee8
DIFF: https://github.com/llvm/llvm-project/commit/88610b79510983d7838e4f97411d5a2ac1f8fee8.diff

LOG: [mlir][vector] Disable transpose -> shuffle lowering for scalable vectors (#79979)

vector.shuffle is not supported for scalable vectors (outside of splats)

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
    mlir/test/Dialect/Vector/vector-transpose-lowering.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
index 97f6caca1b25c..792550dcfaf22 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTranspose.cpp
@@ -434,6 +434,10 @@ class TransposeOp2DToShuffleLowering
       return rewriter.notifyMatchFailure(
           op, "not using vector shuffle based lowering");
 
+    if (op.getSourceVectorType().isScalable())
+      return rewriter.notifyMatchFailure(
+          op, "vector shuffle lowering not supported for scalable vectors");
+
     auto srcGtOneDims = isTranspose2DSlice(op);
     if (failed(srcGtOneDims))
       return rewriter.notifyMatchFailure(

diff  --git a/mlir/test/Dialect/Vector/vector-transpose-lowering.mlir b/mlir/test/Dialect/Vector/vector-transpose-lowering.mlir
index c0b44428d5bcf..97b698edeb059 100644
--- a/mlir/test/Dialect/Vector/vector-transpose-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-transpose-lowering.mlir
@@ -849,3 +849,24 @@ module attributes {transform.with_named_sequence} {
     transform.yield
   }
 }
+
+// -----
+
+// Scalable transposes should not be lowered to vector.shuffle.
+
+// CHECK-LABEL: func @transpose_nx8x2xf32
+func.func @transpose_nx8x2xf32(%arg0: vector<[8]x2xf32>) -> vector<2x[8]xf32> {
+  // CHECK-NOT: vector.shuffle
+  // CHECK: vector.transpose %{{.*}} : vector<[8]x2xf32> to vector<2x[8]xf32>
+  %0 = vector.transpose %arg0, [1, 0] : vector<[8]x2xf32> to vector<2x[8]xf32>
+  return %0 : vector<2x[8]xf32>
+}
+
+module attributes {transform.with_named_sequence} {
+  transform.named_sequence @__transform_main(%func_op: !transform.op<"func.func"> {transform.readonly}) {
+    transform.apply_patterns to %func_op {
+      transform.apply_patterns.vector.lower_transpose lowering_strategy = "shuffle_1d"
+    } : !transform.op<"func.func">
+    transform.yield
+  }
+}