Skip to content

Commit c7c1283

Browse files
krzysz00krzysz00
authored
[mlir][MemRef] Remove integer address space builders (#138579)
The forms of the MemRef builder that took an integer argument instead of an attribute have been deprecated for years now, and have almost no upstream uses (the remaining ones are handled in this PR). Therefore, remove them.
1 parent 6fc092f commit c7c1283

File tree

6 files changed

+44
-100
lines changed

6 files changed

+44
-100
lines changed

‎mlir/include/mlir/IR/BuiltinTypes.td‎

Lines changed: 0 additions & 13 deletions
Original file line numberDiff line numberDiff line change
@@ -807,11 +807,6 @@ def Builtin_MemRef : Builtin_Type<"MemRef", "memref", [
807807
"ArrayRef<int64_t>":$shape, "Type":$elementType,
808808
CArg<"AffineMap">:$map,
809809
CArg<"Attribute", "{}">:$memorySpace)>,
810-
/// [deprecated] `Attribute`-based form should be used instead.
811-
TypeBuilderWithInferredContext<(ins
812-
"ArrayRef<int64_t>":$shape, "Type":$elementType,
813-
"AffineMap":$map,
814-
"unsigned":$memorySpaceInd)>
815810
];
816811
let extraClassDeclaration = [{
817812
using BaseMemRefType::clone;
@@ -1180,14 +1175,6 @@ def Builtin_UnrankedMemRef : Builtin_Type<"UnrankedMemRef", "unranked_memref", [
11801175
Attribute nonDefaultMemorySpace = skipDefaultMemorySpace(memorySpace);
11811176
return $_get(elementType.getContext(), elementType, nonDefaultMemorySpace);
11821177
}]>,
1183-
/// [deprecated] `Attribute`-based form should be used instead.
1184-
TypeBuilderWithInferredContext<(ins "Type":$elementType,
1185-
"unsigned":$memorySpace), [{
1186-
// Convert deprecated integer-like memory space to Attribute.
1187-
Attribute memorySpaceAttr =
1188-
wrapIntegerMemorySpace(memorySpace, elementType.getContext());
1189-
return UnrankedMemRefType::get(elementType, memorySpaceAttr);
1190-
}]>
11911178
];
11921179
let extraClassDeclaration = [{
11931180
using BaseMemRefType::clone;

‎mlir/lib/Dialect/Affine/Utils/LoopUtils.cpp‎

Lines changed: 10 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2097,10 +2097,12 @@ static LogicalResult generateCopy(
20972097
// Check if a buffer was already created.
20982098
bool existingBuf = fastBufferMap.count(memref) > 0;
20992099
if (!existingBuf) {
2100-
AffineMap fastBufferLayout = b.getMultiDimIdentityMap(rank);
2100+
Attribute fastMemorySpace;
2101+
if (copyOptions.fastMemorySpace != 0)
2102+
fastMemorySpace = prologue.getI64IntegerAttr(copyOptions.fastMemorySpace);
21012103
auto fastMemRefType =
21022104
MemRefType::get(fastBufferShape, memRefType.getElementType(),
2103-
fastBufferLayout, copyOptions.fastMemorySpace);
2105+
MemRefLayoutAttrInterface{}, fastMemorySpace);
21042106

21052107
// Create the fast memory space buffer just before the 'affine.for'
21062108
// operation.
@@ -2175,8 +2177,12 @@ static LogicalResult generateCopy(
21752177
} else {
21762178
// DMA generation.
21772179
// Create a tag (single element 1-d memref) for the DMA.
2178-
auto tagMemRefType = MemRefType::get({1}, top.getIntegerType(32), {},
2179-
copyOptions.tagMemorySpace);
2180+
Attribute tagMemorySpace;
2181+
if (copyOptions.tagMemorySpace != 0)
2182+
tagMemorySpace = prologue.getI64IntegerAttr(copyOptions.tagMemorySpace);
2183+
auto tagMemRefType =
2184+
MemRefType::get({1}, top.getIntegerType(32),
2185+
MemRefLayoutAttrInterface{}, tagMemorySpace);
21802186
auto tagMemRef = prologue.create<memref::AllocOp>(loc, tagMemRefType);
21812187

21822188
SmallVector<Value, 4> tagIndices({zeroIndex});

‎mlir/lib/IR/BuiltinTypes.cpp‎

Lines changed: 0 additions & 48 deletions
Original file line numberDiff line numberDiff line change
@@ -484,14 +484,6 @@ bool mlir::detail::isSupportedMemorySpace(Attribute memorySpace) {
484484
return false;
485485
}
486486

487-
Attribute mlir::detail::wrapIntegerMemorySpace(unsigned memorySpace,
488-
MLIRContext *ctx) {
489-
if (memorySpace == 0)
490-
return nullptr;
491-
492-
return IntegerAttr::get(IntegerType::get(ctx, 64), memorySpace);
493-
}
494-
495487
Attribute mlir::detail::skipDefaultMemorySpace(Attribute memorySpace) {
496488
IntegerAttr intMemorySpace = llvm::dyn_cast_or_null<IntegerAttr>(memorySpace);
497489
if (intMemorySpace && intMemorySpace.getValue() == 0)
@@ -583,46 +575,6 @@ MemRefType::getChecked(function_ref<InFlightDiagnostic()> emitErrorFn,
583575
elementType, layout, memorySpace);
584576
}
585577

586-
MemRefType MemRefType::get(ArrayRef<int64_t> shape, Type elementType,
587-
AffineMap map, unsigned memorySpaceInd) {
588-
589-
// Use default layout for empty map.
590-
if (!map)
591-
map = AffineMap::getMultiDimIdentityMap(shape.size(),
592-
elementType.getContext());
593-
594-
// Wrap AffineMap into Attribute.
595-
auto layout = AffineMapAttr::get(map);
596-
597-
// Convert deprecated integer-like memory space to Attribute.
598-
Attribute memorySpace =
599-
wrapIntegerMemorySpace(memorySpaceInd, elementType.getContext());
600-
601-
return Base::get(elementType.getContext(), shape, elementType, layout,
602-
memorySpace);
603-
}
604-
605-
MemRefType
606-
MemRefType::getChecked(function_ref<InFlightDiagnostic()> emitErrorFn,
607-
ArrayRef<int64_t> shape, Type elementType, AffineMap map,
608-
unsigned memorySpaceInd) {
609-
610-
// Use default layout for empty map.
611-
if (!map)
612-
map = AffineMap::getMultiDimIdentityMap(shape.size(),
613-
elementType.getContext());
614-
615-
// Wrap AffineMap into Attribute.
616-
auto layout = AffineMapAttr::get(map);
617-
618-
// Convert deprecated integer-like memory space to Attribute.
619-
Attribute memorySpace =
620-
wrapIntegerMemorySpace(memorySpaceInd, elementType.getContext());
621-
622-
return Base::getChecked(emitErrorFn, elementType.getContext(), shape,
623-
elementType, layout, memorySpace);
624-
}
625-
626578
LogicalResult MemRefType::verify(function_ref<InFlightDiagnostic()> emitError,
627579
ArrayRef<int64_t> shape, Type elementType,
628580
MemRefLayoutAttrInterface layout,

‎mlir/lib/IR/TypeDetail.h‎

Lines changed: 0 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -140,9 +140,6 @@ struct TupleTypeStorage final
140140
/// Checks if the memorySpace has supported Attribute type.
141141
bool isSupportedMemorySpace(Attribute memorySpace);
142142

143-
/// Wraps deprecated integer memory space to the new Attribute form.
144-
Attribute wrapIntegerMemorySpace(unsigned memorySpace, MLIRContext *ctx);
145-
146143
/// Replaces default memorySpace (integer == `0`) with empty Attribute.
147144
Attribute skipDefaultMemorySpace(Attribute memorySpace);
148145

‎mlir/test/Dialect/Vector/vector-warp-distribute.mlir‎

Lines changed: 28 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -19,10 +19,10 @@
1919
// CHECK-SCF-IF-DAG: #[[$TIMES2:.*]] = affine_map<()[s0] -> (s0 * 2)>
2020
// CHECK-SCF-IF-DAG: #[[$TIMES4:.*]] = affine_map<()[s0] -> (s0 * 4)>
2121
// CHECK-SCF-IF-DAG: #[[$TIMES8:.*]] = affine_map<()[s0] -> (s0 * 8)>
22-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_32xf32 : memref<32xf32, 3>
23-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_64xf32 : memref<64xf32, 3>
24-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_128xf32 : memref<128xf32, 3>
25-
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_256xf32 : memref<256xf32, 3>
22+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_32xf32 : memref<32xf32, #gpu.address_space<workgroup>>
23+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_64xf32 : memref<64xf32, #gpu.address_space<workgroup>>
24+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_128xf32 : memref<128xf32, #gpu.address_space<workgroup>>
25+
// CHECK-SCF-IF-DAG: memref.global "private" @__shared_256xf32 : memref<256xf32, #gpu.address_space<workgroup>>
2626

2727
// CHECK-SCF-IF-LABEL: func @rewrite_warp_op_to_scf_if(
2828
// CHECK-SCF-IF-SAME: %[[laneid:.*]]: index,
@@ -47,8 +47,8 @@ func.func @rewrite_warp_op_to_scf_if(%laneid: index,
4747
%r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]
4848
args(%v0, %v1 : vector<4xf32>, vector<8xf32>) -> (vector<1xf32>, vector<2xf32>) {
4949
^bb0(%arg0: vector<128xf32>, %arg1: vector<256xf32>):
50-
// CHECK-SCF-IF: %[[arg1:.*]] = vector.transfer_read %[[buffer_v1]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<256xf32, 3>, vector<256xf32>
51-
// CHECK-SCF-IF: %[[arg0:.*]] = vector.transfer_read %[[buffer_v0]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<128xf32, 3>, vector<128xf32>
50+
// CHECK-SCF-IF: %[[arg1:.*]] = vector.transfer_read %[[buffer_v1]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<256xf32, #gpu.address_space<workgroup>>, vector<256xf32>
51+
// CHECK-SCF-IF: %[[arg0:.*]] = vector.transfer_read %[[buffer_v0]][%[[c0]]], %{{.*}} {in_bounds = [true]} : memref<128xf32, #gpu.address_space<workgroup>>, vector<128xf32>
5252
// CHECK-SCF-IF: %[[def_0:.*]] = "some_def"(%[[arg0]]) : (vector<128xf32>) -> vector<32xf32>
5353
// CHECK-SCF-IF: %[[def_1:.*]] = "some_def"(%[[arg1]]) : (vector<256xf32>) -> vector<64xf32>
5454
%2 = "some_def"(%arg0) : (vector<128xf32>) -> vector<32xf32>
@@ -60,8 +60,8 @@ func.func @rewrite_warp_op_to_scf_if(%laneid: index,
6060
// CHECK-SCF-IF: }
6161
// CHECK-SCF-IF: gpu.barrier
6262
// CHECK-SCF-IF: %[[o1:.*]] = affine.apply #[[$TIMES2]]()[%[[laneid]]]
63-
// CHECK-SCF-IF: %[[r1:.*]] = vector.transfer_read %[[buffer_def_1]][%[[o1]]], %{{.*}} {in_bounds = [true]} : memref<64xf32, 3>, vector<2xf32>
64-
// CHECK-SCF-IF: %[[r0:.*]] = vector.transfer_read %[[buffer_def_0]][%[[laneid]]], %{{.*}} {in_bounds = [true]} : memref<32xf32, 3>, vector<1xf32>
63+
// CHECK-SCF-IF: %[[r1:.*]] = vector.transfer_read %[[buffer_def_1]][%[[o1]]], %{{.*}} {in_bounds = [true]} : memref<64xf32, #gpu.address_space<workgroup>>, vector<2xf32>
64+
// CHECK-SCF-IF: %[[r0:.*]] = vector.transfer_read %[[buffer_def_0]][%[[laneid]]], %{{.*}} {in_bounds = [true]} : memref<32xf32, #gpu.address_space<workgroup>>, vector<1xf32>
6565
// CHECK-SCF-IF: "some_use"(%[[r0]]) : (vector<1xf32>) -> ()
6666
// CHECK-SCF-IF: "some_use"(%[[r1]]) : (vector<2xf32>) -> ()
6767
"some_use"(%r#0) : (vector<1xf32>) -> ()
@@ -1065,18 +1065,18 @@ func.func @warp_execute_has_broadcast_semantics(%laneid: index, %s0: f32, %v0: v
10651065
args(%s0, %v0, %v1, %v2 : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>) -> (f32, vector<f32>, vector<1xf32>, vector<1x1xf32>) {
10661066
^bb0(%bs0: f32, %bv0: vector<f32>, %bv1: vector<1xf32>, %bv2: vector<1x1xf32>):
10671067

1068-
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, 3>, vector<1x1xf32>
1069-
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, 3>, vector<1xf32>
1070-
// CHECK-SCF-IF: vector.transfer_read {{.*}}[]{{.*}} : memref<f32, 3>, vector<f32>
1071-
// CHECK-SCF-IF: memref.load {{.*}}[%[[C0]]] : memref<1xf32, 3>
1068+
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, #gpu.address_space<workgroup>>, vector<1x1xf32>
1069+
// CHECK-SCF-IF: vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, #gpu.address_space<workgroup>>, vector<1xf32>
1070+
// CHECK-SCF-IF: vector.transfer_read {{.*}}[]{{.*}} : memref<f32, #gpu.address_space<workgroup>>, vector<f32>
1071+
// CHECK-SCF-IF: memref.load {{.*}}[%[[C0]]] : memref<1xf32, #gpu.address_space<workgroup>>
10721072
// CHECK-SCF-IF: "some_def_0"(%{{.*}}) : (f32) -> f32
10731073
// CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<f32>) -> vector<f32>
10741074
// CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<1xf32>) -> vector<1xf32>
10751075
// CHECK-SCF-IF: "some_def_1"(%{{.*}}) : (vector<1x1xf32>) -> vector<1x1xf32>
1076-
// CHECK-SCF-IF: memref.store {{.*}}[%[[C0]]] : memref<1xf32, 3>
1077-
// CHECK-SCF-IF: vector.transfer_write {{.*}}[] : vector<f32>, memref<f32, 3>
1078-
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]]] {in_bounds = [true]} : vector<1xf32>, memref<1xf32, 3>
1079-
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x1xf32>, memref<1x1xf32, 3>
1076+
// CHECK-SCF-IF: memref.store {{.*}}[%[[C0]]] : memref<1xf32, #gpu.address_space<workgroup>>
1077+
// CHECK-SCF-IF: vector.transfer_write {{.*}}[] : vector<f32>, memref<f32, #gpu.address_space<workgroup>>
1078+
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]]] {in_bounds = [true]} : vector<1xf32>, memref<1xf32, #gpu.address_space<workgroup>>
1079+
// CHECK-SCF-IF: vector.transfer_write {{.*}}[%[[C0]], %[[C0]]] {in_bounds = [true, true]} : vector<1x1xf32>, memref<1x1xf32, #gpu.address_space<workgroup>>
10801080

10811081
%rs0 = "some_def_0"(%bs0) : (f32) -> f32
10821082
%rv0 = "some_def_1"(%bv0) : (vector<f32>) -> vector<f32>
@@ -1088,10 +1088,10 @@ func.func @warp_execute_has_broadcast_semantics(%laneid: index, %s0: f32, %v0: v
10881088
}
10891089

10901090
// CHECK-SCF-IF: gpu.barrier
1091-
// CHECK-SCF-IF: %[[RV2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, 3>, vector<1x1xf32>
1092-
// CHECK-SCF-IF: %[[RV1:.*]] = vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, 3>, vector<1xf32>
1093-
// CHECK-SCF-IF: %[[RV0:.*]] = vector.transfer_read {{.*}}[]{{.*}} : memref<f32, 3>, vector<f32>
1094-
// CHECK-SCF-IF: %[[RS0:.*]] = memref.load {{.*}}[%[[C0]]] : memref<1xf32, 3>
1091+
// CHECK-SCF-IF: %[[RV2:.*]] = vector.transfer_read {{.*}}[%[[C0]], %[[C0]]]{{.*}} {in_bounds = [true, true]} : memref<1x1xf32, #gpu.address_space<workgroup>>, vector<1x1xf32>
1092+
// CHECK-SCF-IF: %[[RV1:.*]] = vector.transfer_read {{.*}}[%[[C0]]]{{.*}} {in_bounds = [true]} : memref<1xf32, #gpu.address_space<workgroup>>, vector<1xf32>
1093+
// CHECK-SCF-IF: %[[RV0:.*]] = vector.transfer_read {{.*}}[]{{.*}} : memref<f32, #gpu.address_space<workgroup>>, vector<f32>
1094+
// CHECK-SCF-IF: %[[RS0:.*]] = memref.load {{.*}}[%[[C0]]] : memref<1xf32, #gpu.address_space<workgroup>>
10951095
// CHECK-SCF-IF: return %[[RS0]], %[[RV0]], %[[RV1]], %[[RV2]] : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>
10961096
return %r#0, %r#1, %r#2, %r#3 : f32, vector<f32>, vector<1xf32>, vector<1x1xf32>
10971097
}
@@ -1106,22 +1106,22 @@ func.func @warp_execute_nd_distribute(%laneid: index, %v0: vector<1x64x1xf32>, %
11061106
-> (vector<1x64x1xf32>, vector<1x2x128xf32>) {
11071107
// CHECK-SCF-IF-DAG: %[[C0:.*]] = arith.constant 0 : index
11081108

1109-
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[LANEID]], %c0, %c0] {in_bounds = [true, true, true]} : vector<1x64x1xf32>, memref<32x64x1xf32, 3>
1109+
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[LANEID]], %c0, %c0] {in_bounds = [true, true, true]} : vector<1x64x1xf32>, memref<32x64x1xf32, #gpu.address_space<workgroup>>
11101110
// CHECK-SCF-IF: %[[RID:.*]] = affine.apply #[[$TIMES2]]()[%[[LANEID]]]
1111-
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[C0]], %[[RID]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x2x128xf32>, memref<1x64x128xf32, 3>
1111+
// CHECK-SCF-IF: vector.transfer_write %{{.*}}, %{{.*}}[%[[C0]], %[[RID]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x2x128xf32>, memref<1x64x128xf32, #gpu.address_space<workgroup>>
11121112
// CHECK-SCF-IF: gpu.barrier
11131113

11141114
// CHECK-SCF-IF: scf.if{{.*}}{
11151115
%r:2 = gpu.warp_execute_on_lane_0(%laneid)[32]
11161116
args(%v0, %v1 : vector<1x64x1xf32>, vector<1x2x128xf32>) -> (vector<1x64x1xf32>, vector<1x2x128xf32>) {
11171117
^bb0(%arg0: vector<32x64x1xf32>, %arg1: vector<1x64x128xf32>):
11181118

1119-
// CHECK-SCF-IF-DAG: %[[SR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<32x64x1xf32, 3>, vector<32x64x1xf32>
1120-
// CHECK-SCF-IF-DAG: %[[SR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<1x64x128xf32, 3>, vector<1x64x128xf32>
1119+
// CHECK-SCF-IF-DAG: %[[SR0:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<32x64x1xf32, #gpu.address_space<workgroup>>, vector<32x64x1xf32>
1120+
// CHECK-SCF-IF-DAG: %[[SR1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %{{.*}} {in_bounds = [true, true, true]} : memref<1x64x128xf32, #gpu.address_space<workgroup>>, vector<1x64x128xf32>
11211121
// CHECK-SCF-IF: %[[W0:.*]] = "some_def_0"(%[[SR0]]) : (vector<32x64x1xf32>) -> vector<32x64x1xf32>
11221122
// CHECK-SCF-IF: %[[W1:.*]] = "some_def_1"(%[[SR1]]) : (vector<1x64x128xf32>) -> vector<1x64x128xf32>
1123-
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<32x64x1xf32>, memref<32x64x1xf32, 3>
1124-
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W1]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x64x128xf32>, memref<1x64x128xf32, 3>
1123+
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W0]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<32x64x1xf32>, memref<32x64x1xf32, #gpu.address_space<workgroup>>
1124+
// CHECK-SCF-IF-DAG: vector.transfer_write %[[W1]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true, true, true]} : vector<1x64x128xf32>, memref<1x64x128xf32, #gpu.address_space<workgroup>>
11251125

11261126
%r0 = "some_def_0"(%arg0) : (vector<32x64x1xf32>) -> vector<32x64x1xf32>
11271127
%r1 = "some_def_1"(%arg1) : (vector<1x64x128xf32>) -> vector<1x64x128xf32>
@@ -1132,8 +1132,8 @@ func.func @warp_execute_nd_distribute(%laneid: index, %v0: vector<1x64x1xf32>, %
11321132

11331133
// CHECK-SCF-IF: gpu.barrier
11341134
// CHECK-SCF-IF: %[[WID:.*]] = affine.apply #[[$TIMES2]]()[%[[LANEID]]]
1135-
// CHECK-SCF-IF-DAG: %[[R0:.*]] = vector.transfer_read %{{.*}}[%[[LANEID]], %[[C0]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<32x64x1xf32, 3>, vector<1x64x1xf32>
1136-
// CHECK-SCF-IF-DAG: %[[R1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[WID]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<1x64x128xf32, 3>, vector<1x2x128xf32>
1135+
// CHECK-SCF-IF-DAG: %[[R0:.*]] = vector.transfer_read %{{.*}}[%[[LANEID]], %[[C0]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<32x64x1xf32, #gpu.address_space<workgroup>>, vector<1x64x1xf32>
1136+
// CHECK-SCF-IF-DAG: %[[R1:.*]] = vector.transfer_read %{{.*}}[%[[C0]], %[[WID]], %[[C0]]], %cst {in_bounds = [true, true, true]} : memref<1x64x128xf32, #gpu.address_space<workgroup>>, vector<1x2x128xf32>
11371137
// CHECK-SCF-IF: return %[[R0]], %[[R1]] : vector<1x64x1xf32>, vector<1x2x128xf32>
11381138
return %r#0, %r#1 : vector<1x64x1xf32>, vector<1x2x128xf32>
11391139
}

‎mlir/test/lib/Dialect/Vector/TestVectorTransforms.cpp‎

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -518,15 +518,17 @@ struct TestVectorScanLowering
518518
static Value allocateGlobalSharedMemory(Location loc, OpBuilder &builder,
519519
gpu::WarpExecuteOnLane0Op warpOp,
520520
Type type) {
521-
static constexpr int64_t kSharedMemorySpace = 3;
521+
Attribute sharedMemorySpaceAttr =
522+
builder.getAttr<gpu::AddressSpaceAttr>(gpu::AddressSpace::Workgroup);
522523
// Compute type of shared memory buffer.
523524
MemRefType memrefType;
524525
if (auto vectorType = dyn_cast<VectorType>(type)) {
525526
memrefType =
526-
MemRefType::get(vectorType.getShape(), vectorType.getElementType(), {},
527-
kSharedMemorySpace);
527+
MemRefType::get(vectorType.getShape(), vectorType.getElementType(),
528+
MemRefLayoutAttrInterface{}, sharedMemorySpaceAttr);
528529
} else {
529-
memrefType = MemRefType::get({1}, type, {}, kSharedMemorySpace);
530+
memrefType = MemRefType::get({1}, type, MemRefLayoutAttrInterface{},
531+
sharedMemorySpaceAttr);
530532
}
531533

532534
// Get symbol table holding all shared memory globals.

0 commit comments

Comments
 (0)