[Mlir-commits] [mlir] b24c91f - [mlir][Vector][Bigfix] Fix vector transfer to store lowering to insert a proper ExtractOp

[Nicolas Vasilache]

Author: Nicolas Vasilache
Date: 2021-10-12T13:28:12Z
New Revision: b24c91fffc3f28f6b1a6c6f25c7f83efcdeda206

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

LOG: [mlir][Vector][Bigfix] Fix vector transfer to  store lowering to insert a proper ExtractOp

Differential Revision: https://reviews.llvm.org/D111641

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/VectorTransforms.cpp
    mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/VectorTransforms.cpp b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
index 46865160e6f9..6942aad9e790 100644
--- a/mlir/lib/Dialect/Vector/VectorTransforms.cpp
+++ b/mlir/lib/Dialect/Vector/VectorTransforms.cpp
@@ -2600,8 +2600,9 @@ struct VectorStoreToMemrefStoreLowering
     auto vecType = storeOp.getVectorType();
     if (vecType.getNumElements() != 1)
       return failure();
+    SmallVector<int64_t> indices(vecType.getRank(), 0);
     Value extracted = rewriter.create<vector::ExtractOp>(
-        storeOp.getLoc(), storeOp.valueToStore(), ArrayRef<int64_t>{1});
+        storeOp.getLoc(), storeOp.valueToStore(), indices);
     rewriter.replaceOpWithNewOp<memref::StoreOp>(
         storeOp, extracted, storeOp.base(), storeOp.indices());
     return success();

diff  --git a/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir b/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
index 4139a80527af..3e8706f6fbc5 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir
@@ -1,8 +1,9 @@
 // RUN: mlir-opt %s -test-vector-transfer-lowering-patterns -canonicalize -split-input-file | FileCheck %s
 
 // CHECK-LABEL: func @vector_transfer_ops_0d(
-//  CHECK-SAME:   %[[MEM:.*]]: memref<f32>) {
-func @vector_transfer_ops_0d(%M: memref<f32>) {
+//  CHECK-SAME:   %[[MEM:.*]]: memref<f32>
+//  CHECK-SAME:   %[[VV:.*]]: vector<1x1x1xf32>
+func @vector_transfer_ops_0d(%M: memref<f32>, %v: vector<1x1x1xf32>) {
     %f0 = constant 0.0 : f32
 
 //  CHECK-NEXT:   %[[V:.*]] = memref.load %[[MEM]][] : memref<f32>
@@ -13,6 +14,10 @@ func @vector_transfer_ops_0d(%M: memref<f32>) {
     vector.transfer_write %0, %M[] {permutation_map = affine_map<()->(0)>} :
       vector<1xf32>, memref<f32>
 
+//  CHECK-NEXT:   %[[VV:.*]] = vector.extract %arg1[0, 0, 0] : vector<1x1x1xf32>
+//  CHECK-NEXT:   memref.store %[[VV]], %[[MEM]][] : memref<f32>
+    vector.store %v, %M[] : memref<f32>, vector<1x1x1xf32>
+
     return
 }