[Mlir-commits] [mlir] [mlir][vector] Bugfix of linearize `vector.extract` (PR #106836)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Longsheng Mou (CoTinker)

<details>
<summary>Changes</summary>

This patch add check for `vector.extract` with scalar type, which is not allowed when linearize `vector.extract`. Fix #<!-- -->106162.

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


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp (+3) 
- (modified) mlir/test/Dialect/Vector/linearize.mlir (+11) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index 868397f2daaae4..38c608ee338346 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -337,6 +337,9 @@ struct LinearizeVectorExtract final
   matchAndRewrite(vector::ExtractOp extractOp, OpAdaptor adaptor,
                   ConversionPatternRewriter &rewriter) const override {
     Type dstTy = getTypeConverter()->convertType(extractOp.getType());
+    if (!dstTy)
+      return rewriter.notifyMatchFailure(extractOp, "expect vector type.");
+
     if (extractOp.getVector().getType().isScalable() ||
         cast<VectorType>(dstTy).isScalable())
       return rewriter.notifyMatchFailure(extractOp,
diff --git a/mlir/test/Dialect/Vector/linearize.mlir b/mlir/test/Dialect/Vector/linearize.mlir
index 916e3e5fd2529d..0816cb52903e04 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -306,3 +306,14 @@ func.func @test_vector_insert_scalable(%arg0: vector<2x8x[4]xf32>, %arg1: vector
   // ALL: return %[[RES]] : vector<2x8x[4]xf32>
   return %0 : vector<2x8x[4]xf32>
 }
+
+// -----
+
+// ALL-LABEL: test_vector_extract_scalar
+func.func @test_vector_extract_scalar() {
+  %cst = arith.constant dense<[1, 2, 3, 4]> : vector<4xi32>
+  // ALL:     vector.extract
+  // ALL-NOT: vector.shuffle
+  %0 = vector.extract %cst[0] : i32 from vector<4xi32>
+  return
+}

``````````

</details>


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