diff --git a/mlir/lib/Dialect/Vector/VectorTransferPermutationMapRewritePatterns.cpp b/mlir/lib/Dialect/Vector/VectorTransferPermutationMapRewritePatterns.cpp index 3f5c3127a286c..a27ebfc8e5c62 100644 --- a/mlir/lib/Dialect/Vector/VectorTransferPermutationMapRewritePatterns.cpp +++ b/mlir/lib/Dialect/Vector/VectorTransferPermutationMapRewritePatterns.cpp @@ -224,9 +224,15 @@ struct TransferOpReduceRank : public OpRewritePattern { // https://llvm.discourse.group/t/should-we-have-0-d-vectors/3097. // In the meantime, lower these to a scalar load when they pop up. if (reducedShapeRank == 0) { - Value newRead = rewriter.create( - op.getLoc(), originalVecType.getElementType(), op.source(), - op.indices()); + Value newRead; + if (op.getShapedType().isa()) { + newRead = rewriter.create(op.getLoc(), op.source(), + op.indices()); + } else { + newRead = rewriter.create( + op.getLoc(), originalVecType.getElementType(), op.source(), + op.indices()); + } rewriter.replaceOpWithNewOp(op, originalVecType, newRead); 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 866d791c7c19f..a5c0cb584b11b 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,9 +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-LABEL: func @vector_transfer_ops_0d_memref( // CHECK-SAME: %[[MEM:.*]]: memref // CHECK-SAME: %[[VV:.*]]: vector<1x1x1xf32> -func @vector_transfer_ops_0d(%M: memref, %v: vector<1x1x1xf32>) { +func @vector_transfer_ops_0d_memref(%M: memref, %v: vector<1x1x1xf32>) { %f0 = arith.constant 0.0 : f32 // CHECK-NEXT: %[[V:.*]] = memref.load %[[MEM]][] : memref @@ -23,6 +23,22 @@ func @vector_transfer_ops_0d(%M: memref, %v: vector<1x1x1xf32>) { // ----- +// CHECK-LABEL: func @vector_transfer_ops_0d_tensor( +// CHECK-SAME: %[[SOURCE:.*]]: tensor +func @vector_transfer_ops_0d_tensor(%M: tensor) -> vector<1xf32> { + %f0 = arith.constant 0.0 : f32 + +// CHECK-NEXT: %[[S:.*]] = tensor.extract %[[SOURCE]][] : tensor +// CHECK-NEXT: %[[V:.*]] = vector.broadcast %[[S]] : f32 to vector<1xf32> + %0 = vector.transfer_read %M[], %f0 {permutation_map = affine_map<()->(0)>} : + tensor, vector<1xf32> + +// CHECK-NEXT: return %[[V]] + return %0: vector<1xf32> +} + +// ----- + // transfer_read/write are lowered to vector.load/store // CHECK-LABEL: func @transfer_to_load( // CHECK-SAME: %[[MEM:.*]]: memref<8x8xf32>,