[Mlir-commits] [mlir] 55d5c03 - [mlir][vector] Fix crash in `vector.extract` folder (#95912)

[llvmlistbot at llvm.org]

Author: Matthias Springer
Date: 2024-06-18T14:54:31+02:00
New Revision: 55d5c032fc9fd3aeba82cdc5dfc82b2a32f34d4f

URL: https://github.com/llvm/llvm-project/commit/55d5c032fc9fd3aeba82cdc5dfc82b2a32f34d4f
DIFF: https://github.com/llvm/llvm-project/commit/55d5c032fc9fd3aeba82cdc5dfc82b2a32f34d4f.diff

LOG: [mlir][vector] Fix crash in `vector.extract` folder (#95912)

Fix a bug in the `vector.extract` folder when the vector type is 0-d.

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/IR/VectorOps.cpp
    mlir/test/Dialect/Vector/canonicalize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 58951641d33ce..5e5d3e002086a 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -1883,7 +1883,10 @@ static Value foldExtractStridedOpFromInsertChain(ExtractOp extractOp) {
 }
 
 OpFoldResult ExtractOp::fold(FoldAdaptor) {
-  if (getNumIndices() == 0)
+  // Fold "vector.extract %v[] : vector<2x2xf32> from vector<2x2xf32>" to %v.
+  // Note: Do not fold "vector.extract %v[] : f32 from vector<f32>" (type
+  // mismatch).
+  if (getNumIndices() == 0 && getVector().getType() == getResult().getType())
     return getVector();
   if (succeeded(foldExtractOpFromExtractChain(*this)))
     return getResult();

diff  --git a/mlir/test/Dialect/Vector/canonicalize.mlir b/mlir/test/Dialect/Vector/canonicalize.mlir
index 22af91e0eb327..61269e3687ab3 100644
--- a/mlir/test/Dialect/Vector/canonicalize.mlir
+++ b/mlir/test/Dialect/Vector/canonicalize.mlir
@@ -2593,3 +2593,14 @@ func.func @rank_1_shuffle_to_interleave(%arg0: vector<6xi32>, %arg1: vector<6xi3
   %0 = vector.shuffle %arg0, %arg1 [0, 6, 1, 7, 2, 8, 3, 9, 4, 10, 5, 11] : vector<6xi32>, vector<6xi32>
   return %0 : vector<12xi32>
 }
+
+// -----
+
+// CHECK-LABEL: func @extract_from_0d_regression(
+//  CHECK-SAME:     %[[v:.*]]: vector<f32>)
+//       CHECK:   %[[extract:.*]] = vector.extract %[[v]][] : f32 from vector<f32>
+//       CHECK:   return %[[extract]]
+func.func @extract_from_0d_regression(%v: vector<f32>) -> f32 {
+  %0 = vector.extract %v[] : f32 from vector<f32>
+  return %0 : f32
+}