Skip to content

Commit 099fd01

Browse files
meltqmeltq
authored
[mlir][Vector] Remove Vector{Load|Store}ToMemrefLoadLowering (#121454)
0-d vectors are supported now and so these patterns are no longer required. This covers a part of this issue #112913 . Additionally this removes %arg2 in mlir/test/Conversion/GPUCommon/transfer_write.mlir and renames %arg3 to %arg2 as %arg2 was originally not required.
1 parent c745ece commit 099fd01

File tree

4 files changed

+41
-89
lines changed

4 files changed

+41
-89
lines changed

‎mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp

Lines changed: 0 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -492,60 +492,6 @@ struct TransferReadToVectorLoadLowering
492492
std::optional<unsigned> maxTransferRank;
493493
};
494494

495-
/// Replace a 0-d vector.load with a memref.load + vector.broadcast.
496-
// TODO: we shouldn't cross the vector/scalar domains just for this
497-
// but atm we lack the infra to avoid it. Possible solutions include:
498-
// - go directly to LLVM + bitcast
499-
// - introduce a bitcast op and likely a new pointer dialect
500-
// - let memref.load/store additionally support the 0-d vector case
501-
// There are still deeper data layout issues lingering even in this
502-
// trivial case (for architectures for which this matters).
503-
struct VectorLoadToMemrefLoadLowering
504-
: public OpRewritePattern<vector::LoadOp> {
505-
using OpRewritePattern::OpRewritePattern;
506-
507-
LogicalResult matchAndRewrite(vector::LoadOp loadOp,
508-
PatternRewriter &rewriter) const override {
509-
auto vecType = loadOp.getVectorType();
510-
if (vecType.getNumElements() != 1)
511-
return rewriter.notifyMatchFailure(loadOp, "not a single element vector");
512-
513-
auto memrefLoad = rewriter.create<memref::LoadOp>(
514-
loadOp.getLoc(), loadOp.getBase(), loadOp.getIndices());
515-
rewriter.replaceOpWithNewOp<vector::BroadcastOp>(loadOp, vecType,
516-
memrefLoad);
517-
return success();
518-
}
519-
};
520-
521-
/// Replace a 0-d vector.store with a vector.extractelement + memref.store.
522-
struct VectorStoreToMemrefStoreLowering
523-
: public OpRewritePattern<vector::StoreOp> {
524-
using OpRewritePattern::OpRewritePattern;
525-
526-
LogicalResult matchAndRewrite(vector::StoreOp storeOp,
527-
PatternRewriter &rewriter) const override {
528-
auto vecType = storeOp.getVectorType();
529-
if (vecType.getNumElements() != 1)
530-
return rewriter.notifyMatchFailure(storeOp, "not single element vector");
531-
532-
Value extracted;
533-
if (vecType.getRank() == 0) {
534-
// TODO: Unifiy once ExtractOp supports 0-d vectors.
535-
extracted = rewriter.create<vector::ExtractElementOp>(
536-
storeOp.getLoc(), storeOp.getValueToStore());
537-
} else {
538-
SmallVector<int64_t> indices(vecType.getRank(), 0);
539-
extracted = rewriter.create<vector::ExtractOp>(
540-
storeOp.getLoc(), storeOp.getValueToStore(), indices);
541-
}
542-
543-
rewriter.replaceOpWithNewOp<memref::StoreOp>(
544-
storeOp, extracted, storeOp.getBase(), storeOp.getIndices());
545-
return success();
546-
}
547-
};
548-
549495
/// Progressive lowering of transfer_write. This pattern supports lowering of
550496
/// `vector.transfer_write` to `vector.store` if all of the following hold:
551497
/// - Stride of most minor memref dimension must be 1.
@@ -645,7 +591,4 @@ void mlir::vector::populateVectorTransferLoweringPatterns(
645591
patterns.add<TransferReadToVectorLoadLowering,
646592
TransferWriteToVectorStoreLowering>(patterns.getContext(),
647593
maxTransferRank, benefit);
648-
patterns
649-
.add<VectorLoadToMemrefLoadLowering, VectorStoreToMemrefStoreLowering>(
650-
patterns.getContext(), benefit);
651594
}

‎mlir/test/Conversion/GPUCommon/transfer_write.mlir

Lines changed: 8 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,13 +1,15 @@
11
// RUN: mlir-opt %s --gpu-to-llvm | FileCheck %s
22

3-
func.func @warp_extract(%arg0: index, %arg1: memref<1024x1024xf32>, %arg2: index, %arg3: vector<1xf32>) {
3+
// CHECK-LABEL: @warp_extract
4+
// CHECK-SAME: %[[VEC:[a-zA-Z0-9_]+]]: vector<1xf32>
5+
// CHECK:%[[BASE:[0-9]+]] = llvm.extractvalue
6+
// CHECK:%[[PTR:[0-9]+]] = llvm.getelementptr %[[BASE]]
7+
// CHECK:llvm.store %[[VEC]], %[[PTR]] {alignment = 4 : i64} : vector<1xf32>, !llvm.ptr
8+
9+
func.func @warp_extract(%arg0: index, %arg1: memref<1024x1024xf32>, %arg2: vector<1xf32>) {
410
%c0 = arith.constant 0 : index
511
gpu.warp_execute_on_lane_0(%arg0)[32] {
6-
// CHECK:%[[val:[0-9]+]] = llvm.extractelement
7-
// CHECK:%[[base:[0-9]+]] = llvm.extractvalue
8-
// CHECK:%[[ptr:[0-9]+]] = llvm.getelementptr %[[base]]
9-
// CHECK:llvm.store %[[val]], %[[ptr]]
10-
vector.transfer_write %arg3, %arg1[%c0, %c0] {in_bounds = [true]} : vector<1xf32>, memref<1024x1024xf32>
12+
vector.transfer_write %arg2, %arg1[%c0, %c0] {in_bounds = [true]} : vector<1xf32>, memref<1024x1024xf32>
1113
}
1214
return
1315
}

‎mlir/test/Conversion/VectorToLLVM/vector-to-llvm.mlir

Lines changed: 23 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -3292,13 +3292,17 @@ func.func @load_0d(%memref : memref<200x100xf32>, %i : index, %j : index) -> vec
32923292
}
32933293

32943294
// CHECK-LABEL: func @load_0d
3295-
// CHECK: %[[LOAD:.*]] = memref.load %{{.*}}[%{{.*}}, %{{.*}}]
3296-
// CHECK: %[[VEC:.*]] = llvm.mlir.undef : vector<1xf32>
3297-
// CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : i32) : i32
3298-
// CHECK: %[[INSERTED:.*]] = llvm.insertelement %[[LOAD]], %[[VEC]][%[[C0]] : i32] : vector<1xf32>
3299-
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[INSERTED]] : vector<1xf32> to vector<f32>
3300-
// CHECK: return %[[CAST]] : vector<f32>
3301-
3295+
// CHECK: %[[J:.*]] = builtin.unrealized_conversion_cast %{{.*}} : index to i64
3296+
// CHECK: %[[I:.*]] = builtin.unrealized_conversion_cast %{{.*}} : index to i64
3297+
// CHECK: %[[CAST_MEMREF:.*]] = builtin.unrealized_conversion_cast %{{.*}} : memref<200x100xf32> to !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
3298+
// CHECK: %[[REF:.*]] = llvm.extractvalue %[[CAST_MEMREF]][1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
3299+
// CHECK: %[[C100:.*]] = llvm.mlir.constant(100 : index) : i64
3300+
// CHECK: %[[MUL:.*]] = llvm.mul %[[I]], %[[C100]] : i64
3301+
// CHECK: %[[ADD:.*]] = llvm.add %[[MUL]], %[[J]] : i64
3302+
// CHECK: %[[ADDR:.*]] = llvm.getelementptr %[[REF]][%[[ADD]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
3303+
// CHECK: %[[LOAD:.*]] = llvm.load %[[ADDR]] {alignment = 4 : i64} : !llvm.ptr -> vector<1xf32>
3304+
// CHECK: %[[RES:.*]] = builtin.unrealized_conversion_cast %[[LOAD]] : vector<1xf32> to vector<f32>
3305+
// CHECK: return %[[RES]] : vector<f32>
33023306
// -----
33033307

33043308
//===----------------------------------------------------------------------===//
@@ -3392,11 +3396,18 @@ func.func @store_0d(%memref : memref<200x100xf32>, %i : index, %j : index) {
33923396
}
33933397

33943398
// CHECK-LABEL: func @store_0d
3395-
// CHECK: %[[VAL:.*]] = arith.constant dense<1.100000e+01> : vector<f32>
3396-
// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[VAL]] : vector<f32> to vector<1xf32>
3397-
// CHECK: %[[C0:.*]] = llvm.mlir.constant(0 : index) : i64
3398-
// CHECK: %[[EXTRACTED:.*]] = llvm.extractelement %[[CAST]][%[[C0]] : i64] : vector<1xf32>
3399-
// CHECK: memref.store %[[EXTRACTED]], %{{.*}}[%{{.*}}, %{{.*}}]
3399+
// CHECK: %[[J:.*]] = builtin.unrealized_conversion_cast %{{.*}} : index to i64
3400+
// CHECK: %[[I:.*]] = builtin.unrealized_conversion_cast %{{.*}} : index to i64
3401+
// CHECK: %[[CAST_MEMREF:.*]] = builtin.unrealized_conversion_cast %{{.*}} : memref<200x100xf32> to !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
3402+
// CHECK: %[[CST:.*]] = arith.constant dense<1.100000e+01> : vector<f32>
3403+
// CHECK: %[[VAL:.*]] = builtin.unrealized_conversion_cast %[[CST]] : vector<f32> to vector<1xf32>
3404+
// CHECK: %[[REF:.*]] = llvm.extractvalue %[[CAST_MEMREF]][1] : !llvm.struct<(ptr, ptr, i64, array<2 x i64>, array<2 x i64>)>
3405+
// CHECK: %[[C100:.*]] = llvm.mlir.constant(100 : index) : i64
3406+
// CHECK: %[[MUL:.*]] = llvm.mul %[[I]], %[[C100]] : i64
3407+
// CHECK: %[[ADD:.*]] = llvm.add %[[MUL]], %[[J]] : i64
3408+
// CHECK: %[[ADDR:.*]] = llvm.getelementptr %[[REF]][%[[ADD]]] : (!llvm.ptr, i64) -> !llvm.ptr, f32
3409+
// CHECK: llvm.store %[[VAL]], %[[ADDR]] {alignment = 4 : i64} : vector<1xf32>, !llvm.ptr
3410+
// CHECK: return
34003411

34013412
// -----
34023413

‎mlir/test/Dialect/Vector/vector-transfer-to-vector-load-store.mlir

Lines changed: 10 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -6,16 +6,13 @@
66
func.func @vector_transfer_ops_0d_memref(%mem: memref<f32>, %vec: vector<1x1x1xf32>) {
77
%f0 = arith.constant 0.0 : f32
88

9-
// CHECK-NEXT: %[[S:.*]] = memref.load %[[MEM]][] : memref<f32>
10-
// CHECK-NEXT: %[[V:.*]] = vector.broadcast %[[S]] : f32 to vector<f32>
9+
// CHECK-NEXT: %[[S:.*]] = vector.load %[[MEM]][] : memref<f32>, vector<f32>
1110
%0 = vector.transfer_read %mem[], %f0 : memref<f32>, vector<f32>
1211

13-
// CHECK-NEXT: %[[SS:.*]] = vector.extractelement %[[V]][] : vector<f32>
14-
// CHECK-NEXT: memref.store %[[SS]], %[[MEM]][] : memref<f32>
12+
// CHECK-NEXT: vector.store %[[S]], %[[MEM]][] : memref<f32>, vector<f32>
1513
vector.transfer_write %0, %mem[] : vector<f32>, memref<f32>
1614

17-
// CHECK-NEXT: %[[VV:.*]] = vector.extract %arg1[0, 0, 0] : f32 from vector<1x1x1xf32>
18-
// CHECK-NEXT: memref.store %[[VV]], %[[MEM]][] : memref<f32>
15+
// CHECK-NEXT: vector.store %[[VEC]], %[[MEM]][] : memref<f32>, vector<1x1x1xf32>
1916
vector.store %vec, %mem[] : memref<f32>, vector<1x1x1xf32>
2017

2118
return
@@ -191,8 +188,8 @@ func.func @transfer_perm_map(%mem : memref<8x8xf32>, %idx : index) -> vector<4xf
191188
// CHECK-LABEL: func @transfer_broadcasting(
192189
// CHECK-SAME: %[[MEM:.*]]: memref<8x8xf32>,
193190
// CHECK-SAME: %[[IDX:.*]]: index) -> vector<4xf32> {
194-
// CHECK-NEXT: %[[LOAD:.*]] = memref.load %[[MEM]][%[[IDX]], %[[IDX]]] : memref<8x8xf32>
195-
// CHECK-NEXT: %[[RES:.*]] = vector.broadcast %[[LOAD]] : f32 to vector<4xf32>
191+
// CHECK-NEXT: %[[LOAD:.*]] = vector.load %[[MEM]][%[[IDX]], %[[IDX]]] : memref<8x8xf32>, vector<1xf32>
192+
// CHECK-NEXT: %[[RES:.*]] = vector.broadcast %[[LOAD]] : vector<1xf32> to vector<4xf32>
196193
// CHECK-NEXT: return %[[RES]] : vector<4xf32>
197194
// CHECK-NEXT: }
198195

@@ -208,8 +205,7 @@ func.func @transfer_broadcasting(%mem : memref<8x8xf32>, %idx : index) -> vector
208205
// CHECK-LABEL: func @transfer_scalar(
209206
// CHECK-SAME: %[[MEM:.*]]: memref<?x?xf32>,
210207
// CHECK-SAME: %[[IDX:.*]]: index) -> vector<1xf32> {
211-
// CHECK-NEXT: %[[LOAD:.*]] = memref.load %[[MEM]][%[[IDX]], %[[IDX]]] : memref<?x?xf32>
212-
// CHECK-NEXT: %[[RES:.*]] = vector.broadcast %[[LOAD]] : f32 to vector<1xf32>
208+
// CHECK-NEXT: %[[RES:.*]] = vector.load %[[MEM]][%[[IDX]], %[[IDX]]] : memref<?x?xf32>, vector<1xf32>
213209
// CHECK-NEXT: return %[[RES]] : vector<1xf32>
214210
// CHECK-NEXT: }
215211
func.func @transfer_scalar(%mem : memref<?x?xf32>, %idx : index) -> vector<1xf32> {
@@ -222,8 +218,8 @@ func.func @transfer_scalar(%mem : memref<?x?xf32>, %idx : index) -> vector<1xf32
222218
// CHECK-LABEL: func @transfer_broadcasting_2D(
223219
// CHECK-SAME: %[[MEM:.*]]: memref<8x8xf32>,
224220
// CHECK-SAME: %[[IDX:.*]]: index) -> vector<4x4xf32> {
225-
// CHECK-NEXT: %[[LOAD:.*]] = memref.load %[[MEM]][%[[IDX]], %[[IDX]]] : memref<8x8xf32>
226-
// CHECK-NEXT: %[[RES:.*]] = vector.broadcast %[[LOAD]] : f32 to vector<4x4xf32>
221+
// CHECK-NEXT: %[[LOAD:.*]] = vector.load %[[MEM]][%[[IDX]], %[[IDX]]] : memref<8x8xf32>, vector<1x1xf32>
222+
// CHECK-NEXT: %[[RES:.*]] = vector.broadcast %[[LOAD]] : vector<1x1xf32> to vector<4x4xf32>
227223
// CHECK-NEXT: return %[[RES]] : vector<4x4xf32>
228224
// CHECK-NEXT: }
229225

@@ -322,8 +318,8 @@ func.func @transfer_read_permutations(%mem_0 : memref<?x?xf32>, %mem_1 : memref<
322318
// CHECK: vector.transpose %{{.*}}, [2, 1, 3, 0] : vector<16x14x7x8xf32> to vector<7x14x8x16xf32>
323319

324320
%6 = vector.transfer_read %mem_0[%c0, %c0], %cst {in_bounds = [true], permutation_map = #map6} : memref<?x?xf32>, vector<8xf32>
325-
// CHECK: memref.load %{{.*}}[%[[C0]], %[[C0]]] : memref<?x?xf32>
326-
// CHECK: vector.broadcast %{{.*}} : f32 to vector<8xf32>
321+
// CHECK: vector.load %{{.*}}[%[[C0]], %[[C0]]] : memref<?x?xf32>, vector<1xf32>
322+
// CHECK: vector.broadcast %{{.*}} : vector<1xf32> to vector<8xf32>
327323

328324
return %0, %1, %2, %3, %4, %5, %6 : vector<7x14x8x16xf32>, vector<7x14x8x16xf32>,
329325
vector<7x14x8x16xf32>, vector<7x14x8x16xf32>, vector<7x14x8x16xf32>,

0 commit comments

Comments
 (0)