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

[llvmlistbot at llvm.org]

Author: Longsheng Mou
Date: 2024-09-04T16:41:56+08:00
New Revision: 50febdeb64fce345b0fb669e9688d34c8ffe7912

URL: https://github.com/llvm/llvm-project/commit/50febdeb64fce345b0fb669e9688d34c8ffe7912
DIFF: https://github.com/llvm/llvm-project/commit/50febdeb64fce345b0fb669e9688d34c8ffe7912.diff

LOG: [mlir][vector] Bugfix of linearize `vector.extract` (#106836)

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

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
    mlir/test/Dialect/Vector/linearize.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
index 868397f2daaae4..11917ac1e40226 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorLinearize.cpp
@@ -337,6 +337,10 @@ 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,
+                                         "expected n-D 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..543e76b5b26e0c 100644
--- a/mlir/test/Dialect/Vector/linearize.mlir
+++ b/mlir/test/Dialect/Vector/linearize.mlir
@@ -306,3 +306,15 @@ 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-NOT: vector.shuffle
+  // ALL:     vector.extract
+  // ALL-NOT: vector.shuffle
+  %0 = vector.extract %cst[0] : i32 from vector<4xi32>
+  return
+}