[Mlir-commits] [mlir] [mlir][vector] Fix a crash in `VectorExtractOpConversion` (PR #115717)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

This PR fixes a crash when `vector.extract` extract a scalar and the size of `position` smaller than rank of vector. E.g., `vector.extract %arg0[0]: f32 from vector<4x1xf32>`. Fixes #<!-- -->115294.

---
Full diff: https://github.com/llvm/llvm-project/pull/115717.diff


2 Files Affected:

- (modified) mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp (+1-2) 
- (modified) mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir (+20) 


``````````diff
diff --git a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
index 58ca84c8d7bca6..5b501744afd37d 100644
--- a/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
+++ b/mlir/lib/Conversion/VectorToLLVM/ConvertVectorToLLVM.cpp
@@ -1105,8 +1105,7 @@ class VectorExtractOpConversion
 
     // One-shot extraction of vector from array (only requires extractvalue).
     // Except for extracting 1-element vectors.
-    if (isa<VectorType>(resultType) &&
-        position.size() !=
+    if (position.size() <
             static_cast<size_t>(extractOp.getSourceVectorType().getRank())) {
       if (extractOp.hasDynamicPosition())
         return failure();
diff --git a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
index 03bcb341efea2f..26f1f3e5190df5 100644
--- a/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
+++ b/mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir
@@ -1216,6 +1216,26 @@ func.func @extract_vec_1d_from_vec_3d_f32_scalable(%arg0: vector<4x3x[16]xf32>)
 
 // -----
 
+func.func @extract_scalar_from_vec_2d_f32(%arg0: vector<4x1xf32>) -> f32 {
+  %0 = vector.extract %arg0[0]: f32 from vector<4x1xf32>
+  return %0 : f32
+}
+// CHECK-LABEL: @extract_scalar_from_vec_2d_f32
+//       CHECK:   %[[T0:.*]] = llvm.extractvalue {{.*}}[0] : !llvm.array<4 x vector<1xf32>>
+//       CHECK:   %[[T1:.*]] = builtin.unrealized_conversion_cast %[[T0]] : vector<1xf32> to f32
+//       CHECK:   return %[[T1]] : f32
+
+func.func @extract_scalar_from_vec_2d_f32_scalable(%arg0: vector<4x[1]xf32>) -> f32 {
+  %0 = vector.extract %arg0[0]: f32 from vector<4x[1]xf32>
+  return %0 : f32
+}
+// CHECK-LABEL: @extract_scalar_from_vec_2d_f32_scalable
+//       CHECK:   %[[T0:.*]] = llvm.extractvalue {{.*}}[0] : !llvm.array<4 x vector<[1]xf32>>
+//       CHECK:   %[[T1:.*]] = builtin.unrealized_conversion_cast %[[T0]] : vector<[1]xf32> to f32
+//       CHECK:   return %[[T1]] : f32
+
+// -----
+
 func.func @extract_scalar_from_vec_3d_f32(%arg0: vector<4x3x16xf32>) -> f32 {
   %0 = vector.extract %arg0[0, 0, 0]: f32 from vector<4x3x16xf32>
   return %0 : f32

``````````

</details>


https://github.com/llvm/llvm-project/pull/115717